[ai学习笔记]分布式训练原理：DeepSeek千卡集群通信优化策略

# 通信热点分析工具示例
from torch.profiler import profile

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    model.train()
    
print(prof.key_averages().table(
    sort_by="cuda_time_total", 
    row_limit=10
))

class TopologyAwareScheduler:
    def __init__(self, num_gpus, bandwidth_matrix):
        self.bw_matrix = bandwidth_matrix  # 设备间带宽矩阵
        self.groups = self._detect_affinity()
        
    def _detect_affinity(self):
        # 基于带宽检测最优通信组
        groups = []
        visited = set()
        for i in range(len(self.bw_matrix)):
            if i not in visited:
                group = [i]
                for j in range(i+1, len(self.bw_matrix)):
                    if self.bw_matrix[i][j] > 180:  # 200G IB链路
                        group.append(j)
                        visited.add(j)
                groups.append(group)
        return groups

class DynamicSparseCompressor:
    def __init__(self, compression_ratio=0.05):
        self.ratio = compression_ratio
        
    def compress(self, tensor):
        values, indices = torch.topk(
            torch.abs(tensor.flatten()), 
            k=int(self.ratio * tensor.numel())
        return SparseTensor(values, indices, tensor.size())
    
    def decompress(self, sparse_tensor):
        output = torch.zeros(sparse_tensor.size)
        output.data[sparse_tensor.indices] = sparse_tensor.values
        return output

// 自定义AllReduce内核（基于NVIDIA NCCL）
ncclResult_t deepseekAllReduce(
    const void* sendbuff, void* recvbuff, size_t count,
    ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
    cudaStream_t stream
) {
    // 分阶段流水线
    splitToChunks(sendbuff, chunks); 
    for (auto& chunk : chunks) {
        ncclSend(chunk, comm, stream);
        ncclRecv(chunk, comm, stream); 
        overlapCompression(chunk);  // 压缩与传输重叠
    }
    assembleChunks(recvbuff, chunks);
    return ncclSuccess;
}

class GradientAccumulator:
    def __init__(self, model, accum_steps=4):
        self.model = model
        self.accum_steps = accum_steps
        self._init_grad_buffer()
        
    def _init_grad_buffer(self):
        self.grad_buffers = [
            torch.zeros_like(p) for p in self.model.parameters()
        ]
    
    def step(self, closure=None):
        # 累积梯度
        for p, buf in zip(self.model.parameters(), self.grad_buffers):
            if p.grad is not None:
                buf.add_(p.grad / self.accum_steps)
        
        # 条件触发通信
        if self._current_step % self.accum_steps == 0:
            self._sync_gradients()
            self.optimizer.step()
            self.zero_grad()

class HierarchicalDataLoader:
    def __init__(self, dataset, hot_ratio=0.2):
        self.hot_cache = SSDcache(dataset[:hot_ratio])  # 热点数据
        self.cold_store = HDDstore(dataset[hot_ratio:])  # 冷数据
        self.prefetcher = DataPrefetcher()
        
    def __iter__(self):
        for batch in self.prefetcher:
            # 异步预取下一批次
            next_batch = self._fetch_next()
            self.prefetcher.load(next_batch)
            yield batch
            
    def _fetch_next(self):
        if np.random.rand() < 0.8:  # 80%概率访问热数据
            return self.hot_cache.sample()
        else:
            return self.cold_store.fetch()

import torch
import torch.nn as nn
import torch.distributed as dist
import torch.nn.functional as F

class MLAAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_c=128, d_h=64):
        super().__init__()
        self.n_heads = n_heads
        self.d_h = d_h
        self.d_c = d_c
        self.W_DKV = nn.Linear(d_model, d_c)  # KV压缩层
        self.W_UK = nn.Linear(d_c, n_heads * d_h)  # 键扩展
        self.W_UV = nn.Linear(d_c, n_heads * d_h)  # 值扩展
        self.W_Q = nn.Linear(d_model, n_heads * d_h)  # 查询
        self.W_O = nn.Linear(n_heads * d_h, d_model)  # 输出

    def forward(self, x):
        batch_size, seq_len, _ = x.size()
        c = self.W_DKV(x)  # 压缩至潜空间
        k = self.W_UK(c).view(batch_size, seq_len, self.n_heads, self.d_h)
        v = self.W_UV(c).view(batch_size, seq_len, self.n_heads, self.d_h)
        q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_h)

        scores = torch.einsum("bihd,bjhd->bhij", q, k) / (self.d_h ** 0.5)
        weights = F.softmax(scores, dim=-1)
        attn = torch.einsum("bhij,bjhd->bihd", weights, v)
        attn = attn.view(batch_size, seq_len, -1)
        return self.W_O(attn)

class DualPipeModel(nn.Module):
    def __init__(self, layers_per_stage, num_stages):
        super().__init__()
        self.stages = nn.ModuleList([nn.ModuleList([MLAAttention() for _ in range(layers_per_stage)]) for _ in range(num_stages)])
        self.rank = dist.get_rank()
        self.num_stages = num_stages

    def forward(self, x):
        stage_out = x
        if self.rank == 0:  # 前向传播
            for layer in self.stages[0]:
                stage_out = layer(stage_out)
            dist.send(tensor=stage_out, dst=1)
        elif self.rank == self.num_stages - 1:  # 最后阶段
            dist.recv(tensor=stage_out, src=self.rank-1)
            for layer in self.stages[-1]:
                stage_out = layer(stage_out)
        else:  # 中间阶段
            dist.recv(tensor=stage_out, src=self.rank-1)
            for layer in self.stages[self.rank]:
                stage_out = layer(stage_out)
            dist.send(tensor=stage_out, dst=self.rank+1)
        return stage_out

def train(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    model = DualPipeModel(layers_per_stage=2, num_stages=world_size).cuda(rank)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    data = torch.randn(32, 10, 512).cuda(rank)

    for _ in range(10):
        optimizer.zero_grad()
        out = model(data)
        loss = out.sum()  # 模拟损失
        loss.backward()
        optimizer.step()
    dist.destroy_process_group()

python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=0 train.py
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=1 train.py

#!/bin/bash
# DeepSeek集群启动脚本
NUM_NODES=128
GPUS_PER_NODE=8

torchrun \
    --nnodes=$NUM_NODES \
    --nproc_per_node=$GPUS_PER_NODE \
    --rdzv_id=deepseek_train \
    --rdzv_backend=c10d \
    --rdzv_endpoint=coordinator:29500 \
    --max_restarts=3 \
    \
    train.py \
    --batch_size 4096 \
    --use_gradient_accumulation \
    --topology_aware \
    --compression_ratio 0.05

class CheckpointManager:
    def __init__(self, save_interval=3600):
        self.save_interval = save_interval  # 每1小时保存
        self.last_save = time.time()
        
    def maybe_save(self, model):
        if time.time() - self.last_save > self.save_interval:
            # 异步写入分布式存储
            threading.Thread(target=self._async_save, args=(model.state_dict(),)).start()
            self.last_save = time.time()
    
    def _async_save(self, state_dict):
        # 使用Zstandard压缩
        buffer = pickle.dumps(state_dict)
        compressed = zstd.compress(buffer)
        distributed_store.write(checkpoint_path, compressed)