12-框架源码解析

章节概述

本章深入解析主流AI训练和推理框架的源码实现，包括PyTorch分布式、DeepSpeed、Megatron-LM、vLLM和HuggingFace Transformers，帮助AI Infra工程师理解框架内部机制，为性能调优和二次开发奠定基础。

知识体系

┌─────────────────────────────────────────────────────────────────────────────┐
│                        框架源码解析知识体系                                   │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  ┌────────────────────────────────────────────────────────────────────┐     │
│  │                    PyTorch 分布式                                   │     │
│  │                                                                     │     │
│  │   ProcessGroup    │   DDP实现        │   FSDP实现     │   RPC框架  │     │
│  │   ├─ NCCL        │   ├─ Reducer     │   ├─ 分片策略  │   ├─ Agent │     │
│  │   ├─ Gloo        │   ├─ Bucket      │   ├─ 全收集    │   ├─ RRef  │     │
│  │   └─ MPI         │   └─ Hook机制    │   └─ 混合精度  │   └─ Future│     │
│  └────────────────────────────────────────────────────────────────────┘     │
│                                │                                             │
│  ┌────────────────────────────────────────────────────────────────────┐     │
│  │                      DeepSpeed                                      │     │
│  │                                                                     │     │
│  │   ZeRO Stage 1-3  │   Pipeline引擎  │   CPU Offload  │   Activation│     │
│  │   ├─ 优化器分片   │   ├─ 1F1B调度   │   ├─ 参数卸载  │   ├─ 检查点 │     │
│  │   ├─ 梯度分片     │   ├─ 微批次     │   ├─ Adam卸载  │   └─ 分区   │     │
│  │   └─ 参数分片     │   └─ 梯度累积   │   └─ NVMe卸载  │            │     │
│  └────────────────────────────────────────────────────────────────────┘     │
│                                │                                             │
│  ┌────────────────────────────────────────────────────────────────────┐     │
│  │                      Megatron-LM                                    │     │
│  │                                                                     │     │
│  │   MPU并行组       │   张量并行       │   序列并行     │   混合并行  │     │
│  │   ├─ 初始化      │   ├─ Column     │   ├─ QKV分片   │   ├─ 3D并行 │     │
│  │   ├─ 组管理      │   ├─ Row        │   ├─ Dropout   │   ├─ 通信组 │     │
│  │   └─ 进程映射    │   └─ All-Reduce │   └─ LayerNorm │   └─ 调度   │     │
│  └────────────────────────────────────────────────────────────────────┘     │
│                                │                                             │
│  ┌────────────────────────────────────────────────────────────────────┐     │
│  │                        vLLM                                         │     │
│  │                                                                     │     │
│  │   PagedAttention  │   Block管理     │   调度器       │   Continuous│     │
│  │   ├─ 分页机制    │   ├─ BlockTable │   ├─ FCFS      │   Batching  │     │
│  │   ├─ KV缓存      │   ├─ 分配器     │   ├─ 抢占策略  │   ├─ 动态   │     │
│  │   └─ 内存池      │   └─ 碎片管理   │   └─ Swap管理  │   └─ 迭代级 │     │
│  └────────────────────────────────────────────────────────────────────┘     │
│                                │                                             │
│  ┌────────────────────────────────────────────────────────────────────┐     │
│  │                   HuggingFace Transformers                          │     │
│  │                                                                     │     │
│  │   模型抽象层      │   Tokenizer     │   Trainer      │   Generation│     │
│  │   ├─ PreTrained  │   ├─ BPE        │   ├─ 训练循环  │   ├─ Greedy │     │
│  │   ├─ Auto类      │   ├─ Fast       │   ├─ 回调系统  │   ├─ Beam   │     │
│  │   └─ Config      │   └─ 编解码     │   └─ 集成      │   └─ Sample │     │
│  └────────────────────────────────────────────────────────────────────┘     │
│                                                                              │
└─────────────────────────────────────────────────────────────────────────────┘

章节内容

01-PyTorch分布式源码解析

torch.distributed核心

  • ProcessGroup抽象与实现
  • 通信后端: NCCL, Gloo, MPI
  • 初始化与进程组管理

DistributedDataParallel (DDP)

  • Reducer梯度归约
  • Bucket分桶策略
  • 梯度同步Hook机制
  • find_unused_parameters原理

FullyShardedDataParallel (FSDP)

  • 参数分片策略
  • 全收集与分散归约
  • 混合精度与梯度检查点
  • 与DDP对比

RPC框架

  • 分布式自动求导
  • Remote Reference (RRef)
  • Future异步执行

02-DeepSpeed源码深度解析

ZeRO优化器

  • Stage 1: 优化器状态分片
  • Stage 2: 梯度分片
  • Stage 3: 参数分片
  • 通信优化与融合

Pipeline并行引擎

  • 1F1B调度算法
  • 微批次管理
  • 梯度累积与同步
  • 流水线气泡分析

CPU/NVMe Offload

  • 参数卸载机制
  • Adam状态卸载
  • 异步数据传输
  • 分级存储策略

Activation Checkpointing

  • 分区激活检查点
  • CPU激活卸载
  • 内存与计算权衡

03-Megatron-LM源码解析

MPU (Model Parallel Unit)

  • 并行组初始化
  • 进程组管理
  • 全局状态管理

张量并行实现

  • ColumnParallelLinear
  • RowParallelLinear
  • 跨设备All-Reduce

序列并行

  • QKV投影分片
  • Dropout序列分片
  • LayerNorm序列分片
  • 通信优化

3D混合并行

  • TP + PP + DP组合
  • 通信组划分
  • 调度策略

04-vLLM推理引擎源码解析

