02 - 训练基础设施¶

import os
import torch
import torch.nn as nn
from torch.amp import autocast, GradScaler


class MixedPrecisionTrainer:
    """
    混合精度训练器

    核心原理：
    1. 前向传播使用低精度（FP16/BF16）加速计算
    2. 使用 GradScaler 放大损失，防止梯度下溢
    3. 参数更新在 FP32 下进行，保证精度

    FP16 训练需要 GradScaler，BF16 不需要（因为动态范围足够大）
    """

    def __init__(self, model, optimizer, device='cuda', dtype='fp16'):
        self.model = model.to(device)
        self.optimizer = optimizer
        self.device = device
        self.dtype = dtype

        # FP16 需要 GradScaler 防止梯度下溢
        # BF16 动态范围与 FP32 相同，不需要 scaler
        self.scaler = GradScaler() if dtype == 'fp16' else None

    def train_step(self, batch):
        """
        单步混合精度训练
        """
        self.model.train()

        input_ids = batch['input_ids'].to(self.device)
        labels = batch['labels'].to(self.device)

        self.optimizer.zero_grad()

        # 选择精度类型
        amp_dtype = torch.float16 if self.dtype == 'fp16' else torch.bfloat16

        # 自动混合精度上下文
        with autocast(device_type='cuda', dtype=amp_dtype):
            outputs = self.model(input_ids, labels=labels)
            loss = outputs.loss

        if self.scaler is not None:
            # FP16 路径：缩放梯度
            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            # BF16 路径：直接反向传播
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.optimizer.step()

        return loss.item()


# BF16 训练完整示例（推荐，需要 Ampere+ GPU）
def bf16_training_example():
    """
    BF16 训练示例

    BF16 相比 FP16 的优势：
    - 不需要 GradScaler（简化代码）
    - 不容易梯度下溢（训练更稳定）
    - 动态范围与 FP32 相同
    """
    model = nn.Linear(1024, 1024).cuda()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    dummy_input = torch.randn(32, 1024, device='cuda')

    for step in range(100):
        optimizer.zero_grad()

        with autocast(device_type='cuda', dtype=torch.bfloat16):
            output = model(dummy_input)
            loss = output.sum()

        loss.backward()
        optimizer.step()

        if step % 10 == 0:
            print(f"Step {step}, Loss: {loss.item():.4f}")

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
from torch.utils.data.distributed import DistributedSampler


def setup_distributed(rank, world_size):
    """
    初始化分布式训练环境

    Args:
        rank: 当前进程编号（0 到 world_size-1）
        world_size: 总进程数（通常等于 GPU 数量）
    """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    dist.init_process_group(
        backend='nccl',  # NVIDIA GPU 使用 NCCL 后端
        init_method='env://',
        world_size=world_size,
        rank=rank
    )
    torch.cuda.set_device(rank)


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


def train_ddp(rank, world_size, dataset, model_fn, num_epochs=10):
    """
    DDP 训练示例

    DDP 工作原理：
    1. 每个 GPU 持有模型的完整副本
    2. 每个 GPU 处理不同的数据子集（数据并行）
    3. 反向传播时，通过 AllReduce 同步梯度
    4. 所有 GPU 的模型参数保持一致

    通信开销：每个 step 一次 AllReduce（梯度同步）
    内存开销：每卡存完整模型 + 优化器状态
    """
    setup_distributed(rank, world_size)

    # 创建模型并包装为 DDP
    model = model_fn().to(rank)
    model = DDP(model, device_ids=[rank])

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    # 创建分布式采样器（确保每个 GPU 看到不同数据）
    train_sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    dataloader = DataLoader(
        dataset,
        batch_size=32,
        sampler=train_sampler,
        num_workers=4,
        pin_memory=True,
    )

    for epoch in range(num_epochs):
        # 每个 epoch 设置不同的 shuffle 种子
        train_sampler.set_epoch(epoch)

        for batch in dataloader:
            optimizer.zero_grad()
            loss = model(**batch)
            loss.backward()
            optimizer.step()

        if rank == 0:
            print(f"Epoch {epoch}/{num_epochs} completed")

    cleanup_distributed()


# 启动多进程训练
if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    mp.spawn(
        train_ddp,
        args=(world_size, dataset, model_fn, 10),
        nprocs=world_size,
        join=True
    )

