TorchTitan: One-stop PyTorch native solution for production ready LLM pretraining
Authors: Wanchao Liang
Tianyu Liu
Less Wright
Will Constable
Andrew Gu
Chien-Chin Huang
Iris Zhang
Wei Feng
Howard Huang
Junjie Wang
Sanket Purandare
Gokul Nadathur
Stratos Idreos
Venue: ICLR 2025 Poster (2025)
Keywords: large language models
distributed training
pre-training
data parallel
tensor parallel
pipeline parallel
pytorch
llama
distributed checkpointing
3D parallel
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: TorchTitan is an open-source and customizable, PyTorch-native system that enables composable and modular 4D parallel pre-training for LLMs at an elastic scale, achieves significant performance gains, and offers optimized training recipes.

Abstract: The development of large language models (LLMs) has been instrumental in advancing state-of-the-art natural language processing applications. Training LLMs with billions of parameters and trillions of tokens requires sophisticated distributed systems that enable composing and comparing several state-of-the-art techniques in order to efficiently scale across thousands of accelerators. However, existing solutions are complex, scattered across multiple libraries/repositories, lack interoperability, and are cumbersome to maintain. Thus, curating and empirically comparing training recipes requires non-trivial engineering effort.

This paper introduces **TORCHTITAN**$^1$, a PyTorch-native distributed training system that unifies and advances state-of-the-art techniques, streamlining integration and reducing engineering overhead. TORCHTITAN enables seamless application of 4D parallelism in a modular and composable manner, while featuring elastic scaling to adapt to changing computational requirements. The system provides comprehensive logging, efficient checkpointing, and debugging tools, ensuring production-ready training. Moreover, TORCHTITAN incorporates innovative hardware-software co-designed solutions, leveraging cutting-edge features like Float8 training and SymmetricMemory to maximize hardware utilization.

As a flexible experimental test bed, TORCHTITAN facilitates the curation and comparison of custom recipes for diverse training contexts. By leveraging TORCHTITAN, we developed optimized training recipes for the Llama 3.1 family and provide actionable guidance on selecting and combining distributed training techniques to maximize training efficiency, based on our hands-on experiences.

We thoroughly assess TORCHTITAN on the Llama 3.1 family of LLMs, spanning 8 billion to 405 billion parameters, and showcase its exceptional performance, modular composability, and elastic scalability. By stacking training optimizations, we demonstrate accelerations ranging from 65.08% on Llama 3.1 8B at 128 GPU scale (1D), 12.59% on Llama 3.1 70B at 256 GPU scale (2D), to 30% on Llama 3.1 405B at 512 GPU scale (3D) on NVIDIA H100 GPUs over optimized baselines. We also demonstrate the effectiveness of 4D parallelism in enabling long context training.

$^1$ GitHub: [https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan)

Published as a conference paper at ICLR 2025
TORCHTITAN: ONE-STOP PYTORCH NATIVE SOLU- TION FOR PRODUCTION READY LLM PRETRAINING
Wanchao Liang', Tianyu Liu Less Wright', Will Constable', Andrew Gu' Chien-Chin Huang'. Iris Zhang', Wei Feng', Howard Huang', Junjie Wang' Sanket Purandare Gokul Nadathur', Stratos Idreos² Meta, Haavvrd University
ABSTRACT
The development of large language models (LLMs) has been instrumental in ad- vancing state-of-the-art natural language processing applications. Training L.LMs with billions of parameters and trillions of tokens require sophisticated distributed systems that enable composing and comparing several state-of-the-an the-art techniques in order to efficiently scale across thousands of accelerators. However, existing solutions are complex, scattered across multiple libraries/repositories,, lack inter- operability, and cumbersome maintain. Thus, curating and empirically com- paring training recipes require non-trivial engineering effort.
This paper introduces TORCHTITAN, an open-soured Py PyTorch-native distributed training system that unifies and advances state-of-the-ar techniques, streamlining integration and reducing engineering overhead. TORCHTITAN enables seamless application 4D parallelism
modular and composable manner, while featur- ing elastic scaling adapt to changing computational requirements. The system provides comprehensive logging, efficient checkpointing, and debugging tools, ensuring production-reade training. Moreover, TORCHTITAN incorporates inno- vative cordware-softwar co-designed solutions, leveraging cutting-edge features like Float8 training and SymmetricMemorr to maximize hardware utilization. As a flexible experimental test bed, TORCHTITAN facilitates the curation and compar- ison of custom recipes for diverse training contexts. By leveraging TORCHTITAN, we developed optimized training recipes for the Llama 3.1 family and provide ac- tionable guidance on selecting and combining distributed training techniques maximize training efficiency, based on our hands-on experiences.
We thoroughly assess TORCHTITAN on the Ilama 3.1 family of LLMMs, spanning 8 billion to 405 billion parameters, and showcase its exceptional performance, modular composability, and elastic scalability. By stacking training optimizations, we demonstrate accelerations ranging from 65.08% on Llama 3.1 8B at 128 GPU scale (ID), 12.59% on Llama 3.1 70B at 256 GPU scale (2D), to 30% on Llama 3.1 405B at 512 GPU scale (3D) on NVIDIA H100 GPUs over optimized baselines. We also demonstrate the effectiveness of 4D parallelism in enabling long context training.
INTRODUCTION
Large Language Models (LLMs) (Devlin 2018 Liu et al. 2019: Radford et al. 2019 Chowdhery
et al. 2023: Anil al. 2023 Achiam al. 2023: Dubey et al. 2024 Jiang al.] 2024 Abdin
et al. 2024 have been the driving force behind the advancement of natural language processing
(NLP) applications spanning language translation, content/code generation, conversational AI, text data analysis, creative writing and art, education, and research, etc.
Achieving state-of-the-an LLM performance requires massive scale, exemplified by tor-perfogning models like Llama 3.1 (405B parameters, 15T tokens, 30.84M GPU hours, 16K H100 GPUs) (Dubey
"Corresponding author: Tianyu Liu (1ty8meta.com
Work done at Meta
Githhu: http::/gtthbb..c///yacc//oorrrne

---

Published as a conference paper at ICLR 2025
al. 2024) and Google's PaLM (540B parameters, 0.8T tokens, 9.4M TPU hours, 6144 TPUv4 chips) (Chowdhery et al. 2023). These models demonstrate exceptiooal natural language under- standing and generation capabilities, but at the same time necessitate substantial computational resources, memory, and time to train, highlighting the significant investment required to advance natural language processing.
Training large language models (LLMs) at scale is a daunting task that requires a delicate balance of parallelism, computation, and communication, all while navigating intricate memory and com- putation trade-offs. The massive resources required for training make it prone to GPU failures,
underscoring the need for efficient recovery mechanisms and checkpointing strategies to minimize downtime (Eisenman et al.) [022]] Wang et al. [2023] Gupta et al. [2024; Maurya et al.) 2024 Wan
et al. 2024). To optimize resource utilization and achieve elastic scalability, it crucial to combine multiple parallelism techniques, including Data Parallel (L.i et al. 2020, Rajbhandari et al. 2020
Zhang et al.
2022 Zhao et al.] 2023), Tensor Parallel (Narayanan al., 2021] Wang et al. [0.2 Ko-
rthikanti et al. 2023), Context Parallel (Liu et al.) 2023] Liu & Abbeel |2024; NVIDIA 2023 Fang
& Zhao 2024), and Pipeline Parallel (Huang et al. 2019] Narayanan et al. 2019; 2021 1 QQ et al,
2023). By stacking these parallelisms with memory and computation optimization techniques, such as activation recomputation (Chen et al. 2016; Korthikanti et al. 20233 He & Yu 2023] Purandare al., 2023), mixed precision training (Micikevicius et al. 2018: 2022). and deep learning com-
pilers (Bradbury et al.] 2018) Yu et al. 2023 Li et al. 2024 Ansel et al. 2024), it is possible to maximize hardware utilization.
While state-of-the-art distributed training techniques have significantly advanced the field, exist- ing systems that incorporate them usability, adoption and effectiveness for researchers and industry practitioners.
still fall short in addressing critical challenges that hinder their
1. Non-composable: Existing systems struggle integrate and stack parallelism techniques, limiting multi-dimensional exploration and integration with memory and computation op- timizations, thereby reducing training efficiency.
2. Inflexible Architecture: Lack of modularity and extensibility hampers the integration of new techniques, optimizations, and hardware, limiting adaptability to evolving ML landscapes.
3. Inefficient Hardware Utilization: Poor leverage of advanced hardware features results in sub-optimal GPU efficiency and lack of customizable checkpointing strategies for memory- computation trade-offs.
4. Insufficient Support for Production Training: Limited distributed checkpointing scalabil- ity, cumbersome failure recovery, and inadequate debugging tools hinder production- -radee workflows.
5. Framework Limitations: Dependence on external, poorly maintained dependencies and failure to harness PyTorch's optimized kernels, new features, and compiler support lead to inefficiencies and compatibility issues.
The non-composability and inflexibility of distributed systems stem from the absence of unified tensor and device abstractions applied consistently across the stack. Without these foundational components. parallelism strategies, checkpointing, and efficiency optimizations remain fragmented, limiting modularity, scalability, and extensibility.
TORCHTITAN 's primary research contribution lies in identifying and unifying the core principles of parallelism and optimization techniques into cohesive framework. By leveraging and extending PyTorch's 's Distributed Tensor (DTensor) and DeviceMesh (Py Community, 2023a). TORCHTI- TAN provides unified abstraction that simplifies the composition of parallelism strategies, and en- sures correct single device semantics with its sharding primitives. Unlike existing systems that often rely on rigid ad-hoc designs, TORCHTITAN introduces unified template for distributed training, enabling researchers to systematically explore configurations, rigorously evaluate existing methods, and uncover novel techniques within the design space.
ToRCH TITAN represents a complete distributed training system for large language models (LLMs), rather than merely a collection of individual techniques. Its modular, extensible architecture supports seamless composition of 4D parallelism, advanced training optimizations. and scalable distributed checkpoint save/load, all while harnessing PyTorch's native capabilities. The system not only en-
2

---

