Efficient Pre-Training of LLMs via Topology-Aware Communication Alignment on More Than 9600 GPUs

w1 + q1 + q2: Presentation issues. But the understanding is fully formula-based and fragmented to me; it could be better if the authors use a concrete paragraph to explain intuitions and summarise the differences.

We apologize for the presentation issues, and will clarify in the text.

Multi-core CPU parallelism (as in BriskStream) and inter-GPU parallelism (as in LLM training) both aim to use multiple processing units to speed up computation. However, they are fundamentally different due to their underlying memory architecture, communication mechanisms, and architectural granularity.

Memory architecture: Multi-core Parallelism is primarily concerned with efficiently managing a shared memory space among cores, with the main challenge being the performance penalty of Non-Uniform Memory Access (NUMA).

Inter-GPU Parallelism is primarily concerned with orchestrating computation across multiple devices that have their own separate and distributed memory, with the main challenge being the high cost of explicitly transferring data between them.

Communication mechanisms: Multi-core Parallelism happens implicitly through memory. When a producer operator writes a tuple to a queue and a consumer operator reads it, they are accessing a shared location in RAM.

Inter-GPU Parallelism is explicitly managed via libraries like NCCL (NVIDIA Collective Communications Library).

Architectural granularity: Multi-core Parallelism is designed for task parallelism on a single tensor computation level. For example, matrix multiplication for tensors operations.

Inter-GPU Parallelism is designed for massive data parallelism on homogeneous, throughput-oriented, and highly regular workloads. For example, a group of communication processes.

w2: Lack of discussion about data compression and emerging networking architectures, which are also impactful in distributed training.

We will add more discussion on data compression pipeline and networking architectures. Reviewer is right about these factors can impact distributed training.

Data compression and processing. Data processing typically takes place on CPU and causes GPU idle. We adopt MegaScale's optimization, by overlapping data processing with gradient synchronization. They can be fully overlapped because they take place in different devices and the data processing overhead can be hidden.

Emerging networking architectures. Emerging network architectures have significant impacts on the future design of scheduling algorithms. We think the presented algorithm is in principle able to handle other netowrk topology.

We observe that production GPU clusters follow the multi-tier design (Alibaba’s HPN, NVIDIA’s DGX). The collective communication is bottlenecked by the highest hierarchy - the inter-minipod communication in our case. In our case, we observe little impact of intra-minipod topology, however, we can also solve the MIP in intra-minipod level if we can observe any performance degradation.

There are more dedicated network architectures exist for LLM training, e.g. rail-only. However, the inability of any-to-any connectivity make it unlikely to be adopted in production training clusters in the near future.

w1: The description of the ILP formulation is unclear in parts.

We apologize for the presentation and will add more clarification in the text.

In Equation (2) each node is a scheduling unit, while in Equation (4) the communication group (e.g. PP) becomes a scheduling unit.

A simple example. Suppose we have $m$ DP group and $n$ PP group, after the coarsening, the scheduling problem becomes “how to put the $n$ unit into $k$ minipods?” - we are allowed to spread one unit across multiple minipods, but we want to minimize the spread, hence introducing the $T$ auxiliary variable. On the other hand, DP groups’ spread is “orthogonal” to PP group – GPUs in the same stage of different PP groups belong to the same DP group, so if each unit (PP) is packed into a different minipod (the optimal $T$), the spread of DP groups is maximized. Thus, we also minimize the overall usage of minipod, which controls the spread of DP groups.

Combining the above, Equation (4) is our optimizing objective function.

Equation (5) constraints the maximum spread of PP groups less or equal to $T$, and $T$ is being minimized.

Equation (6) says that the allocation cannot exceed the capacity of a minipod, e.g. if a minipod has $10$ nodes free, we cannot allocate more than that. The capacity is normalized by the number of node in a PP group.

Equation (7) says that while we can spread a PP group across minipods, the summation of the allocation percentage must be equal to $1$.

Equation (8) says that if we allocate to a minipod, that it is considered used in the solution, i.e. $y_j = 1$.

Finally, Equation (9) and (10) define the range of the binary variable $y_j$ and $s_{ij}$, and continuous variable $p_{ij}$.

w2: No example of the search output is presented, which would help readers better understand the solution.

