分布式训练：Data Parallel, Model Parallel

import torch

# 估算模型参数量和显存需求
def estimate_memory(model_params, batch_size, seq_len, hidden_dim, dtype_bytes=4):
    """估算训练时所需显存（GB）
    model_params: 模型总参数量
    batch_size * seq_len * hidden_dim: activation memory 的关键变量
    dtype_bytes: 精度字节数 (FP32=4, FP16=2, BF16=2)
    """
    param_memory = model_params * dtype_bytes / 1e9          # 参数本身
    optimizer_memory = model_params * dtype_bytes * 3 / 1e9  # Adam: param + moment1 + moment2
    grad_memory = model_params * dtype_bytes / 1e9           # 梯度
    activation_memory = batch_size * seq_len * hidden_dim * dtype_bytes / 1e9
    
    total = param_memory + optimizer_memory + grad_memory + activation_memory
    return total

# GPT-3 估算：175B 参数，batch=32, seq_len=2048, hidden=12288
gpt3_mem = estimate_memory(175e9, 32, 2048, 12288)
print(f"GPT-3 训练估计显存: {gpt3_mem:.1f} GB")
# 输出: GPT-3 训练估计显存: 4200.0 GB (远超单卡 80GB)

import torch
import torch.distributed as dist

def get_world_info():
    """获取分布式环境信息"""
    if dist.is_available() and dist.is_initialized():
        world_size = dist.get_world_size()
        rank = dist.get_rank()
        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        return world_size, rank, local_rank
    return 1, 0, 0

def print_gpu_stats():
    """打印所有 GPU 的显存使用统计"""
    for i in range(torch.cuda.device_count()):
        mem_alloc = torch.cuda.max_memory_allocated(i) / 1e9
        mem_reserved = torch.cuda.max_memory_reserved(i) / 1e9
        print(f"GPU {i}: 已分配 {mem_alloc:.2f} GB, 已预留 {mem_reserved:.2f} GB")
    
    # 检查 NCCL 通信后端
    print(f"NCCL 可用: {torch.distributed.is_nccl_available()}")
    print(f"后端: {dist.get_backend() if dist.is_initialized() else '未初始化'}")

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

def setup_distributed(rank, world_size):
    """初始化分布式进程组"""
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

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

def train_loop(rank, world_size):
    setup_distributed(rank, world_size)
    
    # 创建模型并移动到当前 GPU
    model = MyModel().to(rank)
    # 关键: 用 DDP 包裹模型，自动处理梯度同步
    model = DDP(model, device_ids=[rank], output_device=rank)
    
    # 使用 DistributedSampler 确保每张卡拿到不同的数据子集
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 每轮重新 shuffle
        for batch in dataloader:
            optimizer.zero_grad()
            loss = model(batch)
            loss.backward()        # 自动 AllReduce 梯度
            optimizer.step()
    
    cleanup_distributed()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(train_loop, args=(world_size,), nprocs=world_size)

# DDP 底层通信原语详解
import torch.distributed as dist
import torch

# AllReduce: 所有进程的 tensor 归约求和/平均
tensor = torch.randn(4, device="cuda")
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
# tensor 现在在所有 GPU 上相同，值为各 GPU 原始值之和

# Broadcast: 从 rank 0 广播到所有进程
if dist.get_rank() == 0:
    data = torch.randn(4, device="cuda")
else:
    data = torch.zeros(4, device="cuda")
dist.broadcast(data, src=0)
# 现在所有 GPU 上的 data 都等于 rank 0 的原始值

# Reduce: 归约到指定进程（非 AllReduce，只有 dst 拿到结果）
tensor = torch.randn(4, device="cuda")
if dist.get_rank() == 0:
    dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)
    
# Gather: 收集所有进程的 tensor 到指定进程
local_tensor = torch.randn(2, device="cuda")
if dist.get_rank() == 0:
    gather_list = [torch.zeros(2, device="cuda") for _ in range(dist.get_world_size())]
    dist.gather(local_tensor, gather_list, dst=0)
# rank 0 的 gather_list 包含所有 GPU 的 tensor

# DeepSpeed ZeRO 配置示例 (deepspeed_config.json)
import json

deepspeed_config = {
    "train_batch_size": 256,
    "gradient_accumulation_steps": 4,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 3e-4,
            "betas": [0.9, 0.95],
            "eps": 1e-8,
            "weight_decay": 0.1
        }
    },
    "zero_optimization": {
        "stage": 3,            # ZeRO-3: 分区参数+梯度+优化器状态
        "offload_optimizer": {
            "device": "cpu",   # 优化器状态卸载到 CPU 内存
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",   # 参数卸载到 CPU 内存
            "pin_memory": True
        },
        "overlap_comm": True,   # 通信与计算重叠
        "contiguous_gradients": True,
        "reduce_bucket_size": 5e8,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e5
    },
    "fp16": {
        "enabled": True,
        "loss_scale": 0,         # 动态 loss scaling
        "initial_scale_power": 16
    }
}