PagedAttention

  • 分页KV缓存机制
  • Block与BlockTable
  • 内存碎片消除
  • CUDA内核实现

BlockManager

  • 物理块分配
  • 虚拟到物理映射
  • Copy-on-Write优化
  • 内存池管理

Scheduler调度器

  • 请求状态机
  • FCFS调度
  • 抢占与恢复
  • Swap管理

Continuous Batching

  • 迭代级批处理
  • 动态序列加入/退出
  • 吞吐量优化

05-HuggingFace Transformers源码解析

模型抽象层

  • PreTrainedModel基类
  • PretrainedConfig配置
  • Auto类自动推断
  • 权重加载与保存

Tokenizer实现

  • PreTrainedTokenizer
  • Fast Tokenizer (Rust)
  • BPE编解码
  • 批量处理优化

Trainer训练框架

  • 训练循环封装
  • TrainingArguments配置
  • 回调系统设计
  • 分布式训练集成

Generation生成系统

  • GenerationMixin
  • 采样策略: Greedy, Beam, Sampling
  • LogitsProcessor链
  • Top-K/Top-P实现

核心代码示例

PyTorch DDP Reducer

# torch/nn/parallel/distributed.py
class Reducer:
    def __init__(self, params, process_group, bucket_bytes_cap):
        self.process_group = process_group
        self.buckets = self._build_buckets(params, bucket_bytes_cap)

    def _build_buckets(self, params, cap):
        """构建梯度桶，按大小分组"""
        buckets = []
        current_bucket = []
        current_size = 0

        for param in reversed(params):
            param_size = param.numel() * param.element_size()
            if current_size + param_size > cap and current_bucket:
                buckets.append(current_bucket)
                current_bucket = []
                current_size = 0
            current_bucket.append(param)
            current_size += param_size

        if current_bucket:
            buckets.append(current_bucket)
        return buckets

    def reduce_bucket(self, bucket_idx):
        """异步归约一个桶"""
        bucket = self.buckets[bucket_idx]
        flat_grads = torch.cat([p.grad.view(-1) for p in bucket])
        dist.all_reduce(flat_grads, group=self.process_group, async_op=True)

DeepSpeed ZeRO Stage 3

# deepspeed/runtime/zero/stage3.py
class DeepSpeedZeroOptimizer_Stage3:
    def __init__(self, module, optimizer, param_names):
        self.module = module
        self.optimizer = optimizer
        self.world_size = dist.get_world_size()

        # 参数分片
        self.param_partitions = self._partition_parameters()

        # 注册钩子
        self._register_hooks()

    def _partition_parameters(self):
        """将参数分片到各rank"""
        partitions = {}
        for name, param in self.module.named_parameters():
            partition_size = param.numel() // self.world_size
            partitions[name] = partition_size
        return partitions

    def step(self):
        """优化器步骤"""
        # 收集完整梯度
        self._allgather_gradients()
        # 本地更新分片参数
        self.optimizer.step()
        # 广播更新后的参数
        self._allgather_parameters()

vLLM PagedAttention

# vllm/attention/backends/flash_attn.py
def paged_attention_forward(
    query: torch.Tensor,      # [num_tokens, num_heads, head_size]
    key_cache: torch.Tensor,  # [num_blocks, num_heads, head_size, block_size]
    value_cache: torch.Tensor,
    block_tables: torch.Tensor,  # [num_seqs, max_num_blocks]
    context_lens: torch.Tensor,  # [num_seqs]
):
    """分页注意力前向传播"""
    output = torch.empty_like(query)

    # 调用CUDA内核
    paged_attention_v1(
        output,
        query,
        key_cache,
        value_cache,
        block_tables,
        context_lens,
        block_size=16,
        max_context_len=max(context_lens),
    )

    return output

性能对比参考

分布式训练框架对比

特性PyTorch DDPDeepSpeed ZeRO-3Megatron-LM TP
参数规模~10B~1T~1T
通信开销低中高
内存效率低高中
实现复杂度低中高
适用场景中小模型超大模型超大模型

推理引擎对比

特性vLLMTensorRT-LLMText Generation Inference
内存效率高 (PagedAttention)中中
吞吐量高很高中
易用性高中高
定制性高低中

学习路径

Week 1-2: PyTorch分布式基础
├── 阅读 01-PyTorch分布式源码解析
├── 理解ProcessGroup和通信原语
├── 深入DDP Reducer实现
└── 实践FSDP训练

Week 3-4: DeepSpeed深入
├── 阅读 02-DeepSpeed源码深度解析
├── 理解ZeRO三个阶段
├── 分析Pipeline引擎调度
└── 实践ZeRO-Offload

Week 5-6: Megatron-LM
├── 阅读 03-Megatron-LM源码解析
├── 理解MPU设计
├── 实现简单张量并行
└── 配置3D混合并行

Week 7-8: 推理优化
├── 阅读 04-vLLM推理引擎源码解析
├── 理解PagedAttention
├── 分析Continuous Batching
└── 性能调优实践

Week 9-10: 框架集成
├── 阅读 05-HuggingFace Transformers源码解析
├── 理解Auto类和模型加载
├── 深入Trainer实现
└── 定制训练流程

推荐资源

官方文档

源码仓库

经典论文

  • "PyTorch Distributed: Experiences on Accelerating Data Parallel Training"
  • "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"
  • "Megatron-LM: Training Multi-Billion Parameter Language Models"
  • "Efficient Memory Management for Large Language Model Serving with PagedAttention"

下一步学习

完成本章后，建议继续学习：

  1. 13-编译优化与图优化 - XLA、TorchDynamo、算子融合
  2. 11-通信与网络底层 - 如果需要补充通信基础
  3. 10-CUDA编程与算子开发 - 如果需要补充CUDA基础