Chapter 13

源码导读 + PyTorch 分布式路线图

📌 commit af33f76TorchTitan 是 PyTorch 主仓的镜像

TorchTitan 仓库不大（核心 ~5000 行），但隐含的 PyTorch 路线图很重。本章给阅读顺序 + 路线图。

13.1仓库结构

torchtitan/                       https://github.com/pytorch/torchtitan
├── torchtitan/
│   ├── train.py                  ★ 入口（500 行）
│   ├── config_manager.py         ★ toml + argparse
│   ├── components/
│   │   ├── checkpoint.py
│   │   ├── optimizer.py
│   │   ├── lr_scheduler.py
│   │   └── tokenizer/
│   ├── distributed/
│   │   ├── parallel_dims.py      ★ 5D mesh
│   │   ├── parallelize_llama.py  ★ TP/SP/CP/FSDP plan
│   │   └── pipeline.py            PP wrap
│   ├── models/
│   │   ├── llama/                ★ Llama 3 实现
│   │   └── ...
│   ├── experiments/              experimental 模型
│   │   ├── deepseek_v3/
│   │   ├── mixtral/
│   │   └── ...
│   ├── datasets/
│   ├── utils/
│   └── ...
├── train_configs/
├── docs/                         设计文档
└── tests/

13.2建议阅读顺序

#	文件	读什么
1	`torchtitan/train.py`	主循环
2	`torchtitan/config_manager.py`	配置系统
3	`torchtitan/distributed/parallel_dims.py`	mesh 构造
4	`torchtitan/models/llama/model.py`	模型实现（看怎么写普通 nn.Module）
5	`torchtitan/distributed/parallelize_llama.py`	怎么挂 TP/FSDP/PP
6	`torchtitan/components/checkpoint.py`	dcp 用法
7	`torchtitan/distributed/pipeline.py`	PP schedule
8	`torchtitan/experiments/deepseek_v3/`	MoE + 多 expert 实例

13.3设计哲学

原则	含义
"PyTorch native first"	不引入新抽象，全用 PyTorch 自带
"少胜于多"	核心 5000 行，能不写就不写
"展示而非黑盒"	每个文件都该能被用户读懂、抄走
"实验场地"	新 PyTorch 特性优先在 TorchTitan 验证

13.4PyTorch 分布式路线图（TorchTitan 透露）

状态	特性
已稳定	FSDP2 / DTensor / ParallelStyle / dcp / SDPA Flash / Context Parallel / async TP / float8 (rowwise/tensorwise)
2025 主线	fp8 MX format / 更智能 compile 与 FSDP / CP 整合 / SymmetricMemory（NVL2/2048 拓扑）
experimental	fp4 / int4 训练 / sparse expert routing / 跨节点 NVLink fabrics

13.5对照其他框架

	TorchTitan	Megatron-LM	nanotron
代码量	~5K	~50K	~5K
分布式抽象	DTensor	自家	自家小型
性能	★★★★（追平 Megatron）	★★★★★	★★★
易读性	★★★★★	★★	★★★★★
跟 PyTorch 同步	★★★★★（同公司）	★★★	★★★
多模态	—	—	—

13.6社区入口

GitHub: https://github.com/pytorch/torchtitan
论坛: PyTorch dev-discuss
设计文档: docs/ 目录（PP / FSDP2 / CP 等）
PyTorch 主仓的 distributed 相关 PR 同步跟进

13.7真实源码树（v2026）

