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

We are grateful to the reviewer for acknowledging, appreciating, validating, and providing critical feedback on our work. Please consider our response labeled as Wi, which aims to address weaknesses i, while Qi aims to address Question i, pointed out by the reviewer.

W1:

We agree with the reviewer that presenting a holistic and comprehensive overview of LEGOSCALE may feel overwhelming to readers without a background in machine learning systems. However, we believe it is crucial to provide a complete picture of the pre-training system’s capabilities, including its parallelisms and memory and compute optimization techniques. To address this, we have kept the descriptions at a high level, focusing on the functionality of each module while avoiding excessive implementation details.

We will carefully review the paper and make adjustments to ensure the right balance between accessibility for general readers and sufficient depth for those with ML systems expertise. This will help make LEGOSCALE’s design and contributions clear to a broader audience while maintaining technical rigor.

W2:

This is a great suggestion that will highlight LEGOSCALE’s model-agnostic and general design. While demonstrated with the Llama model, LEGOSCALE is not model-specific. Its techniques are general-purpose, allowing seamless integration of other models without changes to the system.

We plan to release Llama 3.2 multimodal recipes soon and are working on integrating diffusion models, showcasing LEGOSCALE’s flexibility and extensibility.

W3:

We thank the reviewer for highlighting the importance of context scaling and its integration into LEGOSCALE. We completely agree with the reviewer regarding its significance and are pleased to report that, since submitting the paper, we have successfully incorporated Context Parallelism (CP) into LEGOSCALE. We will include details of this in the camera-ready version of the paper. Below, we provide an outline of how LEGOSCALE addresses context scaling through CP.

Context Parallelism (CP) in LEGOSCALE

To address context scaling, LEGOSCALE extends DTensor’s capabilities to support CP. Following LEGOSCALE’s modular design principles, CP was integrated using two user-friendly APIs:

1. Module Wrapper: Similar to TP, for distributing the Attention module.
2. Context Manager: Dynamically replaces scaled_dot_product_attention calls with context-parallel operators, ensuring no changes to the model code.

Under the hood, CP:

• Shards the DTensor along the sequence dimension across the CP device mesh.
• Extends the DTensor dispatcher to handle context-parallel-specific operations, such as ring-attention and causal attention load balancing, ensuring efficient operation.

By extending DTensor’s design, CP in LEGOSCALE remains fully compatible with other parallelisms (FSDP, TP, PP), optimizations (e.g., activation checkpointing, torch.compile), and Distributed Checkpointing (DCP). This demonstrates LEGOSCALE’s extensibility and future-proof design, accommodating new optimizations seamlessly while maintaining performance and compatibility.

Context Scaling Results (Llama3-8B, 64 H100 GPUs, TP=8, DPxCP=8)

The integration of Context Parallelism highlights LEGOSCALE’s ability to adapt to emerging requirements in LLM training. Its unified DTensor abstraction and modular framework ensure that context scaling and other advancements can be seamlessly incorporated, making it a robust solution for future research and deployment.

The efforts for adding expert parallelism to the framework is underway.

We are grateful to the reviewer for acknowledging the importance of our work, recognizing its performance, appreciating its comprehensiveness, and providing detailed and critical feedback on it. Please consider our response labeled as Wi, which aims to address weaknesses i, while Qi aims to address Question i, pointed out by the reviewer.

W1:

We thank the reviewer for highlighting that LEGOSCALE enables machine learning practitioners and researchers to focus on mathematical problems, architectural exploration, and training dynamics by abstracting away the complexities of distributed systems. Simultaneously, LEGOSCALE provides ML systems researchers with a unified, modular, composable, and extensible testbed, empowering them to invent, extend, and swap newly developed techniques across various parallelism strategies and optimizations. This dual focus ensures that LEGOSCALE serves as a valuable tool for advancing both theoretical and systems-level research in machine learning.

We have observed from 9 active GitHub issues that users are beginning to extend LEGOSCALE for fine-tuning and inference. This strongly indicates that the system’s impact extends beyond pre-training research, benefiting a broader community.

W2 and Q2:

LEGOSCALE’s key research contribution lies in analyzing and identifying a common set of design principles across parallelism and optimization techniques. These principles unify diverse methods into a cohesive system. The engineering contribution is in enforcing these principles across individual techniques, resulting in a modular, composable, and flexible distributed system that enables seamless scaling, efficient deployment, and integration of new techniques. These contributions have shaped PyTorch features like FSDP2, Async-TP, PP, and DTensor advancements.

