Distributed Training: A Complete Guide

import torch.distributed as dist

# Send operation (on rank 0)
if dist.get_rank() == 0:
    tensor = torch.tensor([1.0, 2.0, 3.0], device='cuda')
    dist.send(tensor, dst=1)

# Receive operation (on rank 1)
if dist.get_rank() == 1:
    tensor = torch.zeros(3, device='cuda')
    dist.recv(tensor, src=0)
    # tensor is now [1.0, 2.0, 3.0]

import torch.distributed as dist

# Initialize tensor on all ranks
if dist.get_rank() == 0:
    tensor = torch.tensor([1.0, 2.0, 3.0], device='cuda')
else:
    tensor = torch.zeros(3, device='cuda')

# Broadcast from rank 0 to all ranks
dist.broadcast(tensor, src=0)
# Now all ranks have tensor = [1.0, 2.0, 3.0]

# Each rank has local gradients
local_gradients = compute_gradients(model, batch)

# Reduce sum to rank 0
dist.reduce(local_gradients, dst=0, op=dist.ReduceOp.SUM)

# Only rank 0 has the sum of all gradients
if dist.get_rank() == 0:
    local_gradients /= dist.get_world_size()  # Average

# Root prepares data chunks for each rank
if dist.get_rank() == 0:
    scatter_list = [torch.tensor([i * 10.0], device='cuda') 
                    for i in range(dist.get_world_size())]
else:
    scatter_list = None

# Each rank receives its chunk
output = torch.zeros(1, device='cuda')
dist.scatter(output, scatter_list, src=0)
# Rank 0: [0.0], Rank 1: [10.0], Rank 2: [20.0], ...

# Each rank has local data
local_tensor = torch.tensor([dist.get_rank() * 10.0], device='cuda')

# Gather to rank 0
if dist.get_rank() == 0:
    gather_list = [torch.zeros(1, device='cuda') 
                   for _ in range(dist.get_world_size())]
else:
    gather_list = None

dist.gather(local_tensor, gather_list, dst=0)
# Rank 0: gather_list = [[0.0], [10.0], [20.0], [30.0]]

import torch.distributed as dist

# Each rank computes local gradients
loss = model(inputs).sum()
loss.backward()

# AllReduce to sum gradients across all ranks
for param in model.parameters():
    dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
    param.grad /= dist.get_world_size()  # Average

# Now all ranks have identical averaged gradients
optimizer.step()

def ring_allreduce(tensor, rank, world_size):
    """Ring AllReduce: ReduceScatter + AllGather"""
    chunk_size = tensor.numel() // world_size
    chunks = tensor.chunk(world_size)
    
    # Phase 1: ReduceScatter
    for step in range(world_size - 1):
        send_idx = (rank - step) % world_size
        recv_idx = (rank - step - 1) % world_size
        
        # Send chunk[send_idx] to next rank in ring
        send_to_next = chunks[send_idx].clone()
        recv_from_prev = torch.zeros_like(chunks[recv_idx])
        
        # Simultaneous send/recv
        all_to_all([send_to_next], [recv_from_prev], 
                   src=(rank - 1) % world_size,
                   dst=(rank + 1) % world_size)
        
        # Accumulate received chunk
        chunks[recv_idx] += recv_from_prev
    
    # After ReduceScatter: each rank has one fully reduced chunk
    
    # Phase 2: AllGather
    for step in range(world_size - 1):
        send_idx = (rank - step + 1) % world_size
        recv_idx = (rank - step) % world_size
        
        # Same ring pattern, but just copy (no reduction)
        all_to_all([chunks[send_idx]], [chunks[recv_idx]], ...)
    
    return torch.cat(chunks)

import torch.distributed as dist

# Each rank has a local tensor (e.g., model shard)
local_tensor = torch.randn(1000, device='cuda')  # n = 1000 elements

# Prepare output list for gathered tensors
world_size = dist.get_world_size()
gathered = [torch.zeros_like(local_tensor) for _ in range(world_size)]

# AllGather: collect all shards
dist.all_gather(gathered, local_tensor)

# Result: list of P tensors, each 1000 elements
full_tensor = torch.cat(gathered)  # Shape: [P * 1000]

# Alternative: all_gather_into_tensor (more efficient)
output = torch.zeros(world_size * 1000, device='cuda')
dist.all_gather_into_tensor(output, local_tensor)  # Direct into contiguous buffer

