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

def setup_ddp(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    
    model = MyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    return ddp_model

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp(rank, world_size):
    """初始化分布式环境"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup_ddp():
    """清理分布式环境"""
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    
    def forward(self, x):
        return self.network(x)

def train_ddp(rank, world_size):
    """DDP训练主函数"""
    setup_ddp(rank, world_size)
    
    # 创建模型
    model = SimpleModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    # 创建数据和分布式采样器
    from torchvision import datasets, transforms
    
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)
    
    # 训练
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    for epoch in range(3):
        ddp_model.train()
        total_loss = 0
        
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(rank).view(-1, 784), target.to(rank)
            
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        if rank == 0:  # 只在主进程打印
            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")
    
    cleanup_ddp()

def spawn_ddp():
    """启动多进程DDP训练"""
    world_size = torch.cuda.device_count()
    if world_size < 2:
        print("Need at least 2 GPUs for DDP")
        return
    
    torch.multiprocessing.spawn(
        train_ddp,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

if __name__ == "__main__":
    spawn_ddp()

import torch
import torch.nn as nn

class PipelineStage(nn.Module):
    """流水线的一个阶段"""
    
    def __init__(self, input_size, output_size, num_layers=5):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(input_size if i == 0 else output_size, output_size)
            for i in range(num_layers)
        ])
        self.relu = nn.ReLU()
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
            x = self.relu(x)
        return x

class SimplePipelineParallel:
    """
    简单的流水线并行实现
    
    注意：这是一个教学示例，实际生产应使用torch.distributed.pipelining
    """
    
    def __init__(self, num_stages, input_size, hidden_size, output_size, devices):
        self.num_stages = num_stages
        self.devices = devices
        
        # 创建流水线阶段
        self.stages = nn.ModuleList()
        
        # 第一阶段
        self.stages.append(PipelineStage(input_size, hidden_size).to(devices[0]))
        
        # 中间阶段
        for _ in range(num_stages - 2):
            self.stages.append(PipelineStage(hidden_size, hidden_size).to(devices[_]))
        
        # 最后阶段
        self.stages.append(PipelineStage(hidden_size, output_size).to(devices[-1]))
    
    def forward_pipeline(self, x, micro_batches=4):
        """
        使用微批次的流水线前向传播
        
        Args:
            x: 输入数据
            micro_batches: 微批次数量
        """
        batch_size = x.size(0)
        micro_batch_size = batch_size // micro_batches
        
        outputs = []
        
        for i in range(micro_batches):
            start = i * micro_batch_size
            end = start + micro_batch_size
            micro_x = x[start:end]
            
            # 通过流水线
            current = micro_x
            for stage_idx, stage in enumerate(self.stages):
                device = self.devices[stage_idx]
                current = current.to(device)
                current = stage(current)
            
            outputs.append(current)
        
        return torch.cat(outputs, dim=0)

def demo_pipeline():
    """演示流水线并行"""
    if torch.cuda.device_count() < 2:
        print("Need at least 2 GPUs")
        return
    
    devices = [torch.device(f'cuda:{i}') for i in range(2)]
    
    pipeline = SimplePipelineParallel(
        num_stages=2,
        input_size=784,
        hidden_size=256,
        output_size=10,
        devices=devices
    )
    
    x = torch.randn(32, 784)
    output = pipeline.forward_pipeline(x, micro_batches=4)
    
    print(f"Output shape: {output.shape}")


demo_pipeline()

import torch
import torch.distributed as dist
import time
import os

def benchmark_allreduce(rank, world_size):
    """基准测试AllReduce性能"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12356'
    
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    
    sizes = [1024, 1024*1024, 10*1024*1024, 100*1024*1024]  # 1K, 1M, 10M, 100M elements
    
    for size in sizes:
        tensor = torch.randn(size, device=f'cuda:{rank}')
        
        # 预热
        for _ in range(5):
            dist.all_reduce(tensor)
        
        torch.cuda.synchronize()
        start = time.time()
        
        iterations = 100
        for _ in range(iterations):
            dist.all_reduce(tensor)
        
        torch.cuda.synchronize()
        elapsed = (time.time() - start) / iterations * 1000  # ms
        
        if rank == 0:
            data_size_mb = size * 4 / 1024 / 1024  # FP32 = 4 bytes
            bandwidth_gbps = data_size_mb / elapsed * 1000 * 8 / 1024
            print(f"Size: {size:>12} elements ({data_size_mb:.1f} MB) | "
                  f"Time: {elapsed:.3f} ms | "
                  f"Bandwidth: {bandwidth_gbps:.2f} Gbps")
    
    dist.destroy_process_group()

def run_benchmark():
    """运行通信基准测试"""
    world_size = torch.cuda.device_count()
    if world_size < 2:
        print("Need at least 2 GPUs")
        return
    
    print(f"\n{'='*70}")
    print(f"AllReduce Benchmark ({world_size} GPUs)")
    print(f"{'='*70}")
    
    torch.multiprocessing.spawn(
        benchmark_allreduce,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

if __name__ == "__main__":
    run_benchmark()