distributed-llm-pretraining-torchtitan

# TorchTitan - PyTorch 原生分布式大语言模型预训练 ## 快速开始 TorchTitan 是 PyTorch 官方的大规模大语言模型预训练平台，支持可组合的 4D 并行（FSDP2, TP, PP, CP），在 H100 GPU 上相比基准测试实现了 65% 以上的加速。 **安装**： bash # 从 PyPI 安装 (稳定版) pip install torchtitan # 从源码安装 (最新功能，需要 PyTorch nightly 版本) git clone https://github.com/pytorch/torchtitan cd torchtitan pip install -r requirements.txt **下载分词器 (Tokenizer)**： bash # 从 https://huggingface.co/settings/tokens 获取 HF token python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=... **在 8 个 GPU 上启动训练**： bash CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh ## 常见工作流 ### 工作流 1：在单节点上预训练 Llama 3.1 8B 复制此检查清单： 单节点预训练： - [ ] 第 1 步：下载分词器 - [ ] 第 2 步：配置训练 - [ ] 第 3 步：启动训练 - [ ] 第 4 步：监控与检查点保存 **第 1 步：下载分词器** bash python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=YOUR_HF_TOKEN **第 2 步：配置训练** 编辑或创建 TOML 配置文件： toml # llama3_8b_custom.toml [job] dump_folder = "./outputs" description = "Llama 3.1 8B 训练" [model] name = "llama3" flavor = "8B" hf_assets_path = "./assets/hf/Llama-3.1-8B" [optimizer] name = "AdamW" lr = 3e-4 [lr_scheduler] warmup_steps = 200 [training] local_batch_size = 2 seq_len = 8192 max_norm = 1.0 steps = 1000 dataset = "c4" [parallelism] data_parallel_shard_degree = -1 # 使用所有 GPU 进行 FSDP [activation_checkpoint] mode = "selective" selective_ac_option = "op" [checkpoint] enable = true folder = "checkpoint" interval = 500 **第 3 步：启动训练** bash # 单节点 8 个 GPU 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 **第 4 步：监控与检查点保存** TensorBoard 日志保存至 `./outputs/tb/`： bash tensorboard --logdir ./outputs/tb ### 工作流 2：使用 SLURM 进行多节点训练 多节点训练： - [ ] 第 1 步：配置并行规模 - [ ] 第 2 步：设置 SLURM 脚本 - [ ] 第 3 步：提交任务 - [ ] 第 4 步：从检查点恢复 **第 1 步：配置并行**