With the explanation in w1, the search output is the specific values of each variable, i.e. $y_j$, $s_{ij}$ and $p_{ij}$. For example, $y_j$ indicates whether the $j$-th minipod is used. $p_{0,0}$ tells us how many nodes of the $0$-th PP group should be allocated to the $0$-th minipod. Thus, we use the value of $p_{i,j}$ to know which PP group goes to which minipods.

q1: Figure 4 Analysis

The degradation is due to the increased routing distance and the switch is overwhelmed.

The bandwidth degradation is expected as communication spans more minipods due to the increased physical routing distance. The specific 4-5 minipod degradation may be due to the switch is overwhelmed by the hundreds of bursty flows generated by the communication group over a short period of time, so the ECMP load balancing algorithm is not effective. The flows collide to certain links and cause link saturation, and the saturated link becomes the straggler to slow down the entire communication process.

q2: Bandwidth vs. Latency Bottlenecks

As discussed in the response to q1, the increased physical routing distance and the collision of network flows lead to performance degradation. This is reflected by the measured bandwidth. The theoretical bandwidth is the same for inter-minipods connection.

q3: Scaling Model Size

The training is bottlenecked by communication, and the optimized topology benefits GPU utilizes resources for communication, which is more pronounced when both computation and communication sizes increase.

Computation is much faster than communication for GPUs (TeraFlops vs GB/s). Thus, computation kernels are blocked by the dependent communication kernels. In Figure 10, we observe that an optimized topology benefits communication, which effectively optimizes the bottleneck of training.

q4: ILP Approximation Details

As explained in response to w1, we have clarified the details of each equations:

Equation (4) is the main objective function that minimizes the spread of PP groups and DP groups.

Equation (5) constraints the maximum spread of PP groups less or equal to $T$, and $T$ is being minimized.

Equation (6) says that the allocation cannot exceed the capacity of a minipod, e.g. if a minipod has $10$ nodes free, we cannot allocate more than that. The capacity is normalized by the number of node in a PP group.

Equation (7) says that while we can spread a PP group across minipods, the summation of the allocation percentage must be equal to $1$.

Equation (8) says that if we allocate to a minipod, that it is considered used in the solution, i.e. $y_j = 1$.

Finally, Equation (9) and (10) define the range of the binary variable $y_j$ and $s_{ij}$, and continuous variable $p_{ij}$.

q5: Network Topology Co-design

We think the presented algorithm is in principle able to handle other netowrk topology.

We observe that production GPU clusters follow the multi-tier design (Alibaba’s HPN, NVIDIA’s DGX). The collective communication is bottlenecked by the highest hierarchy - the inter-minipod communication in our case. In our case, we observe little impact of intra-minipod topology, however, we can also solve the MIP at the intra-minipod level if we can observe any performance degradation. In short, at every level of the network hierarchy, we can solve the MIP to find the optimal placement at that level.

q1: Why isn’t equation (4) symmetric? $\alpha$ and $\beta$ are symmetric, but $T$ is the max spread across PP communicator groups as I understand, but for DP communicator groups all the minipods are summed. Wouldn’t it make more sense to have a max spread measure for this as well?

Asymmetry of Equation (4). Equation (2) is symmetric because each node is a scheduling unit. Equation (4) is asymmetric because we coarsen the communication group (e.g. PP) as a scheduling unit. After the coarsening, we lose the granularity to express the spread of DP groups directly, but we can still control that by minimizing the usage of minipods.

A simple example. Suppose we have $m$ DP groups and $n$ PP groups, after the coarsening, the scheduling problem becomes “how to put the $n$ unit into $k$ minipods?” - we are allowed to spread one unit across multiple minipods, but we want to minimize the spread, hence introducing the $T$ auxiliary variable. On the other hand, DP groups’ spread is “orthogonal” to PP group – GPUs in the same stage of different PP groups belong to the same DP group, so if each unit (PP) is packed into a different minipod (the optimal $T$), the spread of DP groups is maximized. Thus, we also minimize the overall usage of minipod, which controls the spread of DP groups. Equation (4) combines the two objectives.

q2: Why isn't $T$ penalized non-linearly?

We acknowledge $T$ could be penalized non-linearly because in Figure 4, it has non-linear bandwidth degradation. We have also thought about constraint $T<=4$. However, it is not feasible in production scenario.