Published as a conference paper at ICLR 2025
able production grade training with thousands of GPUs, but also reduces complexity and fosters innovation, setting a new standard for scalable and flexible distributed training systems.
To develop and evaluate the capabiliities of TORCHTITAN, undertook several key steps, which represent the core contributions of this work, and are summarized as follows:
1. We advance DTTensor by extending its sharding support n-D parallelism, adding compat- ibility with torch. compill le e for compiler optimizations, and enabling efficient check- pointing of n-D models via state dict support. We also resolve critical bugs to bolster DTensor's production readiness.
2. We demonstrate how to compose various parallelism techniques, facilitating the exploration of multi-dimensional parallelism in large language model training ($2.1).
3. We enable novel softwae co-designed solutions exploiting advanced hardware features increase GPU efficiency, offer customizable activation checkpointing strategies
for navigating memory-computatie trade-offs, and utilize torch. compile il to further optimize memory, computation, and communication ($2.2).
4. We offer production grade training by incorporating scalable and efficient distributed checkpoint to facilitate fast failure recovery, integrating debugging tools like Flight Recorder to debug crashed/stuck jobs, and provide extensive logging metrics ($2.3).
5. We extensively evaluate TORCHTITAN on Llama 3.1 family of models, stacking ID to 4D parallelisms (respectively), at the scale from 8 to 512 GPUs to demonstrate elastic scalability while ensuring efficiency, convergence, and accuracy. In summary, we demon- strate training accelerations ranging from 65.08% on Llama 3.1 8B at 128 GPU scale (ID), 12.59% on Llama3.1 70B at 256 GPU scale (2D), to 30% on Llama3.1 405B at 512 GPU scale (3D), and the effectiveness of 4D parallelism in enabling long context training, on latest NVIDIA H100 GPUs over optimized baselines ($3.2).
6. We provide sy systematic training recipes and guidelines that empower users to navigate complexities of distributed training, helping them optimize training efficiency for range of model sizes and cluster configurations ($3.3)
By providing accessible and extensible platform. TORCHTITAN democratizes large language model (LLM) pretraining, empowering a wider range of researchers and developers to tap into the potential of LLMSs and accelerate innovation in the field.
2 ELASTICITY THROUGH COMPOSABILITY
Initiolue
Splt Modal
For each Stage :/Model
Loader
bevice Mesk
Tokenizer
Floxts Madule Swapping
Stage
Tensor/Scquencee Parullel Plan Sharding Row/Cobar Perold
Activation chockppinting Struteyy
torph.compiee block Scdidabyg Coepota
Fully Shorcked Data Paralle!
Loss. Paralel ctx
Contaxt Poralel
Scalin déwee///yranii
Schedules arrehPPa
cbu Officd
Via
dovia.
PP?
Deable TP/ssht
Brobie ACT
Gnable anapde?
Brokiee
Figure 1: Composable and Modular TORCHTITAN initialization workflow.
TORCHTITAN incorporates various parallelisms in modular manner to enable easy, user-selectable combinations of multi-dimensional shardings. This composability enables the tackling of difficult scaling challenges by enhancing the case of exploration for optimizing training efficiencies at scale.
The codebase of TORCUTITAN is organized purposefully to enable composability and extensibility. We intentionally keep three main components separate and as orthogonal possible: (1) the model definition, which is parallelism-agnostt apply parallelisms and training optimizations to particular model, and (3) generalized training loop. All these components are configurable via TOML files with command-lin overrides, and is easy to add new models and parallelism techniques top of the existing codebase.
and designed for readability, (2) parallelism helpers, which

---

Published as a conference paper at ICLR 2025
2.1 COMPOSABLE N-D PARALLELISM TRAINING
In this section, we will walk through the entire regime of scaling model training on large clusters, including meta device initialization and the core composable multi-dimensional parallelisms, to showcase how these techniques can be composed to train LLMs efficiently at increasing scale in TORCHTITAN. The corresponding code snippets in TORCHTITAN can be found in Appendix
2.1.1 LARGE-SCALE MODEL INITIALIZATION USING META DEVICE
As LLMs grow exponentially, scaling challenges arise before training begins, particularly in instantiating large models for sharding without exceeding CPU or GPU memory limits.
To address this, TORCHTITAN enables meta device initialization, where the model is first created on a meta device that stores only metadata, making initialization ultra-fast. The model is then sharded into Distributed Tensors (DTensors), with the local shard of each parameter residing on the meta device. Finally, parameter initialization is performed using user-defined functions, ensuring correct DTensor sharding layouts and proper RNG seed usage.
2.1.2 FULLY SHARDED DATA PARALLEL
The original Fully Sharded Data Parallel (FSDP) (Zaao 2023) is an effective implementation of ZeRO that offers large model training capability in PyTorch. However, the original implementation (FSDP1) in PyTorch suffers from various limitations due to its FlatParameter implementation.
Given these limitations, TORCHTITAN integrates a new version of Fully Sharded Data Parallel (FSDP2), which uses the per-parameter Distributed Tensor sharding representation and thus pro- vides better composability with model parallelism techniques and other features that require the manipulation of individual parameters.
TORCHTITAN integrates and leverages FSDP2 as it's default 1D parallelism, benefiting from the improved memory management (often percent lower per GPU memory requirement FSDP1) and the slight performance gains (average of 1.5 percent gain vs FSDP1). More details on FSDP2 and usage example are shown in Appendix B.T TORCHTITAN makes simple to run with FSDP2 by embedding appropriate defaults, including auto-sharding with your world size automatically.
For scaling to even larger world sizes, TORCHTITAN also integrates Hybrid Sharded Data Parallel (HSDP) which extends FSDP2 by creating 2D DeviceMesh with replica groups Details are shown in Appendix B.2
2.1.3 TENSOR PARALLEI
Tensor Parallel (TP) (Narayanan et al. 2021), together with Sequence Parallel (SP) (Korthikanti et al. 2023), is a key model parallelism technique to enable large model training at scale.
TP is implemented in TORCHTITAN using the PyTorch's RowwiseParalle ColwiseFarallel APIs, where the model parameters are partitioned to DTensors and perform sharded computation with it (Figure 3.. By leveraging DTensor, the TP implementation does not need to touch the model code, which allows faster enablement different models and provides better composability with other features mentioned in this paper.
and
Tensor and Sequence Parallel (TP/SP) While TP partitions the most computationally demanding aspects, Sequence Parallel (SP) performs a sharded computation for the normalization or dropout layers on the sequence dimension, which otherwise generate large replicated activation tensors, and thus can be challenging to memory constraints per GPU. See Appendix B.3 for more details, illustrations, and usage for both TP and FSDP TP.
Due to the synergistic relationship between TP and SP, TORCHTITAN natively bundles these two together, and they are jointly controlled by the TP degree setting.
Loss Parallel When computing the loss function, model outputs typically large, especially with TP/SP, where they are sharded across the vocabulary dimension. Naively computing cross- entropy loss requires gathering all shards, leading high memory usage.
4

---

Published as a conference paper at ICLR 2025
Loss Parallel enables efficient loss computation without fully gathering model outputs, significantly reducing memory consumption and improving training speed by minimizing communication over- head and enabling parallel sharded computation. Due to these advantages, TORCHTITAN imple- ments Loss Parallel by default.
2.1.4 PIPELINE PARALLEL
For large-scale pretraining, TORCHTITAN employs Pipeline Parallelism (PP), which minimizes communication overhead by leveraging P2P communications. PP divides the model into S stages, each running separate group of devices. Typically, each stage represents model layer group adjacent layers, but can include partial layers. During the forward pass, each stage re- ceives input activations (except stage 0), computes locally, and sends output activations (except stage S - 1). The last stage computes the loss and initiates the backward pass, sending gradients in reverse order. To improve efficiency, the input batch is split into microbatches, and the pipeline schedule overlaps computation and communication across microbatches. TORCITITAN supports various pipeline schedules (Narayanan al., 2019 Huang al., 2019] Narayanan et al. 2021| Qi et al. 2023). Recently. TORCHTITAN added support for new schedules including ZuroBubble and FHlexibeeInnterlaaeed FF making use of pipeline IR quickly express new schedules as list of compute actions and rely on compiler passes to insert and optimize communication actions Py Torch
Team 2024d
The PP training loop differs from standard training by creating pipeline stages and executing sched- ules instead directly invoking model. forward (). Since loss is computed per microbatch, TORCHTITAN introduces shared loss fn to unify pipeline and non-pipeline workflows, reduc- ing code divergence.
torch. distribin butt ed. pipe aeppeenniig also simplifies interactions with data parallelism, ensur- ing that reductions occur only after the final microbatch and handling shard/unshard operations (e.g., with ZeRO-3), as well applying gradient scaling transparently within the pipeline schedule ecutor. For more details on TORCHTITAN'S implementation of PP, see Appendix B.4
2.1.5 CONTEXT PARALLELISM
TORCHTITAN has been extended to incorporate Context Parallelism (CP) (Liu et al. 2023
Liu
& Abbeel 2024 NVIDIA 2023). enabling 4D parallelism by adding CP as an additional di- mension to existing DP, TP, and PP. CP scales model training by splitting the context dimension across GPUs, significantly increasing the maximum trainable context length without causing out- of-memory (OOM) errors. For example,
Llama 3.1 8B with 8 H100 GPUs, using CP enabled training at context lengths up to 262, 144 tokens, achieving minor MFU degradation as CP degree increases (PyTorch Team 2025). For more details CP integration please refer to Appendix B.5
2.2
OPTIMIZING TRAINING EFFICIENCIES
2.2.1
NAVIGATING COMPUTE-MEMONT TRADE-OFFS USING ACTIVATION CHECKPOINTING
Activation checkpointing (AC) (Cen et al. 2016 He & Yu 2023 Purandare et al. 2023) and
selective activation checkpointing (SAC) (Korthikanti et al. 2023) are standard training techniques to reduce peak GPU memory usage, by trading activation recomputation during the backward pass for memory savings. It is often needed even after applying multi-dimensional parallelisms.
TorcHTiTAn offers flexible AC and SAC options utilizing torch.util s checkpo ntt ap- plied at the TransformerBloc level. The AC strategies include "full" AC, op-level SAC, and layer-level SAC.
Within TransformerrBooc full AC works by recomputing all activation tensors needed during the backward pass, whereas op-level SAC saves the results from cumputation-intensie Py PoTch operations and only recomputes others. Layer-level SAC works in similar fashion full AC, but the wrapping is applied to every x Trans Transformet formexBlock (where æ is specified by the user) to implement configurable trade-off's between memory and recompute. (Details are in Appendix B.6
5

