# 使用NCCL后端
dist.init_process_group("nccl", rank=rank, world_size=world_size)

ddp_model = DDP(model, bucket_cap_mb=64)  # 增大桶大小

export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

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

def benchmark_communication(rank, world_size):
    """基准测试各种通信原语的性能"""
    
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12357'
    
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    
    sizes = [1024, 1024*1024, 10*1024*1024]  # 1K, 1M, 10M elements
    
    operations = {
        'allreduce': lambda t: dist.all_reduce(t),
        'allgather': lambda t: dist.all_gather([torch.empty_like(t) for _ in range(world_size)], t),
        'broadcast': lambda t: dist.broadcast(t, src=0),
    }
    
    if rank == 0:
        print(f"\n{'='*70}")
        print(f"Communication Benchmark ({world_size} GPUs)")
        print(f"{'='*70}")
    
    for op_name, op_func in operations.items():
        if rank == 0:
            print(f"\n{op_name.upper()}:")
        
        for size in sizes:
            tensor = torch.randn(size, device=f'cuda:{rank}')
            
            # 预热
            for _ in range(5):
                op_func(tensor)
            
            torch.cuda.synchronize()
            start = time.time()
            
            iterations = 50
            for _ in range(iterations):
                op_func(tensor)
            
            torch.cuda.synchronize()
            elapsed = (time.time() - start) / iterations * 1000
            
            if rank == 0:
                data_size_mb = size * 4 / 1024 / 1024
                print(f"  Size: {size:>10} ({data_size_mb:.1f}MB) | Time: {elapsed:.3f} ms")
    
    dist.destroy_process_group()

def run_benchmark():
    world_size = torch.cuda.device_count()
    if world_size < 2:
        print("Need at least 2 GPUs")
        return
    
    torch.multiprocessing.spawn(
        benchmark_communication,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

if __name__ == "__main__":
    run_benchmark()

import torch
import torch.distributed as dist

class TopKCompressor:
    """
    Top-k梯度压缩器
    
    只传输绝对值最大的k%的梯度元素
    """
    
    def __init__(self, compression_ratio=0.1):
        """
        Args:
            compression_ratio: 保留的梯度比例 (0-1)
        """
        self.compression_ratio = compression_ratio
        self.error_buffer = None
    
    def compress(self, gradient):
        """
        压缩梯度
        
        Returns:
            compressed: 压缩后的梯度（稀疏张量）
            indices: 非零元素的索引
        """
        # 加上之前的误差
        if self.error_buffer is not None:
            gradient = gradient + self.error_buffer
        
        # 计算保留的元素数量
        k = max(1, int(gradient.numel() * self.compression_ratio))
        
        # 找到top-k的绝对值
        abs_grad = torch.abs(gradient)
        threshold = torch.topk(abs_grad.view(-1), k).values[-1]
        
        # 创建掩码
        mask = abs_grad >= threshold
        
        # 获取非零元素的索引和值
        indices = torch.nonzero(mask, as_tuple=True)
        values = gradient[indices]
        
        # 计算误差
        reconstructed = torch.zeros_like(gradient)
        reconstructed[indices] = values
        self.error_buffer = gradient - reconstructed
        
        return values, indices
    
    def decompress(self, values, indices, shape, device):
        """解压梯度"""
        decompressed = torch.zeros(shape, device=device)
        decompressed[indices] = values
        return decompressed


class CompressedDDP:
    """
    带梯度压缩的分布式数据并行
    
    注意：这是一个简化示例，实际使用应考虑更多细节
    """
    
    def __init__(self, model, compression_ratio=0.1):
        self.model = model
        self.compressor = TopKCompressor(compression_ratio)
        self.world_size = dist.get_world_size()
        self.rank = dist.get_rank()
    
    def allreduce_compressed(self, param):
        """压缩的AllReduce"""
        if param.grad is None:
            return
        
        # 压缩梯度
        values, indices = self.compressor.compress(param.grad)
        
        # 收集所有进程的索引（需要一致）
        # 这里简化处理，实际需要同步索引
        
        # 传输压缩后的值
        # 实际实现需要考虑不同进程的稀疏模式不同
        
        # 简化的全量AllReduce作为fallback
        dist.all_reduce(param.grad)
    
    def step(self):
        """执行一步参数更新"""
        for param in self.model.parameters():
            if param.requires_grad and param.grad is not None:
                self.allreduce_compressed(param)


# 使用示例
def demo_compression():
    """演示梯度压缩效果"""
    
    # 创建模拟梯度
    gradient = torch.randn(1000000)
    
    compressor = TopKCompressor(compression_ratio=0.1)
    values, indices = compressor.compress(gradient)
    
    original_size = gradient.numel() * 4  # FP32
    compressed_size = values.numel() * 4 + indices[0].numel() * 4
    
    print(f"Original size: {original_size / 1024:.1f} KB")
    print(f"Compressed size: {compressed_size / 1024:.1f} KB")
    print(f"Compression ratio: {compressed_size / original_size * 100:.1f}%")
    
    # 重建并计算误差
    reconstructed = compressor.decompress(values, indices, gradient.shape, gradient.device)
    error = torch.norm(gradient - reconstructed) / torch.norm(gradient)
    print(f"Relative error: {error * 100:.2f}%")


demo_compression()

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

class OverlapOptimizedModel(nn.Module):
    """
    优化通信与计算重叠的模型
    
    通过合理组织层结构，使梯度尽早可用
    """
    
    def __init__(self, input_size, hidden_sizes, output_size):
        super().__init__()
        
        # 将网络分为多个阶段，每个阶段的梯度可以独立同步
        self.stages = nn.ModuleList()
        
        prev_size = input_size
        for hidden_size in hidden_sizes:
            stage = nn.Sequential(
                nn.Linear(prev_size, hidden_size),
                nn.ReLU(),
                nn.LayerNorm(hidden_size),
            )
            self.stages.append(stage)
            prev_size = hidden_size
        
        self.output_layer = nn.Linear(prev_size, output_size)
    
    def forward(self, x):
        for stage in self.stages:
            x = stage(x)
        return self.output_layer(x)


def configure_ddp_overlap(model, bucket_size_mb=25):
    """
    配置DDP以优化通信与计算重叠
    
    Args:
        model: DDP包装的模型
        bucket_size_mb: 梯度桶大小（MB）
    """
    ddp_model = DDP(
        model,
        bucket_cap_mb=bucket_size_mb,
        find_unused_parameters=False,
        gradient_as_bucket_view=True,  # 减少内存拷贝
    )
    return ddp_model


def profile_communication_overlap():
    """
    分析通信与计算的重叠情况
    
    使用PyTorch Profiler
    """
    from torch.profiler import profile, record_function, ProfilerActivity
    
    # 创建模型和数据
    model = OverlapOptimizedModel(784, [512, 256, 128], 10).cuda()
    ddp_model = configure_ddp_overlap(model)
    
    x = torch.randn(64, 784).cuda()
    target = torch.randint(0, 10, (64,)).cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(ddp_model.parameters())
    
    # Profiling
    with profile(
        activities=[ProfilerActivity.CUDA],
        record_shapes=True,
    ) as prof:
        for _ in range(10):
            optimizer.zero_grad()
            output = ddp_model(x)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
    
    # 打印分析结果
    print(prof.key_averages().table(sort_by="cuda_time_total"))


if __name__ == "__main__":
    profile_communication_overlap()