--- title: 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("[%s] %d tensors, %.2f GB total", 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 = []
    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(
            "[%s] avg %.3fs | best %.3fs | throughput %.2f GB/s "
            "(over %d runs, %d warmup, %d GPU(s), node_size=%d)",
            label,
            avg,
            best,
            throughput,
            repeats,
            warmup,
            node_group.world_size,
            node_group.world_size,
        )

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("MASTER_ADDR", "localhost")
    envs.setdefault_external("MASTER_PORT", "29500")
    envs.setdefault_external("RANK", "0")
    envs.setdefault_external("WORLD_SIZE", "1")

    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("Model path: %s", load_path)
        logger.info("Safetensors files: %d", len(files))
        logger.info("World size: %d", 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="to_cpu=True")

    # 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="to_cpu=False (broadcast)")

    cleanup_dist_env_and_memory()