---

Published as a conference paper at ICLR 2025
2.2.2 REGIONAL COMPILATION TO EXPLOIT TORCE. CORCE.COMPIE ILE OPTIMIZATIONS
torch. compile was released in Py PyTorch? 2 (Ansel et al. 2024 with TorchDynamo as the fron- tend to extract PyTorch operations into an FX graph, and TorchInductor as the backend to compile the FX graph into fused Triton code to improve the performance.
In TORCHTITAN, we use regional compilation, which applies torch. compile to each individ-
ual Trans forme Blcck in the Transformer model. This has two main benefits: (1) we get full graph (without graph breaks) for each region, compatible with FSDP2 and TP (and more gen- erally torch. Tensor subclasses such DTensor) and other PyTorch distributed training tech- niques; (2) since the Llama model stacks identical Trans FormerB Bock ock layers one after another, torch. compile ile identify the same structure is being repeatedly compiled and only compile once, thus greatly reducing compilation time.
torch. compile brings efficiency in both throughput and memory (sce Section 3.2 via compu- tation fusions and pepputaiion-ommmiicatt user interface. Below further elaborate how torch. comp i 1 composability helps TORCHTI- TAN unlock hurdware-optimizze performance gain with simple user interface, with the integration of advanced features such as Asynchronous TP and Float8.
reordering, in a model-agnostic way with a simple
2.2.3 ASYNCHRONOUS TENSOR PARALLEI. MAXIMALLY OVERLAP COMMUNICATION
By default, TP incurs blocking communications before/after the sharded computations, causing computation resources to not be effectively utilized. Asynchronous TP (AsyncTP) (Wang et al. 2022) achieves campulatioo-communicce tions within attention and feed-forward modules into smaller chunks, and overlapping communica- tion collectives in between each section. The overlap is achieved by micro- picerooippliiinngggg. optimization, where results are being communicated at the same time that the other chunks of the matmul are being computed.
overlap by fractionalizing the TP matrix multiplica-
PyTorch AsyneTTP is based on a Symmet r cMemrry abstraction, which creates intra-node huffers to write faster communication collectives. This is done by allocating shared memory buffer on each GPU in order to provide direct P2P access (PyTorch Team 2024a).
With TORCUTITAN's integration of torch. compileoomlee. AsyncTP can be easily configured in TORCHTITAN to achieve meaningful end-to-end speedups (see Section 3.2 for details) on newer hardware (H100 or newer GPUs with NVSwitch within node). Usage details are in Appendix B.7
2.2.4 BOOSTING THROUGHPUT WITH MIXED PRECISION TRAINING AND FLOAT8 SUPPORT
Mixed precision training Miciieevicius et al. 2018 provides both memory and computational sav- ings while ensuring training stability. FSDP2 has built-in support for mixed precision training with basic torch. , dtype. This covers the popular usage of performing FSDP all- all-gather and com- putation in a low precision (e.g. torch. bfcctt16 6), and perform lossless FSDP reduce-scatter (gradient) in high precision (e.g. torch. float32) for better numerical results. See Appendix B.8 for usage details.
TORCHTITAN also supports more advanced mixed precision training with Float8, derived data type, applied selectively to linear layers (available on newer hardware like NVIDIA H100), achiev- ing substantial performance gains while ensuring training stability (reported in Section 3.2) The Float8 feature from torchao float B supports multiple per-tensor scaling strategies. including dynamic, delayed, and static (see |Micikevicius et al. (Py Community (2023b), Section 4.3 for details), while being composable with other key nyTorch-native systems such as autograd, torch. comp i le, FSDP2 and TP (with Float8 all-gather capability) (PyTorch Team 2024c).
2.3 PRODUCTION READY TRAINING
To enable production- grade training, TORCHTITAN offers seamless integration with key features out of the box. These include (1) efficient checkpointing using PyTorch Distributed Checkpointing (DCP), and (2) debugging stuck or crashed jobs through integration with Flight Recorder.

---

Published as a conference paper at ICLR 2025
2.3.1 SCALABLE AND EFFICIENT DISTRIBUTED CHECKPOINTING
Checkpoints are crucial in training large language models for two reasons: they facilitate model reuse in applications like inference and evaluation, and they provide recovery mechanism in case of failures. An optimal checkpointing workflow should ensure ease of reuse across different par- allelisms and maintain high performance without slowing down training. There are two typical checkpointing methods. The first aggregates the state (model parameters and optimizer states) into an unsharded version that is srrallelism-agnostie communication. The second method has each trainer save its local sharded state, which speeds up the process but complicates reuse due to embedded parallelism information.
facilitating easy reuse but requiring expensive
DCP addresses these challenges using DTensor, which encapsulates both global and local tensor information independently of parallelism. DCP converts this information into an internal format for storage. During loading, DCP matches the stored shards with the current D'ensor-based model parameters and optimizer states, fetching the necessary shard from storage. TORCHTITAN effec- tively uses DCP to balance efficiency and usability. Furthermore, DCP enhances efficiency through asynchronous checkpointing by processing storage persistence in a separate thread, allowing this op- eration to overlap with subsequent training iterations. TORCH
utilizes DCP's asynchronous checkpointing to reduce the checkpointing overhead by 5-15x compared synchronous distributed checkpointing for the Llama 3.1 8B model (PyTorch Team 2024b).
2.3.2 FLIGHT RECORDER TO DEBUG JOB CRASHES
Debugging NCCL collective timeouts large scales is challenging due to the asynchronous ture of communication kernels. PyTorch's Flight Recorder addresses this by logging the start, end, and enqueue times for all collective and p2p operations, along with metadata like process groups, source/destination ranks, tensor sizes, and stack traces.
This data invaluable for diagnosing hangs in parallelism code. For PP, can pinpoint the latest send or recv completed on the GPU, helping debug schedule bugs. For FSDP and TP, it identifies ranks that failed to call collectives, aiding in uncovering issues with PP scheduling or TP logic.
3 EXPERIMENTATION
In this section, we demonstrate the effectiveness of elastic distributed training using TORCHTITAN, via experiments on Llama 3.1 8B, 70B, and 405B, from ID parallelism to 4D parallelism, at the seale from 8 GPUs to 512 GPUs. We also share the knowledge and experience gained through TORCITI- TAN experimentation. A walkthrough the codebase on how we apply (up to) 4D parallelism can be found in Appendix
3.1 EXPERIMENTAL SETUP
The experiments conducted NVIDIA H100 GPU with 95 GiB memory, where each host is equipped with 8 GPUs and NVSwitch. Two hosts form a rack connected to a TOR switch. A back- end RDMA network connects the TOR switches. In TORCHTITAN we integrate checkpointable data loader and provide built-in support for the C4 dataset (en variant), colossal, cleaned version of Common Crawl's web crawl corpus (Raffel al., 22220). We use the same dataset for all exper- iments in this section. For the tokenizer, we use the official one (tiktoken) released together with Llama 3.
3.2 PERFORMANCE
To showcase the elasticity and scalability of TORCHTITAN, we experiment on a wide range of GPU scales (from 8 to 512), as the underlying model size increases (8B, 70B, and 405B) with a varying number of parallelism dimensions (up to 4D). To demonstrate the effectiveness of the optimization techniques introduced in Section 2.2 we show how training throughput improves when adding each
2The H100 GPUs used for the experiments non-standard. They have HBM2e and are limited to lower TDP. The actual peak FFLOPs should he hetween SXM and NVT., and don't 't know the exact value.
7

---

Published as a conference paper at ICLR 2025
individual technique appropriate baselines. In particular, when training a higher dimensional parallelism with new features, the baseline is always updated to include all previous techniques.
We note that, throughout experimentation, memory readings stable across the whole training proces whereas throughput numbers (token per second, per GPU) are calculated and logged every 10 iterations, and always read at the (arbitrarily determined) 90th iteration. We do not report Model FLOPS Utilization (MFU) (Chowdhery et al. 2023) because when Float8 is enabled in TORCHTI- TAN, both BFLOATI6 Tensor Core and FP8 Tensor Core involved in model training. but they have different peak FLOPS and the definition of MFU under such scenario is not well-defined. We note that the ID Llama 3.1 8B model training on or 128 H100 GPUs without Float8 achieves 33% to 42% MFU.
Table I: ID parallelism (FSDP) on Llama 3.1 8B model, GPUs. Mixed precision training. Selec- tive activation checkpointing. Local batch size 2, global batch size 16. (Stats per GPU)
Techniques
Throughput (Tok/Sec)
Comparison
Memory (GiB)
FSDP
6,258
100%
81.9
+ torch. compile
6,674
+ 6.64%
77.0
+ torch. compile Float8
9,409
+ 50.335
76.8
Table 2: ID parallelism (FSDP) on Llama 3.1 8B model, 128 GPUs. Mixed precision training.
Selective activation checkpointing, Local batch size 2, global batch size 256. (Stats per GPU)
Techniques
Throughput (Tok/Sec)
Comparison
Memory (GiB)
FSDP
5,645
100%
67.0
+ torch. compilolleooooooo
6,482
+ 14.82%
62.1
+ torcch. compili le Float8
9,319
+ -65.08%
61.8
Table 2D parallelism (FSDP + TP) torch. compile + Float8 on Llama 3.1 70B model, 256
GPUs. Mixed precision training. Full activation checkpointing. FSDP degree 32, TP degree 8. Local batch size 16, global batch size 512. (Stats per GPU)
Techniques
Throughput (Tok/Sec)
Comparison
Memory (GiB)
2D
897
100%
70.3
+ AsyncTP
1,010
+ 12.59%
67.7
Table 3D parallelism (FSDP TP PP) + torch. wwwpocmlleeee...... le Float8 + AsyncTP on Llama 3.1 405B model, 512 GPUs. Mixed precision training. Full activation checkpointing. FSDP degree 4, TP degree PP degree 16. Local batch size 32, global batch size 128. (Stats per GPU)
Schedule
Throughput (Tok/Sec)
Comparison
Memory (GiB)
IFIB Interleaved IFIB
100
100%
78.0
130
30.00%
80.3
Additional experimental details and loss-convergence tests for correctness can be found in Ap-
pendix 10
3.3 SCALING WITH TORCHTITAN 4D PARALLELISM
Scaling large language models (LLMS) requires parallelism strategies to handle increasing model sizes and data on thousands of GPUs. TORCHTITAN enables efficient scaling through composable
Different PP ranks have different peak memory usages. We take the maximum all GPUs.
8

---

Published as a conference paper at ICLR 2025
Table 5: FSDP CP
torch... compi le
Float8 Llama 3.1 8B model, 8 GPUs. Mixed
precision training. Full activation checkpointing. Local batch size 1. (Stats GPU)
Schedule
Sequence Length
Throughput (Tok/Sec)
Memory (GiB)
FSDP 8, CP
32,768
3,890
83.9
FSDP 4, CP
65,536
2,540
84.2
FSDP 2, CP 4
131,072
1,071
84.0
FSDP 1, CP 8
262,144
548
84.5
Table 4D parallelism (FSDP TP +P CP) torch. comp le Float8 AsyncTP IFIB
on Llama 3.1 405B model, 512 GPUs. Mixed precision training. Full activation checkpointing. TP degree 8, PP degree 8. Local batch size 8. (Stats per GPU)
Schedule
Sequence Length
Throughput (Tok/Sec)
Memory (GiB)
FSDP 8, CP
32,768
76
75.3
FSDP 4, CP
65,536
47
75,9
FSDP 2, CP
131,072
77.1
FSDP 1, CP 8
262,144
16
84.9
4D parallelism. This section highlights key observations and motivations for using TORCHTITAN 4D parallelism, focusing specific combination shown in Figure 2
larger wodel
Long Context Training: CP or CP + DP
GPUs
4D Parallelism with long context training PP + FSDP + CP + TP
long context A
long combent
langer undel exposed
larger modal exposed
Parallelism
2D Porallehsm TP FSDP
3D Parallelism TP FSDP
FSDP
easy apply come comptta fuly overae
strong scaling coleen/rowwite sharding Sp..eee
wodel spltt by layers vicrebatching inouts with pieeliee small comers, bubble decided by schadule
GPUs
GPUs
more than 4192 GPUs
Figure 2: Scaling with 4D Parallelism
3.3.1 SCALING WITH FSDP
FSDP (ZeRO) is general technique applicable to any model architecture and often sufficient as the first degree of parallelism when communication is faster than computation (e.g., up to 512 GPUs). However, with larger scales, collective latency increases linearly with the world size, limit- ing efficiency. To overcome this, model parallelism like TP and PP can be combined with FSDP.
3.3.2 2D PARALLELISM: TP WITH FSDP
Tensor Parallelism (TP) reduces collective latency by distributing work across GPUs, enabling smaller effective batch sizes and reducing peak memory usage for large models or sequence lengths. Combining FSDP and TP allows strong scaling with a fixed problem/batch size (Details shown in Figure 4). TP also improves FLOP utilization by optimizing matrix multiplication shapes. However, TP introduces blocking collectives and is typically limited to intra-node scaling (c.g.. NVLink), with degrees usually capped at 8. Scaling beyond 4192 GPUs requires combining TP with PP.

---

Published as a conference paper at ICLR 2025
3.3.3 3D PARALLELISM: PP WITH 2D PARALLELISM
Pipeline Parallelism (PP) reduces communication bandwidth requirements by transmitting only ac- tivations and gradients between stages a peer-to-peer manner. PP is particularly effective for mitigating FSDP communication latency at larger scales in Sandwidth-limite clusters. The ef- ficiency of PP depends on pipeline schedules and microbatch sizes, which influence the size of pipeline "bubbles."
3.3.4 LONG CONTEXT TRAINING AND 4D PARALLELISM
Context Parallelism (CP) allows ultra long context training by splitting the context (sequence) di- mension across GPUs to avoid OOM errors. CP is mainly used for long context training, to give the model capability to capture more correlations for tokens, thus enhancing the overall model quality. For scaling sequence length, CP can be used alone together with DP. When training large models or on large number GPUs, we can combine CP with 3D paralleism, where TP usually keeps the innner-most Devi dimension, and CP applies in the next outer DeviceMesh dimension.
4 RELATED WORK
Libraries such Megatron-LM (Narayanan et al. 2021). DeepSpeed (Rasley al.) 2020). veScale (Inc. 2024) and PyTorch Distributed (Paszke et al. [2019) Meta Platforms, Inc. (2024) pro- vide APIs for distributed workflows. However, these frameworks present challenges in flexibility, integration, and scalability. TORCHTITAN addresses these limitations with native support for key
features absent in existing systems:
Megatron-LM: Requires model modifications for TransformerIngiine lacks seamless FSDP integration with TP and PP, and does support advanced pipeline schedules minimize computation overhead.
Depends on Megatron-LM for TP and CP, with limited support for FSDP and advanced pipeline schedules.
DeepSpeed:
veScale: Does support FSDP, CP, SAC, Float8 training, or torch. compilllleeoooo le, and offers only three pipeline schedules, compared to TORCHTITAN 's six.
We note that each of these libraries has its own strengths, and TORCHTITAN is designed provide foundational components that can be leveraged by all of them. A detailed comparison, includ- ing feature breakdowns and code complexity analysis, available in Appendix B.9 Slapo (Chen et al. (202)) introduces a schedule language to convert a PyTorch model for common model train- ing optimizations such as 3D parallelism, and supports progressive optimization through high-level primitives. In contrast, TORCHTITAN provides modular and composable APIs built on DTeesor and DeviceMesh.
CONCLUSION
TORCHTITAN is powerful and flexible framework for LLM training, enabling seamless com- posability of parallelism techniques (FSDP, TP, PP, CP), memory optimizations (Float8, activation checkpointing). and PyTorch compiler integration for enhanced efficiency. modular design sup- ports evolving architectures and hardware, fostering innovation with multi-axis metrics.
Designed for interpretability and production-grast training, TORCIITITAN offers elastic scalability, comprehensive training recipes, and expert guidance distributed training strategies. As demon- strated in experiments, it accelerates training by 65.08% on Llama 3.1 8B (128 GPUs, 1D), 12.59% on Llama 3.1 70B (256 GPUs, 2D), and 30% on Llama 3.1 405B (512 GPUs, 3D) over optimized baselines, while enabling long-context training with 4D composability. With its robust features and high efficiency, TORCHTITAN is ideal one-stop solution for challenging L.M training tasks.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, et al. Phi-3 technical re- port: A highly capable language model locally on your phone. arXiv preprint arXiv:2404. 14219, 2024.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. GPT-4 technical report. arXiv preprint arXiv:2303.087771 2023.
Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean- Rean-Baptiste Alayrac Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, and Gemini Team. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023.
Jason Ansel, Edward Yang. Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao. Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary Vito,,, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael La- 2os, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Shunting Zhang, Michael Suo, Phil Tillet. Xu Zhao, Eikan Wang. Keren Zhou, Richard Zou, Xiaodong Wang. Ajit Mathews, William Wen, Gregory Chanan, Peng Wu, and Soumith Chintala. Py- Torch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Sup- port Programming Languages and Operating Systems, Volume 2, ASPLOS '24, pp. 929-947, New York, NY, USA, 2024. Association for Computing Machinery, ISBN 9798400703850, doi: 10.1145/362065533646666 URL httlhttts//oo// wtw.:/....gg/1... 01011.1611533366663133 0665 3 364 0366
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
htwiittub.comm//m jax //jxx
Hongzheng Chen, Cody Hao Yu, Shuai Zheng, Zhen Zhang, Zhiru Zhang, and Yida Wang. Slapo: A schedule language for progressive optimization of large deep learning model training, 2023.
URL http:///wwwxxxvv.a/...o////
08005
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training Deep Nets with Sublinear htsg//abb1110.11177
Memory Cost, 2016. https:tts:///rx.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. PaLM: Scaling language modeling with Pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Jacob Devlin. BERT: Pre-training of deep bidirectional Transformers for language understanding. arXiv preprini a 1810.04800 2018
Abhimanyu Dubey, Abhinav Jaubri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
Assaf Eisenman, Kiran Kumar Matam, Steven Ingram, Dheevatsa Mudigere, Raghuraman Krish- namoorthi, Krishnakumar Nair, Misha Smelyanskiy, and Murali Annavaram. Check-N-Run Nunn: checkpointing system for training deep learning recommendation models. In 79th USENIX Sym- posium on Networked Systems Design and Implementation (NSDI 22), pp. 929-943, Renton, WA, April 2022. USENIX Association. ISBN 978-1-999113-27-4 URL | www. usenix.
Jiarui Fang and Shangchun Zhao. USP: A unified sequence parallelism approach for long context generative AI, 2024.

---

Published as a conference paper at ICLR 2025
Tanmaey Gupta, Sanjeev Krishnan, Rituraj Kumar, Abhishek Vijeev, Bhargav Gulavani, Nipun Kwatra, Ramachandran Ramjee, and Muthian Sivathanu. Just-in-time checkpointing: Low cost error recovery from deep learning training failures. In Proceedings of the Nineteenth European Conference on Computer Systems, EuroSys '24, pp. 1110-1121 New York, NY, USA, 2024. As- sociation for Computing Machinery. ISBN 9798400704476. doi: 10.1145/362770333550000
URL httpt::t://det oto/ 01145/3627703 3650085
Horace He and Shangdi Yu. Transcending runtime- -untine-mmmory tradeoff's in checkpointing by being fusion aware. Proceedings of Machine Learning and Systems, 5:414427,, 2023.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V. Le, Yonghui Wu, and Zhifeng Chen. GPipe: efficient training of giant neural networks using pipeline parallelism. Curran Associates Inc., Red Hook, NY, USA, 2019.
ByteDance Inc. veScale: A scalable and efficient distributed training framework. https:// 2024. Accessed: 2024-11-21.
Albert Qjang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04085 2024.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael An- dersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recom- putation in large transformer models. In D. Song, M. Carbin, and T. Chen (eds.). Proceedings of Machine Learning and Systems, volume 5, pp. 341-353. Cu- ran, 2023. URL BRRL//tttss///prreeddngs.../// htwp://ww...........ooooooo.o./r fi Eile/ E1220008512205535000000000-----0
pdf
Jianhui Li, Zhennan Qin, Yijie Mei, Jingze Cui, Yunfei Song. Ciyong Chen, Yifei Zhang, Longsheng Du, Xianhang Cheng, Baihui Jin, Yan Zhang, Jason Ye, Eric Lin, and Dan Lavery. oneDNN
graph compiler: A hybrid approach for high-performance deep learning compilation. In 2024 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pp. 460-
470. 2024. doi: 10.1 09.11999C0055660.2224..0000 10444871.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. PyTorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006. 15704, 2020.
Hao Liu and Pieter Abbeel. Blockwise parallel Transformers for large context models. Advances in Neural Information Processing Systems, 36, 2024.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise Transformers for near- infinite context. arXiv preprint arXiv:2310.0188) 2023.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pre- training approach, 2019. htLpstss::
1199)
Avinash Maurya, Robert Underwood, M. Mustafa Rafique, Franck Cappello, and Bogdan Nicolae. Datastates-IIm: Lazy asynchronous checkpointing for large language models. In Proceedings of the 33rd International Symposium on Prrih-Performance Parallel and Distributed Comput- ing. HPDC '24. pp. 227-239, New York, NY, USA, 2024. Association for Computing Machin- ery. ISBN 9798400704130 30. doi: 10.1145/6625549 36588888. URL http:ttt///o...gg/////
1145 7145/3625549.333 3658685
Meta Platforms, Inc. PyTorch Distributed, 2024. URL httpss//pptorrch.org/oees
html Accessed: 2023-09-26.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training, 2018. http://////aaiooooo/a///////////
03740
12

---

Published as a conference paper at ICLR 2025
Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Comea, Pradeep Dubey, Richard Grisen- thwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, Naveen Mellempudi, Stuart Oberman, Mohammad Shoeybi. Michael Siu, and Hao Wu. FP8 formats for deep learning, 2022. /tttttsss//arxii abss 2200. 05433
Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gre- gory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. PipeDream: generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, SOSP '19, pp. 1-15, New York, NY, USA, 2019. Association for Com- puting Machinery. ISBN 9781 9781450368735. doi: 10.1145/33413011 .335966......
URL https:
//doi.org/10: 1...111111331000 3359645
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clus- ters using megatron-Im. In Proceedings of the International Conference for High Performance Computing. Networking. Storage and Analysis, SC '21, New York, NY, USA, 2021. Associa- tion for Computing Machinery. ISBN 450384444. doi: 10.14453458817734466000 URL.
https://dos woogrol...............ooo....... 1 0
NVIDIA. Megatron Core API Guide:
Context Parallel, 2023.
URL
https:
7/does wtt.dia.com/mmceeee
gui de http://xxxt_rrralleeeeeeee html
Accessed: 2023-09-25.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Yang. Zach De Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: an imperative style, high-performance deep learning library. Curran Associates Inc., Red Hook, NY, USA, 2019.
Sanket Purandare, Abdul Wasay, Stratos Idreos, and Animesh Jain. J--WW: 3 Faster Multi- Model Training with Orchestration and Memory Optimization. In D. Song. M. Carbin, and T. Chen (eds.), Proceedings of Machine Learning and Systems, volume 5, pp. 541-562. Cu- ran, 2023. URL http/tttpps//rroeeenggg.... http:///w.........oopaa...p.eeee http:///rrppppppeeee//////rr///a
PyTorch Community. PyTorch DTensor 2023a. BRt.a/ttt.///gattu.oooo//re/eaea
pytorch//ssues 88838
GitHub Issue.
PyTorch Community. Float8 in PyTorch 1.x, 2023b. UK/nttps:/derrddoouss.pppo/ PyTorch Discussion Thread.
org/t/float8-in 1
PyTorch Team. Introducing Asyne Tensor Parallelism in PyTorch. htttps://discuss.
hew.aleeliinnnnnnptttooocttt/tee 2024a44480..0.2...m PyTorch Forum Post.
PyTorch Team.
Optimizing checkpointing efficiency with PyTorch DCP. https:
opt imiziig
2024b.
PyTorch
Forum Post.
PyTorch Team. Enabling Float8 all-gather in FSDP2. https:/
aaaan--nabling aooatt-all-gatiin
Esdp2/209323, 2024c. PoToh Forum Post.
PyTorch Team. Training with zero-bubble Pipeline Parallelism. https://discuss:
zero-
2024d. PyTorch Forum Post.
PyTorh Team. Breaking barriers: Training long context Ilms with IM sequence length in
Parallel. http://www..oooaaaaaaaoooooaaaaa
1m- www.seuueeeeeeeeeeeteeeeeee-eeaa otw.aaae.....tttttaaaaoooooaanna
2025.
PyTorch Forum Post.
13

---

Published as a conference paper at ICLR 2025
Penghui Qi, Xinyi Wan, Guangxing Huang, and Min Lin. Zero bubble pipeline parallelism, 2023. URL inttps://azen org/
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text Transformer. Mach. Learn. Res., 21(1), January 2020, ISSN 1532-4435.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. SC '20. IEEE Press, 2020. ISBN
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase and Yuxiong He. DeepSpeed: System op- timizations enable training deep learning models with over 100 billion parameters. KDD '20, pp. 3505-3506, New York, NY, USA, 2020, Association for Computing Machinery. ISBN URL
9781450379984. doi: 10.1145/39948863400700
11445
3394486. 3406703
Borui Wan, Mingji Han, Yiyao Sheng. Zhichao Lai, Mofan Zhang, Junda Zhang, Yanghua Peng, Haibin Lin, Xin Liu, and Chuan Wu. Bytecheckpoint: A unified checkpointing system for 11m development, 2024. URL http::ttp:
wbt/ /2007 .21143
Shibo Wang. Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, pp. 23-106, 2022.
Zhuang Wang, Zhen Jia, Shuai Zheng, Zhen Zhang, Xinwei Fu, S. Eugene Ng, and Yida Wang. Gemini: Fast failure recovery in distributed training with in-memory checkpoints. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP '23, pp. 364-381, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400702297. doi: 10.1145/ hRRphtttp////o..rrr 111.111..1133661 /tt.://.....1//66000066666611111
3600006.3613144
Cody Hao Yu, Haozheng Fan, Guangtai Huang, Zhen Jia Yizhi Liu, Jie Wang, Zach Zheng, Yuan Zhou, Haichen Shen, Junru Shao, Mu Li, and Yida Wang. Raf: Holistic compilation for deep learning model training, 2023. URL /trttps://axxv...e wtt.oorgg//bb///eeee 04759)
Buyun Zhang, Liang Luo, Xi Liu, Jay Li, Zeliang Chen, Weilin Zhang, Xiaohan Wei, Yuchen Hao, Michael Tsang, Wenjun Wang, Yang Liu, Huayu Li, Yasmine Badr, Jongsoo Park, Jiyan Yang, Dheevatsa Mudigere, and Ellie Wen. DHEN: A deep and hierarchical ensemble network for large- scale click-through rate prediction, 2022.
B wwt.t:tts.... /tappp:://aaxv.../// org/ abs/ 2200 11014
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. PyTorch FSDP: Experiences scaling Fully Sharded Data Parallel. Proc. VLDB Endow, 16(12)3848-3860, aug 2023. ISSN 8097. doi: 10.14778/361 154033611566 URL 10. 10.14778/3611542 36115669
A COMPOSABLE 4D PARALLELISM WALKTHROUGH
We have discussed the scaling with TORCHTITAN 4D parallelism and the motivations to apply different parallelisms to scale training to thousands of GPUs. In this section we will walk through the 4D parallelism code in TORCHTITAN.
The first step is to create an instance of the model (e.g. the transformes for Llama models)
on the meta device. We then apply PP by splitting the model into multiple PP stages according to the pi pipeling satpeeepaaaaeeaaaaaeaaeeeeeeeeee __splisspt______ appii_poontts points config. Note that for PP with looped schedules, may obtain multiple model parrts rts from PP splitting, where each item in mode paarts is
14

---

Published as a conference paper at ICLR 2025
monee--oel--cunk. Next apply SPMD-style distributed training techniques including TP, activation checkpointing, torch. corch.compile. ESDP, and mixed precision training for each model part, before actually initializing the sharded model on GPU.
meta init with torch.device device ("meta"):
model = model_cls. from_model_argee (model_configg
apply PP sp_schedule, model_parts =
[model_name]]
model, pp_mesh, perallel_dims, jeb_config, device, medel_config, Loss_fn
for m in medel_partes
apply SPMD-style distributed training techniques quess modelssparallelllee_mn [model_name] (m, world_mesh, parallel_dime job sob_config!
move sharded model to GPU and initialize weights via DTensor m. to empty (device="cuda"
inin__weighhs()
To apply PP the model,
the following code the high level. manually given _wama__amual_aarraeeeeaaaaaaa manual_split splits the model into multiple stages according to the wttpee_praaal___eteeeeeeeeeeeeee
pipeline 1aama_mmmaaa
config. by removing the unused model components from complete model (on the meta device). Then build pipelipee ine_schedule make the pipeline schedule with various options from torch. distributed pipelining, including IFIB (Narayanan et al. 2019) GPipe (Huang al. 2019), interleaved IFIB (Narayanan
al. 2021). etc. instructed by the
schhuuus config.
stages, models pipeline _allmnne_llaaaaae 11ama haanal_sppiit model, pp_mesh, parallel_dims, job_config, device, model_config
sc_hcccc = return pp_schodule, models
stages, loss_fn)
TP and FSDP are applied in the SPMD-style models hetelllsspaaallle..eomm
function. To apply TP, we utilize the DTTensor pa walllllzeemmer pellelize_mddull API, by providing TP "plan" as the struction of how model parameters should be sharded. In the example below, we showcase the (incomplete) code for sharding the repeated FramsBoore...ock.
for layer_id, transformer_ble
in model. layers. tems ()
layer_tp_plan
"attent battention_norm SequenceParallel1 (),
"attention": PreparaNoduleInpat/
input Lopptt__ayouts
None),
desired_ingun input layouts= l None),
atttention.wgi wq": colwiscParallel
parallleliz__moulee
module=trannffomee__bbl
device Dovice_mesh=tp_nnee parallelize htw.:/iwe..paaallaa..pp...ooooee
Then, we apply the FSDP by wrapping each individual model. Note that the FSDP2 implementation in PyTorch comes with mixed precision training sup- port. By default, we use torch. bfloat 6 on parameters all- all--rtherr and activation computations, and use torch. float 32 on gradient reduce-scatter communication and optimizer updates.
and then the whole
15

