{"id":"torchtitan","name":"distributed-llm-pretraining-torchtitan","summary":"4D並列処理(FSDP2、TP、PP、CP)を用いたtorchtitanを用いたPyTorchネイティブの分散型LLM事前学習を提供します。","body":"# TorchTitan - PyTorch Native Distributed LLM Pretraining\n\n## Quick start\n\nTorchTitan is PyTorch's official platform for large-scale LLM pretraining with composable 4D parallelism (FSDP2, TP, PP, CP), achieving 65%+ speedups over baselines on H100 GPUs.\n\n**Installation**:\n```bash\n# From PyPI (stable)\npip install torchtitan\n\n# From source (latest features, requires PyTorch nightly)\ngit clone https://github.com/pytorch/torchtitan\ncd torchtitan\npip install -r requirements.txt\n```\n\n**Download tokenizer**:\n```bash\n# Get HF token from https://huggingface.co/settings/tokens\npython scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...\n```\n\n**Start training on 8 GPUs**:\n```bash\nCONFIG_FILE=\"./torchtitan/models/llama3/train_configs/llama3_8b.toml\" ./run_train.sh\n```\n\n## Common workflows\n\n### Workflow 1: Pretrain Llama 3.1 8B on single node\n\nCopy this checklist:\n\n```\nSingle Node Pretraining:\n- [ ] Step 1: Download tokenizer\n- [ ] Step 2: Configure training\n- [ ] Step 3: Launch training\n- [ ] Step 4: Monitor and checkpoint\n```\n\n**Step 1: Download tokenizer**\n\n```bash\npython scripts/download_hf_assets.py \\\n  --repo_id meta-llama/Llama-3.1-8B \\\n  --assets tokenizer \\\n  --hf_token=YOUR_HF_TOKEN\n```\n\n**Step 2: Configure training**\n\nEdit or create a TOML config file:\n\n```toml\n# llama3_8b_custom.toml\n[job]\ndump_folder = \"./outputs\"\ndescription = \"Llama 3.1 8B training\"\n\n[model]\nname = \"llama3\"\nflavor = \"8B\"\nhf_assets_path = \"./assets/hf/Llama-3.1-8B\"\n\n[optimizer]\nname = \"AdamW\"\nlr = 3e-4\n\n[lr_scheduler]\nwarmup_steps = 200\n\n[training]\nlocal_batch_size = 2\nseq_len = 8192\nmax_norm = 1.0\nsteps = 1000\ndataset = \"c4\"\n\n[parallelism]\ndata_parallel_shard_degree = -1  # Use all GPUs for FSDP\n\n[activation_checkpoint]\nmode = \"selective\"\nselective_ac_option = \"op\"\n\n[checkpoint]\nenable = true\nfolder = \"checkpoint\"\ninterval = 500\n```\n\n**Step 3: Launch training**\n\n```bash\n# 8 GPUs on single node\nCONFIG_FILE=\"./llama3_8b_custom.toml\" ./run_train.sh\n\n# Or explicitly with torchrun\ntorchrun --nproc_per_node=8 \\\n  -m torchtitan.train \\\n  --job.config_file ./llama3_8b_custom.toml\n```\n\n**Step 4: Monitor and checkpoint**\n\nTensorBoard logs are saved to `./outputs/tb/`:\n```bash\ntensorboard --logdir ./outputs/tb\n```\n\n### Workflow 2: Multi-node training with SLURM\n\n```\nMulti-Node Training:\n- [ ] Step 1: Configure parallelism for scale\n- [ ] Step 2: Set up SLURM script\n- [ ] Step 3: Submit job\n- [ ] Step 4: Resume from checkpoint\n```\n\n**Step 1: Configure parallelism for scale**\n\nFor 70B model on 256 GPUs (32 nodes):\n```toml\n[parallelism]\ndata_parallel_shard_degree = 32  # FSDP across 32 ranks\ntensor_parallel_degree = 8        # TP within node\npipeline_parallel_degree = 1      # No PP for 70B\ncontext_parallel_degree = 1       # Increase for long sequences\n```\n\n**Step 2: Set up SLURM script**\n\n```bash\n#!/bin/bash\n#SBATCH --job-name=llama70b\n#SBATCH --nodes=32\n#SBATCH --ntasks-per-node=8\n#SBATCH --gpus-per-node=8\n\nsrun torchrun \\\n  --nnodes=32 \\\n  --nproc_per_node=8 \\\n  --rdzv_backend=c10d \\\n  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \\\n  -m torchtitan.train \\\n  --job.config_file ./llama3_70b.toml\n```\n\n**Step 3: Submit job**\n\n```bash\nsbatch multinode_trainer.slurm\n```\n\n**Step 4: Resume from checkpoint**\n\nTraining auto-resumes if checkpoint exists in configured folder.\n\n### Workflow 3: Enable Float8 training for H100s\n\nFloat8 provides 30-50% speedup on H100 GPUs.\n\n```\nFloat8 Training:\n- [ ] Step 1: Install torchao\n- [ ] Step 2: Configure Float8\n- [ ] Step 3: Launch with compile\n```\n\n**Step 1: Install torchao**\n\n```bash\nUSE_CPP=0 pip install git+https://github.com/pytorch/ao.git\n```\n\n**Step 2: Configure Float8**\n\nAdd to your TOML config:\n```toml\n[model]\nconverters = [\"quantize.linear.float8\"]\n\n[quantize.linear.float8]\nenable_fsdp_float8_all_gather = true\nprecompute_float8_dynamic_scale_for_fsdp = true\nfilter_fqns = [\"output\"]  # Exclude output layer\n\n[compile]\nenable = true\ncomponents = [\"model\", \"loss\"]\n```\n\n**Step 3: Launch with compile**\n\n```bash\nCONFIG_FILE=\"./llama3_8b.toml\" ./run_train.sh \\\n  --model.converters=\"quantize.linear.float8\" \\\n  --quantize.linear.float8.enable_fsdp_float8_all_gather \\\n  --compile.enable\n```\n\n### Workflow 4: 4D parallelism for 405B models\n\n```\n4D Parallelism (FSDP + TP + PP + CP):\n- [ ] Step 1: Create seed checkpoint\n- [ ] Step 2: Configure 4D parallelism\n- [ ] Step 3: Launch on 512 GPUs\n```\n\n**Step 1: Create seed checkpoint**\n\nRequired for consistent initialization across PP stages:\n```bash\nNGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \\\n  --checkpoint.enable \\\n  --checkpoint.create_seed_checkpoint \\\n  --parallelism.data_parallel_shard_degree 1 \\\n  --parallelism.tensor_parallel_degree 1 \\\n  --parallelism.pipeline_parallel_degree 1\n```\n\n**Step 2: Configure 4D parallelism**\n\n```toml\n[parallelism]\ndata_parallel_shard_degree = 8   # FSDP\ntensor_parallel_degree = 8       # TP within node\npipeline_parallel_degree = 8     # PP across nodes\ncontext_parallel_degree = 1      # CP for long sequences\n\n[training]\nlocal_batch_size = 32\nseq_len = 8192\n```\n\n**Step 3: Launch on 512 GPUs**\n\n```bash\n# 64 nodes x 8 GPUs = 512 GPUs\nsrun torchrun --nnodes=64 --nproc_per_node=8 \\\n  -m torchtitan.train \\\n  --job.config_file ./llama3_405b.toml\n```\n\n## When to use vs alternatives\n\n**Use TorchTitan when:**\n- Pretraining LLMs from scratch (8B to 405B+)\n- Need PyTorch-native solution without third-party dependencies\n- Require composable 4D parallelism (FSDP2, TP, PP, CP)\n- Training on H100s with Float8 support\n- Want interoperable checkpoints with torchtune/HuggingFace\n\n**Use alternatives instead:**\n- **Megatron-LM**: Maximum performance for NVIDIA-only deployments\n- **DeepSpeed**: Broader ZeRO optimization ecosystem, inference support\n- **Axolotl/TRL**: Fine-tuning rather than pretraining\n- **LitGPT**: Educational, smaller-scale training\n\n## Common issues\n\n**Issue: Out of memory on large models**\n\nEnable activation checkpointing and reduce batch size:\n```toml\n[activation_checkpoint]\nmode = \"full\"  # Instead of \"selective\"\n\n[training]\nlocal_batch_size = 1\n```\n\nOr use gradient accumulation:\n```toml\n[training]\nlocal_batch_size = 1\nglobal_batch_size = 32  # Accumulates gradients\n```\n\n**Issue: TP causes high memory with async collectives**\n\nSet environment variable:\n```bash\nexport TORCH_NCCL_AVOID_RECORD_STREAMS=1\n```\n\n**Issue: Float8 training not faster**\n\nFloat8 only benefits large GEMMs. Filter small layers:\n```toml\n[quantize.linear.float8]\nfilter_fqns = [\"attention.wk\", \"attention.wv\", \"output\", \"auto_filter_small_kn\"]\n```\n\n**Issue: Checkpoint loading fails after parallelism change**\n\nUse DCP's resharding capability:\n```bash\n# Convert sharded checkpoint to single file\npython -m torch.distributed.checkpoint.format_utils \\\n  dcp_to_torch checkpoint/step-1000 checkpoint.pt\n```\n\n**Issue: Pipeline parallelism initialization**\n\nCreate seed checkpoint first (see Workflow 4, Step 1).\n\n## Supported models\n\n| Model | Sizes | Status |\n|-------|-------|--------|\n| Llama 3.1 | 8B, 70B, 405B | Production |\n| Llama 4 | Various | Experimental |\n| DeepSeek V3 | 16B, 236B, 671B (MoE) | Experimental |\n| GPT-OSS | 20B, 120B (MoE) | Experimental |\n| Qwen 3 | Various | Experimental |\n| Flux | Diffusion | Experimental |\n\n## Performance benchmarks (H100)\n\n| Model | GPUs | Parallelism | TPS/GPU | Techniques |\n|-------|------|-------------|---------|------------|\n| Llama 8B | 8 | FSDP | 5,762 | Baseline |\n| Llama 8B | 8 | FSDP+compile+FP8 | 8,532 | +48% |\n| Llama 70B | 256 | FSDP+TP+AsyncTP | 876 | 2D parallel |\n| Llama 405B | 512 | FSDP+TP+PP | 128 | 3D parallel |\n\n## Advanced topics\n\n**FSDP2 configuration**: See [references/fsdp.md](references/fsdp.md) for detailed FSDP2 vs FSDP1 comparison and ZeRO equivalents.\n\n**Float8 training**: See [references/float8.md](references/float8.md) for tensorwise vs rowwise scaling recipes.\n\n**Checkpointing**: See [references/checkpoint.md](references/checkpoint.md) for HuggingFace conversion and async checkpointing.\n\n**Adding custom models**: See [references/custom-models.md](references/custom-models.md) for TrainSpec protocol.\n\n## Resources\n\n- GitHub: https://github.com/pytorch/torchtitan\n- Paper: https://arxiv.org/abs/2410.06511\n- ICLR 2025: https://iclr.cc/virtual/2025/poster/29620\n- PyTorch Forum: https://discuss.pytorch.org/c/distributed/torchtitan/44","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/01-model-architecture/torchtitan","license":"MIT","category":"writing","lang":"en","tokens":2461,"stars":0,"calls30d":2,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/checkpoint.md","size":4216,"sha256":"0cacfc3cf87a3fa8cc40eaac4516113afc2bc68ef17186d249e3ed8a25a55b5d"},{"path":"references/custom-models.md","size":7383,"sha256":"05136b49e9f53a3e331c89f59de46be5e26664d8eaa0884d10c8dbf927f8ca76"},{"path":"references/float8.md","size":4055,"sha256":"33f05eacea84e571e4fc77fa39a1ec751c740645af14832322f80f63cdbe6ab8"},{"path":"references/fsdp.md","size":3888,"sha256":"d0743e131ff129485e17f5ab01c604c12fdf682bd2e8a637fa338e5ee40f48a4"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["arxiv.org","discuss.pytorch.org","huggingface.co","iclr.cc"]}}