import torch.distributed as dist

# Each rank has gradients for ALL parameters (before sharding)
full_gradients = torch.randn(4000, device='cuda')  # n*P elements

# Output: each rank gets 1/P of the reduced gradients
world_size = dist.get_world_size()
output_size = full_gradients.size(0) // world_size
reduced_shard = torch.zeros(output_size, device='cuda')

# ReduceScatter: sum gradients across ranks, distribute shards
dist.reduce_scatter_tensor(reduced_shard, full_gradients, op=dist.ReduceOp.SUM)

# Rank 0 gets sum of elements [0:1000] from all ranks
# Rank 1 gets sum of elements [1000:2000] from all ranks
# etc.

import torch.distributed as dist

# Each rank has P chunks to send (one to each rank)
world_size = dist.get_world_size()
rank = dist.get_rank()

# Input: chunks destined for each rank
send_chunks = [torch.randn(100, device='cuda') for _ in range(world_size)]

# Output: chunks received from each rank
recv_chunks = [torch.zeros(100, device='cuda') for _ in range(world_size)]

# All-to-All: exchange chunks
dist.all_to_all(recv_chunks, send_chunks)

# recv_chunks[j] now contains what rank j sent to us (their chunk for us)

# More efficient: all_to_all_single for contiguous tensors
input_tensor = torch.randn(world_size * 100, device='cuda')  # [P * chunk_size]
output_tensor = torch.zeros_like(input_tensor)
dist.all_to_all_single(output_tensor, input_tensor)

def calculate_crossover(alpha_us, bandwidth_gbps):
    """
    Calculate the crossover point where latency = bandwidth term.
    
    Args:
        alpha_us: Latency in microseconds
        bandwidth_gbps: Bandwidth in GB/s
    
    Returns:
        Crossover point in bytes
    """
    alpha_s = alpha_us * 1e-6  # Convert to seconds
    beta = 1 / (bandwidth_gbps * 1e9)  # seconds per byte
    n_star = alpha_s / beta  # crossover in bytes
    return n_star

# Example calculations
interconnects = {
    "NVLink 4.0": (1.0, 450),      # 1 μs latency, 450 GB/s
    "InfiniBand HDR": (1.5, 200),  # 1.5 μs, 200 GB/s
    "100GbE RoCE": (5.0, 12.5),    # 5 μs, 12.5 GB/s
    "TCP/IP Ethernet": (50, 12.5), # 50 μs (with kernel), 12.5 GB/s
}

print("Interconnect Crossover Points:")
print("-" * 45)
for name, (alpha, bw) in interconnects.items():
    n_star = calculate_crossover(alpha, bw)
    if n_star < 1024:
        print(f"{name:20}: n* = {n_star:8.0f} B")
    elif n_star < 1024**2:
        print(f"{name:20}: n* = {n_star/1024:8.1f} KB")
    else:
        print(f"{name:20}: n* = {n_star/(1024**2):8.1f} MB")

# Output:
# Interconnect Crossover Points:
# ---------------------------------------------
# NVLink 4.0          : n* =    450.0 KB
# InfiniBand HDR      : n* =    300.0 KB
# 100GbE RoCE         : n* =     62.5 KB
# TCP/IP Ethernet     : n* =    625.0 KB

import torch
import torch.distributed as dist
import time
import numpy as np

def measure_allreduce_alpha_beta(sizes_bytes, num_warmup=5, num_iters=20):
    """
    Measure α and β for AllReduce by fitting T = α + n*β to timing data.
    
    Run with: torchrun --nproc_per_node=N measure_alpha_beta.py
    """
    times = []
    
    for size in sizes_bytes:
        numel = size // 4  # FP32 elements
        tensor = torch.randn(numel, device='cuda')
        
        # Warmup
        for _ in range(num_warmup):
            dist.all_reduce(tensor)
            torch.cuda.synchronize()
        
        # Timed runs
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(num_iters):
            dist.all_reduce(tensor)
        torch.cuda.synchronize()
        elapsed = (time.perf_counter() - start) / num_iters
        
        times.append(elapsed)
    
    # Linear regression: T = α + n*β
    # For AllReduce: T = 2*log2(P)*α + 2*(P-1)/P * n * β
    sizes = np.array(sizes_bytes)
    times = np.array(times)
    
    # Simple linear fit (ignoring the exact formula for illustration)
    # slope = β, intercept ≈ α
    coeffs = np.polyfit(sizes, times, 1)
    beta = coeffs[0]
    alpha = coeffs[1]
    
    bandwidth_gbps = 1 / beta / 1e9
    
    return alpha, beta, bandwidth_gbps

