hao-ai-lab-FastVideo/fastvideo/models/loader/benchmarks

Weight Loading Benchmarks

Benchmarks for measuring model weight loading speed from safetensors files.

Scripts

benchmark_weight_loading.py

Measures loading throughput for two modes:

• to_cpu=True, broadcast=False: load weights to CPU (memory-mapped, no broadcast)
• to_cpu=False, broadcast=True: load weights to GPU (rank 0 reads from disk, broadcasts to other ranks)

# Single GPU (no torchrun needed)
python fastvideo/models/loader/benchmarks/benchmark_weight_loading.py \
    --model-path /path/to/model --subfolder transformer

# Multi-GPU
torchrun --nproc_per_node=4 \
    fastvideo/models/loader/benchmarks/benchmark_weight_loading.py \
    --model-path /path/to/model --subfolder transformer

benchmark_weight_loading_comparison.py

A/B comparison of independent read vs rank-0 broadcast, with both sync and async broadcast variants. Use this to measure the speedup from distributed weight loading.

torchrun --nproc_per_node=4 \
    fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py \
    --model-path /path/to/model --subfolder transformer

Options

Flag	Default	Description
--model-path	required	Local directory containing .safetensors files
--subfolder	none	Subfolder within model-path (e.g. transformer, text_encoder)
--warmup	1	Warmup iterations before timing
--repeats	3	Timed iterations

Example results

HunyuanVideo transformer (25.64 GB) on NVIDIA B200:

GPUs	Independent read	Sync broadcast	Speedup
1	4.12s (6.22 GB/s)	4.12s (6.22 GB/s)	1.00x
2	4.32s (5.94 GB/s)	4.32s (5.94 GB/s)	1.00x
4	5.62s (4.56 GB/s)	4.54s (5.65 GB/s)	1.24x

The broadcast approach eliminates disk I/O contention as GPU count increases.