Kati8/ICLR_2023_2025 huggingface dataset:
authors: Wanchao Liang
authors: Tianyu Liu
authors: Will Constable
authors: Andrew Gu
authors: Chien-Chin Huang
authors: Iris Zhang
authors: Wei Feng
authors: Howard Huang
authors: Junjie Wang
authors: Sanket Purandare
authors: Gokul Nadathur
authors: Stratos Idreos
reviews: 
Comments: 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.
Strengths: Thank you for submitting this paper to ICLR! I loved reading this paper.

1. Important problem. LLM training is challenging and resource-consuming.

2. 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.

2. Limited model support. I would like to see more application on MoE model, Multi-modal LLM, etc.

3. 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.
Recommendation: 10
Reviewer confidence: 5
Soundness/correctness: 4
Clarity: 3
reviews: 
Comments: 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.
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. 
- 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.
- 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. 
- 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.
Recommendation: 6
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: This paper describes a feature-rich PyTorch native hardware-software system for efficiently training of LLMs.
Strengths: The paper describes a useful tool that addresses an important problem.

The tool substantially improves training time.

The system appears to have sufficient features for production use.

The paper has some tutorial value in that it briefly explains a wide range of features appropriate for pre-training systems.

I 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.

Interest 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.

There 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.

With the exception of the sensitivity analysis in the tables on page 8, the paper does not contain scientific findings.

Figures 1 and 2 are essentially impossible to read without magnification, which circumvents conference page limits.
Recommendation: 6
Reviewer confidence: 4
Soundness/correctness: 4
Clarity: 3
reviews: 
Comments: 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.
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. 
- By incorporating pytorch-native implementations of training optimizations, LegoScale lowers the bar for integration into existing pytorch workflows, likely fostering adoption. 
- 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. 
- 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. 
- 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.
- 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.
Recommendation: 6
Reviewer confidence: 5
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: 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. 
Authors 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.  
The 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.
Strengths: 1. Using the latest Llama 3.1 family of models is promising
2. 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.
2. 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.
3. 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.
4. 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. 
5. 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.
Recommendation: 5
Reviewer confidence: 4
Soundness/correctness: 2
Clarity: 2
reviews: 
Comments: This work presents LEGOSCALE, a PyTorch-integrated open-source framework for distributed training that consolidates existing best practices while introducing novel capabilities. The system provides modular and composable support for three-dimensional parallelism alongside elastic scaling mechanisms to handle variable computational loads. Through its design as an accessible and extensible toolkit, LEGOSCALE makes large language model pre-training more available to diverse research communities, enabling broader participation and fostering continued progress in LLM development.
Strengths: 1. Addresses a critical challenge. The computational and resource demands of LLM training represent a significant barrier in the field.

2. Well-motivated and relevant technical approach. The PyTorch-native design aligns with community requirements and distinguishes itself from comparable systems such as Megatron.
Weaknesses: 1. The paper covers extensive technical components within its page constraints, potentially creating comprehension difficulties for readers without specialized knowledge in machine learning systems architecture.

2. Restricted scope of model coverage. Broader evaluation and implementation across diverse model families including Mixture-of-Experts architectures, multimodal language models, and other emerging paradigms would strengthen the contribution.

3. The parallelism strategy lacks comprehensive coverage relative to contemporary standards in 2024. Enhanced incorporation of additional techniques such as context parallelism (referenced within the manuscript), expert parallelism, and other emerging methods would better position the work.
Recommendation: 10
Reviewer confidence: 5
Soundness/correctness: 4
Clarity: 3