@baimingtao/lbs-agent 1.0.0

package/extensions/skills/optional/mlops/torchtitan/SKILL.md

@@ -0,0 +1,361 @@
+---
+name: distributed-llm-pretraining-torchtitan
+description: 提供 PyTorch 原生的分布式 LLM pretraining，基于 torchtitan 实现 4D parallelism (FSDP2, TP, PP, CP)。适用于 pretraining Llama 3.1、DeepSeek V3 或自定义模型，可在 8 到 512+ 个 GPU 上进行大规模训练，并支持 Float8、torch.compile 和分布式 checkpointing。
+version: 1.0.0
+author: Orchestra Research
+license: MIT
+dependencies: [torch>=2.6.0, torchtitan>=0.2.0, torchao>=0.5.0]
+platforms: [linux, macos]
+metadata:
+  hermes:
+    tags: [Model Architecture, Distributed Training, TorchTitan, FSDP2, Tensor Parallel, Pipeline Parallel, Context Parallel, Float8, Llama, Pretraining]
+
+---
+
+# TorchTitan - PyTorch 原生分布式 LLM Pretraining
+
+## 快速 start
+
+TorchTitan 是 PyTorch 官方的大规模 LLM pretraining 平台，支持可组合的 4D parallelism (FSDP2, TP, PP, CP)，在 H100 GPU 上比基线实现了 65% 以上的速度提升。
+
+**安装**:
+```bash
+# From PyPI (stable)
+pip install torchtitan
+
+# From source (latest features, requires PyTorch nightly)
+git clone https://github.com/pytorch/torchtitan
+cd torchtitan
+pip install -r requirements.txt
+```
+
+**下载 tokenizer**:
+```bash
+# Get HF token from https://huggingface.co/settings/tokens
+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: 在单节点上 Pretrain Llama 3.1 8B
+
+复制此清单：
+
+```
+Single Node Pretraining:
+- [ ] Step 1: Download tokenizer
+- [ ] Step 2: Configure training
+- [ ] Step 3: Launch training
+- [ ] Step 4: Monitor and checkpoint
+```
+
+**Step 1: 下载 tokenizer**
+
+```bash
+python scripts/download_hf_assets.py \
+  --repo_id meta-llama/Llama-3.1-8B \
+  --assets tokenizer \
+  --hf_token=YOUR_HF_TOKEN
+```
+
+**Step 2: 配置训练**
+
+编辑或创建 TOML 配置文件：
+
+```toml
+# llama3_8b_custom.toml
+[job]
+dump_folder = "./outputs"
+description = "Llama 3.1 8B training"
+
+[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  # Use all GPUs for FSDP
+
+[activation_checkpoint]
+mode = "selective"
+selective_ac_option = "op"
+
+[checkpoint]
+enable = true
+folder = "checkpoint"
+interval = 500
+```
+
+**Step 3: 启动训练**
+
+```bash
+# 8 GPUs on single node
+CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh
+
+# Or explicitly with torchrun
+torchrun --nproc_per_node=8 \
+  -m torchtitan.train \
+  --job.config_file ./llama3_8b_custom.toml
+```
+
+**Step 4: 监控与 checkpoint**
+
+TensorBoard 日志保存在 `./outputs/tb/`：
+```bash
+tensorboard --logdir ./outputs/tb
+```
+
+### 工作流 2: 使用 SLURM 进行多节点训练
+
+```
+Multi-Node Training:
+- [ ] Step 1: Configure parallelism for scale
+- [ ] Step 2: Set up SLURM script
+- [ ] Step 3: Submit job
+- [ ] Step 4: Resume from checkpoint
+```
+
+**Step 1: 针对规模配置 parallelism**
+
+在 256 个 GPU（32 个节点）上运行 70B 模型：
+```toml
+[parallelism]
+data_parallel_shard_degree = 32  # FSDP across 32 ranks
+tensor_parallel_degree = 8        # TP within node
+pipeline_parallel_degree = 1      # No PP for 70B
+context_parallel_degree = 1       # Increase for long sequences
+```
+
+**Step 2: 设置 SLURM 脚本**
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=llama70b
+#SBATCH --nodes=32
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+
+srun torchrun \
+  --nnodes=32 \
+  --nproc_per_node=8 \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
+  -m torchtitan.train \
+  --job.config_file ./llama3_70b.toml
+```
+
+**Step 3: 提交作业**
+
+```bash
+sbatch multinode_trainer.slurm
+```
+
+**Step 4: 从 checkpoint 恢复**
+
+如果配置的文件夹中存在 checkpoint，训练将自动恢复。
+
+### 工作流 3: 为 H100 启用 Float8 训练
+
+Float8 在 H100 GPU 上可提供 30-50% 的速度提升。
+
+```
+Float8 Training:
+- [ ] Step 1: Install torchao
+- [ ] Step 2: Configure Float8
+- [ ] Step 3: Launch with compile
+```
+
+**Step 1: 安装 torchao**
+
+```bash
+USE_CPP=0 pip install git+https://github.com/pytorch/ao.git
+```
+
+**Step 2: 配置 Float8**
+
+将以下内容添加到你的 TOML 配置中：
+```toml
+[model]
+converters = ["quantize.linear.float8"]
+
+[quantize.linear.float8]
+enable_fsdp_float8_all_gather = true
+precompute_float8_dynamic_scale_for_fsdp = true
+filter_fqns = ["output"]  # Exclude output layer
+
+[compile]
+enable = true
+components = ["model", "loss"]
+```
+
+**Step 3: 启用 compile 启动**
+
+```bash
+CONFIG_FILE="./llama3_8b.toml" ./run_train.sh \
+  --model.converters="quantize.linear.float8" \
+  --quantize.linear.float8.enable_fsdp_float8_all_gather \
+  --compile.enable
+```
+
+### 工作流 4: 针对 405B 模型的 4D parallelism
+
+```
+4D Parallelism (FSDP + TP + PP + CP):
+- [ ] Step 1: Create seed checkpoint
+- [ ] Step 2: Configure 4D parallelism
+- [ ] Step 3: Launch on 512 GPUs
+```
+
+**Step 1: 创建 seed checkpoint**
+
+这是在 PP 各阶段实现一致性初始化所必需的：
+```bash
+NGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \
+  --checkpoint.enable \
+  --checkpoint.create_seed_checkpoint \
+  --parallelism.data_parallel_shard_degree 1 \
+  --parallelism.tensor_parallel_degree 1 \
+  --parallelism.pipeline_parallel_degree 1
+```
+
+**Step 2: 配置 4D parallelism**
+
+```toml
+[parallelism]
+data_parallel_shard_degree = 8   # FSDP
+tensor_parallel_degree = 8       # TP within node
+pipeline_parallel_degree = 8     # PP across nodes
+context_parallel_degree = 1      # CP for long sequences
+
+[training]
+local_batch_size = 32
+seq_len = 8192
+```
+
+**Step 3: 在 512 个 GPU 上启动**
+
+```bash
+# 64 nodes x 8 GPUs = 512 GPUs
+srun torchrun --nnodes=64 --nproc_per_node=8 \
+  -m torchtitan.train \
+  --job.config_file ./llama3_405b.toml
+```
+
+## 何时使用及替代方案
+
+**在以下情况使用 TorchTitan：**
+- 从头开始 pretraining LLM (8B 到 405B+)
+- 需要无需第三方依赖的 PyTorch 原生解决方案
+- 要求可组合的 4D parallelism (FSDP2, TP, PP, CP)
+- 在 H100 上进行支持 Float8 的训练
+- 希望使用与 torchtune/HuggingFace 互通的 checkpoint
+
+**在以下情况使用替代方案：**
+- **Megatron-LM**：适用于仅限 NVIDIA 部署的最高性能需求
+- **DeepSpeed**：更广泛的 ZeRO 优化生态系统，支持推理
+- **Axolotl/TRL**：侧重于微调而非 pretraining
+- **LitGPT**：用于教育目的和较小规模的训练
+
+## 常见问题
+
+**问题：大模型出现显存溢出 (OOM)**
+
+启用 activation checkpointing 并减小 batch size：
+```toml
+[activation_checkpoint]
+mode = "full"  # Instead of "selective"
+
+[training]
+local_batch_size = 1
+```
+
+或者使用 gradient accumulation：
+```toml
+[training]
+local_batch_size = 1
+global_batch_size = 32  # Accumulates gradients
+```
+
+**问题：TP 配合异步 collectives 导致高内存占用**
+
+设置环境变量：
+```bash
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+```
+
+**问题：Float8 训练没有变快**
+
+Float8 仅对大型 GEMM 有益。过滤掉小层：
+```toml
+[quantize.linear.float8]
+filter_fqns = ["attention.wk", "attention.wv", "output", "auto_filter_small_kn"]
+```
+
+**问题：更改 parallelism 后 checkpoint 加载失败**
+
+使用 DCP 的 resharding 功能：
+```bash
+# Convert sharded checkpoint to single file
+python -m torch.distributed.checkpoint.format_utils \
+  dcp_to_torch checkpoint/step-1000 checkpoint.pt
+```
+
+**问题：Pipeline parallelism 初始化**
+
+请先创建 seed checkpoint（参见工作流 4，Step 1）。
+
+## 支持的 models
+
+| Model | Sizes | Status |
+|-------|-------|--------|
+| Llama 3.1 | 8B, 70B, 405B | Production |
+| Llama 4 | Various | Experimental |
+| DeepSeek V3 | 16B, 236B, 671B (MoE) | Experimental |
+| GPT-OSS | 20B, 120B (MoE) | Experimental |
+| Qwen 3 | Various | Experimental |
+| Flux | Diffusion | Experimental |
+
+## 性能基准测试 (H100)
+
+| Model | GPUs | Parallelism | TPS/GPU | Techniques |
+|-------|------|-------------|---------|------------|
+| Llama 8B | 8 | FSDP | 5,762 | Baseline |
+| Llama 8B | 8 | FSDP+compile+FP8 | 8,532 | +48% |
+| Llama 70B | 256 | FSDP+TP+AsyncTP | 876 | 2D parallel |
+| Llama 405B | 512 | FSDP+TP+PP | 128 | 3D parallel |
+
+## 高级主题
+
+**FSDP2 配置**：有关 FSDP2 与 FSDP1 的详细对比以及 ZeRO 的等价物，请参阅 [references/fsdp.md](references/fsdp.md)。
+
+**Float8 训练**：有关 tensorwise 与 rowwise 缩放的方案，请参阅 [references/float8.md](references/float8.md)。
+
+**Checkpointing**：有关 HuggingFace 转换和异步 checkpointing，请参阅 [references/checkpoint.md](references/checkpoint.md)。
+
+**添加自定义模型**：有关 TrainSpec 协议，请参阅 [references/custom-models.md](references/custom-models.md)。
+
+## 资源
+
+- 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

package/extensions/skills/optional/mlops/torchtitan/references/checkpoint.md

@@ -0,0 +1,181 @@
+# TorchTitan 中的 Checkpointing
+
+TorchTitan 使用 PyTorch Distributed Checkpoint (DCP) 来实现容错且可互操作的 checkpointing。
+
+## 基础配置
+
+```toml
+[checkpoint]
+enable = true
+folder = "checkpoint"
+interval = 500
+```
+
+## 仅保存模型（更小的 Checkpoint）
+
+排除 optimizer state 和训练元数据：
+
+```toml
+[checkpoint]
+enable = true
+last_save_model_only = true
+export_dtype = "bfloat16"  # Optional: export in lower precision
+```
+
+## 从加载中排除键
+
+为修改后的设置进行部分 checkpoint 加载：
+
+```toml
+[checkpoint]
+enable = true
+exclude_from_loading = ["data_loader", "lr_scheduler"]
+```
+
+等效的 CLI 命令：
+```bash
+--checkpoint.exclude_from_loading data_loader,lr_scheduler
+```
+
+## 创建 Seed Checkpoint
+
+Pipeline Parallelism 需要此操作以确保初始化的一致性：
+
+```bash
+NGPU=1 CONFIG_FILE=<path_to_config> ./run_train.sh \
+  --checkpoint.enable \
+  --checkpoint.create_seed_checkpoint \
+  --parallelism.data_parallel_replicate_degree 1 \
+  --parallelism.data_parallel_shard_degree 1 \
+  --parallelism.tensor_parallel_degree 1 \
+  --parallelism.pipeline_parallel_degree 1 \
+  --parallelism.context_parallel_degree 1 \
+  --parallelism.expert_parallel_degree 1
+```
+
+这会在单个 CPU 上进行初始化，从而确保在任意数量的 GPU 上都能实现可复现的初始化。
+
+## 异步 Checkpointing
+
+通过异步写入减少 checkpoint 开销：
+
+```toml
+[checkpoint]
+enable = true
+async_mode = "async"  # Options: "disabled", "async", "async_with_pinned_mem"
+```
+
+## HuggingFace 转换
+
+### 训练期间
+
+直接保存为 HuggingFace 格式：
+
+```toml
+[checkpoint]
+last_save_in_hf = true
+last_save_model_only = true
+```
+
+从 HuggingFace 加载：
+
+```toml
+[checkpoint]
+initial_load_in_hf = true
+
+[model]
+hf_assets_path = "./path/to/hf/checkpoint"
+```
+
+### 离线转换
+
+无需运行训练即可进行转换：
+
+```bash
+# HuggingFace -> TorchTitan
+python ./scripts/checkpoint_conversion/convert_from_hf.py \
+  <input_dir> <output_dir> \
+  --model_name llama3 \
+  --model_flavor 8B
+
+# TorchTitan -> HuggingFace
+python ./scripts/checkpoint_conversion/convert_to_hf.py \
+  <input_dir> <output_dir> \
+  --hf_assets_path ./assets/hf/Llama3.1-8B \
+  --model_name llama3 \
+  --model_flavor 8B
+```
+
+### 示例
+
+```bash
+python ./scripts/convert_from_hf.py \
+  ~/.cache/huggingface/hub/models--meta-llama--Meta-Llama-3-8B/snapshots/8cde5ca8380496c9a6cc7ef3a8b46a0372a1d920/ \
+  ./initial_load_path/ \
+  --model_name llama3 \
+  --model_flavor 8B
+```
+
+## 转换为单一 .pt 文件
+
+将 DCP 分片 checkpoint 转换为单个 PyTorch 文件：
+
+```bash
+python -m torch.distributed.checkpoint.format_utils \
+  dcp_to_torch \
+  torchtitan/outputs/checkpoint/step-1000 \
+  checkpoint.pt
+```
+
+## Checkpoint 结构
+
+DCP 保存的分片 checkpoint 可以通过 resharding 用于不同的并行配置：
+
+```
+checkpoint/
+├── step-500/
+│   ├── .metadata
+│   ├── __0_0.distcp
+│   ├── __0_1.distcp
+│   └── ...
+└── step-1000/
+    └── ...
+```
+
+## 恢复训练
+
+训练会自动从配置文件夹中最新的 checkpoint 恢复。要从特定步骤恢复：
+
+```toml
+[checkpoint]
+load_step = 500  # Resume from step 500
+```
+
+## 与 TorchTune 的互操作性
+
+使用 `last_save_model_only = true` 保存的 checkpoint 可以直接加载到 [torchtune](https://github.com/pytorch/torchtune) 中进行微调。
+
+## 完整配置示例
+
+```toml
+[checkpoint]
+enable = true
+folder = "checkpoint"
+interval = 500
+load_step = -1  # -1 = latest, or specify step number
+last_save_model_only = true
+export_dtype = "bfloat16"
+async_mode = "async"
+exclude_from_loading = []
+last_save_in_hf = false
+initial_load_in_hf = false
+create_seed_checkpoint = false
+```
+
+## 最佳实践
+
+1. **大型模型**：使用 `async_mode = "async"` 让 checkpoint 的保存与训练重叠进行
+2. **微调导出**：启用 `last_save_model_only` 和 `export_dtype = "bfloat16"` 以获得更小的文件
+3. **Pipeline parallelism**：务必先创建 seed checkpoint
+4. **调试**：在开发期间频繁保存 checkpoint，在生产环境中则减少保存频率
+5. **HF 互操作性**：使用转换脚本进行离线转换，在训练工作流中则使用直接保存/加载