def load_independent(files: list[str], device: str):
    """Before-PR behavior: every rank reads every tensor from disk to GPU."""
    for st_file in files:
        with safe_open(st_file, framework="pt", device=device) as f:
            for name in f.keys():  # noqa: SIM118
                param = f.get_tensor(name)
                yield name, param

def load_broadcast(files: list[str], device: str, node_group, async_op: bool = False):
    """After-PR behavior: rank 0 reads from disk, broadcasts to other ranks."""
    # The broadcast source is node-group rank 0, which is not always on CUDA device 0.
    is_leader = node_group.rank_in_group == 0
    handles = []
    for st_file in files:
        with safe_open(st_file, framework="pt", device=device) as f:
            for name in f.keys():  # noqa: SIM118
                if is_leader:
                    param = f.get_tensor(name)
                else:
                    sl = f.get_slice(name)
                    shape = sl.get_shape()
                    dtype = SAFETENSORS_TO_TORCH_DTYPE[sl.get_dtype()]
                    param = torch.empty(shape, device=device, dtype=dtype)
                if node_group.world_size > 1:
                    group = node_group.device_group
                    if async_op:
                        handle = dist.broadcast(
                            param,
                            src=dist.get_global_rank(group, 0),
                            async_op=True,
                            group=group,
                        )
                        handles.append(handle)
                    else:
                        dist.broadcast(
                            param,
                            src=dist.get_global_rank(group, 0),
                            group=group,
                        )
                yield name, param
        if async_op:
            for handle in handles:
                handle.wait()
            handles.clear()

def run_benchmark(label, loader_fn, files, warmup, repeats):
    rank = dist.get_rank()
# Count tensors/bytes on first pass
total_params = 0
total_bytes = 0
for name, tensor in loader_fn(files):
    total_params += 1
    total_bytes += tensor.nelement() * tensor.element_size()
dist.barrier()

# Warmup
for _ in range(warmup):
    for _ in loader_fn(files):
        pass
    dist.barrier()

# Timed runs
times = []
for _ in range(repeats):
    dist.barrier()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in loader_fn(files):
        pass
    torch.cuda.synchronize()
    dist.barrier()
    times.append(time.perf_counter() - t0)

if rank == 0:
    avg = sum(times) / len(times)
    best = min(times)
    throughput = total_bytes / avg / 1e9
    print(f"  {label:30s}  "
          f"avg {avg:.3f}s  best {best:.3f}s  "
          f"throughput {throughput:.2f} GB/s  "
          f"({total_params} tensors, {total_bytes/1e9:.2f} GB)")
    return avg, best, throughput, total_bytes
return None, None, None, None</code></pre></details>
fastvideo.models.loader.benchmarks.benchmark_weight_loading_comparison.main
main()
Source
def main():
    parser = argparse.ArgumentParser(description="A/B benchmark: independent vs broadcast weight loading")
    parser.add_argument("--model-path", type=str, required=True)
    parser.add_argument("--subfolder", type=str, default=None)
    parser.add_argument("--warmup", type=int, default=1)
    parser.add_argument("--repeats", type=int, 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)

maybe_init_distributed_environment_and_model_parallel(1, 1)
rank = dist.get_rank()
world_size = dist.get_world_size()
node_group = get_node_group()
local_rank = node_group.local_rank
device = f&quot;cuda:&#123;local_rank&#125;&quot;

if rank == 0:
    print(f&quot;\n&#123;&#x27;=&#x27;*72&#125;&quot;)
    print(f&quot;Weight Loading Benchmark: &#123;load_path&#125;&quot;)
    print(f&quot;GPUs: &#123;world_size&#125;  |  safetensors files: &#123;len(files)&#125;  &quot;
          f&quot;|  warmup: &#123;args.warmup&#125;  |  repeats: &#123;args.repeats&#125;&quot;)
    print(f&quot;&#123;&#x27;=&#x27;*72&#125;&quot;)

# --- Before PR: every rank reads independently ---
before_avg, before_best, before_tp, total_bytes = run_benchmark(
    &quot;before (independent read)&quot;,
    lambda f: load_independent(f, device),
    files,
    args.warmup,
    args.repeats,
)

# Clear GPU memory between benchmarks
torch.cuda.empty_cache()
dist.barrier()

# --- After PR: rank 0 reads, sync broadcasts ---
after_sync_avg, after_sync_best, after_sync_tp, _ = run_benchmark(
    &quot;after  (sync broadcast)&quot;,
    lambda f: load_broadcast(f, device, node_group, async_op=False),
    files,
    args.warmup,
    args.repeats,
)

# Clear GPU memory between benchmarks
torch.cuda.empty_cache()
dist.barrier()

# --- After PR: rank 0 reads, async broadcasts ---
after_async_avg, after_async_best, after_async_tp, _ = run_benchmark(
    &quot;after  (async broadcast)&quot;,
    lambda f: load_broadcast(f, device, node_group, async_op=True),
    files,
    args.warmup,
    args.repeats,
)

if rank == 0:
    print(f&quot;&#123;&#x27;─&#x27;*72&#125;&quot;)
    if before_avg and after_sync_avg:
        sync_speedup = before_avg / after_sync_avg
        print(f&quot;  sync  speedup: &#123;sync_speedup:.2f&#125;x  &quot;
              f&quot;(saved &#123;before_avg - after_sync_avg:.3f&#125;s)&quot;)
    if before_avg and after_async_avg:
        async_speedup = before_avg / after_async_avg
        print(f&quot;  async speedup: &#123;async_speedup:.2f&#125;x  &quot;
              f&quot;(saved &#123;before_avg - after_async_avg:.3f&#125;s)&quot;)
    print(f&quot;&#123;&#x27;=&#x27;*72&#125;\n&quot;)

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