Lumos: Efficient Performance Modeling and Estimation for Large-scale LLM Training

Lumos: Efficient Performance Modeling and Estimation for Large-scale LLM Training

Speakers: Mingyu Liang, PhD Student, Cornell University (presenting on behalf of Google and Meta collaborators); Hiwot Tadese Kassa; Wenyin Fu; Louis Feng; Christina Delimitrou

Conference: MLSys 2025

YouTube: https://slideslive.com/39042997

Overview

The rapid advancement and increasing scale of Large Language Models (LLMs) necessitate highly efficient training processes. Optimizing the performance of these colossal models is paramount, but it first requires a profound understanding of their intricate execution behavior. This talk introduces Lumos, a novel trace-driven performance modeling and estimation framework designed specifically for large-scale LLM training. Developed through a collaboration between Cornell, Google, and Meta, Lumos addresses critical limitations in existing performance modeling approaches, particularly their inability to accurately capture the complex interplay of compute, communication, and overlap in modern LLMs.

Lumos distinguishes itself by constructing a highly fine-grained execution graph from runtime traces, capturing detailed operator-level activities on both CPU and GPU, along with their intricate dependencies. Unlike prior work that often falls short in accuracy, Lumos achieves an impressive 3% error in replaying real-world LLM training performance. Crucially, it extends beyond mere replay, offering a powerful capability for "what-if" analysis: users can modify an existing execution graph to reflect new model configurations or parallelism strategies, and Lumos can then predict the performance of these unseen scenarios without the need for costly and time-consuming real hardware runs.

The significance of Lumos lies in its potential to revolutionize how LLM developers and infrastructure teams optimize training. By providing accurate performance predictions for various configurations—such as scaling parallelism, adjusting model parameters, or exploring sharding strategies—Lumos enables informed decision-making, significantly reducing development cycles, resource consumption, and the overall cost associated with training state-of-the-art LLMs. This work lays a robust foundation for future downstream optimization studies, making the pursuit of ever-larger and more efficient language models a more tractable endeavor.

Background

▶ Watch: Introduction and limitations of current LLM performance models (0:00)

The training of Large Language Models (LLMs) represents one of the most computationally intensive tasks in modern machine learning. As models like GPT-3 and beyond scale to hundreds of billions or even trillions of parameters, their efficient execution becomes a grand challenge in systems design. Optimizing this performance requires a deep understanding of how these models behave across distributed hardware, a task complicated by the intricate interplay of computation, data movement, and synchronization across hundreds or thousands of accelerators.

Historically, two general approaches have emerged for performance modeling in this domain. The first involves developing analytical models, which offer quick, high-level estimations. While useful for initial sanity checks, these models often lack the fidelity to capture the fine-grained nuances of complex distributed execution, leading to rough and potentially misleading predictions. The second, more accurate approach involves building fine-grained operator-level execution graphs. These graphs aim to precisely map out every computational operation, communication primitive, and their dependencies, providing a detailed blueprint of the model's execution flow. This method, while more complex to construct and simulate, promises superior accuracy.

Despite the promise of fine-grained execution graphs, existing state-of-the-art solutions have struggled to adequately capture the full complexity inherent in today's LLM training. For instance, tools like DeepProf, a prominent work in this field, attempt to model performance through execution graphs. However, as demonstrated by the presented work, when DeepProf's replay results for a full-size GPT-3 training iteration are compared against real runs, significant "gaps" emerge in the breakdown of compute, communication, and their overlap. These discrepancies indicate that even sophisticated prior methods leave substantial room for improvement in accurately characterizing LLM behavior. The limitations of these prior works stem from their inability to fully disentangle and correctly model the myriad of dependencies—both within and between CPU and GPU tasks, and across multiple parallel streams—that dictate the true execution profile of large-scale distributed LLM training. Lumos directly addresses these shortcomings by meticulously capturing these intricate dependencies, thereby achieving a level of accuracy critical for effective optimization.

Key Findings

▶ Watch: Detailed explanation of task dependencies (CPU, GPU, inter-task) (3:00)

Lumos presents several key findings that significantly advance the state of the art in LLM performance modeling and estimation:

