Bf16 Optimizer States Copy Section

BF16 Optimizer States Copy Chapter

In the default fp32 training configuration (training.dtype="float32"), Adam/AdamW keep momentum (exp_avg) and variance (exp_avg_sq) in float32, which roughly doubles optimizer-state memory versus storing those buffers in bfloat16.

Set optimizer.implementation to fused_opt_states_bf16 to use the fused Adam/AdamW CUDA kernel with bf16 optimizer states ,fp32 parameters and fp32 grads (mixed precision). That is the main scenario this option targets: lower optimizer memory while keeping params and grads in full precision.

If you use training.dtype="bfloat16" (params and grads in bf16), you typically keep implementation="fused" (default). PyTorch then aligns optimizer state dtypes with training; you do not need fused_opt_states_bf16 unless you explicitly want the pre-hook initialization path (behavior should match fused training in practice).

This is useful for memory-constrained training where slightly lower precision in moment estimates is acceptable.

Background Copy Section

This technique was notably used by DeepSeek-V3 to train their 671B-parameter MoE model on 14.8 trillion tokens with reduced memory overhead. Their approach demonstrated that both momentum and variance buffers can be stored in bfloat16 without convergence issues, particularly for MoE architectures where expert gradients are smaller in magnitude. The effort to add native bf16 AdamW support to PyTorch is tracked in pytorch/pytorch#146542.

Usage Copy Section

In your config registry function:

from torchtitan.components.optimizer import OptimizersContainer
optimizer=OptimizersContainer.Config(
    name="AdamW",
    implementation="fused_opt_states_bf16",
),

Or via CLI override:

--optimizer.name AdamW --optimizer.implementation fused_opt_states_bf16

Requirements Copy Section

- Optimizer: Must be Adam or AdamW.
- Implementation: Must be fused_opt_states_bf16. The fused CUDA kernel (FusedAdamMathFunctorMP) handles mixed-precision updates (fp32 parameters + bf16 states).

These constraints are validated at config time.

How it works Copy Section

A step pre-hook is registered on each optimizer instance. Before Adam's lazy state initialization runs on the first step, the hook pre-populates exp_avg and exp_avg_sq as bfloat16 tensors. When _init_group finds non-empty state, it skips its own fp32 allocation. The fused kernel detects the dtype mismatch between fp32 parameters and bf16 states and dispatches to the mixed-precision code path.

Interaction with other features Copy Section

- training.dtype: Primary use case is float32 training with fused_opt_states_bf16 for optimizer-state memory savings. With bfloat16 training, default implementation="fused" is usually enough; see the introduction above.
- Checkpointing: Optimizer states are saved in bfloat16 when this option is enabled. On resume, use the same implementation="fused_opt_states_bf16" so checkpoint state matches. The pre-hook only creates bf16 tensors for parameters with empty state; if a checkpoint already populated state, those dtypes are preserved. Mixing checkpoint dtype with a different implementation across save/load is unsupported and can result in dtype-mismatch.
- FSDP: Compatible with FSDP2. The optimizer sees DTensor parameters; the bf16 state hook operates on the local shards.

Limitations Copy Section

- Only supported with OptimizersContainer (standard forward/backward training). Not supported with OptimizersInBackwardContainer (optimizer-step-in-backward); that combination is rejected in OptimizersInBackwardContainer.Config.__post_init__.
- Only Adam and AdamW with fused_opt_states_bf16 are supported.
- Lower precision in moment estimates may affect convergence for some models or hyperparameter settings. Users should verify loss convergence for their specific use case.

---

Checkpoint Copy Section

How to use checkpointing in torchtitan

You may want to enable checkpointing in torchtitan for better fault tolerance during training, or to enable easier importing and exporting of weights between torchtitan and other libraries. torchtitan offers varying degrees of support for other checkpoint formats which are listed further below.

A general guide to use checkpoints during training Copy Chapter

1. ENABLE CHECKPOINTING
In your config_registry function, configure the checkpoint settings:

checkpoint=CheckpointManager.Config(
    interval=500,
),

2. SAVE MODEL ONLY
By setting last_save_model_only to True, the checkpoint will only contain the model and exclude the optimizer state and extra train states, resulting in a smaller checkpoint size.

checkpoint=CheckpointManager.Config(
    interval=500,
    last_save_model_only=True,
),

3. CHOOSE DESIRED EXPORT PRECISION
The default model states are in float32. You can choose to export the checkpoint in a lower precision format such as bfloat16.

checkpoint=CheckpointManager.Config(
    interval=500,
    last_save_model_only=True,
    export_dtype="bfloat16",
),

4. EXCLUDING SPECIFIC KEYS FROM CHECKPOINT LOADING
In some cases, you may want to partially load from a previous-trained checkpoint and modify certain settings, such as the number of GPUs or the current step. To achieve this, you can use the exclude_from_loading parameter to specify which keys should be excluded from loading.

checkpoint=CheckpointManager.Config(
    exclude_from_loading=["data_loader", "lr_scheduler"],
),

5. EXAMPLE CHECKPOINT CONFIGURATION

checkpoint=CheckpointManager.Config(
    interval=10,
    load_step=5,
    last_save_model_only=True,
    export_dtype="bfloat16",
),

A more exhaustive and up-to-date list of checkpoint config options can be found in torchtitan/components/checkpoint.py (CheckpointManager.Config).

Creating a seed checkpoint Copy Chapter

To create a seed checkpoint, use the same model config as you use for training.
e.g.

NGPU=1 ./run_train.sh --module <module_name> --config <config_name> --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

Conversion support Copy Chapter

HuggingFace Copy Section

1. You can directly save huggingface model weights during training by using the --checkpoint.last_save_in_hf and --checkpoint.last_save_model_only options together. To directly load a torchtitan training session from a huggingface safetensors file, enable --checkpoint.initial_load_in_hf, and set either --hf_assets_path or --checkpoint.initial_load_path to the directory containing the huggingface checkpoint. --checkpoint.initial_load_path overrides --hf_assets_path if both are set.

2. To directly reformat the weights without the need to run a training loop, run the corresponding conversion script. The naming scheme is torchtitan-centric, e.g. convert_from_hf means convert hf->tt.

python ./scripts/checkpoint_conversion/convert_from_hf.py <input_dir> <output_dir> --model_name <model_name> --model_flavor <model_flavor>
python ./scripts/checkpoint_conversion/convert_to_hf.py <input_dir> <output_dir> --hf_assets_path ./assets/hf/Llama3.1-8B --model_name <model_name> --model_flavor <model_flavor>
e.g.

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

Torch Copy Section

This guide will walk you through the steps required to convert a checkpoint from torchtitan so that it can be loaded into pt format.

