Optimize distributed training communication and scaling efficiency.

Communication Optimization

Gradient Compression

from torch.distributed.algorithms.ddp_comm_hooks import default_hooks

# PowerSGD gradient compression
from torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook import PowerSGDState, powerSGD_hook

state = PowerSGDState(
    process_group=dist.group.WORLD,
    matrix_approximation_rank=4,  # Compression factor
    warm_start=True
)

model = DDP(model)
model.register_comm_hook(state, powerSGD_hook)

# Reduces gradient communication by 4-8x
# Small accuracy cost on some models

Overlap Communication and Computation

# Bucket-based gradient communication
model = DDP(
    model,
    bucket_cap_mb=25,              # Larger buckets = fewer allreduce ops
    gradient_as_bucket_view=True,  # Reduce memory allocation
    static_graph=True              # Optimize for same-graph repeated execution
)

# pipeline_parallel: compute next micro-batch while syncing current
from torch.distributed.pipeline.sync import Pipe
chunks = 8  # Micro-batches for pipeline parallelism
model = Pipe(model, chunks=chunks)

Scaling Efficiency Metrics

def calculate_scaling_efficiency(baseline_throughput, n_gpus, actual_throughput):
    perfect_throughput = baseline_throughput * n_gpus
    efficiency = actual_throughput / perfect_throughput
    return efficiency

# Example measurements
# 1 GPU: 100 samples/sec
# 8 GPUs: 720 samples/sec (vs ideal 800)
efficiency = calculate_scaling_efficiency(100, 8, 720)
print(f"Scaling efficiency: {efficiency:.1%}")  # 90%

# Good: >85% scaling efficiency
# Acceptable: 70-85%
# Poor: <70% (investigate communication overhead)

Training Recipe for Large Models

# Optimal settings for H100 training
config = {
    "precision": "bf16-true",
    "optimizer": "Adam",
    "learning_rate": 3e-4,
    "weight_decay": 0.1,
    "gradient_clip": 1.0,
    "batch_size_per_gpu": 16,
    "gradient_accumulation": 8,  # Effective batch = 128 per GPU
    "num_gpus": 128,
    "flash_attention": True,
    "gradient_checkpointing": True,
    "compile": True
}