TorchTitan Skill: PyTorch 4D Parallel LLM Pretraining

TorchTitan is PyTorch's official platform for pretraining large language models at scale, and this Hermes Agent skill packages it as a working reference for the agent. You would reach for it when you are training a model from scratch rather than fine-tuning one, when you need composable 4D parallelism (FSDP2, tensor, pipeline, and context parallel), or when you want a PyTorch-native stack without pulling in a third-party training framework.

What it does

The skill gives the agent a structured playbook for TorchTitan runs: installing the package, fetching tokenizer assets, selecting a run configuration from the Python config registry, launching on a single node or across SLURM nodes, enabling Float8 on H100s, and composing all four parallelism dimensions for very large models. It also carries the operational knowledge that usually gets learned the hard way: how to recover from out-of-memory errors, why Float8 sometimes does not speed anything up, and what to do when a checkpoint will not load after you change the parallelism layout.

In practice the agent uses this to draft launch commands, write config registry entries, and diagnose failures in a pretraining job. The source reports 65%+ speedups over baselines on H100 GPUs, which is the headline reason teams pick it over a hand-rolled training loop.

Before you start

This is an optional skill, installed on demand rather than shipped by default. Its path is optional-skills/mlops/torchtitan, version 1.0.1, authored by Orchestra Research under the MIT license. It runs on linux and macos.

The skill assumes you already have a working PyTorch environment and GPUs to train on. TorchTitan itself installs from PyPI for the stable release, or from source when you need the latest features, in which case the source notes that PyTorch nightly is required. Downloading gated model assets such as the Llama tokenizer needs a Hugging Face token, which you generate from your Hugging Face account settings.

Installing TorchTitan

Two install paths exist. The PyPI route gives you the stable release. The source route clones the repository and installs its requirements, and is the one to use when you need features that have not been released yet.

# 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

Downloading the tokenizer

TorchTitan does not bundle tokenizers. You pull them from the Hugging Face Hub with the repository's download script, passing the repo id, the asset type, and your token.

# 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=...

Starting a training run

Configurations are not selected by pointing at a TOML file. They are registered by name in a Python config registry, and you choose one with the CONFIG environment variable alongside the MODULE that names the model family.

# Configs are selected by name from the Python config registry
# (torchtitan/models/llama3/config_registry.py), not by TOML path
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

Workflow 1: Pretrain Llama 3.1 8B on a single node

The canonical first run. Work through the checklist in order.

Single Node Pretraining:
- [ ] Step 1: Download tokenizer
- [ ] Step 2: Configure training
- [ ] Step 3: Launch training
- [ ] Step 4: Monitor and checkpoint

Step 1: Download the tokenizer

Same download script as above, with your token substituted in.

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

Step 2: Configure training

Run configs live in a Python config registry at torchtitan/models/llama3/config_registry.py and are picked by name through CONFIG= or --config. To customize, either register your own entry in that registry or override individual fields on the command line, for example --optimizer.lr 3e-4 --training.steps 1000.

The source shows the equivalent settings for an 8B run as fields. You set them in the registry entry or pass them as --section.key value overrides.

