PyTorch 分布式训练底层原理与 DDP 实战指南

 import torch  

# Let us define a tensor with requires_grad = True  
x = torch.tensor(4.0, requires_grad=True)  
# Creating a function y=x^2  
y = x*x  
# Calculating dy/dx  
y.backward(retain_graph=True)  
# retain_graph = True because pytorch automatically removes the computation  
# graph and intermediate tensors once backward is called to save memory  
# If we want to call backward again, we need to tell pytorch not to delete  
# the computation graph and intermediate tensors  
print(f"Gradient of y w.r.t x after the first backward: {x.grad}")  
# Output: 8.0 as dy/dx = 2*x = 2*4  
# Now let us try to call backward again  
y.backward()  
print(f"Gradient of y w.r.t x after the second backward: {x.grad}")  
 # Output = 16 because 8 from the previous backward is added up here

 import torch  
import torch.nn as nn  
from torch.utils.data import Dataset  

class SimpleDataset(Dataset):  
    def __init__(self, size=100):  
        self.size = size  
        self.data = torch.randn(size, 10) # generate 100 samples each having dimension 10  
        self.labels = torch.randn(size, 1) # generate labels corresponding to each sample  

    def __len__(self):  
        return self.size  

    def __getitem__(self, idx):  
        return self.data[idx], self.labels[idx]  

class SimpleModel(nn.Module):  
    def __init__(self):  
        super().__init__()  
        self.fc1 = nn.Linear(10, 5)  
        self.fc2 = nn.Linear(5, 1)  

    def forward(self, x):  
        x = torch.relu(self.fc1(x))  
         return self.fc2(x)

 import os  
import torch.distributed as dist  

def setup(rank, world_size):  
    os.environ['MASTER_ADDR'] = 'localhost' # IP of the "master" node  
    os.environ['MASTER_PORT'] = '12355' # Port of communication  

    # If we have 4 processes, each process independently calls setup() with  
    # its own rank  
    dist.init_process_group(backend='gloo', rank=rank, world_size=world_size)  
    # 'gloo' is the collective communication backend used for CPU  
    # nccl is used for CUDA

    print(f"\n{'='*60}")  
    print(f"[Rank {rank}] Process initialized!")  
    print(f"[Rank {rank}] Backend: {dist.get_backend()}")  
    print(f"[Rank {rank}] World Size: {dist.get_world_size()}")  
     print(f"{'='*60}\n")

 # Example usage conceptual
# Create DistributedSampler for each rank  
# sampler_rank0 = DistributedSampler(dataset, num_replicas=4, rank=0)  
# ...

# Loop
for epoch in range(num_epochs):  
    train_sampler.set_epoch(epoch)  # Different shuffle each epoch  
    for batch in train_loader:  
         # Training code

 from torch.nn.parallel import DistributedDataParallel as DDP  

def print_separator(rank, message):  
   print(f"\n[Rank {rank}] {'-'*40}")  
   print(f"[Rank {rank}] {message}")  
   print(f"[Rank {rank}] {'-'*40}")  

def train_worker(rank, world_size, num_epochs=2, batch_size=8):  
  # setup the distributed environment  
  setup(rank, world_size)  

  model = SimpleModel()

  # wrap the model with DDP  
  # This is where weights are synchronized across ranks
  ddp_model = DDP(model)  

  dataset = SimpleDataset(size=32)  
  sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False)  
  dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)  

  optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)  
  criterion = nn.MSELoss()  

   for epoch in range(num_epochs):  
       sampler.set_epoch(epoch)  # Ensure different shuffling per epoch  
       epoch_loss = 0.0  
       for batch_idx, (data, target) in enumerate(dataloader):  
           optimizer.zero_grad()  
           output = ddp_model(data)  
           loss = criterion(output, target)  
           
           # Backward pass - THIS IS WHERE DDP MAGIC HAPPENS  
           loss.backward()  
           # Gradients are synchronized (All-Reduce) here automatically
           
           optimizer.step()  
           epoch_loss += loss.item()  
           
  dist.destroy_process_group()  
   print(f"[Rank {rank}] Training completed and cleaned up!\n")

 def validate(model, val_loader, criterion, rank, epoch):  
    model.eval()  
    val_loss = 0.0  
    num_samples = 0  

    with torch.no_grad():  
        for batch_idx, (data, target) in enumerate(val_loader):  
            output = model(data)  
            loss = criterion(output, target)  
            val_loss += loss.item() * data.size(0)  
            num_samples += data.size(0)  

    # Aggregate metrics across all processes  
    total_loss_tensor = torch.tensor([val_loss])  
    total_samples_tensor = torch.tensor([num_samples])  

    # Sum across all processes  
    dist.all_reduce(total_loss_tensor, op=dist.ReduceOp.SUM)  
    dist.all_reduce(total_samples_tensor, op=dist.ReduceOp.SUM)  
      
    global_avg_loss = total_loss_tensor.item() / total_samples_tensor.item()  
     return global_avg_loss

 defmain():  
    world_size=2    
    mp.spawn(  
        demo_worker,  
        args=(world_size,),  
        nprocs=world_size,  
        join=True  
    )

 # Code expects env vars
 rank = int(os.environ["RANK"])  
 local_rank = int(os.environ["LOCAL_RANK"])   
 world_size = int(os.environ["WORLD_SIZE"])  
 demo_worker(rank, world_size)  

 torchrun --nproc_per_node=4 train.py