We choose the current solution due to multi-tenancy. A running cluster has dynamic job submission, hardware failure and bursty network flows. A hard constraint $T<=4$ may fail to solve at real time and it is unacceptable. A soft constraint requires quantification of bandwidth degradation, which is often interfered by other jobs, as shown in Figure 11. Besides, soft constraints are hard to express for MIP solvers.

Another issue is the resource contention in Figure 10 – computation and communication kernels are overlapped for maximized efficiency during training, but they also interfere with each other, which exacerbates the difficulty of quantifying the bandwidth decay.

q3: Are there any experiments against an external baseline (MegaScale) with experimental settings that can be disclosed?

Yes. We disclose the following information:

The experiment settings at framework level are the same, i.e.,

• overlapping DP, TP and PP.
• developping high-performance operators.
• applying asynchronous data pipeline.
• tuning the network performance.

The difference is the scheduling algorithm, and it is the goal of this paper to present a scheduling solution for LLM pre-training. Framework-level optimization is orthogonal.

q4: Is there any explanation for why all gather and send/recv bandwidths degrade over multiple minipods? Is this due to NCCL using a different algorithm? Perhaps the algorithm isn’t being selected intelligently going from 4 to 5 minipods.

Degradation reasons. The degradation is due to the increased routing distance and the switch is overwhelmed. The bandwidth degradation is expected as communication spans more minipods due to the increased physical routing distance. The specific 4-5 minipod degradation may be due to the switch is overwhelmed by the hundreds of bursty flows generated by the communication group over a short period of time, so the ECMP load balancing algorithm is not effective. The flows collide to certain links and cause link saturation, and the saturated link becomes the straggler to slow down the entire communication process. This means the scheduling algorithm needs careful design, not just relying on the network tuning.

NCCL algorithms and network tuning. NCCL uses ring algorithm for the highest bandwdith. The algorithm is tuned by NCCL's cost model. Our network software stack has been tuned meticulously, as observed from the S-curves over the message sizes in Figure 4 (a), which is detailed in NVIDIA’s NCCL performance tuning technical blog.

q5: Lines 175-177 discuss how training is insensitive to intra-minipod topology. But Figure 3 shows an intra-minipod example. Why is this?

Figure 3 is an illustrative example of the alignment of communication and topology. The network architecture is hierarchical, so the illustrative example can be naturally extended to higher layer in the hierarchy. We will refine the figure to be clear about the network topology.

con1: In line 82 (DP), it can be made more clear that this is different from pure DP and is in fact FSDP

It is DP, because FSDP is zero $3$ strategy that shards weights, gradients and optimizer states across all processes in DP groups. It introduces large communication overhead during forward pass. Our DP implementation is zero $1$.

con2: In figure 4(a), it is not clear where the bandwidth is measured between, i.e. is this within a minipod.

This is measured between minipods. We will clarify this in the text.

con3: I think the biggest issue with sections 3 and 4 is that it would be good to show what the alignment pattern is with SoTA communication frameworks

Existing SoTA frameworks use the bin-pack strategy, and it has poor alignment as shown in Figure 3. The misalignment issue of a small training job is less significant due to the small data volume. In Figure $5 (b)$, we empirically observe smaller jobs suffer less from misaligned performance degradation.

However, due to the scale of LLM training, the communication volume increases and the communication patterns become very sparse, as discussed in Figure 1. Thus, the misalignment problem is much more severe, and LLM pre-training needs a more specific scheduling algorithm.

con4: this won’t generalize easily to other large training jobs (for example, different size, parameter shape, GPU) without running a characterization experiment.

Our scheduling algorithm is LLM-specific and LLMs have unified model architecture (decoder-only architecture), increased model sizes by stacking more layers (7B, 33B, 70B), and hybrid parallelism. This makes it easy to pre-characterized automatically and generate database entries, i.e. launch cron-jobs and record the throughput.

con5: Another issue here is the lack of discussion regarding related work. I see this is in Appendix A, but should be part of the main paper, especially since the experimental results are only against 1 external baseline. In addition, it is hard to make sense of the results when the specific number of GPUs and model size is hidden.

We will move the discussion of related work to main body in the final version with an extra page. We will add more discussion of exising DL scheduling, and why they are not sufficient for LLM pre-training jobs - mainly due to ignoring the sparse communication patterns as shown in Figure 1.