print(json.dumps(deepspeed_config, indent=2))

# PyTorch FSDP (Fully Sharded Data Parallel) - ZeRO-3 的 PyTorch 原生实现
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

def setup_fsdp(model, rank, world_size):
    """使用 FSDP 包装模型，等效于 ZeRO-3"""
    
    # 混合精度策略
    mp_policy = MixedPrecision(
        param_dtype=torch.float16,      # 参数以 FP16 存储
        reduce_dtype=torch.float32,     # 梯度归约用 FP32
        buffer_dtype=torch.float16      # buffer 用 FP16
    )
    
    # 自动包装策略：大于 100M 参数的子模块独立分片
    auto_wrap_policy = size_based_auto_wrap_policy(min_num_params=1e8)
    
    fsdp_model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
        mixed_precision=mp_policy,
        auto_wrap_policy=auto_wrap_policy,
        device_id=rank,
        sync_module_states=True,       # 初始化时同步参数
        use_orig_params=True           # 保留原始参数名
    )
    
    return fsdp_model

# 训练循环与 DDP 几乎相同，FSDP 内部自动处理参数的 gather/shard
for batch in dataloader:
    optimizer.zero_grad()
    loss = fsdp_model(batch)
    loss.backward()
    optimizer.step()

# 张量并行：Megatron-LM 风格的列并行 + 行并行
import torch
import torch.nn as nn
from torch.distributed import tensor as dist_tensor
from torch.distributed.device_mesh import init_device_mesh

class ColumnParallelLinear(nn.Module):
    """列并行线性层：将权重矩阵按列切分
    原始: Y = X @ [W1, W2] + [b1, b2]
    拆分: GPU0 计算 Y1 = X @ W1 + b1, GPU1 计算 Y2 = X @ W2 + b2
    """
    def __init__(self, in_features, out_features, world_size, bias=True):
        super().__init__()
        assert out_features % world_size == 0, "输出维度必须能被 GPU 数整除"
        self.local_out = out_features // world_size
        self.weight = nn.Parameter(torch.randn(in_features, self.local_out))
        if bias:
            self.bias = nn.Parameter(torch.zeros(self.local_out))
        else:
            self.register_parameter("bias", None)
    
    def forward(self, x):
        # x 在所有 GPU 上相同
        out = x @ self.weight   # (batch, in) @ (in, local_out) -> (batch, local_out)
        if self.bias is not None:
            out = out + self.bias
        return out  # 需要后续层的行并行来 AllGather

class RowParallelLinear(nn.Module):
    """行并行线性层：将权重矩阵按行切分，输出需要 AllReduce"""
    def __init__(self, in_features, out_features, world_size, bias=True):
        super().__init__()
        assert in_features % world_size == 0
        self.local_in = in_features // world_size
        self.weight = nn.Parameter(torch.randn(self.local_in, out_features))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
    
    def forward(self, x):
        # x 是局部输入 (batch, local_in)
        out = x @ self.weight   # (batch, local_in) @ (local_in, out) -> (batch, out)
        # AllReduce 将各 GPU 的 partial sum 累加
        torch.distributed.all_reduce(out, op=torch.distributed.ReduceOp.SUM)
        if self.bias is not None:
            out = out + self.bias
        return out

# 流水线并行：PyTorch Pipeline Parallel
import torch
from torch.distributed.pipelining import pipeline, SplitPoint
from torch.distributed.pipelining import ScheduleGPipe

class TransformerBlock(nn.Module):
    def __init__(self, dim, n_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.ffn(self.norm2(x))
        return x

class LargeModel(nn.Module):
    def __init__(self, vocab_size=50257, dim=768, n_layers=24, n_heads=12):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, dim)
        self.layers = nn.ModuleList([
            TransformerBlock(dim, n_heads) for _ in range(n_layers)
        ])
        self.head = nn.Linear(dim, vocab_size)
    
    def forward(self, input_ids):
        x = self.embedding(input_ids)
        for layer in self.layers:
            x = layer(x)
        return self.head(x)

# 使用 torch.distributed.pipelining 自动切分模型
model = LargeModel()
# 在每 6 层后设置切分点
split_spec = {
    f"layers.{5}": SplitPoint.BEGINNING,
    f"layers.{11}": SplitPoint.BEGINNING,
    f"layers.{17}": SplitPoint.BEGINNING,
}
pipe = pipeline(model, num_chunks=8, split_spec=split_spec)