---

Published as a conference paper at ICLR 2025
mp polpoiccy MixedPrecisionPolicy (paran reduce _ttpe)
fsdp_ _edp_coniig = I"mesh: dp me_mes, "mp poplooo"""" mp policy]
for layer_id, transformer_blee block in model. layers........rm items () not reshard_after_forrer
for the last TransfozmerBlock since FSDP would prefetch immediately Feshard_after affee_forwaad = int (layer_id) id) len (model.
As an optimization,
milly_shard(
transformer_bles
fsdpp_onffi,
fully_shard (model, **fsdp_config)
Independently, can apply CP by running each training iteration under Python context manager.
opt
opp..ooaa_cccttct_paaalllle.._ee =
<p_mcchr=worddmmmshh"""""'"
cp_bufferrs-input_idd, labels] [m. freqs_sscci for in
model_partss,,
cp_seq_dims-11, 1] [0 for in model_ model_parts],
cp_nc_ pe_no_restore__emooo botoooreeuuteesstrttae ids, labels),
cp cp_rotate_mne jot ctnpeiiiooooooioeoeeeeeee
pawaleel rotaat__etaeeettttt
if parallel_dims
letellddimsccpeeaallle ep_enabled
else None
with train_centext context hptp://pooooccotttteeeeeerllllle pred model (input_ids
loss = lcss_fn (pred, labels)
B SUPPLEMENTARY MATERIALS
B. FULLY SHARDED DATA PARALLEL
FSDP2 advances the tensor sharding approach by replacing the original FSDP1 FlatParameter shard- ing. Specifically, parameters are now represented as DTTensors sharded on the tensor dimension 0. This provides better composability with model parallelism techniques and other features that quires the manipulation of individual parameters, allowing sharded state dict to be represented by DTensor without any communication, and provides for simpler meta-device initialization flow via DTeesor. For example, FSDP2 unlocks finer grained tensor level quantization, especially Float8 tensor quantization, which we will showcase in the results section.
As part of the rewrite from FSDPI to FSDP2, FSDP2 implements an improved memory management system by avoiding using record stream. This enables deterministic memory release, and as result provides lower memory requirements per GPU relative to FSDP1. For example on Llama 2 7B, FSDP2 records an average of 7% lower GPU memory versus FSDP1.
In addition, by writing efficient kernels to perform multi-tensor allgather and reduce scatter, FSDP2 shows on-par performance compared to FSDPI, with even slight performance gains using the Llama 2 7B, FSDP2 shows average gain of 1.5% faster throughput.
The performance gains are the result of employing two small performance improvements. First, only single division kernel is run for the FP32 reduce scatter (pre-dividing the local FP32 reduce- scatter gradient by world size, instead of two step pre and post divide by square root of world size). Secondly, in TORCHTITAN, FSDP2 is integrated with default of not re-sharding the final block in a transformer layer during the forward pass, since it will be immediately re- gathered at the start of the backward pass.
16

