SuperAPI AI 工具集
合作ENAI 行业资讯收藏
Skill 广场distributed-llm-pretraining-torchtitan

distributed-llm-pretraining-torchtitan

智能体与模型

查看项目

更新于 10-21 16:24

技能介绍

用途

TorchTitan 是 PyTorch 官方的大规模 LLM 预训练平台，采用 PyTorch 原生实现，支持可组合的 4D 并行（FSDP2、TP、PP、CP），在 H100 GPU 上相比基线可实现 65% 以上的加速。适用于从零开始预训练 8B 到 405B 及以上规模的 LLM，需要 PyTorch 原生方案、可组合 4D 并行、H100 上的 Float8 支持，或希望获得与 torchtune/HuggingFace 互操作的检查点。

主要能力

  • 可组合 4D 并行：FSDP2、张量并行（TP）、流水线并行（PP）、上下文并行（CP）。
  • Float8 训练：在 H100 GPU 上提供 30-50% 加速。
  • 激活检查点、梯度累积、编译（compile）等内存与性能优化。
  • 检查点保存与自动恢复，支持 DCP 重分片及 HuggingFace 转换。
  • 支持模型：Llama 3.1（8B、70B、405B，生产状态）、Llama 4（实验）、DeepSeek V3（16B、236B、671B MoE，实验）、GPT-OSS（20B、120B MoE，实验）、Qwen 3（实验）、Flux 扩散模型（实验）。

使用方式

安装：

# 从 PyPI（稳定版）
pip install torchtitan

# 从源码（最新特性，需要 PyTorch nightly）
git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt

下载分词器：

python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

在 8 张 GPU 上启动训练：

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh

典型工作流包括：单节点预训练 Llama 3.1 8B（下载分词器、配置 TOML、启动训练、监控与检查点）；使用 SLURM 进行多节点训练；在 H100 上启用 Float8 训练；为 405B 模型配置 4D 并行。配置通过 TOML 文件完成，涵盖 job、model、optimizer、lr_scheduler、training、parallelism、activation_checkpoint、checkpoint 等段落。

适用场景

  • 从零预训练 8B 至 405B+ 的 LLM。
  • 需要 PyTorch 原生、无第三方依赖的方案。
  • 需要可组合 4D 并行（FSDP2、TP、PP、CP）。
  • 在 H100 上使用 Float8 训练。
  • 需要与 torchtune/HuggingFace 互操作的检查点。

限制与替代方案

  • 若追求 NVIDIA 专属部署的极致性能，可考虑 Megatron-LM。
  • 若需要更广泛的 ZeRO 优化生态与推理支持，可考虑 DeepSpeed。
  • 若为微调而非预训练，可考虑 Axolotl/TRL。
  • 若为教育或小规模训练，可考虑 LitGPT。

常见问题包括：大模型显存不足（启用激活检查点、减小 batch size 或梯度累积）、TP 异步集合通信导致高内存（设置 TORCH_NCCL_AVOID_RECORD_STREAMS=1）、Float8 未加速（过滤小层）、并行度变更后检查点加载失败（使用 DCP 重分片）、流水线并行初始化（先创建种子检查点）。

性能基准（H100）

  • Llama 8B，8 GPU，FSDP：5,762 TPS/GPU（基线）。
  • Llama 8B，8 GPU，FSDP+compile+FP8：8,532 TPS/GPU（+48%）。
  • Llama 70B，256 GPU，FSDP+TP+AsyncTP：876 TPS/GPU（2D 并行）。
  • Llama 405B，512 GPU，FSDP+TP+PP：128 TPS/GPU（3D 并行）。

资源

  • GitHub：https://github.com/pytorch/torchtitan
  • 论文：https://arxiv.org/abs/2410.06511
  • ICLR 2025：https://iclr.cc/virtual/2025/poster/29620
  • PyTorch 论坛：https://discuss.pytorch.org/c/distributed/torchtitan/44

操作规范

检查环境（GPU 类型、显存、驱动/CUDA 版本、框架版本、数据集位置与大小）；确认输入格式、标识符与参数；先运行小规模冒烟测试确认 loss 下降与检查点保存；再执行完整任务；在留出数据上验证指标并与基线对比，记录种子、配置与硬件；报告运行内容、检查项与不确定项。遇到 CUDA 显存不足、loss 为 NaN 或发散、版本中缺少函数/标志/端点、输入参数不明确等情况时，按对应方式处理，不得编造 API 或结果。