# GPipe 调度：将 batch 分为多个 micro-batch
micro_batches = 8
schedule = ScheduleGPipe(pipe, micro_batches)
# 执行流水线
loss = schedule.step(input_ids)

# PyTorch 原生 AMP (Automatic Mixed Precision)
import torch
from torch.cuda.amp import autocast, GradScaler

model = MyLargeModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # 自动管理 loss scaling

for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        
        # autocast 自动选择精度：矩阵乘法用 FP16，其他用 FP32
        with autocast(dtype=torch.float16):
            output = model(batch["input_ids"])
            loss = criterion(output, batch["labels"])
        
        # scaler.scale 放大 loss 以避免梯度下溢
        scaler.scale(loss).backward()
        
        # scaler.unscale 反缩放梯度，然后 clip 避免溢出
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        # scaler.step 先 unscale 再 optimizer.step
        scaler.step(optimizer)
        scaler.update()  # 动态调整 scaling factor

# BF16 更简单：不需要 GradScaler（指数范围足够大）
for batch in dataloader:
    optimizer.zero_grad()
    with autocast(dtype=torch.bfloat16):
        output = model(batch["input_ids"])
        loss = criterion(output, batch["labels"])
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

# FP16 vs BF16 vs FP32 精度对比
import struct

def float_info(name, exp_bits, mantissa_bits):
    """打印浮点数格式的精度信息"""
    max_val = (2 - 2(-mantissa_bits)) * 2(2**(exp_bits-1) - 1)
    min_normal = 2(-(2(exp_bits-1) - 2))
    print(f"{name}:")
    print(f"  最大可表示值: {max_val:.2e}")
    print(f"  最小正规数: {min_normal:.2e}")
    print(f"  精度(尾数): {mantissa_bits} bits")

float_info("FP32", 8, 23)
float_info("FP16", 5, 10)
float_info("BF16", 8, 7)
# 输出:
#   FP32: 最大可表示值: 3.40e+38, 最小正规数: 1.18e-38, 精度: 23 bits
#   FP16: 最大可表示值: 6.55e+04, 最小正规数: 6.10e-05, 精度: 10 bits
#   BF16: 最大可表示值: 3.39e+38, 最小正规数: 1.18e-38, 精度: 7 bits

# Tensor Core 加速测试
def benchmark_matmul(size=4096, dtype=torch.float16, device="cuda"):
    A = torch.randn(size, size, dtype=dtype, device=device)
    B = torch.randn(size, size, dtype=dtype, device=device)
    
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    
    start.record()
    C = torch.matmul(A, B)
    end.record()
    torch.cuda.synchronize()
    
    elapsed_ms = start.elapsed_time(end)
    print(f"{dtype} {size}x{size}: {elapsed_ms:.2f} ms")

benchmark_matmul(dtype=torch.float32)  # FP32: 较慢
benchmark_matmul(dtype=torch.float16)  # FP16: Tensor Core 加速
benchmark_matmul(dtype=torch.bfloat16) # BF16: Tensor Core 加速

# 梯度累积的标准实现
import torch

def train_with_gradient_accumulation(model, dataloader, optimizer, criterion,
                                      accumulation_steps=4, max_grad_norm=1.0):
    """使用梯度累积训练模型
    
    Args:
        accumulation_steps: 累积多少个 micro-batch 后才更新
    有效 batch_size = per_gpu_batch_size * accumulation_steps * world_size
    """
    model.train()
    optimizer.zero_grad()
    
    for step, batch in enumerate(dataloader):
        # 前向传播
        output = model(batch["input_ids"])
        loss = criterion(output, batch["labels"])
        
        # 将 loss 除以累积步数，确保梯度尺度正确
        loss = loss / accumulation_steps
        loss.backward()
        
        # 每 accumulation_steps 步执行一次参数更新
        if (step + 1) % accumulation_steps == 0:
            # 梯度裁剪（在全局梯度上操作）
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            optimizer.step()
            optimizer.zero_grad()
    
    # 处理最后不足 accumulation_steps 的剩余 batch
    if (step + 1) % accumulation_steps != 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
        optimizer.step()
        optimizer.zero_grad()

# 使用示例：per-GPU batch_size=8, accumulation=4, 8张GPU
# 有效 batch_size = 8 * 4 * 8 = 256
print("有效 batch_size = 8 * 4 * 8 = 256")

# DDP + 梯度累积 + 混合精度的完整训练循环
import torch
from torch.cuda.amp import autocast, GradScaler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler

class Trainer:
    def __init__(self, model, dataloader, rank, world_size,
                 accumulation_steps=4, lr=1e-4):
        self.model = DDP(model.cuda(rank), device_ids=[rank])
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
        self.scaler = GradScaler()
        self.accumulation_steps = accumulation_steps
        
        sampler = DistributedSampler(dataloader.dataset, 
                                      num_replicas=world_size, rank=rank)
        self.dataloader = torch.utils.data.DataLoader(
            dataloader.dataset, batch_size=8, sampler=sampler)
    
    def train_epoch(self, epoch):
        self.model.train()
        self.dataloader.sampler.set_epoch(epoch)
        self.optimizer.zero_grad()
        
        for step, batch in enumerate(self.dataloader):
            with autocast(dtype=torch.bfloat16):
                output = self.model(batch["input_ids"])
                loss = self.compute_loss(output, batch["labels"])
            
            # 标准化 loss 以支持梯度累积
            loss = loss / self.accumulation_steps
            self.scaler.scale(loss).backward()
            
            # 累积够步数后才更新
            if (step + 1) % self.accumulation_steps == 0:
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()
        
        # 处理剩余 batch
        if (step + 1) % self.accumulation_steps != 0:
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.scaler.step(self.optimizer)
            self.scaler.update()
            self.optimizer.zero_grad()

#!/usr/bin/env python
"""多 GPU 训练入口 - 使用 torchrun 启动
torchrun --nproc_per_node=8 --master_port=29500 train_ddp.py
"""
import os
import argparse
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torch.cuda.amp import autocast, GradScaler
from transformers import AutoModelForCausalLM, AutoTokenizer

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--accum_steps", type=int, default=4)
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--max_grad_norm", type=float, default=1.0)
    parser.add_argument("--checkpoint_dir", type=str, default="./checkpoints")
    args = parser.parse_args()
    
    # torchrun 自动设置这些环境变量
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    rank = int(os.environ["RANK"])
    
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    
    # 加载模型
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
    model.gradient_checkpointing_enable()  # 用计算换显存
    model = model.to(local_rank)
    model = DDP(model, device_ids=[local_rank],
                gradient_as_bucket_view=True,  # 减少显存峰值
                static_graph=True)  # 静态图优化
    
    # 使用 SyncBatchNorm（如果模型中有 BN 层）
    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, 
                                   betas=(0.9, 0.95), weight_decay=0.1)
    scaler = GradScaler()
    
    train_dataset = MyDataset()
    sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(train_dataset, batch_size=args.batch_size, 
                            sampler=sampler, num_workers=4, pin_memory=True)
    
    print(f"Rank {rank}: 开始训练，有效 batch_size = {args.batch_size * args.accum_steps * world_size}")
    # Rank 0: 开始训练，有效 batch_size = 8 * 4 * 8 = 256
    
    for epoch in range(args.epochs):
        train_one_epoch(model, dataloader, optimizer, scaler, sampler,
                        args, rank, epoch)
    
    dist.destroy_process_group()

# 检查点保存与断点续训
import os
import torch
import torch.distributed as dist

def save_checkpoint(model, optimizer, scaler, epoch, step, checkpoint_dir, rank):
    """保存分布式训练检查点"""
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    # 只让 rank 0 保存，避免多进程写同一文件
    if rank == 0:
        checkpoint = {
            "epoch": epoch,
            "step": step,
            "model_state": model.module.state_dict(),  # DDP 包装的模型
            "optimizer_state": optimizer.state_dict(),
            "scaler_state": scaler.state_dict(),
        }
        path = os.path.join(checkpoint_dir, f"checkpoint_epoch{epoch}_step{step}.pt")
        torch.save(checkpoint, path)
        print(f"Rank 0: 检查点已保存: {path}")
    
    dist.barrier()  # 等待所有进程到达此处

def load_checkpoint(model, optimizer, scaler, checkpoint_path, map_location):
    """加载检查点并恢复训练状态"""
    checkpoint = torch.load(checkpoint_path, map_location=map_location)
    
    model.module.load_state_dict(checkpoint["model_state"])
    optimizer.load_state_dict(checkpoint["optimizer_state"])
    scaler.load_state_dict(checkpoint["scaler_state"])
    
    return checkpoint["epoch"], checkpoint["step"]

# 断点续训使用示例
def train_one_epoch(model, dataloader, optimizer, scaler, sampler,
                    args, rank, epoch):
    checkpoint_path = os.path.join(args.checkpoint_dir, 
                                    f"checkpoint_epoch{epoch - 1}_latest.pt")
    start_step = 0
    
    if os.path.exists(checkpoint_path):
        start_epoch, start_step = load_checkpoint(
            model, optimizer, scaler, checkpoint_path, 
            map_location=f"cuda:{rank}")
        print(f"Rank {rank}: 从 epoch {start_epoch} step {start_step} 恢复")
    
    model.train()
    optimizer.zero_grad()
    
    for step, batch in enumerate(dataloader):
        if step <= start_step:
            continue
        # ... 正常训练循环
        pass