LEGOSCALE’s Design Principles

LEGOSCALE adopts a top-down approach, enforcing modularity, non-intrusiveness, and extensibility. By leveraging DTensor and DeviceMesh, it provides a unified framework for diverse parallelisms and optimizations, ensuring scalability and compatibility with PyTorch.

Advancements to DTensor

DTensor sharding was extended to support n-D parallelism, integrating TP, CP, and PP. Compatibility with torch.compile was added for compiler optimizations, and checkpointing efficiency improved with robust state dict support. Production readiness was bolstered by resolving critical bugs.

Contribution to PyTorch FSDP2

FSDP2 replaced FlatParameter with DTensors sharded on dim-0, enabling flexible parameter handling (e.g., freezing, fp8 all-gather). It integrates with torch.compile, reduces GPU memory usage by eliminating CPU synchronization, and simplifies checkpointing with sharded state dicts.

Contribution to PyTorch Async-TP/CP

Async-TP introduced fused operators like fused_all_gather_matmul to reduce synchronization overheads, enabling fused compute and communication. This design facilitates new parallelisms like CP and minimizes bottlenecks with CUDA streams and SymmetricMemory.

Contribution to PyTorch PP

Pipeline Parallelism incorporates LEGOSCALE’s principles, enabling dynamic schedules, fine-grained stage control, and the addition of new schedules while integrating with TP and DP.

A Unified Template for Distributed Training

LEGOSCALE provides a unified design template, simplifying exploration of configurations and enabling rigorous evaluation of techniques. It unifies, modularizes, and extends distributed methods, shaping PyTorch’s library and advancing scalable, flexible systems.

W4:

We thank the reviewer for pointing this out. In the camera-ready version, we will make the text more concise to allocate space for enlarging Figures 1 and 2, ensuring better visibility and readability for the readers.

Q1:

Thank you for the question. We have fully open-sourced the LEGOSCALE system on GitHub (anonymized to comply with ICLR's double-anonymous requirements). The framework has garnered significant interest, with over 2,600 stars, 204 forks, 469 closed pull requests, and 35 active contributors. LEGOSCALE enables researchers and practitioners to explore, develop, and scale advanced parallelism and optimization techniques in a production-grade pre-training system.

LEGOSCALE has been validated on multiple platforms, including the public cloud (AWS), on-premises clusters of a Fortune 500 company, a leading private organization, and a world-renowned university. Results are designed to be easily replicable across various cluster configurations. We are committed to actively maintaining and supporting LEGOSCALE for the foreseeable future.

Since submitting the paper, we have incorporated Context Parallelism (CP) and added three additional Pipeline Parallelism (PP) schedules (Looped-BFS, Flexible-Interleaved-1F1B, ZeroBubble). The camera-ready version will include case studies showcasing how to extend LEGOSCALE and the design principles driving its modularity.

Another key feature is LEGOSCALE’s model-agnostic and general design. While demonstrated with the Llama model, LEGOSCALE is not model-specific. Its techniques are general-purpose, allowing seamless integration of other models without changes to the system.

We plan to release Llama 3.2 multimodal recipes soon and are working on integrating diffusion models. This would also enable research to extend beyond NLP to vision models, highlighting LEGOSCALE's flexibility and extensibility.

Our goal in submitting this paper is to raise awareness within the research community about a system that empowers advanced research on pre-training large language models (LLMs) at scale, fostering innovation in this critical area.

W3:

Our experiments serve multiple objectives:

1. Establish composability and modularity: LEGOSCALE demonstrates seamless integration of various parallelisms and optimization techniques.
2. Showcase performance improvements: Significant speed-ups are observed across parallelisms and optimizations.
3. Validate elastic scalability: LEGOSCALE scales effectively with both model size and the number of GPUs.
4. Ablation studies: Detailed performance gains for individual techniques are presented:
   • Table 1: Highlights improvements from compiler over eager execution, followed by gains with Float8 training.
   • Table 2: Demonstrates how earlier gains scale as the number of GPUs increases.
   • Table 3: Shows speed-up achieved by Async-TP over 2D training combined with torch.compile and Float8 training.
   • Table 4: Quantifies the benefits of Interleaved 1F1B over 1F1B on top of Async-TP, torch.compile, and Float8 training.

We acknowledge the reviewer’s concern and will address it by adding experiments on training convergence and Distributed Checkpointing (DCP) performance in the appendix to provide further scientific findings. We have already done those experiments, the graphs can be found here.

We sincerely thank the reviewer for underscoring the importance of our work, recognizing our contributions, appreciating the design of our system, and providing detailed and critical feedback on our work. Please consider our response labeled as Wi, which aims to address weaknesses i, while Qi aims to address Question i, pointed out by the reviewer.

W1:

We thank the reviewer for highlighting the importance of context scaling and its integration into LEGOSCALE. We are pleased to report that, since submitting the paper, we have successfully incorporated Context Parallelism (CP) into LEGOSCALE. We will include details of this in the camera-ready version of the paper. Below, we provide an outline of how LEGOSCALE addresses context scaling through CP.

Tensor Parallelism (TP) with DTensor

LEGOSCALE leverages DTensor to shard not only model parameters but also inputs and activations along multiple dimensions. For example, in an MLP layer:

• The first linear layer’s parameters are column-sharded, while the second layer’s parameters are row-sharded.
• Inputs are sharded across the batch and feature dimensions, enabling flexible, efficient handling of sharded computations.

This design ensures that TP is not limited to parameter sharding and can effectively scale across different workloads.

Context Parallelism (CP) in LEGOSCALE

To address context scaling, LEGOSCALE extends DTensor’s capabilities to support CP. Following LEGOSCALE’s modular design principles, CP was integrated using two user-friendly APIs:

1. Module Wrapper: Similar to TP, for distributing the Attention module.
2. Context Manager: Dynamically replaces scaled_dot_product_attention calls with context-parallel operators, ensuring no changes to the model code.

Under the hood, CP:

• Shards the DTensor along the sequence dimension across the CP device mesh.
• Extends the DTensor dispatcher to handle context-parallel-specific operations, such as ring-attention and causal attention load balancing, ensuring efficient operation.

By extending DTensor’s design, CP in LEGOSCALE remains fully compatible with other parallelisms (FSDP, TP, PP), optimizations (e.g., activation checkpointing, torch.compile), and Distributed Checkpointing (DCP). This demonstrates LEGOSCALE’s extensibility and future-proof design, accommodating new optimizations seamlessly while maintaining performance and compatibility.

Context Scaling Results (Llama3-8B, 64 H100 GPUs, TP=8, DPxCP=8)

The integration of Context Parallelism highlights LEGOSCALE’s ability to adapt to emerging requirements in LLM training. Its unified DTensor abstraction and modular framework ensure that context scaling and other advancements can be seamlessly incorporated, making it a robust solution for future research and deployment.

W2:

We appreciate the reviewer’s insightful question. Identifying the best configuration parameters for a 4D parallel system, along with compute and memory optimizations, requires a modular, composable, and extensible framework as a foundation. This is the primary contribution of LEGOSCALE—providing a unified system that enables users to systematically explore and evaluate complex configurations.

Optimizing performance involves efficiently navigating a design space influenced by factors such as checkpointing strategies (AC vs. SAC), parallelism degrees (DP/CP/TP/PP), memory layouts (HSDP vs. FSDP), and pipeline schedules (e.g., 1F1B, Zero-Bubble). This requires estimating GPU compute times, communication latencies, and memory footprints, followed by algorithms that recommend a configuration, which presents a new research direction.

To assist users, we provide practical guidelines on choosing the parallelism strategies in Section 3.3. Additionally, we are prototyping tools like AutoSAC and AutoFSDP to automate configuration tuning.

While fully solving the optimization problem remains an open challenge, LEGOSCALE offers the foundation researchers need to experiment and develop solutions for distributed training.

We sincerely thank the reviewer for acknowledging the strengths of our work and highlighting them. We appreciate the detailed and critical feedback on our work. Please consider our response labeled as Wi, which aims to address weaknesses i, while Qi aims to address Question i, pointed out by the reviewer.

W1:

As the reviewer points out, many techniques used in LEGOSCALE are well-documented. However, LEGOSCALE’s key research contribution lies in analyzing and identifying a common set of design principles across parallelism and optimization techniques. These principles unify diverse methods into a cohesive system. The engineering contribution is in enforcing these principles across individual techniques, resulting in a modular, composable, and flexible distributed system that enables seamless scaling, efficient deployment, and integration of new techniques. These contributions have shaped PyTorch features like FSDP2, Async-TP, PP, and DTensor advancements.

LEGOSCALE’s Design Principles

LEGOSCALE adopts a top-down approach, enforcing modularity, non-intrusiveness, and extensibility. By leveraging DTensor and DeviceMesh, it provides a unified framework for diverse parallelisms and optimizations, ensuring scalability and compatibility with PyTorch.

Advancements to DTensor

DTensor sharding was extended to support n-D parallelism, integrating TP, CP, and PP. Compatibility with torch.compile was added for compiler optimizations, and checkpointing efficiency improved with robust state dict support. Production readiness was bolstered by resolving critical bugs.

Contribution to PyTorch FSDP2

FSDP2 replaced FlatParameter with DTensors sharded on dim-0, enabling flexible parameter handling (e.g., freezing, fp8 all-gather). It integrates with torch.compile, reduces GPU memory usage by eliminating CPU synchronization, and simplifies checkpointing with sharded state dicts.

Contribution to PyTorch Async-TP/CP

Async-TP introduced fused operators like fused_all_gather_matmul to reduce synchronization overheads, enabling fused compute and communication. This design facilitates new parallelisms like CP and minimizes bottlenecks with CUDA streams and SymmetricMemory.

Contribution to PyTorch PP

Pipeline Parallelism incorporates LEGOSCALE’s principles, enabling dynamic schedules, fine-grained stage control, and the addition of new schedules while integrating with TP and DP.

A Unified Template for Distributed Training

LEGOSCALE provides a unified design template, simplifying exploration of configurations and enabling rigorous evaluation of techniques. It unifies, modularizes, and extends distributed methods, shaping PyTorch’s library and advancing scalable, flexible systems.

W2:

We thank the reviewer for their insightful comment on how LEGOSCALE’s modular design supports future optimizations and parallelism strategies without compromising performance. Since submitting the paper, we have integrated Context Parallelism (CP) and added three new Pipeline Parallelism (PP) schedules: Looped-BFS, Flexible-Interleaved-1F1B, and ZeroBubble. The camera-ready version will include case studies demonstrating how to extend LEGOSCALE and its design principles. Below, we summarize how CP was enabled.

Context Parallelism (CP) Integration

LEGOSCALE integrates CP through:

1. Module Wrapper: Distributes the Attention module, similar to TP.
2. Context Manager: Dynamically replaces scaled_dot_product_attention with context-parallel operators, requiring no model code changes.

Under the hood, CP shards DTensor along the sequence dimension across the CP device mesh. It extends DTensor’s dispatcher to support context-specific operations like ring-attention and causal attention load balancing, ensuring efficient operation. LEGOSCALE ensures full compatibility of CP with FSDP, TP, PP, optimizations (e.g., torch.compile), and DCP, showcasing its extensibility and seamless integration of new techniques.

Context Scaling Results (Llama3-8B, 64 H100 GPUs, TP=8, DPxCP=8)

We thank the reviewer for underscoring the importance of our work, recognizing its performance, appreciating our design, and providing detailed and critical feedback on our work. Please consider our response labeled as Wi, which aims to address weaknesses i, while Qi aims to address Question i, pointed out by the reviewer.

W1 and Q1:

LEGOSCALE’s key research contribution lies in analyzing and identifying a common set of design principles across parallelism and optimization techniques. These principles unify diverse methods into a cohesive system. The engineering contribution is in enforcing these principles across individual techniques, resulting in a modular, composable, and flexible distributed system that enables seamless scaling, efficient deployment, and integration of new techniques. These contributions have shaped PyTorch features like FSDP2, Async-TP, PP, and DTensor advancements.

LEGOSCALE’s Design Principles

LEGOSCALE adopts a top-down approach, enforcing modularity, non-intrusiveness, and extensibility. By leveraging DTensor and DeviceMesh, it provides a unified framework for diverse parallelisms and optimizations, ensuring scalability and compatibility with PyTorch.

Advancements to DTensor

DTensor sharding was extended to support n-D parallelism, integrating TP, CP, and PP. Compatibility with torch.compile was added for compiler optimizations, and checkpointing efficiency improved with robust state dict support. Production readiness was bolstered by resolving critical bugs.

Contribution to PyTorch FSDP2

FSDP2 replaced FlatParameter with DTensors sharded on dim-0, enabling flexible parameter handling (e.g., freezing, fp8 all-gather). It integrates with torch.compile, reduces GPU memory usage by eliminating CPU synchronization, and simplifies checkpointing with sharded state dicts.

Contribution to PyTorch Async-TP/CP

Async-TP introduced fused operators like fused_all_gather_matmul to reduce synchronization overheads, enabling fused compute and communication. This design facilitates new parallelisms like CP and minimizes bottlenecks with CUDA streams and SymmetricMemory.

Contribution to PyTorch PP

Pipeline Parallelism incorporates LEGOSCALE’s principles, enabling dynamic schedules, fine-grained stage control, and the addition of new schedules while integrating with TP and DP.

A Unified Template for Distributed Training

LEGOSCALE provides a unified design template, simplifying exploration of configurations and enabling rigorous evaluation of techniques. It unifies, modularizes, and extends distributed methods, shaping PyTorch’s library and advancing scalable, flexible systems.

W2:

We have been actively involved with PyTorch developers and have significantly contributed to the design and enhancement of the PyTorch distributed library, as discussed earlier. By making LEGOSCALE fully PyTorch-native, we ensure seamless compatibility with upstream PyTorch developments, enabling users to benefit from the latest features and support for new GPU models. Despite offering a full range of 4D parallel training capabilities and optimizations, LEGOSCALE’s codebase is only 9K lines, making it highly maintainable and adaptable for future advancements.

Q2:

Thank you for the question. We have fully open-sourced the LEGOSCALE system on GitHub (anonymized to comply with ICLR's double-anonymous requirements). The framework has garnered significant interest, with over 2,600 stars, 204 forks, 469 closed pull requests, and 35 active contributors. LEGOSCALE enables researchers and practitioners to explore, develop, and scale advanced parallelism and optimization techniques in a production-grade pre-training system.

LEGOSCALE has been validated on multiple platforms, including the public cloud (AWS), on-premises clusters of a Fortune 500 company, a leading private organization, and a world-renowned university. Results are designed to be easily replicable across various cluster configurations. The convergence graphs can be found here.

We are committed to actively maintaining and supporting LEGOSCALE for the foreseeable future.

Since submitting the paper, we have incorporated Context Parallelism (CP) and added three additional Pipeline Parallelism (PP) schedules (Looped-BFS, Flexible-Interleaved-1F1B, ZeroBubble).

Another key feature is LEGOSCALE’s model-agnostic and general design. While demonstrated with the Llama model, LEGOSCALE is not model-specific. Its techniques are general-purpose, allowing seamless integration of other models without changes to the system.

Our goal in submitting this paper is to raise awareness within the research community about a system that empowers advanced research on pre-training large language models (LLMs) at scale, fostering innovation in this critical area.

We thank the reviewer for providing detailed and critical feedback on our work. Please consider our response labeled as Wi, which aims to address weaknesses i, pointed out by the reviewer.

W1:

As the reviewer points out, many techniques used in LEGOSCALE are well-documented. However, LEGOSCALE’s key research contribution lies in analyzing and identifying a common set of design principles across parallelism and optimization techniques. These principles unify diverse methods into a cohesive system. The engineering contribution is in enforcing these principles across individual techniques, resulting in a modular, composable, and flexible distributed system that enables seamless scaling, efficient deployment, and integration of new techniques. These contributions have shaped PyTorch features like FSDP2, Async-TP, PP, and DTensor advancements.

LEGOSCALE’s Design Principles

LEGOSCALE adopts a top-down approach, enforcing modularity, non-intrusiveness, and extensibility. By leveraging DTensor and DeviceMesh, it provides a unified framework for diverse parallelisms and optimizations, ensuring scalability and compatibility with PyTorch.

Advancements to DTensor

DTensor sharding was extended to support n-D parallelism, integrating TP, CP, and PP. Compatibility with torch.compile was added for compiler optimizations, and checkpointing efficiency improved with robust state dict support. Production readiness was bolstered by resolving critical bugs.

Contribution to PyTorch FSDP2

FSDP2 replaced FlatParameter with DTensors sharded on dim-0, enabling flexible parameter handling (e.g., freezing, fp8 all-gather). It integrates with torch.compile, reduces GPU memory usage by eliminating CPU synchronization, and simplifies checkpointing with sharded state dicts.

Contribution to PyTorch Async-TP/CP

Async-TP introduced fused operators like fused_all_gather_matmul to reduce synchronization overheads, enabling fused compute and communication. This design facilitates new parallelisms like CP and minimizes bottlenecks with CUDA streams and SymmetricMemory.

Contribution to PyTorch PP

Pipeline Parallelism incorporates LEGOSCALE’s principles, enabling dynamic schedules, fine-grained stage control, and the addition of new schedules while integrating with TP and DP.

A Unified Template for Distributed Training

LEGOSCALE provides a unified design template, simplifying exploration of configurations and enabling rigorous evaluation of techniques. It unifies, modularizes, and extends distributed methods, shaping PyTorch’s library and advancing scalable, flexible systems.

W3:

We thank the reviewer for their insightful observation. Although recent pre-trained LLM checkpoints primarily use FP16 or BF16, results with FP8 have been promising, and LEGOSCALE is designed to be future-ready. In LEGOSCALE, we adopt a cautious approach by applying FP8 selectively to linear layers, ensuring significant performance improvements while maintaining training stability.

Our experiments show a 49% improvement in throughput with FP8 compared to BF16, alongside identical loss convergence curves. To demonstrate FP8’s efficacy, we compare training speed and loss convergence under equivalent configurations, highlighting its feasibility and benefits. While LEGOSCALE is primarily designed for pre-training, it also enables future exploration of fine-tuning on FP8-pretrained models.

We will include detailed loss convergence curves in the appendix to support our findings further. Our graphs can be found here: FP-8-Loss-Convergence

W4:

We thank the reviewer for their observation regarding LEGOSCALE’s use of PyTorch’s Distributed Checkpointing (DCP) and will include additional details about its design and performance in the appendix. LEGOSCALE addresses the challenge of production-critical features being dispersed across third-party libraries by building on the unified abstraction of DTensor, simplifying state dicts and enabling an efficient DCP implementation that is compatible across diverse parallelisms.

Comparison with Other Systems

• Gemini: Tightly coupled with DeepSpeed, unavailable to native PyTorch users, supports only ZeRO-3 (FSDP), and lacks compatibility with TP, CP, and PP.
• ByteCheckpoint: Adds asynchronous operations and disaggregated storage but relies on DCP’s foundation without significant enhancements for advanced parallelism.

LEGOSCALE’s DCP Integration

LEGOSCALE enhances DCP by integrating it with DTensor to standardize shard representations across TP, CP, PP, and FSDP. Asynchronous checkpointing overlaps storage operations with training, reducing overhead by 19x compared to synchronous methods. A zero-overhead prototype further overlaps GPU-to-CPU transfers with forward and backward passes, achieving a 5x additional reduction and reducing total checkpoint time to under one second. This unified, efficient design ensures LEGOSCALE’s checkpointing is robust, scalable, and production-ready for large-scale distributed training.

Performance Graphs can be found here: DCP-Performance

W5:

We acknowledge the reviewer’s observation that frameworks like JAX and ONNX leverage compiler optimizations to enhance training throughput and resource utilization. PyTorch and JAX, however, differ fundamentally in their programming paradigms:

• PyTorch uses an imperative style with dynamic computation graphs, enabling immediate execution. This approach is intuitive, user-friendly, and facilitates debugging and fast prototyping, making it highly accessible.
• JAX, in contrast, employs a functional paradigm with static computation graphs, achieving performance gains via XLA but introducing a steeper learning curve for users unfamiliar with functional programming.

PyTorch’s imperative style, combined with its Pythonic syntax, active community, and versatility, has made it the leading framework for deep learning, with over 20K papers published in the last 12 months using PyTorch.

LEGOSCALE bridges the gap between these paradigms by introducing composability and modularity within PyTorch’s imperative environment. This approach retains the ease of programming and fast prototyping of PyTorch while achieving performance benefits akin to compiler-driven frameworks like JAX.

To address the reviewer’s point, torch.compile, introduced in PyTorch 2.0, enables compiler optimizations without compromising accessibility. LEGOSCALE enhances DTensor compatibility with torch.compile, allowing parallelism strategies like TP, CP, and FSDP to fully leverage these optimizations. This ensures state-of-the-art memory and computational efficiency while maintaining PyTorch’s flexibility. LEGOSCALE delivers a performant and modular system for pretraining large language models, balancing usability with significant efficiency gains.

Great question. No, the optimizer is still used in FP32. The optimizer states are sharded across all GPUs, so as the scale of training increases the cost of optimizer updates decreases linearly. For completeness, we use float8 linear layers, and float8 all gather in conjunction with FSDP2. For this training, we use float8 per tensor (tensorwise) scaling granularity rather than rowwise. We leverage torch.compile to ensure that we get maximum performance gains.