benchmark_process_group() — pytorch Function Reference

Architecture documentation for the benchmark_process_group() function in benchmark.py from the pytorch codebase.

Function python CoreTensor Boxing calls 3 called by 1

Entity Profile

CoreTensor→ Boxing→ benchmark_process_group() — pytorch Function Reference

Dependency Diagram

graph TD
  320a5772_c0de_3469_ef0c_57110d2440fd["benchmark_process_group()"]
  5a01be28_6429_ebf7_fc47_02232b445a1f["run_benchmark()"]
  5a01be28_6429_ebf7_fc47_02232b445a1f -->|calls| 320a5772_c0de_3469_ef0c_57110d2440fd
  72e5e449_391d_8039_3669_200d59371962["create_model()"]
  320a5772_c0de_3469_ef0c_57110d2440fd -->|calls| 72e5e449_391d_8039_3669_200d59371962
  9f38f69b_2e5e_c7b6_7ca9_8d3405e67d05["generate_inputs()"]
  320a5772_c0de_3469_ef0c_57110d2440fd -->|calls| 9f38f69b_2e5e_c7b6_7ca9_8d3405e67d05
  a4e63779_c312_ebbe_ab8e_11bfec3beef4["generate_target()"]
  320a5772_c0de_3469_ef0c_57110d2440fd -->|calls| a4e63779_c312_ebbe_ab8e_11bfec3beef4
  style 320a5772_c0de_3469_ef0c_57110d2440fd fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/distributed/ddp/benchmark.py lines 53–83

def benchmark_process_group(pg, benchmark, use_ddp_for_single_rank=True):
    torch.manual_seed(pg.rank())
    torch.cuda.manual_seed(pg.rank())

    model = benchmark.create_model()
    data = [(benchmark.generate_inputs(), benchmark.generate_target())]
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), 0.001, momentum=0.9, weight_decay=1e-4)
    if use_ddp_for_single_rank or pg.size() > 1:
        model = torch.nn.parallel.DistributedDataParallel(
            model,
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False,
            process_group=pg,
            bucket_cap_mb=benchmark.bucket_size,
        )

    measurements = []
    warmup_iterations = 5
    measured_iterations = 10
    for inputs, target in data * (warmup_iterations + measured_iterations):
        start = time.time()
        output = model(*inputs)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        torch.cuda.synchronize()
        measurements.append(time.time() - start)

    # Throw away measurements for warmup iterations
    return measurements[warmup_iterations:]