if __name__ == "__main__":
    dist.init_process_group()
    
    # Test various message sizes from 1KB to 1GB
    sizes = [1024 * (2 ** i) for i in range(20)]  # 1KB to 1GB
    
    alpha, beta, bw = measure_allreduce_alpha_beta(sizes)
    
    if dist.get_rank() == 0:
        print(f"Measured Parameters:")
        print(f"  α (latency)   = {alpha*1e6:.2f} μs")
        print(f"  β (inv. bw)   = {beta*1e12:.2f} ps/byte")
        print(f"  Bandwidth     = {bw:.1f} GB/s")
        print(f"  Crossover n*  = {alpha/beta/1024:.1f} KB")
    
    dist.destroy_process_group()

def calculate_ccr_ddp(model_params, batch_size, seq_len, dtype_bytes=2):
    """
    Calculate computation-to-communication ratio for DDP training.
    
    For transformers: FLOPs ≈ 6 * N * tokens (forward + backward)
    Communication: 2 * N * dtype_bytes (AllReduce gradients)
    """
    tokens = batch_size * seq_len
    flops = 6 * model_params * tokens  # Approximate transformer FLOPs
    comm_bytes = 2 * model_params * dtype_bytes  # AllReduce volume
    
    ccr = flops / comm_bytes
    return ccr

# Example: GPT-3 175B with different batch sizes
model_params = 175e9
seq_len = 2048

print("GPT-3 175B CCR Analysis:")
print("-" * 50)
for batch_size in [1, 4, 16, 64, 256]:
    ccr = calculate_ccr_ddp(model_params, batch_size, seq_len)
    print(f"Batch size {batch_size:4}: CCR = {ccr:8.1f} FLOPS/byte")

# Output:
# GPT-3 175B CCR Analysis:
# --------------------------------------------------
# Batch size    1: CCR =   3072.0 FLOPS/byte
# Batch size    4: CCR =  12288.0 FLOPS/byte
# Batch size   16: CCR =  49152.0 FLOPS/byte
# Batch size   64: CCR = 196608.0 FLOPS/byte
# Batch size  256: CCR = 786432.0 FLOPS/byte

def analyze_roofline_position(
    model_flops_per_sample: float,
    model_params: float,
    batch_size: int,
    gpu_peak_tflops: float,
    network_bw_gbps: float,
    dtype_bytes: int = 2,
    comm_factor: float = 2.0,  # AllReduce ≈ 2x model size
):
    """
    Determine where a workload sits on the communication roofline.
    
    Returns analysis with optimization recommendations.
    """
    # Calculate CCR
    total_flops = model_flops_per_sample * batch_size
    comm_bytes = comm_factor * model_params * dtype_bytes
    ccr = total_flops / comm_bytes
    
    # Calculate ridge point
    ccr_star = (gpu_peak_tflops * 1e12) / (network_bw_gbps * 1e9)
    
    # Theoretical throughput limits
    compute_limit = (gpu_peak_tflops * 1e12) / model_flops_per_sample  # samples/s
    comm_limit = (network_bw_gbps * 1e9) / (comm_bytes / batch_size)  # samples/s
    
    # Effective throughput (min of both limits)
    effective_throughput = min(compute_limit, comm_limit)
    efficiency = effective_throughput / compute_limit * 100
    
    # Determine regime
    if ccr < ccr_star * 0.2:
        regime = "SEVERELY_COMM_BOUND"
        recommendation = "Increase batch size by 4x+ or use gradient compression"
    elif ccr < ccr_star * 0.8:
        regime = "COMM_BOUND"
        recommendation = "Increase batch size or overlap communication"
    elif ccr < ccr_star * 1.2:
        regime = "NEAR_RIDGE"
        recommendation = "Near optimal! Small batch increase may help"
    else:
        regime = "COMPUTE_BOUND"
        recommendation = "Excellent! Focus on GPU optimizations"
    
    return {
        "ccr": ccr,
        "ccr_star": ccr_star,
        "regime": regime,
        "efficiency": efficiency,
        "recommendation": recommendation,
        "critical_batch": batch_size * (ccr_star / ccr),  # Batch to reach ridge
    }