---

Published as a conference paper at ICLR 2025
Usage:
TORCHTITAN has fully integrated FSDP2 the default parallelism when training, and the data wwwpaaalllasaaaaeeeeeelllaaaaaaa
degree is the controlling dimension in the command line or TOML file. Note that for ease of use, the default data parallel shard shard_degree is -1, means to simply use all GPUs available, so user do not need to specify the actual world size.
B.2 HYBRID SHARDED DATA PARALLEI.
Hybrid Sharded Data Parallel (HSDP) is an extension of FSDP (Zhang et al. 2022). In FSDP, communication occurs between all devices within the FSDP group. However, at some point, the FSDP communication overhead exceeds its corresponding computation because the latency of allgather/reduce -catter communications increases linearly with the number of devices. This results in low MFU and becomes worthless to add more GPUs for scaling.
HSDP obviates this to some degree by creating a 2-D DeviceMesh that contains replica groups on dimension and shard groups the other dimension, where each shard group runs FSDP and the replica group runs normal data parallel. This ensures the FSDP communications happen a fraction of the original world size, with the addition of backward gradient allreduce across replica groups. HSDP reduces FSDP communication overhead and allows further scaling with data parallel.
Usage: TORCHTITAN makes it easy experiment with HSDP by using the configurable and data_paral perrraalleerepllleeeeeeeeeeeelle cecateeeeeeeeeee
settings: data_parallel_s sealel_ssrrrdderreee which controls the degree of the shard and replica groups we are creating. The product of both replicate and shard degree is the actual data parallel world size.
B.3 TENSOR PARALLEL
TP partitions the attention and feed forward network (MLP) modules of transformer layer across multiple devices, where the number of devices used is the TP degree. This allows for multiple GPUs cooperatively process the same batch by using the local sharded model parameters, the cost of adding all-redduee aell-reduue/alll paauuce/allggae perduue/alllgatheeeerrddeee wwwceeeeeeruueeeeeeeeeeeeeeeeeee scatter operations to synchronize intermediate activations.
Weight Sharding 16
Weigt Charding
4 12
Weights
Yow (P:1)
1813
Mestal
6'12
6-12
All Reduce
GPU
MatMall
Weight
acteadoo
12
Find Dut
Weights
Welghta
Figure 3: Tensor Parallel in detail (2 GPUs, data moves from left to right).
Due to the additional collectives introduced by TP. it needs to happen within a fast network (i.e NVLink), When training LLMs, TP is usually combined with FSDP, where TP shards within nodes and FSDP shards across nodes to create the 2D hierarchical sharding on different DeviceMesh di- mensions.
Usage: Because of the synergistic relationship between TP and SP, TORCHTITAN natively bundles these together and they are jointly controlled by the TP degree setting in the command line
17

