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

distributed-llm-pretraining-torchtitan

智能体与模型

查看项目

更新于 10-15 13:09

技能介绍

TorchTitan 是 PyTorch 官方的大规模 LLM 预训练平台，提供 PyTorch 原生的分布式预训练能力，支持可组合的 4D 并行（FSDP2、TP、PP、CP），在 H100 GPU 上相比基线可实现 65% 以上的加速。

用途

TorchTitan 主要用于从零开始预训练大语言模型，覆盖 8B 到 405B+ 参数规模。它适合需要 PyTorch 原生解决方案、不希望引入第三方依赖的场景，也适合需要可组合 4D 并行（FSDP2、TP、PP、CP）的分布式训练任务。此外，它支持在 H100 上进行 Float8 训练，并希望检查点能与 torchtune/HuggingFace 互操作。

主要能力

  • 支持 4D 并行：FSDP2、张量并行（TP）、流水线并行（PP）、上下文并行（CP）。
  • 支持 Float8 训练，在 H100 GPU 上可提供 30-50% 的加速。
  • 支持多节点训练，可通过 SLURM 提交作业。
  • 支持检查点保存与恢复，训练可在配置文件夹中存在检查点时自动恢复。
  • 支持与 torchtune/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

常见工作流

  1. 单节点预训练 Llama 3.1 8B：下载分词器、配置 TOML 文件、启动训练、监控与检查点。配置示例包括模型、优化器、学习率调度、训练参数、并行度、激活检查点和检查点设置。启动命令可使用 ./run_train.sh 或 torchrun --nproc_per_node=8 -m torchtitan.train --job.config_file ./llama3_8b_custom.toml。TensorBoard 日志保存在 ./outputs/tb/。
  2. 多节点训练（SLURM）：配置并行度（如 70B 模型在 256 GPU 上设置 data_parallel_shard_degree = 32、tensor_parallel_degree = 8），编写 SLURM 脚本，提交作业，检查点自动恢复。
  3. H100 上的 Float8 训练：安装 torchao（USE_CPP=0 pip install git+https://github.com/pytorch/ao.git），配置 Float8（如 converters = ["quantize.linear.float8"]），启用编译后启动。
  4. 405B 模型的 4D 并行：创建种子检查点，配置 4D 并行（如 data_parallel_shard_degree = 8、tensor_parallel_degree = 8、pipeline_parallel_degree = 8），在 512 GPU 上启动。

适用场景

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

限制与替代方案

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

常见问题

  • 大模型内存不足：启用激活检查点（mode = "full"）并减小批次大小，或使用梯度累积。
  • TP 异步集合通信导致高内存：设置环境变量 export TORCH_NCCL_AVOID_RECORD_STREAMS=1。
  • Float8 训练未加速：Float8 仅对大 GEMM 有益，可过滤小层（如 filter_fqns = ["attention.wk", "attention.wv", "output", "auto_filter_small_kn"]）。
  • 并行度变更后检查点加载失败：使用 DCP 的重分片能力，如 python -m torch.distributed.checkpoint.format_utils dcp_to_torch checkpoint/step-1000 checkpoint.pt。
  • 流水线并行初始化：需先创建种子检查点。

性能基准（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