FSDP

# DDP vs FSDP 显存占用对比def calc_memory(params_billion, n_gpus, mode='ddp'):    params = params_billion * 1e9    param_mem = params * 4 / 1e9    grad_mem = params * 4 / 1e9    optim_mem = params * 8 / 1e9    if mode == 'ddp':        total = param_mem + grad_mem + optim_mem    elif mode == 'fsdp':        total = (param_mem + grad_mem + optim_mem) / n_gpus    return total for params in [7, 13, 70]:    ddp = calc_memory(params, 8, 'ddp')    fsdp = calc_memory(params, 8, 'fsdp')    print(f'{params}B参数: DDP={ddp:.0f}GB, FSDP={fsdp:.0f}GB (8卡)')

import torchclass FlatParamShard:    def __init__(self, params, world_size, rank):        flat = torch.cat([p.data.flatten() for p in params])        self.total_size = flat.numel()        self.shard_size = self.total_size // world_size        if self.total_size % world_size != 0:            pad = world_size - (self.total_size % world_size)            flat = torch.nn.functional.pad(flat, (0, pad))            self.total_size = flat.numel()            self.shard_size = self.total_size // world_size        start = rank * self.shard_size        self.shard = flat[start:start + self.shard_size].clone()    def all_gather(self, shards_list):        return torch.cat(shards_list) params = [torch.randn(100), torch.randn(200), torch.randn(300)]shard_0 = FlatParamShard(params, world_size=2, rank=0)shard_1 = FlatParamShard(params, world_size=2, rank=1)print(f'总参数量: {shard_0.total_size}')print(f'GPU0分片: {shard_0.shard.numel()}, GPU1分片: {shard_1.shard.numel()}')full = shard_0.all_gather([shard_0.shard, shard_1.shard])print(f'All-Gather恢复: {full.numel()} 个参数')

import torchclass FSDPCommunicationSim:    def __init__(self, param_size, world_size, rank):        self.param_size = param_size        self.world_size = world_size        self.rank = rank        self.shard_size = param_size // world_size        self.param_shard = torch.randn(self.shard_size)        self.grad_shard = torch.zeros(self.shard_size)    def forward_step(self):        full_param = self.param_shard.repeat(self.world_size)        print(f'GPU{self.rank} 前向All-Gather: 获取{full_param.numel()}个参数')        output = full_param.sum()        del full_param        print(f'GPU{self.rank} 释放完整参数，保留{self.shard_size}个分片')        return output    def backward_step(self):        full_param = self.param_shard.repeat(self.world_size)        full_grad = torch.ones_like(full_param)        print(f'GPU{self.rank} 反向All-Gather: 计算完整梯度')        self.grad_shard = full_grad[self.rank*self.shard_size:(self.rank+1)*self.shard_size]        del full_param, full_grad        print(f'GPU{self.rank} Reduce-Scatter: 保留{self.grad_shard.numel()}个梯度分片') sim = FSDPCommunicationSim(param_size=1000, world_size=4, rank=0)sim.forward_step()sim.backward_step()

import torchimport torch.nn as nnfrom torch.distributed.fsdp.wrap import ModuleWrapPolicy class SimpleTransformer(nn.Module):    def __init__(self, vocab_size, d_model, n_layers):        super().__init__()        self.embedding = nn.Embedding(vocab_size, d_model)        self.layers = nn.ModuleList([            nn.TransformerDecoderLayer(d_model, nhead=8)            for _ in range(n_layers)        ])        self.lm_head = nn.Linear(d_model, vocab_size) model = SimpleTransformer(vocab_size=50000, d_model=1024, n_layers=12)wrap_policy = ModuleWrapPolicy({nn.TransformerDecoderLayer})print(f'模型总参数量: {sum(p.numel() for p in model.parameters())/1e6:.1f}M')for name, module in model.named_modules():    n_params = sum(p.numel() for p in module.parameters(recurse=False))    if n_params > 0:        print(f'  {name}: {n_params/1e6:.1f}M params')

class PrefetchSimulator:    def __init__(self, n_units, compute_time=0.1, comm_time=0.08):        self.n_units = n_units        self.compute_time = compute_time        self.comm_time = comm_time    def run_without_prefetch(self):        total = 0        for i in range(self.n_units):            total += self.comm_time + self.compute_time        return total    def run_with_prefetch(self):        total = self.comm_time + self.compute_time        for i in range(1, self.n_units):            step_time = max(self.comm_time, self.compute_time)            total += step_time        return total sim = PrefetchSimulator(n_units=12)t_no = sim.run_without_prefetch()t_yes = sim.run_with_prefetch()print(f'无预取: {t_no:.2f}s, 有预取: {t_yes:.2f}s, 加速比: {t_no/t_yes:.2f}x')