# fields for a llama3 8B run (register in config_registry.py or pass as --overrides)
[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: Launch training

Three ways to start, from the convenience wrapper to the explicit torchrun invocation. run_train.sh wraps torchrun, so the third form is what the script ultimately runs.

# 8 GPUs on single node (config selected by name from the registry)
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

# Override individual fields on the command line
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh --optimizer.lr 3e-4 --training.steps 1000

# Or explicitly with torchrun (run_train.sh wraps this)
torchrun --nproc_per_node=8 \
  -m torchtitan.train \
  --module llama3 --config llama3_8b

Step 4: Monitor and checkpoint

TensorBoard logs are written under ./outputs/tb/. Point TensorBoard at that directory to watch the run.

tensorboard --logdir ./outputs/tb

Workflow 2: Multi-node training with SLURM

Scaling past one node means configuring the parallelism dimensions deliberately and handing the launch to 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: Configure parallelism for scale

The source's example is a 70B model on 256 GPUs across 32 nodes. FSDP shards across 32 ranks, tensor parallelism stays inside a node at degree 8, pipeline parallelism is off, and context parallelism is left at 1 unless you have long sequences.

[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: Set up the SLURM script

The batch script requests 32 nodes with 8 tasks and 8 GPUs each, then launches torchrun under srun with the c10d rendezvous backend pointed at the master address and port.

#!/usr/bin/env 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 \
  --module llama3 --config llama3_70b

Step 3: Submit the job

sbatch multinode_trainer.slurm

Step 4: Resume from checkpoint

Training resumes automatically when a checkpoint already exists in the configured folder. You do not pass a resume flag; the presence of the checkpoint is what triggers it.

Workflow 3: Enable Float8 training on H100s

Float8 gives a 30-50% speedup on H100 GPUs according to the source. It requires torchao and torch.compile to pay off.

Float8 Training:
- [ ] Step 1: Install torchao
- [ ] Step 2: Configure Float8
- [ ] Step 3: Launch with compile

Step 1: Install torchao

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

Step 2: Configure Float8

Float8 is not a TOML section in current TorchTitan. It is applied at config time through the quantization parameter of your model_registry() call inside the config registry. You add a Float8LinearConverter.Config there.

# in torchtitan/models/llama3/config_registry.py (your model_registry(...) call)
from torchtitan.components.quantization import Float8LinearConverter

model_spec = model_registry(
    "8B",
    quantization=[
        Float8LinearConverter.Config(
            recipe_name="rowwise",          # or "rowwise_with_gw_hp"
            filter_fqns=["output"],          # skip layers too small to benefit
            model_compile_enabled=True,      # requires torch.compile for competitive perf
        ),
    ],
)

You also need to turn on torch.compile in the run config.

[compile]
enable = true
components = ["model", "loss"]

Step 3: Launch with compile

The Float8 configuration is baked into the registered config, so you select that config and enable compile on the command line.

# Float8 config is baked into the registered config; just select it and enable compile
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh --compile.enable

Workflow 4: 4D parallelism for 405B models

At 405B you combine all four dimensions. The seed checkpoint step is not optional here: pipeline stages need a consistent starting point.

4D Parallelism (FSDP + TP + PP + CP):
- [ ] Step 1: Create seed checkpoint
- [ ] Step 2: Configure 4D parallelism
- [ ] Step 3: Launch on 512 GPUs

Step 1: Create a seed checkpoint

This is required for consistent initialization across pipeline parallel stages. You run with a single GPU and all parallelism degrees forced to 1, asking the run to create a seed checkpoint.

NGPU=1 MODULE=llama3 CONFIG=llama3_405b ./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: Configure 4D parallelism

FSDP at degree 8, tensor parallel at 8 within a node, pipeline parallel at 8 across nodes, and context parallel left at 1 unless you need it for long sequences.

[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: Launch on 512 GPUs

64 nodes of 8 GPUs each, launched through srun and torchrun.

# 64 nodes x 8 GPUs = 512 GPUs
srun torchrun --nnodes=64 --nproc_per_node=8 \
  -m torchtitan.train \
  --module llama3 --config llama3_405b

When not to use it

The source is explicit about where TorchTitan fits and where it does not. Choose it when you are pretraining from scratch at 8B to 405B+, when you want a PyTorch-native solution with no third-party dependencies, when you need composable 4D parallelism, when you are on H100s and want Float8, or when you need checkpoints that interoperate with torchtune and HuggingFace.

Reach for something else in these cases:

• Megatron-LM when you want maximum performance and are committed to NVIDIA-only deployments.
• DeepSpeed when you want a broader ZeRO optimization ecosystem and inference support.
• Axolotl or TRL when the job is fine-tuning rather than pretraining.
• LitGPT for educational or smaller-scale training.

Limits and gotchas

The source collects the failures people actually hit. Each has a documented remedy.

Out of memory on large models. Switch activation checkpointing from selective to full and drop the local batch size to 1.

[activation_checkpoint]
mode = "full"  # Instead of "selective"

[training]
local_batch_size = 1

Alternatively, use gradient accumulation by setting a global batch size larger than the local one.

[training]
local_batch_size = 1
global_batch_size = 32  # Accumulates gradients

Tensor parallelism causing high memory with async collectives. Set an environment variable before launching.

export TORCH_NCCL_AVOID_RECORD_STREAMS=1

Float8 training not faster. Float8 only helps large GEMMs. Filter out small layers through the converter's filter_fqns, and add auto_filter_small_kn to skip layers too small to benefit automatically.

from torchtitan.components.quantization import Float8LinearConverter

Float8LinearConverter.Config(
    # add "auto_filter_small_kn" to auto-skip layers too small to benefit
    filter_fqns=["attention.wk", "attention.wv", "output", "auto_filter_small_kn"],
    model_compile_enabled=True,
)

Checkpoint loading fails after a parallelism change. Use DCP's resharding capability to convert the sharded checkpoint into a single file.

# Convert sharded checkpoint to single file
python -m torch.distributed.checkpoint.format_utils \
  dcp_to_torch checkpoint/step-1000 checkpoint.pt

Pipeline parallelism initialization. Create the seed checkpoint first, as in Workflow 4 Step 1.

Supported models

Llama 3.1 at 8B, 70B, and 405B is the production-ready path. Everything else is marked experimental: Llama 4, DeepSeek V3 (16B, 236B, 671B MoE), GPT-OSS (20B, 120B MoE), Qwen 3, and Flux for diffusion. Treat experimental entries as usable but not something to build a production schedule around without your own validation.

Performance benchmarks on H100

The source's reference numbers, useful for sanity-checking your own throughput. Llama 8B on 8 GPUs with plain FSDP gives 5,762 TPS/GPU; adding compile and FP8 raises it to 8,532, a 48% gain. Llama 70B on 256 GPUs with FSDP, TP, and AsyncTP reaches 876 TPS/GPU. Llama 405B on 512 GPUs with FSDP, TP, and PP reaches 128 TPS/GPU.

Going deeper

The skill links to four reference documents in the Hermes Agent repository: references/fsdp.md for the FSDP2 versus FSDP1 comparison and ZeRO equivalents, references/float8.md for tensorwise versus rowwise scaling recipes, references/checkpoint.md for HuggingFace conversion and async checkpointing, and references/custom-models.md for the TrainSpec protocol when adding your own architectures.

What pairs with this

The upstream project lives at github.com/pytorch/torchtitan. The design is described in arXiv paper 2410.06511 and was presented at ICLR 2025. For questions, the PyTorch Forum has a dedicated TorchTitan category under discuss.pytorch.org/c/distributed/torchtitan/44. If you are working inside Hermes Agent, pair this skill with the checkpoint and FSDP reference files above before attempting a multi-node run.