TorchTitan Distributed LLM Pretraining

A skill that walks Claude through configuring and launching large-scale LLM pretraining with torchtitan — 4D parallelism (FSDP2/TP/PP/CP), Float8, and distributed checkpointing from 8 to 512+ GPUs.

Dev & CodingAdvanced★ 12,666⑂ 912AI score 8/10Last updated: Jun 16, 2026

What it does

Covers installing torchtitan, PyTorch's official large-scale pretraining platform, and getting a first run going.
Supplies ready-to-edit TOML configs for 4D parallelism: FSDP2 sharding, tensor parallel, pipeline parallel and context parallel, sized to your model and GPU count.
Shows how to turn on Float8 + torch.compile for 30–50% throughput gains on H100s, submit multi-node jobs via SLURM, and manage distributed (DCP) checkpoints including resharding and conversion.
Includes a practical troubleshooting section: OOM, TP memory blowups with async collectives, Float8 not helping, checkpoint load failures after parallelism changes, and PP seed-checkpoint setup.

Who it's for

ML engineers and researchers pretraining from scratch: Llama 3.1 (8B/70B/405B), DeepSeek V3, Qwen 3, GPT-OSS MoE models.
Infra owners planning a scale-up path from one 8-GPU node to a 512-GPU cluster.
Teams that want a PyTorch-native stack instead of Megatron-LM or DeepSpeed.

Note: for fine-tuning (LoRA/SFT), Axolotl or TRL is the better fit.

Example uses

Pretrain Llama 3.1 8B on a single 8×H100 node: download the tokenizer, write llama3_8b_custom.toml, run run_train.sh, monitor via TensorBoard.
Train 70B on 32 nodes / 256 GPUs: FSDP 32 × TP 8 config plus a SLURM sbatch script, with automatic resume from checkpoint.
4D parallelism for 405B: create a seed checkpoint first, then run FSDP 8 × TP 8 × PP 8 across 64 nodes (512 GPUs).

· · · Install guide · · ·

Try it now, no install

Paste this into Claude to use the skill without installing anything.

Read the instructions in this file and follow them to help me:
https://raw.githubusercontent.com/Orchestra-Research/AI-Research-SKILLs/HEAD/01-model-architecture/torchtitan/SKILL.md

What I want: (describe your task here)

If Claude can't open the link, open it yourself and paste the contents instead.

↓ If it works for you, download the ZIP below and install it. Then it runs on its own — no pasting each time.

Install in the Claude app (no terminal)

Download the ZIP with the button below.
In Claude, open Settings → Capabilities and turn on 'Code execution and file creation'. (one time)
Go to Customize → Skills → + → 'Upload a skill' and upload the ZIP.

↓ Download ZIP

Install in Claude Code

Let Claude do it — paste this into Claude Code

Install the skill I found on Claude Skill Mart.
Copy the 01-model-architecture/torchtitan folder from the GitHub repo Orchestra-Research/AI-Research-SKILLs into my ~/.claude/skills/distributed-llm-pretraining-torchtitan/.
When it's done, tell me in one line what this skill can do.

Install with a command instead

git clone https://github.com/Orchestra-Research/AI-Research-SKILLs.git && mkdir -p ~/.claude/skills && cp -r AI-Research-SKILLs/01-model-architecture/torchtitan ~/.claude/skills/distributed-llm-pretraining-torchtitan

⚠ This is a third-party skill. Check the source repository before installing.

Open a terminal and clone the skills repo: git clone https://github.com/Orchestra-Research/AI-Research-SKILLs.git
Create the Claude skills directory: mkdir -p ~/.claude/skills
Copy just this skill: cp -r AI-Research-SKILLs/01-model-architecture/torchtitan ~/.claude/skills/distributed-llm-pretraining-torchtitan
Restart Claude Code and confirm the skill appears in your skills list.
Prepare the runtime for real training: pip install torchtitan (or install from source) with torch>=2.6 and torchao>=0.5.
Get a Hugging Face token and fetch a tokenizer: python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...
Trigger the skill with a prompt like "Set up Llama 3.1 8B pretraining on one 8×H100 node with Float8 enabled."

View source on GitHub ↗License: MIT