---

Published as a conference paper at ICLR 2025
Shanding vit 40 16'4
Feniiie Hegh
Cottl Din
Stastin
Seuning
Stariig
Swarg
GPU
GPU4
GPU
GPU
PSOP Daher
PEDP Dather
Figure 4: FSDP2 Tensor Parallel (TP degree 4) sharding layout, with 2 nodes of 4 GPUs.
or the TOML entry of tensor cetesor_prraaeele _ellllee__ggee...ee Setting this to for example would mean that 2 GPUs within the node will share the computational load for each transformer layers attention and MLP modules via TP, and normalization/dropol layers via Sequence Parallel. Loss Parallel is implemented via context manager as needs control the loss computation outside of the model's forward computation. It can be enabled via enable loss loss_parallel.
B.4 PIPELINE PARALLEL
We expose several parameters configure PP. pipel controls the number of ranks participating in PP. of strings, representing layer fully-qualffied-names before which split will be performed. Thus, the total number of pipeline stages V will be determined by the length of this list. pipeline sateeeeppraeee____heeeeeeee schecule accepts the name of the schedule to be used. If the schedule is multi-stage, there should be V > stages assigned to each pipeline rank, otherwise V == 1. pipel iie wweeeeeppaaleeeeeeeeeeeaaeaaaeee www.mbirabba.m. ches controls the number of microbatches to split a data batch into.
accepts a list
B.5 ENABLING 4D PARALLEL TRAINING: CONTEXT-PARALLER (CP)
To address context scaling, we have incorporated Context Parallelism (CP) into TORCUTI- TAN. Following the principles of modular design of TORCHTITAN, CP was integrated context manager that dynamically replaces calls to attention operators (namely, with CP operations, ensuring no changes to the model
scaled dot code are required.
Under the hood, CP shards the DTensorr along the sequence dimension across the CP device mesh. extends the DTensor dispatcher to handle CP-specific operations, such as Ring Attention and causal attention load balancing, ensuring efficient operation. By extending DTensor's capabilities to support CP, TORCHTITAN ensures that CP fully compatible with all other parallelisms (FSDP, TP, PP), optimizations (e.g., activation checkpointing, torch. compile) and DCP. This demonstrates the extensibility of TORCHTITAN 's modular design. which accommodates future optimizations seamlessly while maintaining performance and compatibility.
18

---

Published as a conference paper at ICLR 2025
B.6 ACTIVATION CHECKPOINTING
TORCHTITAN offers two types of Selective Activation Checkpointing which allow for more nu- anced tradcoff between memory and recomputation. Specifically, we offer the option to selectively checkpoint "per layer" or "per operation". The goal for per operation is to free memory used by op- erations that are faster to recompute and save intermediates (memory) for operations that are slower to recompute and thus deliver more effective throughput/memory trade-off.
Usage: AC is enabled via a two-line setting in the command line or TOML file. Specifically, mode be either none, selective, or full. When selective set, then the next config of se sele seleettin sestivee___ypp is used which can be either a positive integer to enable selective layer checkpointing, op enable selective operation checkpointing. Per layer takes integer input guide the checkpointing policy, where 1 = checkpoint every layer (same as full), 2 checkpoint ev- cry other layer. checkpoint every third layer, etc. Per op(eration) driven by the save _ave_list policy in para lellll1l (ellin ize_ 1____lamo lamaa. . py which flags high arithmetic intensity operations such mat- mul (matrix multiplication) and SPDA (Scaled Dot Product Attention) for saving the intermediate results, while allowing other lower intensity operations to be recomputed. Note that for balancing total throughput, only every other matmul is flagged for saving.
B.7 ASYNCTP
The Symmetri cMemory collectives used in AsyncTP are faster than standard NCCL collectives and operate by having each GPU allocate an identical memory buffer in order to provide direct P2P access. Symmet t r cMemory relies on having NVSwitch within the node, and is thus generally only available for H100 or newer GPUs.
Usage: AsyncTP enabled within the experimental section of the TORCHTITAN TOML config file and turned on or off via the satplll_asnn__teeseeeee pap.lllllllllllll.lllllllmmlllll boolean setting.
B.8 CUSTOMIZING FSDP2 MIXED PRECISION IN TORCHTITAN
Mixed Precision is controlled by the wiieePrecisionnPo tion, which is then customized with param_dtte FP32 by default in TORCHTITAN. The reduce dt ype in FP32 means that the reduce-scatter in backwards pass for gradient computation will take place in FP32 to help maximize both stability and precision of the gradient updates.
class in the apply_fssee func- as BF16, and reduce siduce_dtypee defaulting to
B.9 TORCHTITAN: COMPREHENSIVE FEATURE SET AND REDUCED COMPLEXITY
B.9.1 TORCHTITAN ENABLES NEW DESIGNS
TORCHTITAN extensive feature set and broad design space coverage are driven by its unified design principles i.e. modularity, composability, and extensibility. Leveraging these principles, TORCHTITAN seamlessly integrates diverse parallelism strategies (FSDP, TP, PP, and CP) and opti- mizations (e.g., SAC, Float8 training). This unified framework not only supports advanced pipeline schedules and multi-dimensional parallelism but also simplifies the integration of new techniques, making it highly adaptable for cutting-edge research and production-grade deployments.
The following table highlights TORCHTITAN capabilities in context of parallelism, checkpointing and compiler support offerings compared to Megatron-LM, DeepSpeed, and veScale:
B.9.2 CODE COMPLEXITY AND MAINTAINABILITY
TorcHTiTan 's design principles also contribute to its significantly reduced code complexity. De- spite offering a rich feature set, TORCHTITAN maintains a compact and modular codebase, making it casier to extend, maintain, and evolve while ensuring high performance. The following table compares the lines code (LOC) for TORCHTITAN with Megatron- and DeepSpeed:
*Custom Fusion Kernels
19

---

Published as a conference paper at ICLR 2025
Table Comparison of TORCHTITAN with Megatron-L.M, DeepSpeed, and veScale with respect parallelism, compiler support, activation checkpointing, and model checkpoining.
Features
TORCHTITAN
Megatron-I.M
DeepSpeed
veScale
FSDP-Zero2
Yes
Yes
Yes
No
FSDP-Zem3
Yes
Yes
Yes
No
HSDP
Yes
Yes
No
No
TP
Yes
Yes
No
Yes
Asyne TP (Micro-pipelininn)
Yes
Yes
No
Yes
CP
Yes
Yes
No
No
PP-Gpipe
Yes
Yes
Yes
No
PP- Interleaveed (IFIB)
Yes
Yes
Yes
Yes
PP-Looped-BFS
Yes
No
No
No
PP-1FIB
Yes
Yes
Yes
Yes
I.-FlxibbeeImtteilaaee---
Yes
No
No
No
PP-ZeroBubble
Yes
No
No
Yes
(TP+SP)+PP
Yes
Yes
No
Yes
DDP+(TP+SP)+PP
Yes
Yes
No
Yes
FSDP+(TP+SP)
Yes
No
No
No
(TDPS((PP
Yes
No
No
No
(SDD=<(PPSS)*PPPCCP
Yes
No
No
No
MoE
Ongoing
Yes
No
No
Full AC
Yes
Yes
Yes
Yes
Flexible SAC
Yes
No
No
No
DCP
Yes
Yes
Yes
Yes
Float8 Training
Yes
Yes
No
No
torch. compile
Yes
No
Partial
No
Table 8: Lines of Code (LOC) comparison across systems.
Lines of Code (LOCC
TORCHTITAN Megatron-I.M
DeepSpeed
Core Codebase
7K
93K
94K
Total Codebase (Including Utils)
9K
269K
194K
B.10
EXTENDED EXPERIMENTS ANALYSIS: PERFORMANCE AND Loss CONVERGING
1.10.1
PERFORMANCE
Our experiments in Section 3.2 serve multiple objectives:
Establish composability and modularity: TORCITITTA demonstrates seamless integra- tion of various parallelisms and optimization techniques.
Showcase performance improvements: Significant speed-ups are observed across paral- lelisms and optimizations.
Validate elastic scalability: TORCHTITAN scales effectively with both the model size and the number of GPUs.
Ablation studies: Detailed performance gains for individual techniques are presented.
In particular
Table I Highlights improvements from compiler support over eager execution, followed by further gains with Float8 training.
Table 2
Demonstrates how earlier gains scale the number of GPUs increases.
Table Shows speed-up achieved by AsyncTP (a HW/SW co-designed technique) over 2D training combined with torch. comp 1 and Float8 training.
20

---

Published as a conference paper at ICLR 2025
Table 4
Quantifies the benefits of Interleaved IFIB scheduling over IFIB top of Async T.. torch. compile, and Float8 training.
Table 5
Demonstrates the effectiveness CP on enabling long context training, even at small scale.
6 Demonstrate the composability of 4D parallelism, and the effectiveness of CP on enabling long context training large scale.
For FSDP, the ZeRO-3 variant is used for all experiments except for those involving PP where the ZeRO-2 variant used. This distinction due to the inefficiency ZeRO-3 in PP. where it incurs additional all-gather calls for each microbatch. In contrast, 7eRO-2 gathers parameters only once for the first microbatch and reshards after the last microbatch's backward pass.
B.10.2 Loss CONVERGING
TORCHTITAN 's design principles have influenced the development of advanced distributed training features such as FSDP2, As,neTI,, PP, and CP in PyTorch's 's distributed library. Throughout these contributions, we have ensured the loss converging of individual techniques as well as their various combinations of parallelisms and optimizations.
For example, below is a series of loss-converging tests covering both parallelisms and training op- timizations. We notations of "FSDP 8" for experiment in which the degree of FSDP is 8, "FSDP 8, CP 8" for an experiment on 64 GPUs where FSDP degree is 8 and CP degree is 8, etc. We assume the correctness of FSDP, which can be further verified by comparing it with DDP or even single-device jobs.
Table 9: Loss-convergingg tests setup.
Parallelism
Techniques
FSDP 8 (ground truth)
default
FSDP 8, TP 2, PP FSDP 8, TP 2, CP 2, PP FSDP 8, CP 8
torch.compile Float8, asyne TP, Interleaved 1FIB torch.compile, Float8, async TP, Interleaved IFIB default
PERPLCS
1990
Figure 5: Loss converging tests on Llama 3.1 8B. C4 dataset. Local batch size 4, global batch size 32. 3000 steps, 600 warmup steps.
21


Reviews:
[{"review_id": "cNnsxjkguo", "reviewer": "Reviewer_3H4A", "summary": "The proposed LEGOSCALE framework aims to enhance the training efficiency of the Llama 3.1 model in distributed training by integrating several established techniques, including n-D parallelism support, efficient checkpoint, compiler enhancement, debugging tools, and Float8 support. \nAuthors evaluate the scaling performance of LEGOSCALE using the Llama 3.1 family of models (8B, 70B, and 405B) on scales from 8 to 512 H100 GPUs.  \nThe paper claims 65.08%, 12.59%, and 30% improvement on training speed on the 8B model with data parallel, 70B model with data- and tensor-parallel, and 405B model with 3D parallel, repsectively.", "questions": "Please address the points in weakness.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Using the latest Llama 3.1 family of models is promising\n2. The scale with up to 512 H100 GPUs, though about two orders of magnitude smaller than the original Llama 3.1 paper, is sufficient.", "weaknesses": "1. The primary concern is that none of the proposed techniques is new or invented by the authors, significantly limiting LEGOSCALE's originality. Most of the performance improvement in Table 1-4 is attributed to the incorporation of existing solutions, not the new design of LEGOSCALE.\n2. LEGOSCALE employs a composable multi-dimensional parallelism approach by combining Fully Sharded Data Parallel (FSDP), Tensor Parallel (TP), and Pipeline Parallel (PP) methods, which have been extensively validated in systems like DeepSpeed (https://github.com/microsoft/DeepSpeed), Megatron-LM, and GPipe for efficient model scaling.\n3. Additionally, LEGOSCALE leverages hardware-software co-designed solutions such as Float8 mixed precision and optimized activation checkpointing to improve hardware utilization and memory efficiency. Most of the recent LLM training was done with BF16 and FP32. FP8 was rarely used due to practical concerns in optimization stability. To justify the innovation of FP8, the authors need to provide end-to-end training results and downstream task performance.\n4. Furthermore, LEGOSCALE incorporates PyTorch’s Distributed Checkpointing (DCP) to enable efficient failure recovery and state persistence, building upon checkpointing strategies already prevalent in systems like ByteCheckpoint and Gemini. \n5. LEGOSCALE optimizes memory and computation efficiency through `torch.compile` regional compilation, a technique similarly found in frameworks like JAX and ONNX, which use compiler optimizations to enhance training throughput and resource utilization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zPyYeXJYQU", "reviewer": "Reviewer_izCA", "summary": "This paper introduces `LegoScale`, an open-source, PyTorch-native distributed training system designed to unify and enhance current engineering efforts in LLM pre-training. `LegoScale` improves training efficiency through a modular, composable design that supports various parallelism techniques, including Fully Sharded Data Parallel (FSDP), Tensor Parallel (TP), and Pipeline Parallel (PP), as well as memory optimization methods such as FP8 precision and activation checkpointing. It also integrates with `torch.compile` for further optimizations. Experiments on the Llama 3.1 family of LLMs demonstrate that `LegoScale` accelerates training by 65.08% on the Llama 3.1 8B model at a 128-GPU scale (1D), 12.59% on the Llama 3.1 70B model at a 256-GPU scale (2D), and 30% on the Llama 3.1 405B model at a 512-GPU scale (3D), all on NVIDIA H100 GPUs, compared to optimized baselines.", "questions": "-   Could you elaborate on the improvements from FSDP1 to FSDP2, and how these enhancements impact different parallelism settings?\n-   How does each parallelism strategy scale with the number of GPUs?\n-   What does GPU utilization look like during training? Is GPU efficiency maximized?\n-   Is this method extensible to other platforms, such as TPU clusters?\n-   Could you provide more details on the system's fault tolerance and crash recovery design?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is well-organized and clearly presents the design of  `LegoScale`  across parallelism strategies, memory optimization techniques, and integration with  `torch.compile`. While prior research has addressed various aspects of LLM training, this paper delivers a comprehensive, feature-complete open-source solution that focuses on key challenges, including:\n\n1.  Composability of parallelism techniques,\n2.  Extensibility for incorporating new optimizations,\n3.  Efficient hardware utilization, and\n4.  Production-ready features.\n\nThe figures, particularly Figures 1 and 2, effectively illustrate how the system supports parallelism and customization features.\n\nNotably, the authors provide clear abstractions for N-D parallelism in training and analyze strategies for maximizing efficiency when combining parallelism methods. The experiments demonstrate promising performance improvements compared to strong baselines. Beyond these contributions, the engineering effort invested in this project is substantial and valuable for future work in this area.", "weaknesses": "The proposed training system builds extensively on prior work, especially in parallelism design and memory optimization. While the tensor abstraction is novel and more expressive, the primary contribution is from an engineering perspective. Additionally, possibly due to page limitations, much of the detail is relegated to the appendix. There are a few areas where the paper could strengthen its credibility.\n\nFirst, the paper could benefit from more discussion on how the modular design facilitates the integration of future optimizations and enables further parallelism without compromising performance.\n\nFor the PyTorch integration, although  `torch.compile`  is powerful, the paper lacks specifics on how it sustains high GPU efficiency in a highly parallel training environment—for example, insights on GPU usage.\n\nAn ablation study on the effect of individual features, such as transitioning from FSDP1 to FSDP2, and on performance changes as GPU scale increases within 1D, 2D, and 3D parallelism settings, would also help readers better understand the system's effectiveness.\n\n**Minor issue:**  In line 315, \"torch\" is misspelled as \"torchao.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UG8T4Q5f3s", "reviewer": "Reviewer_Xmiy", "summary": "The paper presents LegoScale, a framework that enables easy comparison and composition of different LLM training functionalities and optimizations. LegoScale can significantly improve the productivity of model scientists and practitioners by reducing the engineering burden of model development. LegoScale incorporates Pytorch-native implementation of popular techniques like mixed precision, 3D parallelism, ZeRO, activation checkpointing, model checkpointing, logging, etc. The evaluation studies usage of LegoScale for training Llama-3 model family and demonstrates model and hardware scaling scenarios by composing 3D parallelism, mixed precision, and compiler optimizations.", "questions": "1. Can you clarify whether the throughput results in Tables 1-4 are per GPU or aggerate.\n2. How are the optimization configurations (e.g., PP vs FSDP, AC vs SAC) used in Tables 1-4 obtained? Was there some exploration of the configuration space to discover optimal settings?\n3. Can you clarify that LegoScale applies to other model families besides Llama? In particular, line 279 refers to Llama model when discussing the torch.compile integration (2.2.2). \n4. What ZeRO stages are the FSDP configurations in Tables 1-4?\n5. HSDP is not mentioned in Figure 1 or 2, although mentioned as an important FSDP option. Is HSDP not available in LegoScale?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- LegoScale tackles an important problem of making it easier for model practitioners and researchers to compare and compose SOTA training techniques. Advanced LLM training techniques are typically scattered across different frameworks making it difficult if not impractical for practitioners and researchers to leverage these techniques. LegoScale could be a very valuable tool for the community. \n- By incorporating pytorch-native implementations of training optimizations, LegoScale lowers the bar for integration into existing pytorch workflows, likely fostering adoption. \n- The evaluation results highlight intuitive compositions of relevant optimizations for model and hardware scaling, which helps the user appreciate the developer benefits of LegoScale.", "weaknesses": "- Considering the growing importance of context scaling, to millions of tokens, LegoScale seems to be missing an important scenario. Although, this is discussed as ongoing work, it is unclear how context scaling fits into LegoScale. This is because the current tensor abstraction focuses on model partitioning, whereas SOTA context scaling optimizations such as Ulyssess involve input partitioning. \n- While LegoScale allows composition of various optimizations, it is unclear how users will be guided to make optimal choices in terms of performance metrics such as throughput, latency, or hardware utilization. \n- While the evaluation does a good presenting the scaling performance results, it is missing results of the production-ready-training features (2.3). For example, it would be useful to show convergence plots demonstrating that effectiveness of parallelism changes using DCP.\n- The performance evaluation lacks a common baseline metric across Tables 1-4, which makes it difficult to understand LegoScale efficiency across model and hardware scaling dimensions. While the authors rightly identified how different precisions hampers MFU usage, I think a good compromise is to report bfloat16 and float8 MFUs separately.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kSE7cVYaD5", "reviewer": "Reviewer_5X5L", "summary": "This paper describes a feature-rich PyTorch native hardware-software system for efficiently training of LLMs.", "questions": "Is there a plan to make the software components of the pre-training system available to other researchers?\n\nWhat do you view as the most novel new ideas within the design of your pre-training system?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper describes a useful tool that addresses an important problem.\n\nThe tool substantially improves training time.\n\nThe system appears to have sufficient features for production use.\n\nThe paper has some tutorial value in that it briefly explains a wide range of features appropriate for pre-training systems.\n\nI think this paper's greatest value is in advertizing the pre-training system to researchers and developers who may want to use it.", "weaknesses": "There are few new ideas. The focus is on  making a powerful tool by building on existing ideas. However, this weakness is partially offset by the fact that the work may enable others to develop new ideas by reducing the time required for pre-training.\n\nInterest will be bimodal: those who need to pre-train LLMs will be very interested, especially if they have access to the pre-training system, and those who do not will probably find little of interest in the paper.\n\nThere is not a strong distinction between engineering efforts to add features to the pre-training system and research efforts to develop new system design ideas. The paper explains the design of the tool but does little to justify design decisions. Improvements to writing, alone, might address these problem.\n\nWith the exception of the sensitivity analysis in the tables on page 8, the paper does not contain scientific findings.\n\nFigures 1 and 2 are essentially impossible to read without magnification, which circumvents conference page limits.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WzOdHBsOVJ", "reviewer": "Reviewer_DprR", "summary": "This submission comprises a technical report / white paper about the development of LegoScale, a pytorch-based solution for scalable LLM training in distributed settings. The proposed framework extends the capabilities of vanilla pytorch for LLM training in many aspects, ranging from improved scalability through a principled combination of numerous optimised parallelism techniques and memory management, to improved logging and debugging mechanisms.", "questions": "This review is intentionally short and high-level, due to the technical excellence, but also lack of research novelty, of the proposed system. \nIn my opinion, white papers of successful training frameworks can still be published in ML venues; but will be willing to adjust my score considering the outcome of the upcoming discussion with other reviewers'. \n\nA piece of information that can facilitate this process is any insights that the authors have obtained and can share with us about the usage of their framework from the community (without violating the double-blind process). Has the framework been open-sourced? How many people are actively contributing to it? How many people are currently using it? Are the reported results easily replicated out-of-the-box on different clusters? Is long-term support and maintenance of the framework part  of the plans of the authors?\n\nPost rebuttal edit:\n I am increasing my score from 5 to 6. Although I maintain the position that the novelty in the manuscript is limited, I believe that the proposed framework can have notable impact to the community, and as such it is worth being presented to an ML venue in my opinion.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The proposed framework offers a streamlined approach for LLM training across distributed systems with different scales of compute. This scalability is achieved through a principled and well-engineered framework that allows tunable exploitation of different parallelism dimensions and other optimisations. This abstracts tons of complexities from deep learning practitioners. \n- Evidently, the majority of design choices in the proposed framework have relied on insights and experience from both the LLM training and the GPU deployment communities.\n- The reported experiments demonstrate that the proposed framework can achieve notable speed-ups at training time.", "weaknesses": "- It is difficult to identify much research novelty in the proposed framework, since its contribution can mostly be found in the design and tons of non-trivial engineering efforts behind LegoScale. \n- It is unclear whether this contribution will last or may wind down due to enhancements of the original pytorch framework towards the same direction, or lack of maintenance on the proposed framework which may lead to lack of support from latest pytorch versions or for latest GPU models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i4pMgQ0HMv", "reviewer": "Reviewer_UCk4", "summary": "This paper introduces LEGOSCALE, an open-source, PyTorch-native distributed training system that unifies and advances state-of-the-art techniques, streamlining integration and reducing engineering overhead. LEGOSCALE enables 3D parallelism in a modular and composable featuring elastic scaling to adapt to changing computational requirements. By providing an accessible and extensible platform, LEGOSCALE democratizes large language model pre-training, empowering a wider range of researchers and developers to tap into the potential of LLMs and accelerate innovation in the field.", "questions": "Although this work is more of an engineer's contribution than a research novelty contribution, I still think it is important and should be accepted by ICLR. There are some questions I would like to see more discussion of.\n\n1. For Pytorch native training framework, there is another framework implemented by Bytedance: https://github.com/volcengine/veScale. Could you compare the differences between them? It seems they support more parallisms.\n\n2. I would like to see some results of end-to-end throughput comparison with the latest Megatron, Deepspeed.", "limitations": "", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Thank you for submitting this paper to ICLR! I loved reading this paper.\n\n1. Important problem. LLM training is challenging and resource-consuming.\n\n2. Very graceful, interesting and timely solution. Pytorch-native is what the community needs, which is different with existing systems like Megatron.", "weaknesses": "1. Too many modules are introduced in limited pages, which is not friendly to general readers without machine learning system background.\n\n2. Limited model support. I would like to see more application on MoE model, Multi-modal LLM, etc.\n\n3. 3D-parallelism is not comprehensive / state-of-the-art at the time of 2024. Consider adding more support on context parallel (also mentioned in paper), expert parallel, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)