Data Parallelism › Code: Scaling Analysis

# Model parameters affect gradient size
param_configs = {
    "125M (GPT-2 small)": 125e6,
    "1.3B (GPT-2 XL equiv)": 1.3e9,
    "7B (Llama-7B equiv)": 7e9,
}

# Hardware parameters
nvlink_bandwidth_gb = 600  # NVLink bandwidth per GPU in GB/s (intranode)
infiniband_bandwidth_gb = (
    25  # InfiniBand effective bandwidth per GPU in GB/s (internode)
)
nvlink_step_latency_ms = 0.02
infiniband_step_latency_ms = 0.05
bytes_per_param_bf16 = 2  # BF16 gradients

# Illustrative compute times chosen to expose the compute/communication balance
compute_times_ms = {
    "125M (GPT-2 small)": 3.0,
    "1.3B (GPT-2 XL equiv)": 50.0,
    "7B (Llama-7B equiv)": 400.0,
}


def ring_allreduce_time_ms(num_params, bandwidth_gb_per_s, N, step_latency_ms):
    """Estimate ring all-reduce time in milliseconds."""
    gradient_bytes = num_params * bytes_per_param_bf16
    bandwidth_bytes_per_ms = bandwidth_gb_per_s * 1e9 / 1e3
    # Ring all-reduce: send+receive ~ 2 * (N-1)/N * gradient_bytes
    comm_bytes = 2 * ((N - 1) / N) * gradient_bytes
    bandwidth_time_ms = comm_bytes / bandwidth_bytes_per_ms
    latency_time_ms = 2 * (N - 1) * step_latency_ms
    return bandwidth_time_ms + latency_time_ms


def compute_speedup(
    N_values, num_params, compute_ms, bandwidth_gb, step_latency_ms
):
    """Compute actual vs ideal speedup for a range of GPU counts."""
    speedups = []
    efficiencies = []
    for N in N_values:
        comm_ms = ring_allreduce_time_ms(
            num_params, bandwidth_gb, N, step_latency_ms
        )
        # Step time = compute (overlapped with comm) = max(compute, comm)
        # Plus a small serial fraction for sync, data loading, etc.
        serial_overhead_ms = 0.5  # fixed overhead per step
        step_time = max(compute_ms, comm_ms) + serial_overhead_ms
        single_gpu_step_ms = compute_ms + serial_overhead_ms
        speedup = (
            single_gpu_step_ms * N / step_time
        )  # effective throughput gain
        speedups.append(speedup)
        efficiencies.append(speedup / N * 100)
    return speedups, efficiencies


N_values = [1, 2, 4, 8, 16, 32, 64, 128]

# Compute scaling curves for each model and bandwidth scenario
results = {}
for model_name, num_params in param_configs.items():
    compute_ms = compute_times_ms[model_name]
    speedups_nv, efficiencies_nv = compute_speedup(
        N_values,
        num_params,
        compute_ms,
        nvlink_bandwidth_gb,
        nvlink_step_latency_ms,
    )
    speedups_ib, efficiencies_ib = compute_speedup(
        N_values,
        num_params,
        compute_ms,
        infiniband_bandwidth_gb,
        infiniband_step_latency_ms,
    )
    results[model_name] = {
        "speedup_nvlink": speedups_nv,
        "speedup_infiniband": speedups_ib,
        "efficiency_nvlink": efficiencies_nv,
        "efficiency_infiniband": efficiencies_ib,
    }

# Compute ideal speedup
ideal_speedup = N_values

# Simulate the effect of gradient bucketing on step time
# With perfect overlap: step_time = max(compute, comm)
# With no overlap: step_time = compute + comm
# Reality is between these extremes


def step_time_analysis(compute_ms, comm_ms, overlap_fraction=0.85):
    """
    Estimate step time with partial overlap.
    overlap_fraction: fraction of communication that overlaps with computation.
    """
    comm_serial = comm_ms * (1 - overlap_fraction)
    comm_overlapped = comm_ms * overlap_fraction

    # Overlapped communication hides within compute time
    effective_step = compute_ms + comm_serial
    return effective_step


model_name = "7B (Llama-7B equiv)"
num_params = param_configs[model_name]
compute_ms = compute_times_ms[model_name]

print(f"Analysis for {model_name}:")
print(f"Compute time per step: {compute_ms:.1f} ms")
print()

for N in [1, 2, 4, 8, 16, 32]:
    comm_nv = ring_allreduce_time_ms(
        num_params, nvlink_bandwidth_gb, N, nvlink_step_latency_ms
    )
    comm_ib = ring_allreduce_time_ms(
        num_params, infiniband_bandwidth_gb, N, infiniband_step_latency_ms
    )
    step_nv = step_time_analysis(compute_ms, comm_nv, overlap_fraction=0.85)
    step_ib = step_time_analysis(compute_ms, comm_ib, overlap_fraction=0.85)
    throughput_nv = N * compute_ms / step_nv  # relative throughput
    throughput_ib = N * compute_ms / step_ib
    print(
        f"  N={N:3d}  NVLink comm={comm_nv:5.1f} ms  IB comm={comm_ib:5.1f} ms  "
        f"  NVLink eff={throughput_nv:.1f}x  IB eff={throughput_ib:.1f}x"
    )