• Superior Performance Replay Accuracy: Lumos achieves an impressive 3% average error in replaying the end-to-end training time across a wide range of GPT-3 model configurations and parallelism strategies. This accuracy is a substantial improvement over existing baselines like DeepProf, which exhibit noticeable "gaps" when comparing their replayed execution breakdowns to real-world runs. Lumos's ability to precisely capture the detailed breakdown into compute, communication, and their overlap is a critical differentiator.
• Accurate Prediction for New Model Configurations: A core contribution of Lumos is its capability to accurately predict the performance of LLMs under new configurations that have not been physically executed. This includes scenarios involving:
• Scaling Parallelism Strategies: Lumos can predict performance when scaling data and pipeline parallelism (e.g., from a base 2x2x4 configuration). This allows users to explore different sharding and parallelism choices without incurring the massive computational cost of running each variant on actual hardware.
• Varying Model Parameters: The framework can accurately estimate performance for models with different numbers of layers, hidden dimension sizes, and overall model sizes (e.g., predicting performance for larger GPT-3 variants based on a 15 billion parameter base model).
• Efficient "What-If" Analysis through Graph Manipulation: Lumos enables powerful "what-if" questions to be answered by allowing users to directly modify an existing execution graph (derived from a base configuration) to reflect changes for new model configurations. This graph manipulation, combined with simulation, provides a resource-efficient alternative to costly real-hardware experimentation, accelerating the optimization cycle for LLMs.
• Comprehensive Dependency Modeling: The framework's high accuracy stems from its meticulous capture of a fine-grained execution graph, which includes identifying and modeling all possible relationships among CPU and GPU tasks. This comprehensive approach to dependency modeling, spanning intra-thread, inter-thread, CPU-to-GPU, GPU-to-CPU, intra-stream, and crucial inter-stream dependencies, is fundamental to its predictive power.

Technical Deep Dive

▶ Watch: Evaluation setup: GPT-3, H100 GPUs, DeepProf baseline (8:00)

Lumos's technical prowess lies in its meticulous, trace-driven approach to constructing and manipulating fine-grained execution graphs. The framework's methodology can be broken down into several key stages, each contributing to its high accuracy and predictive capabilities.

The overall approach begins with instrumentation and trace capture. Lumos inserts hooks into the LLM training code to capture runtime traces from a real execution. These traces serve as the raw data from which the model's behavior is distilled. From these traces, the system then extracts key meta-information, which includes identifying all individual operators or tasks executed, along with the precise dependencies between them. This forms the foundation of the execution graph. A crucial capability of Lumos is the ability to modify these graphs to reflect new model configurations or parallelism strategies, enabling "what-if" analysis. Finally, these constructed and potentially modified graphs are fed into simulators to estimate performance.

The essence of LLM execution modeling in Lumos revolves around its detailed execution graph. This graph comprises two primary types of tasks:

1. CPU Tasks: These include high-level PyTorch operators and underlying CUDA runtime events, which orchestrate GPU operations.
2. GPU Tasks: These are the low-level CUDA kernels that perform the actual computation on the accelerator.

A core innovation of Lumos is its comprehensive identification of all possible relationships and dependencies among these tasks:

• Dependencies within CPU Tasks:
• Intra-thread Dependency: Operators assigned to the same CPU thread are inherently sequential, executing one after another.
• Inter-thread Dependency: In complex frameworks like PyTorch, different CPU threads might be dedicated to distinct phases (e.g., forward pass, backward pass). Dependencies exist between these threads, where one thread's progress might block or be contingent on another's.

• Dependency from CPU to GPU Tasks:
• Launching Dependency: GPU kernels are launched by corresponding CUDA runtime calls from the CPU. Lumos captures this control dependency by using correlation IDs that tag both the CUDA runtime event and the launched kernel in the trace, establishing a direct link between the CPU trigger and the GPU execution.

• Dependency from GPU to CPU Tasks:
• Synchronization Dependency: CPU threads often need to wait for GPU kernels to complete before proceeding. This occurs during explicit synchronizations (e.g., cuda.synchronize()) or implicit synchronizations, where the CPU thread is blocked until one or more kernels on specific CUDA streams finish their execution.

• Dependencies among GPU Tasks: These are particularly critical for modern LLMs, which extensively utilize multiple CUDA streams for concurrent execution.
• Intra-stream Dependency: Kernels dispatched to the same CUDA stream must execute sequentially, respecting the order of submission.
• Inter-stream Dependency: This is a highly significant dependency for LLMs. Kernels on one CUDA stream may be blocked by or depend on the completion of kernels from other streams. This dependency directly dictates the degree of parallelism achievable between computation and communication, or between different stages of a pipeline. Accurately modeling these inter-stream relationships is paramount for predicting the end-to-end performance of LLMs that leverage complex overlapping strategies across multiple streams.

Beyond replaying existing performance, Lumos empowers users to answer "what-if" questions by modifying existing execution graphs. This capability is fundamental for exploring new configurations without real hardware. The framework currently supports changing model parameters (e.g., number of layers, hidden dimension) and parallelism strategies (e.g., data parallelism, pipeline parallelism, tensor parallelism).