Additionally to the response to q3, where we release details model configurations and framework-level experiment settings. The largest job uses $1216$ nodes and each node has $8$ GPUs.

con6: I am concerned about whether the theoretical formulation of the MIP makes sense (and hope this will be clarified in the rebuttal) but more importantly, whether there is enough significance/novelty as a research work suitable for NeurIPS

The theoretical formualtion is clarified in the response to q1.

Suitability. The topic of "Infrastructure (e.g., libraries, improved implementation and scalability, distributed solutions)" is explicitly included in the call for paper statement in NeurIPS 2025. The scope of our paper falls in the infrastructure track, which includes improved implementation, scalability and distributed solutions.

Significance/novelty. LLMs' learning is benefited by large-scale training, which necessitates innovative infrastructure support. In section 3 and 4, we share our insights and lessons learned from large-scale training on real-world clusters. In section 5, 6 and 7, we present a scheduling algorithm capable of handling very large-scale training jobs. We hope our work shed light on scalable training solutions and benefit the community.

Other comment:

We thank the reviewer for pointing out the grammar issues and lack of explanation in the caption. We will fix the grammar and add more text in the captions to make the terminology clear.

• What about experiments to justify lines 247 - 250?

Lines 247-250 say that we use proportionally small-scaled models and small numbers of GPUs to approximate the performance of large-scale experiments. This is well-studied in Megatron-LM[1], in section 5.1.1, termed weak scaling. A more up-to-date result is also shown in Megatron-LM's official Github, under the section of "Training Speed and Scalability", where it is clear that throughput linearly increases as both model sizes and the number of GPUs scale up, and therefore the performance characteristics are comparable.

[1] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020.

• Any other baselines?

We add the GPU-packing scheduling algorithm as a baseline and integrate it into the scheduling system, so that we can compare its performance with Arnold in the Llama3 experiment shown later.

• It is still important to provide some idea of the experimental setup for reproducibility purposes. I appreciate that there are business concerns, but then some proxy setup could be used to obtain more experiments.

For the training framework, MegaScale is open source, and it is capable of being used to train large models. For the scheduling layer, we build on top of Kubernetes, which is also open source. By utilizing those open-source systems, one can reproduce the experimental results.

• What about different types of models (different MoE config, or non-MoE), or data center network? Some small result could have been used here.

We add an additional experiment on Llama3 8B, replacing the InfiniBand network with a RoCE network and using 12 nodes equipped with L20 GPUs. Average throughput (PetaFLOPs):

The result is consistent with our findings - our scheduling algorithm achieves better speedups, but the speedup is more prominent in larger-scale runs.

Suitability of work: To be clear, I have no issue at all with the suitability of the work from the point of view of topic. The topic is indeed a fit for the Infrastructure track and the claimed results (10%) is also a great result for such an important/well optimized workload. My issue is with the significance/novelty of the research contributions which are mainly just the formulation of the MIP and observation regarding DP/PP misalignment. The work though has a significant amount of engineering contribution.

We wish to gently point out that the MIP algorithm represents innovative system research contributions. Its motivation stems from our observation of misalignment in real-world clusters, and its effectiveness is tested on production training jobs. The 10.6% speedup of a production training job has a huge impact on cost-saving, swift productionization, and fast iteration.

We appreciate your response. If you have any further concerns, please feel free to let us know. Thank you.

w1: This is a system paper, the authors made not much effort to connect this to learning. I think a system venue would be more suitable than NeurIPS.

Scope of our paper. The topic of "Infrastructure (e.g., libraries, improved implementation and scalability, distributed solutions)" is explicitly included in the call for paper statement in NeurIPS 2025. The scope of our paper falls in the infrastructure track, which includes improved implementation, scalability and distributed solutions. Other innovative solutions such as flash-attention, also have wide impacts on the NeurIPS community.

Connection to learning. LLMs' learning is benefited by large-scale training, which necessitates innovative infrastructure support. In section 3 and 4, we share our insights and lessons learned from large-scale training on real-world clusters. In section 5, 6 and 7, we present a scheduling algorithm capable of handling very large-scale training jobs. We hope our work shed light on scalable training solutions and benefit the community.

w2: The empirical validation is lacking. I think throughput should have been used as a metric throughout (and not just in the last section). The time to schedule is absolutely negligible relative to the total training time.