# Example analysis
result = analyze_roofline_position(
    model_flops_per_sample=6 * 175e9 * 2048,  # GPT-3 175B, seq=2048
    model_params=175e9,
    batch_size=32,
    gpu_peak_tflops=1979,  # H100 FP16
    network_bw_gbps=900,  # NVLink-4
)

print(f"CCR: {result['ccr']:.0f} FLOPS/byte")
print(f"Ridge CCR*: {result['ccr_star']:.0f} FLOPS/byte")
print(f"Regime: {result['regime']}")
print(f"Efficiency: {result['efficiency']:.1f}%")
print(f"Recommendation: {result['recommendation']}")

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp(rank, world_size):
    """Initialize distributed process group."""
    dist.init_process_group(
        backend="nccl",  # NVIDIA Collective Communication Library
        init_method="env://",
        world_size=world_size,
        rank=rank,
    )
    torch.cuda.set_device(rank)

def train_ddp(rank, world_size, model, dataloader, optimizer):
    setup_ddp(rank, world_size)
    
    # Wrap model with DDP
    model = model.to(rank)
    model = DDP(model, device_ids=[rank])
    
    for batch in dataloader:
        inputs, labels = batch
        inputs, labels = inputs.to(rank), labels.to(rank)
        
        optimizer.zero_grad()
        
        # Forward pass (independent per GPU)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # Backward pass (DDP hooks trigger AllReduce)
        loss.backward()  # ← Gradients synchronized here!
        
        # Optimizer step (identical on all GPUs)
        optimizer.step()
    
    dist.destroy_process_group()

# Launch with: torchrun --nproc_per_node=8 basic_ddp.py

from torch.nn.parallel import DistributedDataParallel as DDP

# Advanced DDP configuration
model = DDP(
    model,
    device_ids=[local_rank],
    
    # Bucket size (bytes) - smaller = more overlap, more latency overhead
    bucket_cap_mb=25,  # Default: 25 MB per bucket
    
    # Find unused parameters (for models with conditional execution)
    find_unused_parameters=False,  # Set True if some params don't get gradients
    
    # Gradient as bucket view (memory optimization)
    gradient_as_bucket_view=True,  # Reduces memory by 1 copy
    
    # Static graph optimization (PyTorch 1.11+)
    static_graph=True,  # Enable if computation graph doesn't change
)

# Tune bucket size based on your model and network
# - Smaller buckets: Better overlap, more latency overhead
# - Larger buckets: Less overhead, less overlap opportunity
# - Rule of thumb: bucket_size > α/β (crossover point)

import torch
import torch.distributed as dist

class TopKCompressor:
    """Top-K gradient sparsification with error feedback."""
    
    def __init__(self, ratio=0.01):
        self.ratio = ratio
        self.error_feedback = {}  # Accumulated residuals
    
    def compress(self, name, grad):
        # Add error feedback from previous iteration
        if name in self.error_feedback:
            grad = grad + self.error_feedback[name]
        
        # Select Top-K elements
        k = max(1, int(grad.numel() * self.ratio))
        values, indices = torch.topk(grad.abs().flatten(), k)
        
        # Store residual for next iteration
        mask = torch.zeros_like(grad.flatten())
        mask[indices] = 1
        self.error_feedback[name] = grad.flatten() * (1 - mask)
        self.error_feedback[name] = self.error_feedback[name].view_as(grad)
        
        # Return sparse representation
        return indices, grad.flatten()[indices]
    
    def decompress(self, indices, values, shape):
        grad = torch.zeros(shape, device=values.device).flatten()
        grad[indices] = values
        return grad.view(shape)

# Usage in training loop
compressor = TopKCompressor(ratio=0.01)  # Keep top 1%

for name, param in model.named_parameters():
    if param.grad is not None:
        idx, vals = compressor.compress(name, param.grad)
        
        # AllGather compressed gradients (much smaller!)
        all_idx = [torch.zeros_like(idx) for _ in range(world_size)]
        all_vals = [torch.zeros_like(vals) for _ in range(world_size)]
        dist.all_gather(all_idx, idx)
        dist.all_gather(all_vals, vals)
        
        # Decompress and average
        param.grad.zero_()
        for i, v in zip(all_idx, all_vals):
            param.grad.flatten()[i] += v / world_size