def estimate_memory(params_b, seq_len, n_layers, n_gpus,                    use_bf16=False, use_offload=False, use_ckpt=False):    params = params_b * 1e9    if use_bf16:        base = (params*2 + params*2 + params*4) / n_gpus / 1e9    else:        base = (params*4 + params*4 + params*8) / n_gpus / 1e9    d_model = 4096    activation = seq_len * d_model * n_layers * 4 * 2 / 1e9    if use_ckpt:        activation *= 0.1    offload_save = base * 0.3 if use_offload else 0    return base + activation - offload_save print('70B, seq=2048, 128层, 8卡A100:')for bf16 in [False, True]:    for offload in [False, True]:        for ckpt in [False, True]:            mem = estimate_memory(70, 2048, 128, 8, bf16, offload, ckpt)            cfg = f'BF16={bf16} Offload={offload} CKPT={ckpt}'            print(f'  {cfg}: {mem:.1f}GB [{"OK" if mem<80 else "OOM"}]')

def compare_frameworks():    features = {        '特性': ['参数分片','梯度分片','优化器分片','原生PyTorch集成','CPU/NVMe Offload','与torch.compile兼容','Pipeline Parallel'],        'DDP': ['否','否','否','是','否','是','需手动'],        'FSDP': ['是','是','是','是','CPU Offload','FSDP2支持','需组合'],        'DeepSpeed ZeRO-3': ['是','是','是','否','CPU+NVMe','部分支持','原生支持']    }    header = f"{'特性':<24} {'DDP':<10} {'FSDP':<16} {'DeepSpeed ZeRO-3':<16}"    print(header)    print('-'*70)    for i in range(len(features['特性'])):        print(f"{features['特性'][i]:<24} {features['DDP'][i]:<10} {features['FSDP'][i]:<16} {features['DeepSpeed ZeRO-3'][i]:<16}")compare_frameworks()

import torchimport torch.nn as nn class TransformerBlock(nn.Module):    def __init__(self, d_model, nhead, dim_ff):        super().__init__()        self.attn = nn.MultiheadAttention(d_model, nhead)        self.ffn = nn.Sequential(nn.Linear(d_model, dim_ff), nn.GELU(), nn.Linear(dim_ff, d_model))        self.norm1 = nn.LayerNorm(d_model)        self.norm2 = nn.LayerNorm(d_model)    def forward(self, x):        h = self.norm1(x)        h = x + self.attn(h, h, h)[0]        h = self.norm2(h)        h = h + self.ffn(h)        return h class LargeLM(nn.Module):    def __init__(self, vocab_size, d_model, n_layers, nhead, dim_ff):        super().__init__()        self.embedding = nn.Embedding(vocab_size, d_model)        self.layers = nn.ModuleList([TransformerBlock(d_model, nhead, dim_ff) for _ in range(n_layers)])        self.lm_head = nn.Linear(d_model, vocab_size)    def forward(self, input_ids):        h = self.embedding(input_ids)        for layer in self.layers:            h = layer(h)        return self.lm_head(h) def analyze_memory(model, n_gpus, mode='fsdp', use_bf16=False):    params = sum(p.numel() for p in model.parameters())    bpp = 2 if use_bf16 else 4    param_mem = params * bpp / 1e9    grad_mem = params * bpp / 1e9    optim_mem = params * 8 / 1e9    if mode == 'ddp':        return param_mem + grad_mem + optim_mem    elif mode == 'fsdp':        return (param_mem + grad_mem + optim_mem) / n_gpus model = LargeLM(vocab_size=50000, d_model=4096, n_layers=32, nhead=32, dim_ff=16384)print(f'参数量: {sum(p.numel() for p in model.parameters())/1e9:.2f}B')for n in [8, 16, 32]:    print(f'{n}卡: DDP={analyze_memory(model,n,"ddp"):.1f}GB FSDP={analyze_memory(model,n,"fsdp"):.1f}GB FSDP+BF16={analyze_memory(model,n,"fsdp",True):.1f}GB') print('\n--- 扁平参数分片 ---')block = model.layers[0]flat = torch.cat([p.data.flatten() for p in block.parameters()])print(f'TransformerBlock参数: {flat.numel()/1e6:.1f}M')for r in range(4):    s = flat.numel() // 4    print(f'  GPU{r}: {s/1e6:.1f}M params')