Throughput as the metric throughout the paper. In the simulator described in section 7.1, we use Equation (2) as the score to benchmark scheduling algorithms. This is based on the observation from Figure 4 and 5. Alternatively, we can replace the metric in the simulator by throughput via using the characterization results from section 4 (line 141-155). For example, as communication groups span over more minipods, the bandwidth decreases from the maximum bandwidth, and we use the percentage of the decrease to estimate the decrease of throughput from maximum throughput.

setting (i) throughput (PetaFlops)

setting (ii) throughput (PetaFlops)

setting (iii) throughput (PetaFlops)

The time to schedule is not negligible. We wish to gently point out that the time to schedule is not negligible as shown in Figure 8. The enumeration method has time complexity $O(k^m)$, i.e., for each of the $m$ nodes, we can place it in one of the $k$ minipods. This is not feasible for online scheduling. Moreover, cloud providers price Hopper GPUs at $2.99/h. For our largest job, this is \$28,704/hour or \\$478/min. GPUs can be idle during scheduling time, which directly leads to significant cost. Figure 8 shows that our method scales well for large job configurations and data center topology.

w3: The reported gain when using 9600 GPUs is a bit marginal. If large runs take a month, let's say, Arnold let us save only 3 days. It is unclear to me whether this is because the scheduling is sub-optimal or mega-scale is close to optimal already.

Arnold's impact on cost savings. We wish to highlight that a 10.6% speedup has a huge impact on cost savings. If 9,600 GPUs save 3 days, this saves us \2,066,688, because cloud providers price Hopper GPUs at \2.99/h. This improvement has a significant impact on cost-saving, swift productionization, and fast iteration.

MegaScale's scheduling is suboptimal. MegaScale, which employs the bin-pack algorithm, ignores the inherent sparse communication patterns in LLM pre-training (shown in Figure 1 and 3), while Arnold's MIP algorithm specifically considers these patterns. In section 7, we demonstrate how the improved scheduling algorithm works well in both simulation and real-world environments.

w4: The paper lacks insights a bit. It would have been nice to show how different hardware topologies / cluster configurations yield / benefit from different scheduling.

Benchmarks in different setups. We have shown benchmarks of different hardware topologies (TOR switch, spine switch), cluster configurations (small, medium, large job sizes), and different scheduling algorithms (a list of SoTA schedulers from literacture) in section 7.1 (line 290-306).

Scheduling insights. We have shown that when scheduling small jobs in a simple topology, scheduling algorithms achieve similar scores. As the job size and scale become larger, Arnold's algorithm shows more advantageous, and it is able to scale to very large sizes, as detailed in line 321-325, section 7.1:

We also observe that as $\alpha$ increases, our scheduling algorithm is closer to other baselines. This is because $\alpha$ controls the affinity of the DP group, and if $\alpha$ is $1$, the objective function reduces to a bin-packing formulation and therefore has no difference from other bin-packing algorithms. In practice, we would not set $\alpha$ greater than $0.5$ as observed from our characterization results. As a result, our algorithm usually scores higher than other baselines.

w5: When scaling to very large models, a node may not be sufficient to host the entire model and tensor parallelism might be needed across nodes, but the paper considers only DP and PP.

Clarification. Our model is not hosted by one node. As discussed in section 7.2 (line 354-358), our training systems have trained LLMs with more than 400 billion parameters, and it is sharded by TP and PP across multiple nodes.

Why only PP and DP are considered. TP must stay within high-bandwidth low-latency network domain (e.g. intra-node NVLink), as Megatron-LM [1] has in-depth experiments to show that. PP can scale out to inter-node connections. Therefore, to scale even larger than our models, more PP stages can be added to accommodate the memory pressure.

[1] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clusters using megatron-lm. In SC21: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–14, 2021.

q1: how much better is the enumeration approach in terms of throughput across various cluster configurations?

Compare with enumeration approach on setting (i). We add additional experiments on setting (i) (line 302) comparing the throughput of the enumeration approach with our scheduling algorithm:

setting (i) throughput (PetaFlops)

As shown in the table, our scheduling algorithm can match the performance of the enumeration method with much shorter scheduling time.

Compare on other settings. For setting (ii) and (iii), the algorithm has complexities of $O(5^{96})$ and $O(11^{368})$. They are computationally intractable.