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 DDP	DeepSpeed ZeRO-3	Megatron-LM TP
参数规模	~10B	~1T	~1T
通信开销	低	中	高
内存效率	低	高	中
实现复杂度	低	中	高
适用场景	中小模型	超大模型	超大模型

推理引擎对比

特性	vLLM	TensorRT-LLM	Text 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实现
└── 定制训练流程

下一步学习

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

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