### Bf16 Optimizer States ## BF16 Optimizer States 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 This technique was notably used by [DeepSeek-V3](https://arxiv.org/abs/2412.19437) 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](https://github.com/pytorch/pytorch/issues/146542). ### Usage In your config registry function: ```python from torchtitan.components.optimizer import OptimizersContainer optimizer=OptimizersContainer.Config( name="AdamW", implementation="fused_opt_states_bf16", ), ``` Or via CLI override: ```bash --optimizer.name AdamW --optimizer.implementation fused_opt_states_bf16 ``` ### Requirements - **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 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 - **`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 - 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 # 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 1. ENABLE CHECKPOINTING In your config_registry function, configure the checkpoint settings: ```python checkpoint=CheckpointManager.Config( interval=500, ), ``` Or via CLI: `--checkpoint.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. ```python 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`. ```python 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. ```python checkpoint=CheckpointManager.Config( exclude_from_loading=["data_loader", "lr_scheduler"], ), ``` When used in command line: `--checkpoint.exclude_from_loading data_loader,lr_scheduler`. 5. EXAMPLE CHECKPOINT CONFIGURATION ```python 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 Sometimes one needs to create a seed checkpoint to initialize a model from step 0. E.g. it is hard, if not impossible, for meta initialization on multiple devices to reproduce the initialization on a single device. A seed checkpoint does initialization of the model on a single CPU, and can be loaded from another job on an arbitrary number of GPUs via DCP resharding. To create a seed checkpoint, use the same model config as you use for training. e.g. ```bash NGPU=1 ./run_train.sh --module --config --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 ### HuggingFace `torchtitan` offers two ways to work with Hugging Face models: either by directly saving and loading a Hugging Face checkpoint during training, or by using an example conversion script to directly reformat the model weights on cpu. 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. ```bash python ./scripts/checkpoint_conversion/convert_from_hf.py --model_name --model_flavor python ./scripts/checkpoint_conversion/convert_to_hf.py --hf_assets_path ./assets/hf/Llama3.1-8B --model_name --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 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 ```python 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. ```bash 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 # Building a Clean, Readable Distributed LLM One of the main goals for torchtitan was to provide a version of distributed LLM that was not only high performance, but utilized native PyTorch techniques and readable code. The challenge is how to compose together so many individual library components (FSDP, TP, PP, Float8, Compile, DCP, ..., just to name a few), and avoid having to make too many changes to the model guts in the process. A lot of the work is behind the scenes, designing individual components to make fewer assumptions, use common abstractions (e.g. DTensor) and generally "get along". But we found a few tweaks to the model code invaluable as well, and wanted to share those changes and the rationale for them. ## Making the model "pipeline friendly" When applying Pipeline Parallelism, you will have to construct nn.Module objects representing the portion of the model that runs on a given pipeline stage. Whether you plan to manually edit your model code, or use techniques like tracing to extract model chunks, a few changes to the original model code can go a long way to making this process easier. ### Simplifying the top-level model forward Most likely, you can write your model in such a way that the top-level nn.Module owns a sequence of child modules that it calls during forward, delegating most of the complexity to the child module forwards. If you can reduce your top level forward to mostly a for-loop over child module calls, then you'll simplify the pipeline-partitioning task to choosing the set of submodules to keep per stage. If you have non-trivial logic in the top-level forward, you'll have to find a way to patch that logic back onto the resulting pipeline stage model, which can be annoying. Example ([PR #321](https://github.com/pytorch/torchtitan/pull/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](https://github.com/pytorch/torchtitan/pull/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 Initializing the pipeline-parallel model is challenging because we assume the model could be so large as to not fit on local GPU (or possibly, even on CPU), and we also want to use the (bitwise) same initialization as we use for 1D or 2D parallel models, to ease debugging or comparisons between runs. It's not that easy to rewrite the original model's `init_weights` function to be tolerant of initializing only some layers, and also serializing initialization operations globally for consistent RNG order. 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](../torchtitan/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 We intentionally upcast the final output tensor to fp32 inside the loss function rather in the `Transformer.forward()` so that forward and backward casts can be fused with the loss forward and backward respectively when we `torch.compile()` the loss function. This can improve both throughput and memory usage. ## Setting `TORCH_NCCL_AVOID_RECORD_STREAMS=1` for TP Users should set the environment variable `TORCH_NCCL_AVOID_RECORD_STREAMS=1` when using tensor parallelism (TP) to avoid unexpectedly high memory usage. 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`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) 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 This note clarifies the recommended practices to follow when testing the loss converging of a new feature. ### Disclaimers 1. We assume the vanilla 1D FSDP to be “correct”, and would serve as the baseline for comparisons. The correctness of FSDP can be verified by comparing with DDP on small models, which has been widely adopted and believed to be correct. 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 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. `dp*tp*pp*cp*p`) 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 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 - Base config: `llama3_8b` (from [config_registry.py](../torchtitan/models/llama3/config_registry.py)) - `training.local_batch_size = 4`, which is a minimum for Pipeline Parallel with `pipeline_parallel_degree = 2` and `pipeline_parallel_schedule = "Interleaved1F1B"` - `training.data_parallel_shard_degree = 8`, resulting in global batch size 32 - `training.steps = 3000`, `lr_scheduler.warmup_steps = 600` | 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 [^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 # 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 ``` torchtitan/hf_datasets/text_datasets.py # pre-training and SFT torchtitan/hf_datasets/multimodal/mm_datasets.py # vision ``` --- ## Pre-training datasets ### Adding a custom text dataset You need three components: a loader function, a sample processor, and a registry entry. #### 1. Define a dataset loader ```python 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 ```python 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 ```python DATASETS = { # ... existing datasets ... "wikipedia": DatasetConfig( path="wikipedia", loader=load_wikipedia_dataset, sample_processor=process_wikipedia_text, ), } ``` #### 4. Configure training ```python dataloader=HuggingFaceTextDataLoader.Config( dataset="wikipedia", infinite=True, ), ``` --- ## Instruction-tuning / SFT datasets (chat) 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 ```python 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 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 ```python 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 ```python 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 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 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 | 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 ## Enable Memory Profiling 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 ``` * `--profiler.enable_memory_snapshot`: to enable memory profiling * `--profiler.save_memory_snapshot_folder`: configures the folder which memory snapshots are dumped into (`./outputs/memory_snapshot/` by default) * `--profiler.memory_snapshot_freq`: controls how often regular memory snapshots are taken. When unset, it defaults to `--profiler.profile_freq` for backward compatibility. + In case of OOMs, the snapshots will be in `./outputs/memory_snapshot/iteration_x_exit`. + Regular snapshots will be in `memory_snapshot/iteration_x`. + For example, set `--profiler.memory_snapshot_freq 3` to take a snapshot every three iterations independently of trace profiling. You can find the saved pickle files in your output folder. To visualize a snapshot file, you can drag and drop it to . To learn more details on memory profiling, please visit this [tutorial](https://pytorch.org/blog/understanding-gpu-memory-1/). ## Overriding Boolean Flags from Config via CLI 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: ```python def my_config() -> Trainer.Config: return Trainer.Config( profiler=Profiler.Config(enable_memory_snapshot=True), # ... ) ``` You can override it at runtime via CLI with: ```bash --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 To inspect how configuration values are interpreted—including those from config_registry functions and CLI overrides—run the config manager directly: ```bash python -m torchtitan.config.manager --module llama3 --config llama3_8b [your cli args...] ``` For example, ```bash python -m torchtitan.config.manager --module llama3 --config llama3_8b --profiler.enable_memory_snapshot ``` To list all available CLI flags and usage: ```bash 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 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 #### 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: ```bash 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:** ```bash # 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: ```bash 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:** ```bash # 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 - **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 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](https://pytorch.org/tutorials/prototype/flight_recorder_tutorial.html). ## Reproducibility between Runs 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](https://docs.pytorch.org/docs/stable/notes/randomness.html). ### Seed Configuration Set consistent random seeds across all parallelism dimensions: ```bash ./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 Enable deterministic algorithms to ensure bit-for-bit reproducibility across runs: ```bash ./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 ### 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 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 ```bash 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: ```bash 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 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 # 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](https://github.com/pytorch/torchtitan/blob/main/docs/checkpoint.md#huggingface) for users who want to do evaluation using external tools such as ELeutherAI's `lm_eval`. ## Validation For users who want to perform validation directly during the training loop, we provide the `Validator` class which can be conveniently configured via `Validator.Config` in your config_registry function. The validator class has access to and reuses many of the trainer's functions such as its parallelization, including pipelining. Below is an example validation config: ```python validator=Validator.Config( freq=500, dataset="c4_validation", steps=-1, # consumes the entire validation set ), ``` ## Third-Party Evaluation With `./scripts/checkpoint_conversion/convert_to_hf.py`, `torchtitan` offers support for converting checkpoints from DCP to safetensors format. Using this script, users can perform efficient evaluation separate from their training using external libraries that support HuggingFace e.g. `lm_eval` with `vllm` backend. ### Example usage of `lm_eval` with `vllm`: To use this specific setup make sure to include a HuggingFace `config.json` file which is not provided by conversion script or `last_save_in_hf` option. The HF config file can be downloaded by running `python ./scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets config`. Note that pip installing `lm-eval` may result in breaking `torchtitan` dev environment so we recommend creating a separate env. ```bash 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 ``` | Groups |Version|Filter|n-shot|Metric| |Value | |Stderr| |------------------|------:|------|------|------|---|-----:|---|-----:| |mmlu | 2|none | |acc |↑ |0.6209|± |0.0038| | - humanities | 2|none | |acc |↑ |0.5481|± |0.0066| | - other | 2|none | |acc |↑ |0.7045|± |0.0078| | - social sciences| 2|none | |acc |↑ |0.7351|± |0.0078| | - stem | 2|none | |acc |↑ |0.5357|± |0.0085| --- ### Extension 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`](../torchtitan/protocols/model_spec.py) 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](../torchtitan/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](../torchtitan/models/llama3/) for an example. ### Train script 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](../torchtitan/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` 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: ```python # 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`: ```python # 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: ```bash MODULE=your_folder CONFIG=my_experiment_debugmodel ./run_train.sh ``` --- ### Fsdp # FSDP1 -> FSDP2 ## Why FSDP2? PyTorch's fully sharded data parallelism (FSDP) API, [`FullyShardedDataParallel`](https://pytorch.org/docs/stable/fsdp.html), looks to offer a performant eager-mode implementation, including communication bucketing and communication/computation overlap. It defines a `FlatParameter` by flattening and concatenating a group of parameters to represent a communication bucket. However, this `FlatParameter` complicates applying different behaviors to individual parameters within the `FlatParameter`, e.g. parameter freezing, parameter casting, etc., hurting composability, and it complicates the internal implementation, e.g. making state dict logic thousands of lines and requiring additional communications. 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 `DTensor`s 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](https://github.com/pytorch/torchtitan/pull/165)). 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](https://github.com/pytorch/pytorch/issues/114299). In this README, we try to provide more user-facing info and less system design details. ## FSDP1 <> FSDP2 API Differences We go over some API differences between FSDP1 and FSDP2. Overall, we hope to minimize the API surface (including the number of arguments) to avoid having a monolithic API. ```python @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 `DTensor`s. - `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_group`s/2D `device_mesh` + `HYBRID_SHARD` | 2D `mesh` + `reshard_after_forward=True` | MiCS | | 2 `process_group`s/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](https://github.com/pytorch/pytorch/issues/108190) 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](#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 Before with FSDP1: ```python 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) ``` After with FSDP2: ```python 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)` ``` FSDP1 requires either `reset_parameters` or `param_init_fn` to materialize a module onto GPU immediately before sharding. To do this correctly without re-initializing any tensors requires care and can be unwieldy. However, FSDP2 allows materializing tensors onto GPU _after_ sharding (taking advantage of `DTensor` and a new `swap_tensors` path for `nn.Module._apply` methods). --- ### Metrics 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 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 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](../torchtitan/models/llama3/config_registry.py) Note that if both W&B and Tensorboard are enabled then we will prioritize W&B. --- ### Release ## Stable Releases Currently we follow a lightweight release process. - Update the version number in `assets/version.txt` with a PR. The version numbering should follow https://semver.org/. - E.g. for a pre-release `0.y.z` - if major features are added, increment `y` - if minor fixes are added, increment `z` - Create a new release at https://github.com/pytorch/torchtitan/releases/new - In the tag section, add a new tag for the release. The tag should use the version number with a `v` prefix (for example, `v0.1.0`). Make sure to select the `main` branch as the target. - In the release notes - include proper nightly versions for `torch` and `torchao`, which can be found in [latest CI](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu.yaml) test log "Run script in container" section. E.g. - "Successfully installed ... `torch-2.8.0.dev20250605+cu130`" - "Successfully installed `torchao-0.12.0.dev20250605+cu130`" - describe the release at a high level compared to the last release, e.g. - "added an experiment for multimodal LLM training" - or simply state "this is a regular release" - For now, choose "Set as a pre-release". - As we set up the GitHub workflow [release.yml](/.github/workflows/release.yml), it should trigger a [GitHub action](https://github.com/pytorch/torchtitan/actions/workflows/release.yml) to update the [torchtitan package on PyPI](https://pypi.org/project/torchtitan/), which requires approval from one of the maintainers to run. The general instruction on managing releases can be found [here](https://docs.github.com/en/repositories/releasing-projects-on-github/managing-releases-in-a-repository). ## Nightly Builds Nightly builds are automatically triggered by a [nightly GitHub workflow](/.github/workflows/build_whl_and_publish.yaml) and can be installed by ```bash pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130 ``` You can replace `cu130` with another version of cuda or an AMD GPU (e.g. `rocm6.3`). --- ### CONTRIBUTING # Contributing to torchtitan We want to make contributing to this project as easy and transparent as possible. Contributions should follow the [Contributing Guidelines](#contributing-guidelines) below. ### Setup ``` pip install -r requirements.txt -r requirements-dev.txt ``` ### Pull Requests We actively welcome your pull requests. 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") In order to accept your pull request, we need you to submit a CLA. You only need to do this once to work on any of Meta's open source projects. Complete your CLA here: ### Issues We use GitHub issues to track public bugs. Please ensure your description is clear and has sufficient instructions to be able to reproduce the issue. Meta has a [bounty program](https://www.facebook.com/whitehat/) 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 By contributing to `torchtitan`, you agree that your contributions will be licensed under the LICENSE file in the root directory of this source tree. --- ## Contributing Guidelines Note: To accelerate contributions to and innovations around `torchtitan`, we are adding a new, experimental folder [`torchtitan/experiments`](torchtitan/experiments/), which has its own [Contributing Guidelines](torchtitan/experiments/README.md#contributing-guidelines). The content below is for the core portions of `torchtitan`. ### Principles of contribution - 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](docs/composability.md). - Keep code modularized, especially for [train.py](torchtitan/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](torchtitan/config/configs.py). - Create separate functions/files. ### Proof of Value 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](docs/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 When appropriate, one should consider - Adding CPU/GPU unit/integration tests. - To add a unit test, put it in the [tests](tests/) folder and follow the existing test files. - To add a GPU integration test, create a new `OverrideDefinitions` in [integration_tests](tests/integration_tests/). It will override the default config to run on the Llama 3 debug model (see [config_registry.py](torchtitan/models/llama3/config_registry.py)). - Updating [README](README.md) and writing a new note in the [docs](docs/) folder on installation and usage, similar to [float8.md](torchtitan/components/quantization/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](torchtitan/models/common/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](benchmarks) folder. - Creating GitHub issues for things that cannot be addressed at the moment. - Writing a post on [PyTorch Forums](https://discuss.pytorch.org/c/distributed/torchtitan/44) and linking to it. --- ### README
# 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)
`torchtitan` is under extensive development. To use the latest features of `torchtitan`, we recommend using the most recent PyTorch nightly. ## Latest News - [2026/08] [TitanRL](torchtitan/experiments/rl) is a hackable RL stack for scaling and debugging. It reuses TorchTitan model definitions and kernels across training and vLLM generation and supports batch-invariant mode. - [2025/11] AMD released an [optimized fork](https://github.com/AMD-AGI/torchtitan-amd/tree/main) of `torchtitan` for AMD GPUs. - [2025/10] We released `torchtitan` [v0.2.0](https://github.com/pytorch/torchtitan/releases). - [2025/10] SkyPilot now supports `torchtitan`! See the tutorial [here](https://docs.skypilot.co/en/latest/examples/training/torchtitan.html). - [2025/07] We published [instructions](/torchtitan/models/README.md) on how to add a model to `torchtitan`. - [2025/04] Our paper was accepted by [ICLR 2025](https://iclr.cc/virtual/2025/poster/29620). - [2024/12] GPU MODE [lecture](https://www.youtube.com/watch?v=VYWRjcUqW6w) on torchtitan. - [2024/07] [Presentation](https://pytorch2024.sched.com/event/1fHn3) at PyTorch Conference 2024. ## Overview `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](docs/extension.md), 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 We look forward to your contributions! * To accelerate contributions to and innovations around torchtitan, we host an [`experiments`](torchtitan/experiments) folder. New ideas should start there. To contribute, follow the [`experiments guidelines`](torchtitan/experiments/README.md). * For fixes and contributions to core, follow these [`guidelines`](CONTRIBUTING.md). ## Llama 3.1 training ### Key features available 1. Multi-dimensional composable parallelisms - [FSDP2](docs/fsdp.md) with per-parameter sharding - [Tensor Parallel](https://pytorch.org/docs/stable/distributed.tensor.parallel.html) (including [async TP](https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487)) - [Pipeline Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420) - [Context Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) 2. [Meta device](https://pytorch.org/docs/stable/meta.html) initialization 3. Per-op selective and full activation checkpointing 4. [Distributed checkpointing](https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250) (including async checkpointing) - [Interoperable checkpoints](docs/checkpoint.md) which can be loaded directly into [`torchtune`](https://github.com/pytorch/torchtune) for fine-tuning 5. `torch.compile` support 6. [Float8](https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323) support ([how-to](torchtitan/components/quantization/float8.md)) 7. [MXFP8 training for dense and MoE models](torchtitan/components/quantization/mxfp8.md) on Blackwell GPUs. 8. Supervised Fine-Tuning (SFT) with chat-formatted datasets 9. DDP and HSDP 10. [TorchFT](https://github.com/pytorch/torchft) integration 11. Checkpointable data-loading, with the C4 dataset pre-configured (144M entries) and support for [custom datasets](docs/datasets.md) 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](docs/bf16_optimizer_states.md) for reduced memory usage 15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via [Tensorboard or Weights & Biases](/docs/metrics.md) 16. [Debugging tools](docs/debugging.md) including CPU/GPU profiling, memory profiling, Flight Recorder, etc. 17. All options easily configured via [Python config registry](torchtitan/models/llama3/config_registry.py) with `--module` and `--config` CLI flags 18. Structured logging: per-rank trace of key training phases; (see [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md)) 19. [Helper scripts](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](benchmarks/llama3_h100_202412_torchtitan.md) on up to 512 GPUs, and verify [loss converging](docs/converging.md) correctness of various techniques. ### Dive into the code 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](torchtitan/train.py) - the main training loop and high-level setup code * [torchtitan/models/llama3/model.py](torchtitan/models/llama3/model.py) - the Llama 3.1 model definition * [torchtitan/models/llama3/parallelize.py](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](torchtitan/distributed/pipeline_parallel.py) - helpers for applying Pipeline Parallel to the model * [torchtitan/components/checkpoint.py](torchtitan/components/checkpoint.py) - utils for saving/loading distributed checkpoints * [torchtitan/components/quantization/float8.py](torchtitan/components/quantization/float8.py) - utils for applying Float8 techniques ## Installation One can directly run the source code, or install `torchtitan` from a nightly build, or a stable release. ### From source This method requires the nightly build of PyTorch, or the latest PyTorch built [from source](https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source). ```bash 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 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`). ```sh 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 One can install the latest [stable release](https://github.com/pytorch/torchtitan/releases) of `torchtitan` via `pip` or `conda`. ```sh pip install torchtitan ``` ```sh conda install conda-forge::torchtitan ``` Note that each stable release pins the nightly versions of `torch` and `torchao`. Please see [release.md](docs/release.md) for more details. ### Downloading a tokenizer `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](https://huggingface.co/meta-llama/Llama-3.1-8B) 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. ```bash # 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 Llama 3 8B model locally on 8 GPUs ```bash MODULE=llama3 CONFIG=llama3_8b ./run_train.sh ``` ### Multi-Node Training For training on ParallelCluster/Slurm type configurations, you can use the `multinode_trainer.slurm` file to submit your sbatch job. 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 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](https://openreview.net/forum?id=SFN6Wm7YBI) ``` @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 Source code is made available under a [BSD 3 license](./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. ---