import torch
import torch.nn as nn


class ModelParallelModel(nn.Module):
    """
    简单模型并行：将模型不同层放在不同 GPU 上

    优点：
    - 突破单卡内存限制
    - 实现简单

    缺点：
    - 任何时刻只有一个 GPU 在计算（利用率低）
    - GPU 间通信开销
    - 仅适用于层数较多的模型

    适用场景：
    - 模型太大无法放入单卡
    - 作为流水线并行的简化版本
    """

    def __init__(self):
        super().__init__()
        # 不同层放在不同 GPU
        self.layer1 = nn.Linear(4096, 4096).to('cuda:0')
        self.layer2 = nn.Linear(4096, 4096).to('cuda:1')
        self.layer3 = nn.Linear(4096, 4096).to('cuda:2')

    def forward(self, x):
        # 逐层计算，跨 GPU 传输中间结果
        x = torch.relu(self.layer1(x.to('cuda:0')))
        x = torch.relu(self.layer2(x.to('cuda:1')))
        x = self.layer3(x.to('cuda:2'))
        return x

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy


def train_fsdp(rank, world_size):
    """
    FSDP 训练示例
    """
    setup_distributed(rank, world_size)

    # 混合精度策略
    mp_policy = MixedPrecision(
        param_dtype=torch.bfloat16,    # 参数使用 BF16
        reduce_dtype=torch.bfloat16,   # 梯度归约使用 BF16
        buffer_dtype=torch.bfloat16,   # Buffer 使用 BF16
    )

    # 创建模型
    model = create_model().to(rank)

    # FSDP 包装
    model = FSDP(
        model,
        device_id=rank,
        sharding_strategy=ShardingStrategy.FULL_SHARD,  # 等价于 ZeRO-3
        mixed_precision=mp_policy,
        # 自动按 Transformer 层分片
        auto_wrap_policy=transformer_auto_wrap_policy,
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    # 训练循环（与普通 DDP 相同）
    for batch in dataloader:
        optimizer.zero_grad()
        loss = model(**batch)
        loss.backward()
        optimizer.step()

    cleanup_distributed()


# ShardingStrategy 对比
SHARDING_STRATEGIES = {
    "NO_SHARD": "等价于 DDP，每卡存完整模型",
    "SHARD_GRAD_OP": "等价于 ZeRO-2，分片梯度和优化器状态",
    "FULL_SHARD": "等价于 ZeRO-3，分片参数、梯度和优化器状态",
}

import json
import deepspeed
import torch.nn as nn


class DeepSpeedTrainer:
    """
    DeepSpeed 训练器

    DeepSpeed 的优势：
    - ZeRO 优化器大幅减少显存占用
    - 支持 CPU/NVMe Offload（进一步节省显存）
    - 自动处理混合精度和梯度累积
    - 支持流水线并行和张量并行
    """

    def __init__(self, model, config_path):
        # 加载配置
        with open(config_path, 'r') as f:
            ds_config = json.load(f)

        # 初始化 DeepSpeed
        model_engine, optimizer, _, _ = deepspeed.initialize(
            model=model,
            model_parameters=model.parameters(),
            config=ds_config
        )

        self.model = model_engine
        self.optimizer = optimizer

    def train_step(self, batch):
        """
        DeepSpeed 训练步骤

        DeepSpeed 自动处理：
        - 梯度累积（根据 gradient_accumulation_steps）
        - 混合精度（根据 bf16/fp16 配置）
        - 梯度分片和通信（根据 ZeRO stage）
        """
        # 前向传播
        outputs = self.model(**batch)
        loss = outputs.loss

        # 反向传播
        self.model.backward(loss)

        # 更新参数（自动处理梯度累积）
        self.model.step()

        return loss.item()


# ZeRO Stage 对比
ZERO_STAGES = """
ZeRO Stage 0: 标准DDP，无优化
  每卡存储：完整参数 + 完整梯度 + 完整优化器状态

ZeRO Stage 1: 优化器状态分片（节省 ~4x 内存）
  每卡存储：完整参数 + 完整梯度 + 1/N 优化器状态

ZeRO Stage 2: 优化器状态 + 梯度分片（节省 ~8x 内存）
  每卡存储：完整参数 + 1/N 梯度 + 1/N 优化器状态

ZeRO Stage 3: 全分片（与 GPU 数量线性相关）
  每卡存储：1/N 参数 + 1/N 梯度 + 1/N 优化器状态

例如 8 卡训练 7B 模型（FP16）：
  显存需求 ≈ 14GB 参数 + 14GB 梯度 + 84GB 优化器状态 = 112GB
  - Stage 0: 每卡 112GB（不可能）
  - Stage 1: 每卡 ~28GB（A100 40GB 可行）
  - Stage 2: 每卡 ~18GB（A100 40GB 充裕）
  - Stage 3: 每卡 ~14GB（更小显存也可行）
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class TensorParallelLinear(nn.Module):
    """
    张量并行线性层（教学示例）

    生产中应使用 Megatron-LM 或 Megatron-Core 的实现，
    它们包含了通信优化、异步执行等高级特性。
    """

    def __init__(self, input_size, output_size, rank, world_size):
        super().__init__()
        self.rank = rank
        self.world_size = world_size

        # 每个GPU只存储一部分权重
        self.output_size_per_partition = output_size // world_size
        self.weight = nn.Parameter(
            torch.empty(self.output_size_per_partition, input_size)
        )
        self.bias = nn.Parameter(torch.empty(self.output_size_per_partition))

        # 初始化
        nn.init.kaiming_uniform_(self.weight)
        nn.init.zeros_(self.bias)

    def forward(self, input):
        """
        列并行前向传播
        每个GPU计算部分输出，然后通过 All-Gather 聚合
        """
        # 本地计算
        output_local = F.linear(input, self.weight, self.bias)

        # All-Gather 聚合所有 GPU 的输出
        # 生产中使用 torch.distributed.all_gather
        output_list = [torch.empty_like(output_local) for _ in range(self.world_size)]
        torch.distributed.all_gather(output_list, output_local)
        output = torch.cat(output_list, dim=-1)

        return output

class GradientAccumulationTrainer:
    """
    梯度累积训练器

    核心思想：用多个小 batch 的梯度累积来模拟大 batch 训练。

    等效 batch size = micro_batch_size × accumulation_steps × num_gpus

    示例：
    - 目标 batch size = 256
    - GPU 数量 = 8
    - 每 GPU micro batch = 4
    - accumulation_steps = 256 / (8 × 4) = 8

    注意：loss 需要除以 accumulation_steps 来保持梯度量级一致
    """

    def __init__(self, model, optimizer, accumulation_steps=4):
        self.model = model
        self.optimizer = optimizer
        self.accumulation_steps = accumulation_steps
        self.step_count = 0

    def train_step(self, batch):
        """
        梯度累积训练步骤
        """
        # 前向传播
        outputs = self.model(**batch)

        # 缩放 loss（保持梯度量级一致）
        scaled_loss = outputs.loss / self.accumulation_steps

        # 反向传播（梯度会累积）
        scaled_loss.backward()

        self.step_count += 1

        # 每 accumulation_steps 更新一次参数
        if self.step_count % self.accumulation_steps == 0:
            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), max_norm=1.0
            )
            self.optimizer.step()
            self.optimizer.zero_grad()

        # 返回未缩放的 loss
        return scaled_loss.item() * self.accumulation_steps

def clip_gradients(parameters, max_norm=1.0, norm_type=2.0):
    """
    梯度裁剪

    防止梯度爆炸（gradient explosion），在 Transformer 训练中必不可少。

    原理：
    - 计算 all 梯度的 L2 范数
    - 如果范数 > max_norm，等比缩放所有梯度
    - total_norm = ||∇θ||₂
    - clip_coef = max_norm / total_norm
    - ∇θ = ∇θ × min(clip_coef, 1.0)

    经验值：
    - max_norm = 1.0 是最常用的设置
    - 某些任务可能需要 0.5 或 5.0
    """
    total_norm = torch.nn.utils.clip_grad_norm_(
        parameters,
        max_norm=max_norm,
        norm_type=norm_type
    )
    return total_norm

import torch
from transformers import get_cosine_schedule_with_warmup


def create_scheduler(optimizer, num_warmup_steps, num_training_steps):
    """
    创建余弦退火学习率调度器（LLM 训练最常用的调度策略）

    调度过程：
    1. Warmup 阶段：学习率从 0 线性增加到 lr
    2. Decay 阶段：学习率按余弦函数从 lr 衰减到 0

    为什么需要 Warmup：
    - 训练初期模型参数随机，大学习率可能导致训练不稳定
    - Warmup 让模型先用小学习率找到正确的优化方向
    - 典型 warmup 步数：总步数的 1-10%

    常见学习率调度对比：
    - Cosine Decay: 最常用，平滑衰减
    - Linear Decay: 简单线性衰减
    - Constant with Warmup: warmup 后保持不变
    - WSD (Warmup-Stable-Decay): 2024年新策略
    """
    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
        num_cycles=0.5  # 半个余弦周期
    )
    return scheduler


# 可视化学习率调度
def visualize_lr_schedule():
    """生成学习率调度曲线数据"""
    model = torch.nn.Linear(10, 10)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    num_warmup = 1000
    num_total = 10000

    scheduler = get_cosine_schedule_with_warmup(
        optimizer, num_warmup, num_total
    )

    lrs = []
    for step in range(num_total):
        lrs.append(optimizer.param_groups[0]['lr'])
        scheduler.step()

    return lrs  # 返回学习率列表供绘图使用

import os
import torch
from pathlib import Path


class CheckpointManager:
    """
    检查点管理器

    大模型训练必须具备故障恢复能力：
    - 训练可能持续数天到数周
    - 硬件故障（GPU、网络）时有发生
    - 检查点允许从最近的保存点恢复

    检查点内容：
    - 模型参数（state_dict）
    - 优化器状态（动量、方差等）
    - 学习率调度器状态
    - 当前训练步数和 epoch
    - 随机数生成器状态（确保可复现）
    """

    def __init__(self, checkpoint_dir, max_keep=3):
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.max_keep = max_keep  # 最多保留的检查点数量

    def save_checkpoint(self, model, optimizer, scheduler, step, loss):
        """
        保存检查点
        """
        checkpoint = {
            'step': step,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': (
                scheduler.state_dict() if scheduler else None
            ),
            'loss': loss,
            'rng_state': torch.random.get_rng_state(),
        }

        # 保存到编号文件
        ckpt_path = self.checkpoint_dir / f'checkpoint_step_{step}.pt'
        torch.save(checkpoint, ckpt_path)

        # 更新 latest 指针（兼容 Windows 和 Linux）
        latest_path = self.checkpoint_dir / 'checkpoint_latest.pt'
        # Windows 不支持 symlink，使用 copy 替代
        try:
            if latest_path.exists():
                latest_path.unlink()
            os.symlink(ckpt_path.name, str(latest_path))
        except OSError:
            # Windows 可能不支持 symlink，直接复制
            import shutil
            shutil.copy2(str(ckpt_path), str(latest_path))

        # 清理旧检查点
        self._cleanup_old_checkpoints()

        print(f"Checkpoint saved: {ckpt_path}")

    def load_checkpoint(self, model, optimizer, scheduler=None, path=None):
        """
        加载检查点

        Args:
            path: 指定路径，或 None 加载 latest
        """
        if path is None:
            path = self.checkpoint_dir / 'checkpoint_latest.pt'

        checkpoint = torch.load(
            path, map_location='cpu', weights_only=False
        )

        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

        if scheduler and checkpoint.get('scheduler_state_dict'):
            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

        # 恢复随机数状态
        if 'rng_state' in checkpoint:
            torch.random.set_rng_state(checkpoint['rng_state'])

        step = checkpoint['step']
        loss = checkpoint['loss']

        print(f"Checkpoint loaded: {path}, step={step}, loss={loss:.4f}")
        return step, loss

    def _cleanup_old_checkpoints(self):
        """保留最近 max_keep 个检查点，删除更早的"""
        checkpoints = sorted(
            self.checkpoint_dir.glob('checkpoint_step_*.pt')
        )
        while len(checkpoints) > self.max_keep:
            oldest = checkpoints.pop(0)
            oldest.unlink()
            print(f"Removed old checkpoint: {oldest}")


# DeepSpeed 检查点（自动处理分片保存）
# 保存：model.save_checkpoint(save_dir, tag='step_1000')
# 加载：model.load_checkpoint(load_dir, tag='step_1000')