import torch
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

# Define sharding strategy
sharding_strategy = ShardingStrategy.FULL_SHARD  # ZeRO-3 equivalent
# Options: FULL_SHARD, SHARD_GRAD_OP (ZeRO-2), NO_SHARD (DDP)

# Mixed precision policy
mixed_precision = MixedPrecision(
    param_dtype=torch.float16,
    reduce_dtype=torch.float16,
    buffer_dtype=torch.float16,
)

# Auto-wrap policy for transformers
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},  # Your layer class
)

# Wrap model with FSDP
model = FSDP(
    model,
    sharding_strategy=sharding_strategy,
    mixed_precision=mixed_precision,
    auto_wrap_policy=auto_wrap_policy,
    device_id=torch.cuda.current_device(),
    
    # Performance optimizations
    use_orig_params=True,  # Better compatibility with optimizers
    limit_all_gathers=True,  # Prevent OOM from too many concurrent gathers
)

# Training loop (same as regular PyTorch!)
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch).loss
    loss.backward()
    optimizer.step()

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    BackwardPrefetch,
    ForwardPrefetch,
)

# Configure prefetching
model = FSDP(
    model,
    # Forward prefetching (gather next layer during current layer's compute)
    forward_prefetch=True,
    
    # Backward prefetching strategy
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,  # or BACKWARD_POST
    # BACKWARD_PRE: Prefetch before backward (default, best for most cases)
    # BACKWARD_POST: Prefetch after backward (use if memory-tight)
)

# For maximum overlap, also enable:
# - async_op=True in NCCL operations
# - CUDA streams for concurrent execution

def tensor_parallel_comm_per_layer(
    batch_size: int,
    seq_len: int,
    hidden_dim: int,
    tp_degree: int,
    dtype_bytes: int = 2  # FP16
) -> float:
    """Calculate communication volume per transformer layer with TP."""
    
    # Activation size
    activation_size = batch_size * seq_len * hidden_dim * dtype_bytes
    
    # AllReduce cost (ring algorithm)
    allreduce_factor = 2 * (tp_degree - 1) / tp_degree
    
    # 4 AllReduces per layer (attn fwd, attn bwd, mlp fwd, mlp bwd)
    comm_volume = 4 * allreduce_factor * activation_size
    
    return comm_volume

# Example: GPT-3 175B style layer
comm_bytes = tensor_parallel_comm_per_layer(
    batch_size=1,
    seq_len=2048,
    hidden_dim=12288,
    tp_degree=8
)
print(f"Comm per layer: {comm_bytes / 1e6:.1f} MB")
# Output: ~175 MB per layer

# With 96 layers: ~16.8 GB total TP communication per iteration!

import torch
import torch.nn as nn
import torch.distributed as dist

class ColumnParallelLinear(nn.Module):
    """Linear layer with column-parallel weight distribution.
    
    Input X is replicated, weight W is split along columns (output dim).
    Output Y is split along columns (each GPU has different output features).
    """
    
    def __init__(self, in_features, out_features, tp_group, gather_output=False):
        super().__init__()
        self.tp_group = tp_group
        self.tp_size = dist.get_world_size(tp_group)
        self.gather_output = gather_output
        
        # Each GPU gets out_features/tp_size columns
        assert out_features % self.tp_size == 0
        self.out_features_per_partition = out_features // self.tp_size
        
        self.weight = nn.Parameter(
            torch.empty(self.out_features_per_partition, in_features)
        )
        self.bias = nn.Parameter(torch.empty(self.out_features_per_partition))
        
    def forward(self, x):
        # Local matmul: no communication!
        output = F.linear(x, self.weight, self.bias)
        
        if self.gather_output:
            # Optionally gather to get full output
            output = _gather_along_last_dim(output, self.tp_group)
        
        return output