Consider changing pipeline parallelism as an example. This involves several sophisticated graph manipulations:

1. Update Pipeline Schedule: The first step is to redefine the order of forward and backward passes for each micro-batch across the new pipeline stages.
2. Partition or Aggregate Layers: Based on the new pipeline schedule, existing model layers are either partitioned into smaller segments or aggregated to form new pipeline stages.
3. Insert Collective Operators: To ensure correct data transfer between these newly defined stages, appropriate collective communication operators (e.g., all-reduce, all-gather, send, recv) are strategically inserted into the graph.

A crucial aspect of graph modification is updating the execution duration of individual operations. If model parameters or parallelism strategies change, the shapes and sizes of tensors processed by operators, or the communication patterns, will also change. Lumos addresses this by leveraging auxiliary performance models:

• For compute operations, roofline models are employed to update the estimated execution time based on new input shapes.
• For communication operations, in-house performance models are utilized to predict communication times under the new configurations and network environments. This ensures that the simulated performance reflects the actual cost of operations under the modified settings. If a completely new operator appears due to graph manipulation, Lumos would require an external mechanism (e.g., another performance model or a small benchmark run) to estimate its execution time.

Experimental Setup & Results

▶ Watch: Lumos accuracy for performance replay and execution breakdown (8:40)

Lumos's effectiveness was rigorously evaluated using the GPT-3 model, a cornerstone of the MLPerf training benchmark. The experiments involved creating several variants of GPT-3, with model weights ranging from a manageable 15 billion parameters up to the full 175 billion parameters. This range allowed for comprehensive testing across different scales of LLMs.

The evaluation was conducted on a production cluster utilizing up to 512 NVIDIA H100 GPUs. This high-performance distributed environment provided a realistic setting for assessing Lumos's capabilities in large-scale LLM training scenarios. For comparison, DeepProf, a state-of-the-art performance modeling tool based on execution graphs, was chosen as the primary baseline.

The experimental results highlighted Lumos's superior accuracy in two key areas:

1. Performance Replay Accuracy:

• Lumos demonstrated remarkable accuracy in replaying the end-to-end training time for various model configurations and parallelism strategies. Across all tested scenarios, Lumos consistently achieved an average error of only 3%.
• Beyond just the total time, the framework also accurately captured the underlying execution breakdown, precisely reflecting the proportions of compute, communication, and their overlap. This detailed fidelity is crucial for understanding performance bottlenecks and is where Lumos significantly outperforms baselines like DeepProf, which showed observable "gaps" in their breakdown analysis compared to real runs.

1. Performance Prediction for New Model Configurations:

• Scaling Parallelism: In one scenario, Lumos was tasked with predicting performance when scaling data and pipeline parallelism from a base configuration (e.g., 2x2x4 tensor-pipeline-data parallelism). By modifying the base graph to reflect these parallelism changes and feeding it into the simulator, Lumos achieved "pretty good accuracy" in predicting the performance for these unseen configurations. This demonstrates its utility in exploring optimal parallelism strategies without extensive hardware testing.
• Varying Model Parameters: Another experiment involved predicting performance for GPT-3 variants with different model parameters. Starting with a base model of 15 billion parameters, Lumos predicted the performance for variants with altered numbers of layers, hidden dimension sizes, and overall model sizes. Again, the framework achieved "good accuracy," showcasing its ability to extrapolate performance for models with different architectural dimensions.

The robustness of Lumos's prediction capabilities relies on its auxiliary models for operator timing. When configurations change (e.g., hidden size, number of devices for communication), the execution time for individual compute and communication operations must be updated. Lumos employs roofline models for compute operations to estimate new durations based on changed shapes, and in-house performance models for communication operations to predict times under new network environments and collective patterns. While the current implementation of network models doesn't explicitly account for switch-level congestion, it provides a solid foundation for predicting communication costs, with the possibility of incorporating more advanced network models for greater fidelity in the future.

Practical Implications

▶ Watch: Predicting performance for varying model parameters (layers, dimensions) (10:00)

Lumos offers profound practical implications for various stakeholders involved in the development, deployment, and optimization of large-scale LLMs.

For Practitioners and Infrastructure Teams:

• Reduced Resource Consumption and Time: The most significant implication is the ability to conduct extensive "what-if" analyses without consuming vast amounts of expensive GPU cluster time. Exploring different parallelism strategies, sharding configurations, or hardware setups can be simulated efficiently, drastically cutting down on the trial-and-error cycle that currently plagues LLM development. This translates directly to lower operational costs and faster iteration times.
• Informed Decision-Making: Lumos provides accurate performance predictions, enabling infrastructure teams to make data-driven decisions about hardware provisioning, network topology requirements, and optimal resource allocation for LLM training jobs. It can help answer questions like "How many GPUs do I need for X throughput?" or "Which parallelism strategy scales best for my model size on this cluster?"
• Bottleneck Identification: By accurately replaying and predicting the breakdown of compute, communication, and overlap, Lumos helps pinpoint performance bottlenecks more precisely. This allows teams to focus optimization efforts on the most impactful areas, whether it's optimizing kernel execution, improving communication patterns, or enhancing scheduling.

For Model Builders and Deployers:

• Optimized Model Design and Scaling: Model developers can use Lumos early in the design phase to evaluate how architectural choices (e.g., number of layers, hidden dimensions) will impact training performance on target hardware. This allows for proactive optimization rather than retrospective debugging.
• Efficient Deployment Strategy: When deploying LLMs, understanding how they will perform across different hardware configurations and parallelism schemes is critical. Lumos provides a tool to predict these behaviors, guiding the selection of efficient deployment strategies for inference or fine-tuning, even if the primary focus of the talk is training.

Tradeoffs and Limitations:

• Initial Tracing Requirement: Lumos is a trace-driven framework, meaning it requires an initial run on real hardware to capture the base execution trace. While this is a one-time cost for a given base configuration, it's still a prerequisite.
• Dependency on Auxiliary Models: The accuracy of predictions for new configurations heavily relies on the fidelity of the auxiliary performance models used to estimate individual operator and communication durations (e.g., roofline models for compute, in-house models for networking). If these auxiliary models are inaccurate or incomplete for certain operations or hardware, the overall prediction accuracy can suffer.
• Network Environment Complexity: While Lumos incorporates in-house models for communication timing, the complexity of real-world network environments, especially concerning congestion from other users or dynamic traffic patterns, is not yet fully captured directly within Lumos's core modeling. Its accuracy here depends on how sophisticated the underlying network performance model is.
• Handling Unseen Operators: If a graph modification introduces an operator that was not present in the original trace and for which no auxiliary model exists, Lumos would require an external mechanism (e.g., profiling the new operator in isolation or developing a new specific model for it) to estimate its execution time. This highlights a need for a comprehensive library of operator performance models.
• Scope of "What-If" Analysis: Currently, Lumos supports changes to model parameters and parallelism strategies. Extending it to cover other aspects, such as different optimizers, mixed precision settings, or novel hardware architectures, would require further development and potentially new modeling techniques.

Despite these limitations, Lumos represents a significant step forward in making LLM performance optimization more systematic, efficient, and accessible. Its ability to provide accurate insights without extensive physical experimentation is a game-changer for the field.

Key Takeaways

• High-Accuracy Performance Modeling: Lumos is a trace-driven framework that achieves an impressive 3% error rate in replaying end-to-end training time and detailed execution breakdowns (compute, communication, overlap) for large-scale GPT-3 models.
• Efficient "What-If" Analysis: It enables accurate performance prediction for new model configurations and parallelism strategies through intelligent graph manipulation, eliminating the need for costly and time-consuming real hardware runs.
• Comprehensive Dependency Capture: Its high fidelity stems from meticulously identifying and modeling all types of dependencies—intra-CPU, inter-CPU, CPU-to-GPU, GPU-to-CPU, intra-stream, and critical inter-stream GPU dependencies—that govern LLM execution.
• Graph Manipulation for Parallelism: Lumos can modify execution graphs to reflect changes in parallelism, such as pipeline parallelism, by updating schedules, partitioning layers, and inserting collective communication operators.
• Auxiliary Models for Operator Timing: The framework leverages auxiliary performance models (e.g., roofline models for compute, in-house models for communication) to estimate operator durations when model parameters or configurations change.
• Reduced Optimization Costs: Lumos significantly reduces the time and resource consumption associated with optimizing LLM training, empowering developers and infrastructure teams to make data-driven decisions more efficiently.

About the Speaker(s)

The talk was presented by Mingyu Liang, who is a PhD student at Cornell University. This work was conducted during his time as a PhD student at Cornell and involved a collaborative effort with Meta. Mingyu Liang is also associated with Google, indicating a strong background in industry-leading AI/ML research and development. The list of collaborators on this work includes Hiwot Tadese Kassa, Wenyin Fu, Louis Feng, and Christina Delimitrou, highlighting a team effort from prominent institutions in the field.