benchmark_weight_loading

Benchmark for model weight loading speed.

Measures the time to load model weights from safetensors files using different strategies (CPU vs GPU, broadcast vs independent).

Usage (single GPU): python fastvideo/models/loader/benchmarks/benchmark_weight_loading.py --model-path /path/to/model

Usage (multi-GPU, e.g. 4 GPUs): torchrun --nproc_per_node=4 fastvideo/models/loader/benchmarks/benchmark_weight_loading.py --model-path /path/to/model

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.cleanup_dist_env_and_memory`

cleanup_dist_env_and_memory(shutdown_ray: bool = False)

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.get_node_group`

get_node_group() -> GroupCoordinator

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.maybe_init_distributed_environment_and_model_parallel`

maybe_init_distributed_environment_and_model_parallel(tp_size: int, sp_size: int, distributed_init_method: str = 'env://', timeout: timedelta | None = None, sp_group_ranks: list[list[int]] | None = None, dp_group_ranks: list[list[int]] | None = None)

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.init_logger`

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.resolve_safetensors_files`

resolve_safetensors_files(model_path: str) -> list[str]

Discover safetensors files in a model directory.

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.safetensors_weights_iterator`

safetensors_weights_iterator(hf_weights_files: list[str], to_cpu: bool = False, broadcast: bool = True, async_broadcast: bool = False) -> Generator[tuple[str, torch.Tensor], None, None]

Iterate over the weights in the model safetensor files. Args: hf_weights_files: List of safetensor files to load. to_cpu: Whether to load the weights to CPU. If False, will load to the GPU device bound to the current process. broadcast: Whether node-group rank 0 should read GPU weights and broadcast them to the other ranks of the node. async_broadcast: Whether to overlap loading from disk and broadcasting to other ranks. If True, must iterate over all the weights before use. Only used when broadcast is True and to_cpu is False.

`fastvideo.models.loader.benchmarks.benchmark_weight_loading.benchmark_loading`

benchmark_loading(files: list[str], to_cpu: bool, broadcast: bool, warmup: int, repeats: int, label: str) -> None

Run the weight loading benchmark and print results.

Source

def benchmark_loading(
    files: list[str],
    to_cpu: bool,
    broadcast: bool,
    warmup: int,
    repeats: int,
    label: str,
) -> None:
    """Run the weight loading benchmark and print results."""
    rank = dist.get_rank() if dist.is_initialized() else 0
    node_group = get_node_group()
# Count total params and bytes on first pass
total_params = 0
total_bytes = 0
for name, tensor in safetensors_weights_iterator(files, to_cpu=to_cpu, broadcast=broadcast):
    total_params += 1
    total_bytes += tensor.nelement() * tensor.element_size()

if rank == 0:
    logger.info(&quot;&#91;%s&#93; %d tensors, %.2f GB total&quot;, label, total_params, total_bytes / 1e9)

# Warmup
for _ in range(warmup):
    for _ in safetensors_weights_iterator(files, to_cpu=to_cpu, broadcast=broadcast):
        pass
    if dist.is_initialized():
        dist.barrier()

# Timed runs
times = &#91;&#93;
for i in range(repeats):
    if dist.is_initialized():
        dist.barrier()
    torch.cuda.synchronize() if torch.cuda.is_available() else None

    t0 = time.perf_counter()
    for _ in safetensors_weights_iterator(files, to_cpu=to_cpu, broadcast=broadcast):
        pass
    torch.cuda.synchronize() if torch.cuda.is_available() else None

    if dist.is_initialized():
        dist.barrier()
    elapsed = time.perf_counter() - t0
    times.append(elapsed)

if rank == 0:
    avg = sum(times) / len(times)
    best = min(times)
    throughput = total_bytes / avg / 1e9
    logger.info(
        &quot;&#91;%s&#93; avg %.3fs | best %.3fs | throughput %.2f GB/s &quot;
        &quot;(over %d runs, %d warmup, %d GPU(s), node_size=%d)&quot;,
        label,
        avg,
        best,
        throughput,
        repeats,
        warmup,
        node_group.world_size,
        node_group.world_size,
    )</code></pre></details>
fastvideo.models.loader.benchmarks.benchmark_weight_loading.main
main() -> None
Source
def main() -> None:
    parser = argparse.ArgumentParser(description="Benchmark model weight loading speed")
    parser.add_argument("--model-path",
                        type=str,
                        required=True,
                        help="Path to a local model directory containing .safetensors files")
    parser.add_argument("--subfolder",
                        type=str,
                        default=None,
                        help="Subfolder within model-path to load (e.g. 'transformer', "
                        "'text_encoder'). If not set, loads from model-path directly.")
    parser.add_argument("--warmup", type=int, default=1, help="Number of warmup iterations (default: 1)")
    parser.add_argument("--repeats", type=int, default=3, help="Number of timed iterations (default: 3)")
    args = parser.parse_args()
load_path = args.model_path
if args.subfolder:
    load_path = os.path.join(load_path, args.subfolder)

files = resolve_safetensors_files(load_path)

# Set default env vars for single-GPU standalone mode (no torchrun)
envs.setdefault_external(&quot;MASTER_ADDR&quot;, &quot;localhost&quot;)
envs.setdefault_external(&quot;MASTER_PORT&quot;, &quot;29500&quot;)
envs.setdefault_external(&quot;RANK&quot;, &quot;0&quot;)
envs.setdefault_external(&quot;WORLD_SIZE&quot;, &quot;1&quot;)

maybe_init_distributed_environment_and_model_parallel(1, 1)
rank = dist.get_rank() if dist.is_initialized() else 0
world_size = dist.get_world_size() if dist.is_initialized() else 1

if rank == 0:
    logger.info(&quot;Model path: %s&quot;, load_path)
    logger.info(&quot;Safetensors files: %d&quot;, len(files))
    logger.info(&quot;World size: %d&quot;, world_size)

# Benchmark: load to CPU (no broadcast, each rank reads independently)
benchmark_loading(files,
                  to_cpu=True,
                  broadcast=False,
                  warmup=args.warmup,
                  repeats=args.repeats,
                  label=&quot;to_cpu=True&quot;)

# Benchmark: load to GPU (rank 0 reads, broadcasts to others)
if torch.cuda.is_available():
    benchmark_loading(files,
                      to_cpu=False,
                      broadcast=True,
                      warmup=args.warmup,
                      repeats=args.repeats,
                      label=&quot;to_cpu=False (broadcast)&quot;)

cleanup_dist_env_and_memory()</code></pre></details>