class RowParallelLinear(nn.Module):
    """Linear layer with row-parallel weight distribution.
    
    Input X is split along last dim, weight W is split along rows (input dim).
    Outputs are partial sums that need AllReduce.
    """
    
    def __init__(self, in_features, out_features, tp_group, input_is_parallel=True):
        super().__init__()
        self.tp_group = tp_group
        self.tp_size = dist.get_world_size(tp_group)
        self.input_is_parallel = input_is_parallel
        
        # Each GPU gets in_features/tp_size rows
        assert in_features % self.tp_size == 0
        self.in_features_per_partition = in_features // self.tp_size
        
        self.weight = nn.Parameter(
            torch.empty(out_features, self.in_features_per_partition)
        )
        self.bias = nn.Parameter(torch.empty(out_features))
        
    def forward(self, x):
        if not self.input_is_parallel:
            # Split input if it's replicated
            x = _split_along_last_dim(x, self.tp_group)
        
        # Local matmul (partial result)
        output_partial = F.linear(x, self.weight)
        
        # AllReduce to sum partial results
        dist.all_reduce(output_partial, group=self.tp_group)
        
        # Add bias after reduction (only on one rank to avoid duplication)
        if dist.get_rank(self.tp_group) == 0:
            output_partial = output_partial + self.bias
        
        return output_partial


class TensorParallelMLP(nn.Module):
    """Megatron-style tensor parallel MLP.
    
    Column-parallel first layer, row-parallel second layer.
    Only ONE AllReduce for the entire MLP!
    """
    
    def __init__(self, hidden_size, ffn_hidden_size, tp_group):
        super().__init__()
        
        # Column-parallel: split FFN hidden dim
        self.fc1 = ColumnParallelLinear(
            hidden_size, ffn_hidden_size, tp_group, gather_output=False
        )
        
        # Row-parallel: recombine to hidden_size
        self.fc2 = RowParallelLinear(
            ffn_hidden_size, hidden_size, tp_group, input_is_parallel=True
        )
        
    def forward(self, x):
        # x: [B, S, H] - replicated across TP ranks
        x = self.fc1(x)      # [B, S, FFN/TP] - split output
        x = F.gelu(x)        # Local activation
        x = self.fc2(x)      # [B, S, H] - AllReduce inside
        return x

def pipeline_comm_per_iteration(
    batch_size: int,
    seq_len: int,
    hidden_dim: int,
    num_micro_batches: int,
    num_stages: int,
    dtype_bytes: int = 2
) -> dict:
    """Calculate pipeline parallel communication."""
    
    micro_batch_size = batch_size // num_micro_batches
    
    # Activation size per micro-batch
    act_size = micro_batch_size * seq_len * hidden_dim * dtype_bytes
    
    # P2P comm per stage boundary
    # Forward: send activations
    # Backward: send gradients
    p2p_per_boundary = 2 * act_size  # fwd + bwd
    
    # Total boundaries = num_stages - 1
    # Total micro-batches = num_micro_batches
    total_p2p = p2p_per_boundary * (num_stages - 1) * num_micro_batches
    
    # Bubble analysis
    bubble_fraction = (num_stages - 1) / (num_micro_batches + num_stages - 1)
    
    return {
        "p2p_per_microbatch_boundary": p2p_per_boundary,
        "total_p2p_bytes": total_p2p,
        "bubble_fraction": bubble_fraction,
        "efficiency": 1 - bubble_fraction
    }

# Example: 8 stages, 32 micro-batches
result = pipeline_comm_per_iteration(
    batch_size=512,
    seq_len=2048,
    hidden_dim=8192,
    num_micro_batches=32,
    num_stages=8
)
print(f"Efficiency: {result['efficiency']*100:.1f}%")  # ~82%
print(f"Total P2P: {result['total_p2p_bytes']/1e9:.1f} GB")

import torch
import torch.distributed as dist