1. CHECKPOINT CONFIGURATION

checkpoint=CheckpointManager.Config(
    interval=10,
    last_save_model_only=True,
    export_dtype="bfloat16",
),

2. SAVE THE FINAL CHECKPOINT\
Once the above have been set, the final checkpoint at the end of the training step will consist of model only with the desired export dtype. However, if the final step has not been reached yet, full checkpoints will still be saved so that training can be resumed.

3. CONVERT SHARDED CHECKPOINTS TO A SINGLE FILE\
Finally, once you have obtained the last checkpoint, you can use the following command to convert the sharded checkpoints to a single .pt file.

python -m torch.distributed.checkpoint.format_utils dcp_to_torch torchtitan/outputs/checkpoint/step-1000 checkpoint.pt

That's it. You have now successfully converted a sharded torchtitan checkpoint for use with pytorch formats.

---

Composability Copy Section

Building a Clean, Readable Distributed LLM

Making the model "pipeline friendly" Copy Chapter

Simplifying the top-level model forward Copy Section

Example (PR #321):
We used to slice the freqs_cis buffer by seq_len in the top level forward, pass that into child modules, and expect that inside the child modules the seq_len would match up with the size of other local tensors. But we don't know about whether TP was applied or not when we consider PP splitting and could create a mismatch. Its just as easy to perform the freqs_cis slicing inside the child submodule, using the runtime-accurate local seq_len, and this sidesteps the issue at PP slicing time.

Example (PR #322):
We decided to actually reuse the top-level model object on every PP stage, just delete the layers we don't want, and make sure that the top-level forward would do the right thing. This means we don't have to make a separate runtime pp_forward that glues together child modules per stage. The first change was using a moduledict instead of modulelist to store layers. This preserves layer Fully Qualified Names (FQNs) even when deleting some layers - e.g. layers.1 stays layers.1 even if you remove layers.0, which isn't true for a list- this matters for checkpoint save/load. Preserving FQNs is a requirement for using Distributed Checkpointing (DCP) since it uses FQNs as globally unique IDs for sharding metadata. The second change was making the input and output layers optional- if the layer exists, we run it, otherwise we feed the input through to bypass it. With these two changes, we can just (meta)-initialize the whole model, delete the unused parts per stage, then materialize the remaining part on GPU before loading a checkpoint.

Using a seed checkpoint for init Copy Chapter

For now, we sidestep all these problems with a simple but brutal solution: Initialize the whole model on some CPU instance, save a checkpoint file, and then lean on Distributed Checkpointing's "load" functionality to initialize the FQNs that are present on a given PP stage after stage creation. For future work, we consider adding a more elaborate initialization scheme to torch.pipelining.

One issue with seed checkpoints is that we rely on initializing _every_ model state from the checkpoint, which means the model can't have any non-persistent buffers, or else we have to specially initialize those in train.py after pipeline splitting. freqs_cis was originally a non-persistent buffer, and we changed this to persistent in order to load it from the seed checkpoint.

On upcasting the final output to fp32 Copy Chapter

Setting TORCH_NCCL_AVOID_RECORD_STREAMS=1 for TP TORCH_NCCL_AVOID_RECORD_STREAMS=1 for TP')" title="Copy chapter prompt for LLMs"> Copy Chapter

TP uses async collectives (i.e. with async_op=True), such as all-gather, reduce-scatter, and all-reduce, to overlap communication with compute. Under the hood, an async collective runs the NCCL communication kernel in a separate CUDA stream, owned by the process group. Calling wait() on the returned work object has the current stream wait for the process group's stream, allowing the current stream to correctly use the result of the collective.

This represents a producer-consumer pattern across streams: the collective tensors are produced in a compute stream (usually the default stream), and they are consumed in a communication stream (from the process group). Under such producer-consumer patterns across streams, we must ensure that the tensors are not freed before their usage in the consumer stream.

Tensor.record_stream is a legacy approach for ensuring this. The process group will call record_stream(comm_stream) on the collective input and output tensors after issuing the collective kernel in the process group's comm_stream. This records a CUDA event in comm_stream, and the CUDA caching allocator that manages CUDA tensor memory in PyTorch will query this recorded event upon future allocations. Only once the event has completed, meaning that the collective has finished running, can the tensor memory be freed and considered for future reuse. This couples the caching allocator's memory reuse with _GPU kernel timing_, which does not otherwise happen without record_stream. While the collective kernel runs on GPU, any allocations made from the CPU for future ops cannot reuse that memory, even if we know that those future ops must run after the current collective. This inability to reuse leads to unexpected memory stacking.

By setting TORCH_NCCL_AVOID_RECORD_STREAMS=1, the process group avoids calling record_stream on the collective tensors and instead uses a different approach. It simply stashes references to the collective tensors until the user calls wait() on the work object. Holding references ensures that the collective tensors will not be freed by the caching allocator. This can only lead to a memory regression if the user never calls wait(), where with record_stream, the caching allocator would still eventually free the collective tensors once the collective finishes on the GPU. Since this is not common or an expected usage, we recommend setting this environment variable.

---

Converging Copy Section

This note clarifies the recommended practices to follow when testing the loss converging of a new feature.

Disclaimers Copy Section

2. The focus is on the correctness of new distributed training techniques. For a new model size / architecture, the demonstration of loss-converging is not in the scope of this note.

3. Different tests can be run with different random seeds, as
- It’s closer to the production environment.
- One anyway cannot hope to achieve exact reproducibility when working with different distributed settings[^1][^2].

Guidelines Copy Chapter

To validate the correctness of a distributed training technique, one should try to keep the determinism in the input data to minimize the differences it could cause. To make sure the global batch size and in general #tokens per iteration stay the same, one can fix the local batch size (training.local_batch_size) in the config_registry function, and at the same time fix the data parallel degree.

If the technique is a parallelism (TP/PP/CP/etc)
- The control set is a 1D FSDP job on dp GPUs (or any other verified setups), with a trusted training config (e.g. those in config_registry.py).
- The minimal test set is a 2D job on dp*p GPUs, where p >= 2 is the degree of the experimented parallelism.
- For some parallelisms, larger p may cause larger discrepancies in numeric due to various reasons. For example, current implementation of CP uses torch.bfloat16 (under default mixed precision training configs) when accumulating intermediate results. A higher p is desired to ensure the parallelism works properly, at the cost of more hardware resources.
- Certain parallelisms may impose additional requirements on the batch size. For instance, PP requires local batch size to be at least the number of microbatches (or equivalently, the number of pipeline stages) to reduce bubbles. A valid comparison example would be 1D FSDP on N GPUs with local batch size 8, and 2D FSDP + PP on 4N GPUs (DP N, PP 4) with Interleaved 1F1B schedule (also with local batch size 8), where each PP rank gets two pipeline stages.
- N-D parallelisms should be tested on more (e.g. dptpppcpp) GPUs.
- For mutually compatible parallelisms, one can assume the converging of a higher dimensional combination implies that of a lower dimensional subset. For example, in order to verify the converging of FSDP, TP, PP, and CP, in principle one can just run two jobs, the control set being 1D FSDP and the test set being 4D including all.
- A lower dimensional combination can be helpful for enlarging the degrees of some parallelisms and for debugging purposes in general.

If the technique is not a parallelism
- The control set is a 1D FSDP job on dp GPUs (or any other verified setups), without the technique enabled.
- One may argue that to test an optimization technique for other / higher dimensional parallelism, we could directly use that parallelism as the control set. E.g. to test Async TP, the baseline could be the verified 2D FSDP + TP. We use 1D FSDP for two reasons: (1) in general loss-converging correctness should be transitive; and (2) depending on the definition of correctness, small discrepancies are sometimes acceptable (e.g. due to randomness, mixed precision, etc.), but the discrepancies could add up to be unacceptable after many transitive links.
- The minimal test set is a job with the technique, on top of proper parallelisms & training techniques. For example:
- For Async TP, the basic test set is a FSDP + TP + torch.compile job on dp*tp GPUs.
- For a PP schedule, the basic test is a FSDP + PP job on on dp*pp GPUs.
- Similar to parallelism testing, maximal composability should be tested on multi-dimensional parallelism settings, with multiple optimization techniques enabled.

Example Copy Chapter

This is a series of loss-converging tests on Llama 3.1, covering both parallelisms and training optimizations.
Results are obtained on 2025/01/21, with the latest torch, torchao, and torchtitan.

Setup Copy Section

| Parallelism | Techniques | Remarks |
| ------------------------ | ------------------------------------------------- | --------------------------------- |
| FSDP 8 | default | 1D control set |
| FSDP 8, TP 2, PP 2 | torch.compile, Float8, async TP, Interleaved 1F1B | 3D test set |
| FSDP 8, TP 2, CP 2, PP 2 | torch.compile, Float8, async TP, Interleaved 1F1B | 4D test set |
| FSDP 8, CP 8 | default | to verify CP with a larger degree |

Test results Copy Section

[^1]: Model initialization in a sharded setting can hardly match that in a single-device setting (or a differently sharded setting), because each time a random operator is called, the underlying RNG state offset is advanced by a quantized amount, often not aligned with the amount of randomness needed, thus “wasting” different amount of randomness on differently sharded settings.

[^2]: With a seed checkpoint, one can guarantee that the model initialization matches across different distributed jobs. But still, it’s hard to obtain identical results due to (1) other random ops in the model due to similar reasons as above, and (2) subtle differences in accumulation orders in computation.

---

Datasets Copy Section

Custom Datasets in torchtitan

torchtitan is designed to work seamlessly with most HuggingFace datasets. It supports three training flavours — pre-training (plain text), instruction-tuning / SFT (chat), and multimodal (vision) — each with its own dataloader. Both text flavours support single-source and multi-source interleaved configurations.

Dataset file locations Copy Chapter

torchtitan/hf_datasets/text_datasets.py        # pre-training and SFT
torchtitan/hf_datasets/multimodal/mm_datasets.py  # vision

---

Pre-training datasets Copy Chapter

Adding a custom text dataset Copy Section

You need three components: a loader function, a sample processor, and a registry entry.

#### 1. Define a dataset loader

def load_wikipedia_dataset(dataset_path: str, kwargs):
    """Load Wikipedia dataset with specific configuration."""
    return load_dataset(
        dataset_path,
        name="20220301.en",
        split="train",
        streaming=True,
        trust_remote_code=True,
    )

#### 2. Define a sample processor

def process_wikipedia_text(sample: dict[str, Any]) -> str:
    """Process Wikipedia dataset sample text."""
    return f"{sample['title']}\n\n{sample['text']}"

#### 3. Register your dataset

DATASETS = {
    # ... existing datasets ...
    "wikipedia": DatasetConfig(
        path="wikipedia",
        loader=load_wikipedia_dataset,
        sample_processor=process_wikipedia_text,
    ),
}

#### 4. Configure training

dataloader=HuggingFaceTextDataLoader.Config(
    dataset="wikipedia",
    infinite=True,
),

---

Instruction-tuning / SFT datasets (chat) Copy Chapter

The ChatDataLoader handles single-turn [user, assistant] message pairs. It tokenizes samples using the model's chat template, masks prompt tokens in labels so loss is computed on the assistant response only, and packs multiple short samples into each sequence.

Configuring a chat dataloader Copy Section

from torchtitan.hf_datasets.text_datasets import ChatDataLoader
def process_gsm8k(sample: dict) -> list[dict]:
    return [
        {"role": "user",      "content": sample["question"]},
        {"role": "assistant", "content": sample["answer"]},
    ]
dataloader=ChatDataLoader.Config(
    dataset_path="openai/gsm8k",
    load_dataset_kwargs={"name": "main", "split": "train"},
    sample_processor=process_gsm8k,
    infinite=True,
),

---

Multi-source interleaved dataloaders Copy Chapter

Both text flavours support interleaving multiple sources with configurable sampling weights. At each step a source is drawn proportionally to its weight. When a source is drawn, it returns a packed sample, potentially consisting multiple data points from the source. Iteration stops depending on stopping strategy (on_first_exhausted / all_exhausted), defining an epoch boundary — re-looping and shuffling are handled per source exactly as in the single-source case.

All sources must share the same infinite setting.

Interleaved pre-training Copy Section

from torchtitan.hf_datasets.text_datasets import (
    HFDataSource,
    InterleavedHuggingFaceTextDataLoader,
)
dataloader=InterleavedHuggingFaceTextDataLoader.Config(
    sources=[
        HFDataSource(dataset="c4",         weight=7.0, infinite=True),
        HFDataSource(dataset="wikipedia",  weight=2.0, infinite=True),
        HFDataSource(dataset="my_dataset", weight=1.0, infinite=True),
    ],
    seed=42,
),

Interleaved SFT Copy Section

from torchtitan.hf_datasets.text_datasets import (
    ChatDataSource,
    InterleavedChatDataLoader,
)
def process_gsm8k(sample):
    return [
        {"role": "user",      "content": sample["question"]},
        {"role": "assistant", "content": sample["answer"]},
    ]
def process_alpaca(sample):
    return [
        {"role": "user",      "content": sample["instruction"]},
        {"role": "assistant", "content": sample["output"]},
    ]
dataloader=InterleavedChatDataLoader.Config(
    sources=[
        ChatDataSource(
            dataset_path="openai/gsm8k",
            load_dataset_kwargs={"name": "main", "split": "train"},
            sample_processor=process_gsm8k,
            weight=3.0,
            infinite=True,
        ),
        ChatDataSource(
            dataset_path="tatsu-lab/alpaca",
            load_dataset_kwargs={"split": "train"},
            sample_processor=process_alpaca,
            weight=1.0,
            infinite=True,
        ),
    ],
    seed=42,
),

Weight semantics Copy Section

Weights are sampling probabilities, normalised internally. A weight of 3.0 alongside 1.0 means the first source is drawn three times as often on average — it does not mean the source is iterated three times per epoch. The epoch boundary is defined by whichever source exhausts first.

This makes weights easy to reason about as a token mixture ratio: if source A has weight 3 and source B has weight 1, roughly 75 % of training tokens will come from A and 25 % from B, regardless of the absolute dataset sizes.

Checkpointing Copy Section

Interleaved dataloaders are fully stateful. The interleaver RNG and the state of every source are saved together, so resuming from a checkpoint produces byte-identical continuations.

---

Summary Copy Chapter

| Use case | Dataloader |
|---|---|
| Single pre-training source | HuggingFaceTextDataLoader |
| Multiple pre-training sources | InterleavedHuggingFaceTextDataLoader |
| Single SFT source | ChatDataLoader |
| Multiple SFT sources | InterleavedChatDataLoader |
| Multimodal (vision + text) | MMDataLoader |

---

Debugging Copy Section

Enable Memory Profiling Copy Chapter

Launch training job with the following command (or alternatively set configs in your config_registry function)

MODULE=llama3 CONFIG=llama3_debugmodel ./run_train.sh --profiler.enable_memory_snapshot --profiler.save_memory_snapshot_folder memory_snapshot

You can find the saved pickle files in your output folder.
To visualize a snapshot file, you can drag and drop it to <https://pytorch.org/memory_viz>. To learn more details on memory profiling, please visit this tutorial.

Overriding Boolean Flags from Config via CLI Copy Chapter

Boolean flags are treated as actions. To disable a flag from the command line, use the --no prefix.

For example, given the following in your config_registry function:

def my_config() -> Trainer.Config:
    return Trainer.Config(
        profiler=Profiler.Config(enable_memory_snapshot=True),
        # ...
    )

--profiler.no_enable_memory_snapshot
--profiler.no-enable-memory-snapshot  # Equivalent

Note: --enable_memory_snapshot=False will not work. Use --no_enable_memory_snapshot instead.

Debugging Config Values Copy Chapter

To inspect how configuration values are interpreted—including those from config_registry functions and CLI overrides—run the config manager directly:

python -m torchtitan.config.manager --module llama3 --config llama3_8b [your cli args...]

For example,

python -m torchtitan.config.manager --module llama3 --config llama3_8b --profiler.enable_memory_snapshot

To list all available CLI flags and usage:

python -m torchtitan.config.manager --module llama3 --config llama3_debugmodel --help

This will print a structured configuration to stdout, allowing you to verify that overrides are being applied correctly.

Communication Mode (COMM_MODE) for Debugging Copy Chapter

The COMM_MODE environment variable provides specialized debugging modes that allow you to test and validate your training setup without requiring full multi-GPU distributed execution. This is particularly useful for rapid iteration during development and debugging.

Available Modes Copy Section

#### 1. fake_backend - Configuration Validation Mode

This mode enables dry-run validation of your configuration, model setup, and rank-0 program logic without actual distributed communication:

NGPU=32 COMM_MODE="fake_backend" ./run_train.sh

What it does:
- Uses fake process groups that simulate distributed communication without actual data transfer
- Runs on a single GPU without torchrun or NCCL initialization
- Validates configuration parsing, model initialization, and overall training workflow
- Executes only one training step by default

When to use it:
- Quick validation of configuration files before launching expensive multi-GPU jobs
- Debugging training and parallelism logic that doesn't require actual communication. Note that No data-dependent logic should be validated with "fake_backend".

Example use case:

Validate a 128-GPU configuration on a single GPU

NGPU=128 COMM_MODE="fake_backend" MODULE=llama3 CONFIG=llama3_70b ./run_train.sh

#### 2. local_tensor - Single-GPU Distributed Simulation

This mode simulates the full distributed training workflow on a single GPU by executing all communication and computation locally:

NGPU=32 COMM_MODE="local_tensor" ./run_train.sh

What it does:
- Simulates multi-GPU behavior on a single shared GPU
- Executes all collectives (all-reduce, all-gather, etc.) locally without network communication
- Maintains the same code paths as distributed training for accurate debugging
- Runs only one training step by default

When to use it:
- Debugging distributed training logic (FSDP, TP, PP, CP, EP) with data dependencies without multi-GPU setup. Note that local tensor doesn't support FSDP2 but should support SimpleFSDP.
- Verifying correctness of parallelism strategies locally
- Testing gradient synchronization and communication patterns
- Reproducing distributed training bugs in a simplified environment

Example use case:

Debug 8-way TP + 2-way FSDP on a single GPU

NGPU=16 COMM_MODE="local_tensor" ./run_train.sh \
  --parallelism.tensor_parallel_degree 8 \
  --parallelism.data_parallel_shard_degree 2

Limitations Copy Section

- Performance testing: Neither mode provides accurate performance metrics; use actual distributed runs for benchmarking
- Memory requirement: Local tensor runs require more memory on a single GPU than the actual distributed runs

Troubleshooting jobs that timeout Copy Chapter

If you encounter jobs that timeout, you'll need to debug them to identify the root cause. To help with this process, we've enabled Flight Recorder, a tool that continuously collects diagnostic information about your jobs.
When a job times out, Flight Recorder automatically generates dump files on every rank containing valuable debugging data. You can find these dump files in the dump_folder directory.
To learn how to analyze and diagnose issues using these logs, follow our step-by-step tutorial link.

Reproducibility between Runs Copy Chapter

When debugging issues with multi-dimensional parallelism (combinations of FSDP, TP, PP, CP, EP), ensuring reproducible behavior is crucial for isolating and fixing problems. torchtitan provides several mechanisms to achieve deterministic training runs. For more information on ensuring reproducibility and managing randomness in PyTorch, you can refer to the official PyTorch documentation on randomness: PyTorch Randomness Documentation.

Seed Configuration Copy Section

./run_train.sh --debug.seed 42

Seed behavior with parallelism:
- Data Parallel (DP/FSDP), Tensor Parallel (TP), Context Parallel (CP): All ranks use the same seed.
- Note: For FSDP and TP, DTensor will do special RNG management to make sure a Replicate tensor get the same init across ranks, but a Shard tensor get "random"-like init across ranks.
- Pipeline Parallel (PP): Each PP stage gets a different seed to ensure different initialization across layers on different PP ranks.

Deterministic Mode Copy Section

Enable deterministic algorithms to ensure bit-for-bit reproducibility across runs:

./run_train.sh --debug.deterministic

What it does:
- Forces all CUDA operations to use deterministic algorithms
- Disables CuDNN benchmarking and enables deterministic mode
- Sets deterministic workspace configuration for CuBLAS operations
- Note: This will significantly reduce training performance but ensures exact reproducibility

Use --debug.deterministic_warn_only to only warn about (not stop running) kernel without deterministic implementation.

Activation Checkpointing Debugging ### Copy Section

The following debug configs are available for AC.

preserve_rng_state - if deterministic output compared to non-checkpointed passes is required, set to true. Results in stashing and restoring the RNG state during each checkpoint, may be slower.

determinism_check - A string specifying the determinism function

debug - capture ac debug information. Will be slower.

See https://docs.pytorch.org/docs/stable/checkpoint.html for details.

Seed-Checkpoint-based Reproducibility Copy Section

For multiple experimental runs with different parallelism configs, we need to use a "seed" checkpoint to ensure model initializations are the same across runs. This is because in torchtitan/train.py, the model parameters are sharded first, and then have their weights initialized on each rank separately. As a result, it is not equivalent to initialize the model on one rank and then shard it. Using a seed checkpoint helps different runs load the same model weights from checkpoint -- DCP resharding will make sure the loaded weights are sharded correctly according to the parallelism configs.

#### Creating a Seed Checkpoint

NGPU=1 MODULE=llama3 CONFIG=llama3_debugmodel ./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

#### Loading Seed Checkpoints for Debugging

When using seed checkpoints for debugging or validation purposes, you can enable the load_only configuration to load checkpoints without saving any new ones during training. This is particularly useful when you only want to verify model correctness or compare different configurations without cluttering your disk:

MODULE=llama3 CONFIG=llama3_debugmodel ./run_train.sh --checkpoint.enable --checkpoint.load_only

The --checkpoint.load_only flag prevents the training process from saving any checkpoints, allowing you to:
- Run debugging sessions without generating unwanted checkpoint files
- Compare model behaviors using the same initial weights without checkpoint overhead

Note: Using a seed checkpoint will only make sure a model has same initial weights when configs change, but the training process may not be the same even after setting the seed and the deterministic mode, e.g. due to tensor shape change, data precision change, usage of randomness in model code, etc.

Example: Reproducing loss curves with different parallelism configs Copy Section

A common scenario is when you introduce a new parallelism strategy to the model, you need to ensure that the loss curve remains numerically equivalent to the previous parallelism config, thereby confirming the accuracy of your implementation. To achieve consistent behavior across multiple runs with varying parallelism configurations, it's crucial to make sure dataloader behaves consistently. We need to fix the DP degree (dp_replicate * dpshard) to ensure the dataloader operates consistently.

Here's a typical comparison setup (maintaining an overall DP degree of 4):
- Run 1: dp_shard = 4
- Run 2: dp_replicate = 2, dp_shard = 2, TP degree = 2
- Run 3: dp_replicate = 2, dp_shard = 2, CP degree = 2, PP degree = 2

To reproduce loss curves across above runs, you'll need to create a seed checkpoint, and then load the same seed checkpoint for all runs to ensure consistent model initialization on each rank. You might also need to set the deterministic mode to ensure consistent training behavior.

We also provided an example of verifying the numerical consistency across parallelism plans configs on Llama 3 in https://github.com/pytorch/torchtitan/blob/main/docs/converging.md.

---

Evaluation Copy Section

Validation and Evaluation

torchtitan provides direct and indirect support for validation to support user's training goals. Direct support is provided by the Validator class which interacts directly with the training loop, and indirect support is provided through HuggingFace checkpoint conversion for users who want to do evaluation using external tools such as ELeutherAI's lm_eval.

Validation Copy Chapter

Below is an example validation config:

validator=Validator.Config(
    freq=500,
    dataset="c4_validation",
    steps=-1,  # consumes the entire validation set
),

Third-Party Evaluation Copy Chapter

Example usage of lm_eval with vllm: lm_eval with vllm:')" title="Copy section prompt for LLMs"> Copy Section

Note that pip installing lm-eval may result in breaking torchtitan dev environment so we recommend creating a separate env.

pip install "lm-eval[vllm]"
lm_eval --model vllm \
    --model_args pretrained=./outputs/checkpoint/step-1000,tensor_parallel_size=8,dtype=auto,gpu_memory_utilization=0.8, \
    --tasks mmlu \
    --batch_size auto

---

Extension Copy Section

To support rapid experimentation with torchtitan, we provide several extension points. The principle for adding these extension points is to support various use cases with flexible component swapping and reuse, while trying to keep the code clean and minimal.

The extension points and protocols mentioned in this note are subject to change.

ModelSpec ModelSpec')" title="Copy section prompt for LLMs"> Copy Section

ModelSpec supports configuring high-level components in model training, including
- definitions of model config and model class
- model parallelization functions
- loss functions

The coarse level abstraction tries to hit a balance between flexible component swapping and a straightforward train script (train.py).

To register a model, define a model_registry(flavor) function in your model's __init__.py that returns a ModelSpec. Then define training configs in a config_registry.py module. See torchtitan/models/llama3 for an example.

Train script Copy Section

To perform various tasks, from adding a new model (possibly with a new modality), to trying out a new training paradigm (e.g. async training), a single train script cannot handle all the cases, unless customization points are inserted everywhere to make it less readable. Instead of always starting and maintaining a standalone train script, we group code in train.py into functions to allow for reuse.

This is an ongoing effort, and the level of grouping is subject to change.

Extending Trainer.Config Trainer.Config')" title="Copy section prompt for LLMs"> Copy Section

To add custom configuration for an experiment, subclass Trainer.Config (or Trainer itself) and add new fields. Define config_registry functions that return your custom Config type.

#### Example

To add a custom config section for an experiment:

torchtitan/experiments/your_folder/trainer.py

from dataclasses import dataclass, field
from torchtitan.trainer import Trainer
@dataclass
class CustomConfig:
    how_is_your_day: str = "good"
    """Just an example."""
class MyTrainer(Trainer):
    @dataclass(kw_only=True, slots=True)
    class Config(Trainer.Config):
        custom_config: CustomConfig = field(default_factory=CustomConfig)

Then in your config_registry.py:

torchtitan/experiments/your_folder/config_registry.py

from .trainer import MyTrainer, CustomConfig
def my_experiment_debugmodel() -> MyTrainer.Config:
    return MyTrainer.Config(
        custom_config=CustomConfig(how_is_your_day="great"),
        training=TrainingConfig(steps=100),
        # ... other fields
    )

Then run with:

MODULE=your_folder CONFIG=my_experiment_debugmodel ./run_train.sh

---

Fsdp Copy Section

FSDP1 -> FSDP2

Why FSDP2? Copy Chapter

With these limitations in mind, we designed and implemented an FSDP rewrite removing the FlatParameter. We refer to this rewrite as FSDP2 and the original as FSDP1. FSDP2 targets the same use cases as FSDP1 plus more, and FSDP2 still strives for good performance in eager mode, using several of the same techniques.

Compared to FSDP1:
- FSDP2 represents sharded parameters as DTensors sharded on dim-0, allowing for easy manipulation of individual parameters, communication-free sharded state dicts, and a simpler meta-device initialization flow.
- FSDP2 implements an improved memory management system that achieves lower and deterministic GPU memory by avoiding recordStream and does so without any CPU synchronization.

In the future, FSDP2 will offer an extension point to customize the all-gather (e.g. for fp8 all-gather for fp8 linears) and improved torch.compile support.

We have validated FSDP2 numerics and performance using torchtitan (e.g. see this PR). For example, on some Llama-7B runs on 8x H100s, FSDP2 achieves higher MFU with 7% lower peak memory than FSDP1, matching the same loss curve.

For more details on motivation, API, and system design, refer to here. In this README, we try to provide more user-facing info and less system design details.

FSDP1 <> FSDP2 API Differences Copy Chapter

@contract(state_cls=FSDPState)
def fully_shard(
  module: nn.Module,
  *,
  mesh: Optional[DeviceMesh] = None,
  reshard_after_forward: Union[bool, int] = True,
  mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
  offload_policy: OffloadPolicy = OffloadPolicy(),
) -> nn.Module:  # returns module for contract checks

| FSDP1 | FSDP2 |
| ----- | ----- |
| module | module |
| process_group/device_mesh | mesh |
| sharding_strategy | reshard_after_forward |
| cpu_offload | offload_policy |
| auto_wrap_policy | removed |
| backward_prefetch | removed |
| mixed_precision | mp_policy |
| param_init_fn | removed |
| device_id | removed |
| sync_module_states | removed |
| forward_prefetch | not yet implemented |
| limit_all_gathers | removed |
| use_orig_params | removed |
| no_sync | set_requires_gradient_sync |
| ignored_modules, ignored_states | ignored_params |

- fully_shard(module) is similar to FullyShardedDataParallel(module), constructing one communication bucket from module.parameters() except those already assigned to a nested fully_shard/FullyShardedDataParallel call.
- fully_shard(module) adds an FSDPState object on module, accessible via fully_shard.state(module), instead of being an nn.Module wrapper. This is done via the @contract decorator.
- Calling model.named_parameters() for a model with FSDP2 applied returns unchanged parameter names and DTensor sharded parameters. This means that the optimizer and gradient norm clipping see DTensors.
- fully_shard(module) performs a dynamic class swap on module. E.g., if type(module) is Transformer, then FSDP2 constructs a new class FSDPTransformer that inherits from a class FSDPModule and Transformer and sets module.__class__ to be FSDPTransformer. This allows us to add new methods and override methods via FSDPModule without constructing an nn.Module wrapper.
- FSDP1's sharding_strategy and process_group/device_mesh maps to FSDP2's mesh and reshard_after_forward.
- mesh should be 1D for FSDP and 2D for HSDP. For HSDP, we assume replication on the 0th mesh dim and sharding on the 1st mesh dim. If mesh is None, then FSDP2 initializes a 1D global mesh over the default process group.
- reshard_after_forward=True or False determines whether parameters are resharded (freed) after forward. If True, then they are re-all-gathered in backward. This trades off saving memory at the cost of extra communication.
- (Experimental) reshard_after_forward: int means that parameters are resharded to a smaller world size after forward (e.g. reshard_after_forward=8 can mean intra-node) so that the backward all-gather is over a smaller world size.
- | FSDP1 | FSDP2 | DeepSpeed |
| --- | --- | --- |
| 1 process_group + FULL_SHARD | 1D mesh + reshard_after_forward=True | ZeRO-3 |
| 1 process_group + SHARD_GRAD_OP | 1D mesh + reshard_after_forward=False | ZeRO-2 |
| 2 process_groups/2D device_mesh + HYBRID_SHARD | 2D mesh + reshard_after_forward=True | MiCS |
| 2 process_groups/2D device_mesh + _HYBRID_SHARD_ZERO2 | 2D mesh + reshard_after_forward=False | - |
| - | 1D/2D mesh + reshard_after_forward=8 (int) | ZeRO++ hpZ |
- FSDP2 maps mixed_precision to mp_policy and cpu_offload to offload_policy.
- For mp_policy, we remove buffer_dtype, simplify cast_forward_inputs and cast_root_forward_inputs into just cast_forward_inputs, and add an output_dtype.
- For offload_policy, we add a pin_memory option to avoid pinning CPU memory. (This feature may not have landed yet.)
- FSDP2 removes auto_wrap_policy, backward_prefetch, param_init_fn, device_id, sync_module_states, limit_all_gathers, and use_orig_params.
- auto_wrap_policy provides a syntactic sugar for calling FullyShardedDataParallel on modules based on a predicate given by the policy and assigning the wrapped module to its parent. FSDP2 is no longer an nn.Module wrapper, so there is need to assign the module back to its parent. We prefer for this functionality to exist above fully_shard, and we may provide a utility like auto_wrap_policy in the future.
- FSDP2 always follows backward_prefetch=BACKWARD_PRE without option since that is the only way to overlap collectives in backward correctly. BACKWARD_POST can prefetch incorrectly in nested-module cases.
- FSDP2 supports a new meta-device initialization flow that does not require materializing a module on GPU before sharding it, removing the need for param_init_fn. See Meta-Device Initialization for more details.
- FSDP2 always moves managed parameters/buffers to the mesh's corresponding device, removing the need for device_id. For example, if mesh.device_type is "cuda", then FSDP2 uses the current CUDA device.
- FSDP2 uses a new memory management system that preserves communication/computation overlap while achieving deterministic and lower memory usage than FSDP1. This system does not require any CPU synchronization, so there is no need for limit_all_gathers.
- FSDP2 always "uses the original parameters" since there is no more FlatParameter, removing the need for use_orig_params.
- How to implement forward_prefetch in FSDP2 is under discussion.

| FSDP1 | FSDP2 |
| ----- | ----- |
| model.state_dict(): full state dict | model.state_dict(): sharded state dict (no communication) |
| optim.state_dict(): local state dict | optim.state_dict(): sharded state dict (no communication) |
| summon_full_params() | use DTensor APIs like full_tensor() |
| FSDP.clip_grad_norm_() | nn.utils.clip_grad_norm_() |
| ShardedGradScaler | amp.grad_scaler.GradScaler |

Meta-Device Initialization Copy Chapter

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
with torch.device("meta"):
    model = Transformer()
policy = ModuleWrapPolicy({TransformerBlock})
Call reset_parameters() on every module

model = FSDP(model, auto_wrap_policy=policy)
Call param_init_fn on every module

def param_init_fn(module: nn.Module) -> None: ...
model = FSDP(model, auto_wrap_policy=policy, param_init_fn=param_init_fn)

with torch.device("meta"):
    model = Transformer()
for module in model.modules():
    if isinstance(module, TransformerBlock):
        fully_shard(module)
fully_shard(model)
for tensor in itertools.chain(model.parameters(), model.buffers()):
    assert tensor.device == torch.device("meta")
Allocate buffers and sharded parameters on GPU

model.to_empty(device="cuda")
Run user-defined initializers

model.init_weights() # or model.apply(init_weights)

---

Metrics Copy Section

We support automatically collecting metrics such as
1. High level system metrics such as MFU, average loss, max loss and words per second along with some
2. Memory metrics to measure max VRAM consumption and the number of OOMs
3. Timing metrics to measure data loading bottlenecks

Those metrics can then be visualized in either a TensorBoard or WandDB dashboard

TensorBoard Copy Chapter

To visualize TensorBoard metrics of models trained on a remote server via a local web browser:

1. Make sure metrics.enable_tensorboard option is set to true in model training (either from a config_registry function or from CLI).

2. Set up SSH tunneling, by running the following from local CLI

ssh -L 6006:127.0.0.1:6006 [username]@[hostname]

3. Inside the SSH tunnel that logged into the remote server, go to the torchtitan repo, and start the TensorBoard backend

tensorboard --logdir=./outputs/tb

4. In the local web browser, go to the URL it provides OR to http://localhost:6006/.

Weights and Biases Copy Chapter

Weights and Biases will automatically send metrics to a remote server if you login with wandb login

So all you need to do is make sure that metrics.enable_wandb is enabled

For an example you can inspect the Llama 3 config_registry.py

Note that if both W&B and Tensorboard are enabled then we will prioritize W&B.

---

Release Copy Section

Stable Releases Copy Chapter

The general instruction on managing releases can be found here.

Nightly Builds Copy Chapter

pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130

---

CONTRIBUTING Copy Section

Contributing to torchtitan

Setup Copy Section

pip install -r requirements.txt -r requirements-dev.txt

Pull Requests Copy Section

1. Fork the repo and create your branch from main.
2. If you've added code that should be tested, add tests.
3. If you've changed APIs, update the documentation.
4. Ensure the test suite passes.
5. Make sure your code lints (pre-commit run --all-files).
6. If you haven't already, complete the Contributor License Agreement ("CLA").

Contributor License Agreement ("CLA") Copy Section

Complete your CLA here: <https://code.facebook.com/cla>

Issues Copy Section

Meta has a bounty program for the safe
disclosure of security bugs. In those cases, please go through the process
outlined on that page and do not file a public issue.

License Copy Section

---

Contributing Guidelines Copy Chapter

Note: To accelerate contributions to and innovations around torchtitan, we are adding a new, experimental folder torchtitan/experiments, which has its own Contributing Guidelines. The content below is for the core portions of torchtitan.

Principles of contribution Copy Section

- Apply PyTorch-native training techniques.
- The technique should be of general interest for distributed training.
- A technique with moderate to large complexity should be sitting in the proper repo (e.g. pytorch/pytorch for a new parallelism, or pytorch/data for a new data loader) instead of torchtitan.
- The main branch of torchtitan should have minimal dependencies on non-PyTorch libraries. Interesting models/techniques that depend on external libraries can be demonstrated in the experiments folder, or in forks of torchtitan.
- Aim for minimal (if not zero) code change to the model. For the Llama model in torchtitan, if one has to make justifiable model change(s):
- After the model change, it should still load the original checkpoint correctly.
- Document the reasons for the code change, similar to composability.md.
- Keep code modularized, especially for train.py, so that it remains easy to copy-paste into a minimal code example. If necessary:
- Introduce new config options/category in configs.py.
- Create separate functions/files.

Proof of Value Copy Section

It is the contributor’s responsibility to justify the change. The requirements include, but are not limited to

#### Loss

- If a change does not impact computation results, one should see identical loss before vs. after, with fixed random seeds (training.seed) and deterministic algorithms (training.deterministic). An example is activation checkpointing.
- If a change is expected to impact computation results, loss converging should be verified via end-to-end training on representable datasets (e.g. Llama 3 models on the C4 dataset). Please refer to the recommended practices in converging.md.

#### Performance
- Memory and TPS / MFU, which are available from logging, should meet expectations.
- It is worth noting that performance expectations vary from case to case. For example, there are cases when a technique targeting memory reduction may cause throughput regression but still be acceptable (e.g. activation checkpointing). Again, it is the contributor's job to justify the feature, whether by achieving hypothetical performance, or by comparing with existing well-known implementations, etc.
- If necessary, verify the numbers on jobs spanning multiple nodes (e.g. on 64 GPUs). Please reach out to the torchtitan team for help if you are resource-constrained.
- When appropriate, one should show profile traces and/or memory snapshots to prove the effectiveness.

Best practices Copy Section

When appropriate, one should consider

- Adding CPU/GPU unit/integration tests.
- To add a unit test, put it in the tests folder and follow the existing test files.
- To add a GPU integration test, create a new OverrideDefinitions in integration_tests. It will override the default config to run on the Llama 3 debug model (see config_registry.py).
- Updating README and writing a new note in the docs folder on installation and usage, similar to float8.md.
- Following the tensor shape-suffix naming convention for new model code (e.g. x_BLD, q_BLNH, out_TNH), with a per-module legend comment as in attention.py. Capital suffixes name logical tensor dimensions (not sharding layout) and are scoped per file.
- Adding a new file with benchmark results in benchmarks folder.
- Creating GitHub issues for things that cannot be addressed at the moment.
- Writing a post on PyTorch Forums and linking to it.

---

README Copy Section

<div align="center">

torchtitan

#### A PyTorch native platform for training generative AI models

[](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml?query=branch%3Amain)
[](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml?query=branch%3Amain)
[](https://arxiv.org/abs/2410.06511)
[](https://iclr.cc/virtual/2025/poster/29620)
[](https://discuss.pytorch.org/c/distributed/torchtitan/44)
[](./LICENSE)
[](https://pypi.org/project/torchtitan/)
[](https://anaconda.org/conda-forge/torchtitan)

</div>

torchtitan is under extensive development. To use the latest features of torchtitan, we recommend using the most recent PyTorch nightly.

Latest News Copy Chapter

Overview Copy Chapter

torchtitan is a PyTorch native platform designed for rapid experimentation and large-scale training of generative AI models. As a minimal clean-room implementation of PyTorch native scaling techniques, torchtitan provides a flexible foundation for developers to build upon. With torchtitan extension points, one can easily create custom extensions tailored to specific needs.

Our mission is to accelerate innovation in the field of generative AI by empowering researchers and developers to explore new modeling architectures and infrastructure techniques.

The Guiding Principles when building torchtitan
* Designed to be easy to understand, use and extend for different training purposes.
* Minimal changes to the model code when applying multi-dimensional parallelism.
* Bias towards a clean, minimal codebase while providing basic reusable / swappable components.

torchtitan has been showcasing PyTorch's latest distributed training features, via support for pretraining Llama 3.1 LLMs of various sizes.

Contributing Copy Chapter

We look forward to your contributions!

* To accelerate contributions to and innovations around torchtitan, we host an experiments folder. New ideas should start there. To contribute, follow the experiments guidelines.
* For fixes and contributions to core, follow these guidelines.

Llama 3.1 training Copy Chapter

Key features available Copy Section

1. Multi-dimensional composable parallelisms
- FSDP2 with per-parameter sharding
- Tensor Parallel (including async TP)
- Pipeline Parallel
- Context Parallel
2. Meta device initialization
3. Per-op selective and full activation checkpointing
4. Distributed checkpointing (including async checkpointing)
- Interoperable checkpoints which can be loaded directly into torchtune for fine-tuning
5. torch.compile support
6. Float8 support (how-to)
7. MXFP8 training for dense and MoE models on Blackwell GPUs.
8. Supervised Fine-Tuning (SFT) with chat-formatted datasets
9. DDP and HSDP
10. TorchFT integration
11. Checkpointable data-loading, with the C4 dataset pre-configured (144M entries) and support for custom datasets
12. Gradient accumulation, enabled by giving an additional --training.global_batch_size argument on the CLI
13. Flexible learning rate scheduler (warmup-stable-decay)
14. BF16 optimizer states for reduced memory usage
15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via Tensorboard or Weights & Biases
16. Debugging tools including CPU/GPU profiling, memory profiling, Flight Recorder, etc.
17. All options easily configured via Python config registry with --module and --config CLI flags
18. Structured logging: per-rank trace of key training phases; (see torchtitan/observability/structured_logger/README.md)
19. Helper scripts to
- download tokenizers from Hugging Face
- convert original Llama 3 checkpoints into the expected DCP format
- estimate FSDP/HSDP memory usage without materializing the model
- run distributed inference with Tensor Parallel

We report performance on up to 512 GPUs, and verify loss converging correctness of various techniques.

Dive into the code Copy Section

You may want to see how the model is defined or how parallelism techniques are applied. For a guided tour, see these files first:
* torchtitan/train.py - the main training loop and high-level setup code
* torchtitan/models/llama3/model.py - the Llama 3.1 model definition
* torchtitan/models/llama3/parallelize.py - helpers for applying Data Parallel, Tensor Parallel, activation checkpointing, and torch.compile to the model
* torchtitan/distributed/pipeline_parallel.py - helpers for applying Pipeline Parallel to the model
* torchtitan/components/checkpoint.py - utils for saving/loading distributed checkpoints
* torchtitan/components/quantization/float8.py - utils for applying Float8 techniques

Installation Copy Chapter

One can directly run the source code, or install torchtitan from a nightly build, or a stable release.

From source Copy Section

This method requires the nightly build of PyTorch, or the latest PyTorch built from source.

git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt
pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu

Note: The nightly build of torchdata is required when using a PyTorch nightly. Install it from the nightly index as shown above.

Note: You can run directly from the source tree. If you need to import torchtitan as a package from elsewhere, install it in editable mode without re-resolving dependencies: pip install -e . --no-deps.

Nightly builds Copy Section

This method requires the nightly build of PyTorch. You can replace cu130 with another version of cuda or an AMD GPU (e.g. rocm6.3).

pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 --force-reinstall
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130

Stable releases Copy Section

pip install torchtitan

conda install conda-forge::torchtitan

Downloading a tokenizer Copy Section

torchtitan currently supports training Llama 3.1 (8B, 70B, 405B) out of the box. To get started training these models, we need to download the tokenizer. Follow the instructions on the official meta-llama repository to ensure you have access to the Llama model weights.

Once you have confirmed access, you can run the following command to download the Llama 3.1 tokenizer to your local machine.

Get your HF token from https://huggingface.co/settings/tokens
Llama 3.1 tokenizer

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

Start a training run Copy Section

MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

Multi-Node Training Copy Section

To get started adjust the number of nodes and GPUs

#SBATCH --ntasks=2
#SBATCH --nodes=2

Then start a run where nnodes is your total node count, matching the sbatch node count above.

srun torchrun --nnodes 2

If your gpu count per node is not 8, adjust --nproc_per_node in the torchrun command and #SBATCH --gpus-per-task in the SBATCH command section.

Citation Copy Chapter

We provide a detailed look into the parallelisms and optimizations available in torchtitan, along with summary advice on when to use various techniques.

TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

@inproceedings{
   liang2025torchtitan,
   title={TorchTitan: One-stop PyTorch native solution for production ready {LLM} pretraining},
   author={Wanchao Liang and Tianyu Liu and Less Wright and Will Constable and Andrew Gu and Chien-Chin Huang and Iris Zhang and Wei Feng and Howard Huang and Junjie Wang and Sanket Purandare and Gokul Nadathur and Stratos Idreos},
   booktitle={The Thirteenth International Conference on Learning Representations},
   year={2025},
   url={https://openreview.net/forum?id=SFN6Wm7YBI}
}

License Copy Chapter

Source code is made available under a BSD 3 license, however you may have other legal obligations that govern your use of other content linked in this repository, such as the license or terms of service for third-party data and models.

---