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

distributed-llm-pretraining-torchtitan

智能体与模型

查看项目

更新于 10-30 06:48

技能介绍

用途

TorchTitan 是 PyTorch 官方的大规模 LLM 预训练平台，采用 PyTorch 原生实现，提供可组合的 4D 并行能力（FSDP2、TP、PP、CP），据资料称在 H100 GPU 上相较基线可实现 65% 以上加速。该技能面向以下场景：从零预训练 8B 至 405B 以上规模的 LLM；需要不依赖第三方组件的 PyTorch 原生方案；需要可组合的 4D 并行；在 H100 上使用 Float8 支持；以及希望获得可与 torchtune/HuggingFace 互操作的检查点。

安装与准备

从 PyPI 安装稳定版：

pip install torchtitan

从源码安装最新特性（需要 PyTorch nightly）：

git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt

下载 tokenizer（HF token 需从 https://huggingface.co/settings/tokens 获取）：

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

主要工作流

单节点预训练 Llama 3.1 8B：依次下载 tokenizer、编辑或创建 TOML 配置文件（设置 job、model、optimizer、lr_scheduler、training、parallelism、activation_checkpoint、checkpoint 等段落）、启动训练并监控。启动方式为：

CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh

或显式使用 torchrun：

torchrun --nproc_per_node=8 -m torchtitan.train --job.config_file ./llama3_8b_custom.toml

TensorBoard 日志保存在 ./outputs/tb/，可用 tensorboard --logdir ./outputs/tb 查看。

SLURM 多节点训练：针对 70B 模型在 256 GPU（32 节点）上，可在 [parallelism] 中设置 data_parallel_shard_degree = 32、tensor_parallel_degree = 8、pipeline_parallel_degree = 1、context_parallel_degree = 1，编写 SLURM 脚本后通过 sbatch 提交；若配置的文件夹中存在检查点，训练会自动恢复。

H100 上的 Float8 训练：据资料称可带来 30–50% 加速。先安装 torchao：

USE_CPP=0 pip install git+https://github.com/pytorch/ao.git

再在 TOML 中添加 converters = ["quantize.linear.float8"] 及相关 [quantize.linear.float8] 配置，并启用 [compile]。

405B 模型的 4D 并行：先创建 seed 检查点（用于跨 PP 阶段的一致初始化），再配置 FSDP、TP、PP 并行度，最后在 512 GPU（64 节点 × 8）上通过 srun torchrun 启动。

适用与替代

资料建议在需要从零预训练、要求 PyTorch 原生与 4D 并行、使用 H100 Float8 或需要与 torchtune/HuggingFace 互操作检查点时选用 TorchTitan。替代方案包括：Megatron-LM（面向 NVIDIA 部署追求最高性能）、DeepSpeed（更广的 ZeRO 优化生态与推理支持）、Axolotl/TRL（微调而非预训练）、LitGPT（教育用途、较小规模训练）。

常见问题

  • 大模型显存不足：启用完整激活检查点并减小 batch size，或使用梯度累积。
  • TP 配合异步集合通信时显存过高：设置环境变量 TORCH_NCCL_AVOID_RECORD_STREAMS=1。
  • Float8 未加速：Float8 仅对大型 GEMM 有益，可通过 filter_fqns 过滤小层。
  • 改变并行度后检查点加载失败：使用 DCP 的重分片能力转换检查点。
  • 流水线并行初始化：需先创建 seed 检查点。

支持模型与性能

资料列出：Llama 3.1（8B、70B、405B，生产状态）、Llama 4（实验）、DeepSeek V3（16B、236B、671B MoE，实验）、GPT-OSS（20B、120B MoE，实验）、Qwen 3（实验）、Flux 扩散模型（实验）。H100 基准示例：Llama 8B 在 8 GPU 上使用 FSDP 为 5,762 TPS/GPU；叠加 compile 与 FP8 后为 8,532 TPS/GPU（+48%）；Llama 70B 在 256 GPU 上使用 FSDP+TP+AsyncTP 为 876 TPS/GPU；Llama 405B 在 512 GPU 上使用 FSDP+TP+PP 为 128 TPS/GPU。

限制

资料未提供价格、收费、平台授权或 SLA 等信息，相关内容暂未提供。性能数据均为资料中列出的 H100 基准，不保证在其他硬件上复现。Float8 与部分模型的标记为实验性，实际可用性需以项目当前状态为准。