class PipelineStage:
    """Simple 1F1B pipeline stage implementation."""
    
    def __init__(self, module, stage_id, num_stages, process_group):
        self.module = module
        self.stage_id = stage_id
        self.num_stages = num_stages
        self.group = process_group
        
        # Determine neighbors
        self.prev_rank = stage_id - 1 if stage_id > 0 else None
        self.next_rank = stage_id + 1 if stage_id < num_stages - 1 else None
        
    def recv_forward(self, shape, dtype):
        """Receive activations from previous stage."""
        if self.prev_rank is None:
            return None
        
        tensor = torch.empty(shape, dtype=dtype, device='cuda')
        dist.recv(tensor, src=self.prev_rank, group=self.group)
        return tensor.requires_grad_()
    
    def send_forward(self, tensor):
        """Send activations to next stage."""
        if self.next_rank is None:
            return
        dist.send(tensor, dst=self.next_rank, group=self.group)
    
    def recv_backward(self, shape, dtype):
        """Receive gradients from next stage."""
        if self.next_rank is None:
            return None
        
        tensor = torch.empty(shape, dtype=dtype, device='cuda')
        dist.recv(tensor, src=self.next_rank, group=self.group)
        return tensor
    
    def send_backward(self, tensor):
        """Send gradients to previous stage."""
        if self.prev_rank is None:
            return
        dist.send(tensor, dst=self.prev_rank, group=self.group)

    def forward_step(self, input_tensor):
        """Execute one forward micro-batch."""
        output = self.module(input_tensor)
        self.send_forward(output.detach())
        return output
    
    def backward_step(self, output_tensor, output_grad):
        """Execute one backward micro-batch."""
        torch.autograd.backward(output_tensor, output_grad)
        
        if output_tensor.grad is not None:
            self.send_backward(output_tensor.grad)


def schedule_1f1b(stage, num_micro_batches, input_tensors):
    """1F1B schedule for one pipeline stage."""
    
    num_warmup = stage.num_stages - stage.stage_id - 1
    num_microbatches_remaining = num_micro_batches - num_warmup
    
    # Store outputs for backward
    outputs = []
    
    # Warmup: only forward passes
    for i in range(num_warmup):
        inp = stage.recv_forward(...) if stage.prev_rank else input_tensors[i]
        out = stage.forward_step(inp)
        outputs.append((inp, out))
    
    # Steady state: 1F1B
    for i in range(num_microbatches_remaining):
        # Forward
        inp = stage.recv_forward(...) if stage.prev_rank else input_tensors[num_warmup + i]
        out = stage.forward_step(inp)
        outputs.append((inp, out))
        
        # Backward (for micro-batch i)
        old_inp, old_out = outputs.pop(0)
        grad = stage.recv_backward(...) if stage.next_rank else None
        stage.backward_step(old_out, grad)
    
    # Cooldown: only backward passes
    for _ in range(num_warmup):
        old_inp, old_out = outputs.pop(0)
        grad = stage.recv_backward(...) if stage.next_rank else None
        stage.backward_step(old_out, grad)

import torch.distributed as dist