torchtitan/
├── train.py:17                  ★ main()：ConfigManager → Trainer
├── trainer.py:56                 ★ Trainer 主类 (~900 行)
│
├── config/                       ★ 配置系统
│   ├── manager.py:19             ConfigManager
│   ├── configs.py:26-130         TrainingConfig / ParallelismConfig
│   └── job_config.py             顶层 JobConfig
│
├── distributed/                  ★ 并行核心
│   ├── parallel_dims.py:26       ParallelDims (5D + EP)
│   ├── tensor_parallel.py:102    maybe_enable_async_tp / NoParallel
│   ├── context_parallel.py:35    apply_cp_to_forward
│   ├── pipeline_parallel.py      1F1B / looped
│   ├── fsdp.py                    FSDP2 wrapper
│   └── full_dtensor/              ★ 纯 DTensor 模式（新）
│
├── models/                       ★ 6 个架构
│   ├── llama3/
│   │   ├── model.py:58           Llama3Model
│   │   ├── parallelize.py:41     parallelize_llama
│   │   └── config_registry.py    llama3_8b/70b/llama3_debugmodel
│   ├── llama4/model.py:110       Llama4Model
│   ├── deepseek_v3/
│   │   └── model.py:180          DeepSeekV3Model（MLA + MoE）
│   ├── qwen3/ / qwen3_vl/
│   ├── gpt_oss/                  GPT-OSS (attention sink + YaRN)
│   ├── flux/                      Flux 图像
│   └── common/                    attention / decoder / moe / rope / nn_modules
│
├── components/                   ★ 训练组件
│   ├── checkpoint.py             CheckpointManager (DCP + HuggingFaceStorageWriter)
│   ├── optimizer.py:67           OptimizersContainer（AdamW / Muon）
│   ├── dataloader.py
│   ├── loss.py / lr_scheduler.py
│   └── quantization/
│       └── float8.py:24          Float8LinearConverter
│
├── experiments/                  ★ 研究功能
│   ├── graph_trainer/            autoparallel (FX graph 自动切分)
│   ├── ft/                        TorchFT 微调
│   ├── rl/                         RL
│   ├── transformers_modeling_backend/   HF transformers 直接训
│   └── forge/                      （开发中）
│
├── datasets/                     数据集
└── observability/                结构化日志

13.8三条阅读路径

路线 A · "我要跑大模型"：run_train.sh → train.py:17 → trainer.py:56 看 __init__ 和 train_step；
路线 B · "我要懂分布式"：distributed/parallel_dims.py:26 → models/llama3/parallelize.py:41 → distributed/tensor_parallel.py / fsdp.py / pipeline_parallel.py；
路线 C · "我要加新模型"：抄 models/llama3/ 整个目录（5 个文件：model / parallelize / config_registry / args / __init__）→ 在 train.py 加 module 注册。

13.9experiments/ 看什么

experiment	研究主题
`experiments/graph_trainer/`	★ Autoparallel：FX graph + 自动切分（llama3/qwen3/deepseek_v3 各一份）
`experiments/ft/`	TorchFT：弹性容错微调
`experiments/rl/`	纯 PyTorch RL 训练（无 vLLM 依赖）
`experiments/transformers_modeling_backend/`	★ 直接训 HF transformers 模型不用迁移
`experiments/forge/`	（2026 开发中）

规则：experiments 只能依赖 core，不能反向。core 模块稳定，experiments 可大改。

13.10TorchTitan vs Megatron：哲学对比

	TorchTitan	Megatron-LM
核心抽象	DTensor + DeviceMesh（SPMD）	显式 process group + 手写 sharding
算子来源	PyTorch / torchao 上游	自家 CUDA / TE
代码规模	~5000 行核心	~50000+ 行
新特性节奏	跟 PyTorch nightly	跟 NVIDIA H100/H200 硬件
研发理念	"参考实现"，简洁可读	"生产框架"，全场景兜底
FP8 / Async TP	torchao + inductor	TransformerEngine
适合谁	PyTorch 开发者、研究员、想读懂分布式的人	千亿模型训练，工业级生产

13.11这章你需要带走的

TorchTitan 是 PyTorch 分布式的参考实现；
核心 5000 行，按 train.py → parallelize → model 顺序读；
性能追平 Megatron，但代码量是 1/10；
跟踪 TorchTitan = 跟踪 PyTorch 分布式新特性。