def initialize_3d_parallelism(
    tp_size: int,
    pp_size: int, 
    dp_size: int
):
    """Initialize process groups for 3D parallelism.
    
    GPU layout (example: TP=2, PP=2, DP=2 = 8 GPUs):
    
    Node 0: GPU 0,1 (TP group 0, PP stage 0)
    Node 1: GPU 2,3 (TP group 1, PP stage 1)
    Node 2: GPU 4,5 (TP group 2, PP stage 0)  # DP replica
    Node 3: GPU 6,7 (TP group 3, PP stage 1)  # DP replica
    """
    
    world_size = tp_size * pp_size * dp_size
    rank = dist.get_rank()
    
    # Calculate coordinates in 3D grid
    tp_rank = rank % tp_size
    pp_rank = (rank // tp_size) % pp_size
    dp_rank = rank // (tp_size * pp_size)
    
    # Create Tensor Parallel groups
    # GPUs in same TP group: same pp_rank, same dp_rank
    tp_groups = []
    for dp in range(dp_size):
        for pp in range(pp_size):
            ranks = [
                dp * (tp_size * pp_size) + pp * tp_size + tp
                for tp in range(tp_size)
            ]
            group = dist.new_group(ranks)
            if rank in ranks:
                tp_group = group
    
    # Create Pipeline Parallel groups
    # GPUs in same PP group: same tp_rank, same dp_rank
    pp_groups = []
    for dp in range(dp_size):
        for tp in range(tp_size):
            ranks = [
                dp * (tp_size * pp_size) + pp * tp_size + tp
                for pp in range(pp_size)
            ]
            group = dist.new_group(ranks)
            if rank in ranks:
                pp_group = group
    
    # Create Data Parallel groups
    # GPUs in same DP group: same tp_rank, same pp_rank
    for pp in range(pp_size):
        for tp in range(tp_size):
            ranks = [
                dp * (tp_size * pp_size) + pp * tp_size + tp
                for dp in range(dp_size)
            ]
            group = dist.new_group(ranks)
            if rank in ranks:
                dp_group = group
    
    return {
        'tp_group': tp_group,
        'pp_group': pp_group,
        'dp_group': dp_group,
        'tp_rank': tp_rank,
        'pp_rank': pp_rank,
        'dp_rank': dp_rank,
    }

# Example: 128 GPUs with TP=8, PP=4, DP=4
groups = initialize_3d_parallelism(tp_size=8, pp_size=4, dp_size=4)
print(f"GPU {dist.get_rank()}: TP={groups['tp_rank']}, PP={groups['pp_rank']}, DP={groups['dp_rank']}")

import torch
import torch.nn.functional as F

def load_balancing_loss(router_logits, num_experts, top_k=2):
    """Auxiliary loss to encourage balanced expert usage.
    
    From Switch Transformer / GShard papers.
    """
    # router_logits: [batch_size * seq_len, num_experts]
    
    # Get routing probabilities
    routing_probs = F.softmax(router_logits, dim=-1)
    
    # Get top-k expert assignments
    _, top_k_indices = torch.topk(router_logits, k=top_k, dim=-1)
    
    # Create one-hot masks for selected experts
    expert_mask = torch.zeros_like(routing_probs)
    expert_mask.scatter_(-1, top_k_indices, 1.0)
    
    # Fraction of tokens routed to each expert
    tokens_per_expert = expert_mask.mean(dim=0)  # [num_experts]
    
    # Average routing probability per expert
    router_prob_per_expert = routing_probs.mean(dim=0)  # [num_experts]
    
    # Load balancing loss: encourage uniform distribution
    # Minimize product of (fraction routed) × (avg probability)
    aux_loss = num_experts * (tokens_per_expert * router_prob_per_expert).sum()
    
    return aux_loss


class MoELayer(torch.nn.Module):
    def __init__(self, hidden_size, ffn_size, num_experts, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        
        # Router network
        self.router = torch.nn.Linear(hidden_size, num_experts, bias=False)
        
        # Expert FFN networks
        self.experts = torch.nn.ModuleList([
            torch.nn.Sequential(
                torch.nn.Linear(hidden_size, ffn_size),
                torch.nn.GELU(),
                torch.nn.Linear(ffn_size, hidden_size)
            )
            for _ in range(num_experts)
        ])
        
    def forward(self, x):
        batch_seq, hidden = x.shape
        
        # Compute routing scores
        router_logits = self.router(x)  # [B*S, E]
        
        # Get top-k experts and weights
        routing_weights, selected_experts = torch.topk(
            F.softmax(router_logits, dim=-1), 
            self.top_k, 
            dim=-1
        )
        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
        
        # Compute expert outputs (simplified - real impl uses All-to-All)
        output = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            # Find tokens routed to this expert
            expert_mask = (selected_experts == i).any(dim=-1)
            if expert_mask.any():
                expert_input = x[expert_mask]
                expert_output = expert(expert_input)
                
                # Weight by routing probability
                weights = routing_weights[expert_mask]
                weight_mask = (selected_experts[expert_mask] == i)
                expert_weight = (weights * weight_mask).sum(dim=-1, keepdim=True)
                
                output[expert_mask] += expert_weight * expert_output
        
        # Compute auxiliary loss for load balancing
        self.aux_loss = load_balancing_loss(router_logits, self.num_experts, self.top_k)
        
        return output

# DeepSeek-V2 236B Configuration
model:
  total_params: 236B
  active_params: 21B  # Per token
  
  architecture:
    layers: 60
    hidden_dim: 5120
    num_experts: 160
    experts_per_token: 6  # Top-6 routing
    shared_experts: 2     # Always-active experts
    
  innovations:
    - Multi-head Latent Attention (MLA)  # Reduces KV cache
    - DeepSeekMoE with shared experts
    - Fine-grained expert segmentation

training:
  parallelism:
    expert_parallel: 8-16
    tensor_parallel: 1-4
    pipeline_parallel: 1
    data_parallel: FSDP
    
  hardware:
    gpus: ~2000 H800
    interconnect: NVLink + IB
    
  efficiency:
    tokens_trained: 8.1T
    training_cost: ~$5.5M  # Remarkably efficient!