Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models
Authors: Hanxiao Zhang
Lin JU
Chan Wu
Jinjing Huang
Youshao Xiao
Zhenglei Zhou
Zhiming fan
Zhaoxin Huan
Siyuan Li
Fanzhuang Meng
Lei Liang
Xiaolu Zhang
JUN ZHOU
Venue: NeurIPS 2024 poster (2024)
Keywords: Large Language Model
Distributed Training
Communication Topology
Primary area: infrastructure

Abstract: Recently, various strategies for distributed training of large language models (LLMs) have been proposed.
By categorizing them into basic strategies and composite strategies, we have discovered that existing basic strategies provide limited options in specific scenarios, leaving considerable room for optimization in training speed.
In this paper, we rethink the impact of memory and communication costs on the training speed of LLMs, taking into account the impact of intra- and inter-group communication performance disparities, and then propose a new set of basic strategies named the \textbf{Pa}rtial \textbf{R}edundancy \textbf{O}ptimizer (PaRO).
PaRO Data Parallelism (PaRO-DP) accelerates LLM training through refined model state partitioning and tailored training procedures. At the same time, PaRO Collective Communications (PaRO-CC) speeds up collective communication operations by rearranging the topology. We also propose a guideline for choosing different DP strategies based on simple quantitative calculations, which yields minimal ranking errors.
Our experiments demonstrate that PaRO improves the training speed of LLMs by up to 266\% that of ZeRO-3 as basic DP strategies.
Moreover, employing PaRO-CC independently for model parallel strategies, such as Megatron, can also boost the training speed by 17\%.

Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models
Hanxiao Zhang, Lin Ju, Chan Wu, Jinjing Huang, Youshao Xiao' Zhenglei Zhou, Zhiming Fan, Zhaoxin Huan, Siyuan Li, Fanzhuang Meng
Lei Liang, Xiaolu Zhang, Jun Zhou'
Ant Group, Hangzhou, China
(zhanghanxiao zhx, julin. ji, wuchan wu, huangjinjing youshaou xys, zhouzhenglei zzl, zhiming fzm, zhaoxin hzx, lisiyuan. 1i mengfanzhuang mfz, leywar liang, yueyin. 2x1, jun,
com
Abstract
Recently, various strategies for distributed training of large language models (LLMs) have been proposed. By categorizing them into basic strategies and composite strategies, we have discovered that existing basic strategies provide limited options in specific scenarios, leaving considerable room for optimization in training speed. In this paper, we rethink the impact of memory and communication costs the training speed of LLMs when employ data parallelism based tech- niques. We take the impact of intra- and inter-group communication performance disparities into account, and then propose new set of basic strategies named the Partial Redundancy Optimizer (PaRO). PaRO Data Parallelism (PaRO-DP) accelerates LLM training through refined model state partitioning and tailored training procedures. Additionally, PaRO Collective Communications (PaRO-CC) speeds up collective communication operations by rearranging the topology. We also propose guideline for choosing different DP strategies based on simple quan- titative calculations, which yields minimal ranking errors. Our experiments show that PaRO improves the training speed of L.LMs by up to 266% that of ZeRO-3 as basic DP strategies. Moreover, employing PaRO-CC independently for model parallel strategies, such as Megatron, can also boost the training speed by 17%.
Introduction
Large language models (LLMs) have demonstrated extraordinary capabilities across various domains, expanding their parameter sizes into the intricate challenge of training such large models, different distributed training strategies have been proposed.
hundreds of billions 4 29 To tackle the
Data Parallelism (DP) II 18 divides the input data into multiple mini-batches, with each GPU independently processing mini-batch through the forward and backward passes, followed by synchronization of the gradients. ZeRO [25] as a DP strategy, suggests partitioning model states across multiple GPUs. This strategy facilitates trade-off between memory and communication costs, offering three distinct partitioning methods, known as 7eRO-1, ZeRO-2 and ZeRO-3.
Model Parallelism (MP) 223
19] which encompasses Tensor Parallelism (TP) 27 35 2 and
Pipeline Parallelism (PP)
113 10,, distributes the model components across multiple GPUs. For
instance, Megatron-LM 27 achieves ID tensor parallelism by distributing the row or column
"Corresponding Authors
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

dimensions of the Transformer layers across multiple GPUs. However, the output of each layer requires all-reduce operations to aggregate the input, which increases the cost communication.
We classify these strategies as basic strategies for distributed training. We argue that choosing a basic strategy essentially involves making trade-off between memory and communication costs, based on different model sizes and hardware configurations to achieve faster training speed.
Beyond the traditional approach of training on the full set of model parameters, Low-Rank Adaptation (LoRA) [2] is frequently employed the fine-tuning of large models as form of Parameter-Efficient Fine-Tuming (PEFT) 21] methods. Present approaches to speed up PEFT, such as QLORA (7] and LoRA-FA (3),, leverage quantization and fash-attention [6] techniques. However, there a lack of discussion the distributed training strategy for finetuning scenarios including PEFT.
Additionally, mainstream GPU clusters adopt a heterogeneous network architecture that utilizes NVLink within machines and Remote Direct Memory Access (RDMA) between machines. Similarly, a set of GPU machines are interconnected via high-performance switch, whereas training cluster comprises multiple sets of GPU machines that need to communicate across switches [ In this paper, we refer to a multi-GPU machine or a switch as a group, applicable to other heterogeneous network scenarios as well. The constraint of inter-group communication performance being poorer than intra-growe group communication performance can slow down the collective communication operations that commonly used in DP and MP methods. Moreover, for methods like ZeRO, the entire training process needs to be redesigned because there a considerable amount communication operations during the forward pass, backward pass, and gradient update stages.
In the analysis of various mainstream basic distributed training strategies, we consider two key questions:
Are there more optimal basic strategies that offer additional trade-off' options for the DP training, including suitable strategies for scenarios involving full-purameters training and portial-parameterr training, thereby accelerating the training speed?
Is there more optimal collective communication method that can provide more efficient communication for DP and MP methods, thereby accelerating the training speed?
This paper addresses these questions by exploring potential basic strategies for LLM distributed training. The main contributions of this paper are as follows:
We systematically take into account the disparity in intra- and inter- anter-group communication performance when rethinking the trade-off between memory and communication costs. Our in-depth analysis covers various scenarios involving both full-parameter training, partial- parameter training and PEFT.
We propose the PaRO-DP set of strategies by refined model state partitioning and tailored training procedures. These strategies provide options for the trade-off between memory and communication costs, allowing for training speed improvements of up to 266% over ZeRO in scenarios where inter-group communication performance is poor. Notably, our proposal of a tailored training strategy for PEFT constitutes a new attempt within the field our knowledge.
We propose guideline that can be used to quantitatively select specific PaRO-DP strategies in various scenarios where the model size and communication performance vary. The effectiveness of this guideline has been validated via our extensive experiments.
We propose PaRO Collective Communication (PaRO-CC) tailored for collective communi- cation operations clusters with performance disparities between intra- and inter- inter-group communications, PaRO-CC can be incorporated into most basic training strategies, as demonstrated by 17% improvement in Megatron end-to-end training task. Compared to the traditional Ring topology, PaRO-CC exhibits a 36.5% improvement in collective communication efficiency.
Furthermore, we provide the open-source release of our code https com/ hww.aiio/haa///aa///eeee//aa//e

---

2 Related Work
Enhancement of Basic Strategies To enhance communication efficiency when training LLMs on large-scale GPU clusters, especially with heterogeneous networks, several methods have been proposed to employ group partitioning techniques to improve the basic strategy. ZeRO++ [32] extends ZeRO-3 by performing secondary partitioning of parameters while keeping other model states partitioned across all GPUs. This modification shifts the global all-gather operation in the backward pass to intra- group all-gather, reducing the volume of inter- group communication. Additionally, ZeRO++ compresses model parameters and gradients using quantization to decrease communication volume. MiCS 39] introduces group partitioning strategy, whereby the GPU cluster is divided into smaller and location-based subgroups. The model state is then partitioned within and replicated across these subgroups. By configuring appropriate subgroup sizes, MiCS can leverage high-quality intra-group networks and a hierarchical communication strategy to reduce communication volume between groups. PyTorch's 's official Fully Sharded Data Parallel (FSDP) provides a I/YBRID_SI __TAADD (FSDP-hs)) [40] strategy that leverages data center locality like MiCS to accelerate training and reduce inter-group communication. These optimization approaches, which incorporate grouping strategies, have inspired the direction of our subsequent research efforts.
Composite Strategy Combining multiple basic parallel strategies to integrate their benefits is common practice. These combinations include multi-dimensional hybrid parallelism and auto- parallelism. Multi-dimensional hybrid parallelism refers to the combination of multiple parallel strategies for computation simultaneously in distributed training. By combining these parallel strategies judiciously, we can fully exploit the advantages of various resources to improve training efficiency. This method is widely used in the industry for training large language models, such as LLama 29], GLM(8), and Megatron-Turinn NLG(28).
Auto- poll-parlllelism automatically selects a better or optimal parallel strategy for efficient execution based on the given model and machine resources used. Unlike hybrid parallel methods, automatic parallelism typically selects different basic strategies at a finer granularity, such as choosing different basic strategies for different layers or operators of the model. There are two types of auto- autto-prallellsm: semi-automatie and fully automatic. In semi-automatic mode, certain tensor and operator slices need to specified, such as Mesh-TensorFlow (26). GShard and GSPMD 36. In fully automatic mode, the framework adaptively selects all tensors and operators to determine the best slicing strategy, including Alpa and Unity [15 41 30 22].
We argue that through dedicated design efforts, new basic strategies can be leveraged across spectrum of composite strategies. This paper concentrates on the enhancement of basic strategies and refrains from comparing them directly with composite strategies.
3 PaRO Design
In the following discussion, we first delve into an analysis of the trade-off between memory and communication costs in LLM training. Based on this, we propose the PaRO-DP solution offering refined model state partitioning and guideline for selecting an appropriate strategy in various scenarios. Additionally, we introduce PaRO-CC further optimize the collective communication operation. Lastly, we discuss the applications of PaRO.
3.1 PaRO-DP
3.1.1 Analysis and Insights
As aforementioned, substantial performance gap of communication exists between intra- and inter-group networks, creating a bottleneck that impedes training efficiency. However, in data-parallel training tasks, by allowing for an acceptable level of memory redundancy, we can reduce the need for -nter-group communication during the training process while fully leveraging nigh-performance communication within groups, thus reducing the overall communication cost.
Contrasting with ZeRO, our approach augments the model states with an additional intra-group partitioning state. For further illustration, we define three partitioning states: No partitioning (N), Intra- Intra-group group partitioning (!) and Global partitioning (G) from coarse caase-graiied to fine- fine-grained, which
3

---

act on three components of model states: parameter (p), gradient (g), optimizer state (os). More specifically, I means that model states are partitioned within the group, while each group holds a complete replica. N means that each GPU holds replica of model states. G means that each GPU holds a part of model states, with only one complete model state on the global scale. We introduce the following notations to aid in the explanation. Note that, for simplification, we substitute network performance with bandwidth here.
m: The number of GPUs in the group.
71: The number of GPUs in total.
ng:
The number of groups, ng = n/m.
V: The number of parameters.
B:
Inter-group bandwidth between GPUs.
T The number of trainable parameters.
B':
Intra-group bandwidth between GPUs.
st The number of gradient accumulation.
T:
The Time cost of training a mini-batch.
Pp ptgtoos Partitioning strategy of p. 9 and os. The value combines N/T/G for each component of
model states. For example, PNIG means the strategy with no partitioning of Parameter, intra-group partitioning of Gradient and global partitioning of Optimizer State.
In the context of gradient accumulation where one mini-batch step contains micro-batch steps, we analyze the communication cost of model states with different partitioning states.
Parameter partitioning: Parameters are utilized in both forward and backward compu- tations during each step of the micro-batch.
both global partitioning and partitioning states, an all-gather operation is necessary obtain all parameters of the current layer before use. While only an intra-group all-gather is required when partitioning model parameters within group.
Gradient partitioning: Gradients are computed during the backward computation and used in the model update stage. Likewise, in both global partitioning and intra-group partitioning states, the aggregated gradient of the corresponding local partition is obtained through the reduce-scatter operation. Only intra- intra-group reduce-scatter is required when partitioning gradients within group, which is quicker than the global reduce-scatter used in global partitioning. Subsequently, each GPU performs local accumulation of the aggregated gradients in every micro-batcl step.
Optimizer state partitioning: The optimizer state is utilized during the model updating stage. If the partitioning scope of optimizer states differs from that of gradients or parameters, communication operations will be needed before and after the optimizer step. For instance, it requires performing an inter-group reduce-scatter of gradients before the optimizer step and an inter- group all-gather of updated parameters after the optimizer step, when the partitioning strategy is PIGG.
The above three levels of partitioning granularity on p. g and as bring up 27 combinations of model- partitioning strategies, while not all strategies effective. Employing fine- fine-grained partitioning can enhance memory efficiency for larger batch sizes, potentially increasing throughput, yet it also incurs more communication costs, thereby reducing throughput. Therefore, a trade-off between memory and communication costs should be made when selecting the appropriate model partitioning strategy.
In a mainstream training process using mixed precision and Adam optimizer [66, the memory consumption for the parameters, gradients and optimizer states are respectively 2 2 and 12 (25). Note that in PEFT tasks, where ! << v the sizes of and os relatively small compared to p. This suggests that PEFT tasks should be treated differently when making the trade-off between memory and communication costs.
In our findings, we identify a key insight: all scenarios involving various numbers of trainable parameters, the partitioning granularity of os should be the same as or finer than that of p and Otherwise, if os partitioned by a coarser granularity than p or g. more memory will be used without any reduction in communication cost, compared to when os is partitioned by the same granularity of p and y. According to this insight, we can eliminate out of the 27 possible combinations mentioned earlier and Appendix A.I.I provides more details.
3.1.2 Implementations of PaRO-DP
Following the analysis and resulting insights, we have formulated a strategy set of 14 combinations that constitute our proposed PaRO-DP. Specifically, DDP, ZeRO-1, ZeRO-2, ZeRO-3/FSDP, and

---

Panmdaass
Group-O
Group1
Paramesess
Group-
Group-1
Gradients
GPL-0
GRIA
GPl
GPUST
Gindieuts
GPL
GPE
GPI-2
GPU-
Optiminer
Optinizer
stalea
Intro Altisheee
Istra Atra-AlKiatber
00
Micio Nateree
...
Intra AllGnther
Intra AlKGather
Reckwarl
Backwand & Releso
Reckwwrdd& & Release
ReduveScatier
Rola-RolaccSsutts Accmmiiation
ReduceScuter Accumulttion
RediceSiatter
Accemmition
Accummbiioo
Inter ReduecScatter
Inter-ReinceSertter
Apply gralient
Apply graiee
Apply gradieu
Apply gradiemt
Inter AllGather
Gloui-AllGasbe
(a) Schematic of PIIG
(b) Schematic PNIG
The schematic of PaRO-DP illustrated within cluster comprising four GPUs. Each group in the cluster consists of two GPU devices. In particular, we illustrate the computation and communication phases of global step that incorporates gradient accumulation. We detail the intra- and inter-group communication by using specific prefixes for each collective communication primitive.
Figure 1:
MiCS are equivalent or approximately equivalent to PNNN, PNNG. PNGG, PEGGs PIr in the PaRO-DP strategy set.
As previously mentioned, when the model states are re-partitioned different ways, the procedures for Forward, Backward, and Gradient Accumulation and Apply Gradient during the training process require redesigning. Below, we provide a detailed explanation and PNIG examples. To simplify the examples, we use four GPUs and divide them into two groups to demonstrate the training process. These approaches can be easily scaled larger GPU clusters and Appendix A.1 provides other solutions in detail. It includes explanations of various collective communication operations used the diagrams.
their training processes using PIGG
Figure ((a) illustrates the schematic of PIIg, where the model parameters and gradients are intra- group partitioned, while the optimizer states are globally partitioned. Therefore, complete replica of model parameters and gradients are preserved within each group. The memory cost of PnG is less than that of PNNG (ZeRO-1), greater than PGGG (ZeRO-3), and in most scenarios less than PNGG (ZeRO-2), depending on the grouping situation. In the Forward and Backward stages, each GPU gathers a complete replica of model parameters through intra- intra-group all-gather operations. After the backward computation, each GPU aggregates gradients from other GPUs through intra-group reduce-scatter operations for local gradient synchronization. These gradients are temporarily stored in each GPU group through gradient accumulation. Once the gradients of the last micro-batch are accumulated, each GPU performs enter-group reduce-scatter operations to achieve global gradient synchronization. Each GPU utilizes the gradient partition to update the optimizer state maintained by itself and update model parameters. Finally, the updated model parameters of each GPU are obtained from other groups through inter-group all-gather operations. Note that compared to ZeRO-3, PHG employs intra-gromd group collective communication in place of global collective communication in the Forward and Backward stages of each micro-batch step. This makes PLıG faster than ZeRO-3 in context, at the cost of additional but limited memory of parameters and gradients.
Figure (bb) illustrates the schematic of PNIG. where the parameters of the model are not partitioned, while the gradients are gooru partitioned, and the optimizer states are globally partitioned. The memory cost of PNIG is less than that of PNNG (ZeRO-1), greater than PNGg (ZeRO-2). Different
5

---

from PLIG, each GPU retains complete model parameters in PNIG. Therefore, in the Forward and Backward stages, each GPU can directly perform the computation without collecting and releasing model parameters. The subsequent four-step computation process of PNIG is consistent with that of PrIG- Finally, each GPU collects updated parameters via PaRO-CC all-gather operations. Note that compared to ZeRO-2, PG employs intra-group collective communication in place of global collective communication during the gradient accumulation of each micro-batch step. This makes PNIG faster than ZeRO-2 in our context, at the cost of additional but limited memory redundancy of gradients.
Moreover, PaRO-DP strategies show efficiency advantages over both ZeRO++ and FSDP HYBRID SHARD ZERO2 (FSDP-hsz) approaches. Notably. under equal peak memory conditions, the global all all-gather operations in the Forward stage of ZeRO++ are less efficient than Prag's 's intra-group all-gather. Similarly, FSDP-hsz incurs additional communication costs due to the all-gather operations of parameters in the Forward stage compared to PNII.
3.1.3 The Guideline for PaRO-DP Strategy Selection
While we propose a set of effective strategies in PaRO-DP, there remains a gap in applying a suitable strategy from our PaRO-DP across diverse training scenarios with different v and V. Therefore, we present a guideline based on quantitative calculations. We are dedicated to identifying the optimal partitioning states (P²) that minimize the total time cost (T) of mini-batch training in a given scenario, as the throughput is inversely proportional to T. The total time cost T contains the time of communication loomm. and the time of compute (oomp., and Loverlop denoting the overlapping time of communication and computation. It could be formalized as a optimization problem as illustrated in [
min Pytaioo
Leomm. + leomp. Loverlap + min max (loomm.. Leomp.)
Potesos
Assuming that communication and computation can be fully overlapped, T can be approximately regarded as the maximum of t.comm. and feomp. Given a specific batch size, ccmmp can be straight- forwardly determined through a single forward computation. But the time of communication is more complex as shown in Formula 2
fcomm. = (4param + (graddent) # + Lupdate
For each micro-batch step, there is time cost Eporam of all-gather for parameters in both the forward and backward passes and a time cost gradsentt of reduce-scatter for gradients in the backward pass. For the last batro-batch of mini-batch, there is time cost, tupdate, for reduce-scatter//ll-reeduu reeuuce of gradients and all-gather for the updated parameters.
Therefore, the guideline implies that finding the partitioning states (P (op+ in which minimizing T. as defined by Formula [ is the recommended strategy for the given scenario. The process of calculation and selection is as follows, with further details in Appendix A.3 A.3
1. Input: n, m, B, B', V, V, s and tcomp.
2. Calculate the T values for various strategies.
3. Select the strategy with the smallest T value.
3.2 PaRO-CC
Distributed training multiple groups of GPUs using existing methods often requires global collective communication operations. For example, ZeRO-3 uses a global all-gather to obtain parameters of the current layer and global reduce-scatter to synchronize gradients. Megatron-TP uses global all-reduce to synchronize the computation results of a split layer. Although PaRO-DP decreases the frequency of global communication operations by redesigning the training process, some operations remain necessary. For example, Prog needs a global reduce-scatter to synchronize the gradients obtained from the backward pass of each layer.
In the Ring topology, collective communication divides a collective communication into n steps. In each step, each GPU is responsible for sending the currently relevant data block to the next GPU in

---

Node-
Kannk22
Step
Nodoo
Node-1
Baniko
Romb-1
Step
guther
Nodeodl
Nade-l
Intervaate sllepeher
Figure 2: The communication stages of PaRO-CC. C* denotes a data chunk. Only the procedure of all-gather operation Node 0 is shown for brevity. In step 1, inter- and intra-node all-gather executed in parallel. In step 2, only inter-node all-gather is executed.
the ring. Clearly, in this process, each step includes both intra-group and inter-group communication, leading to the speed bottleneck of each step being on the inter-group communication.
Therefore, we introduce PaRO-CC to accelerate global collective communication operations by organizing GPUs into groups based on the network topology and rearranging the communication topology accordingly. Groups are configured to form an outer ring, with each group further organizing an inner ring. The operation of PaRO-CC is segmented into intra-group and inter-group parts, both of which can utilize the Ring topology. Consequently, within a single step, communication confined to either intra-grou aatuu inter-group interactions, avoiding the need to address both simultaneously. is straightforward to ascertain that for the all-gather operation, execution of the inter- part should precede that of the intra-group part. Conversely, in the of the reduce-scatter operation, the intra-group part should be executed before the inter-group part. Additionally, during the all-gather be executed concurrently, assuming that suitable
operation, the two parts of communication blocking points are established to manage the flow of data. Figure [ illustrates an example of performing global all-gather operation using PaRO-CC. More details about PaRO-CC can be found in Appendix A.2
3.3
Applications of PaRO
PaRO can be employed independently as a basic strategy, or be used together with others in a com- posite strategy. For example, PaRO-DP can be alternative method to data parallelism in hybrid parallelism, e.g. 4D parallelism, when training large LLMs in heterogeneous networks [29] can be orthogonally integrated with Sequence Parallelism [ 14 20). At the same time, PaRO-CC can be applied in various distributed training strategies that require global collective communication opera- tions, accelerate the training. Furthermore, set provides more flexibility to complicated machine learning systems, such as distributed RLHF where each sub-model has different memory or communication requirements.
argue that the comprehensive PaRO-DP strategy
systems 24 37,
Experiments and Analysis
4.1 Experiment Environments
Our experimental cluster consists of up to 16 DGX nodes as 16 groups, each containing 8 Ampere A100 SXM4 80GB GPUs The GPUs in each node are interconnected via NVLink/NVSwitch with bidirectional bandwidth of up to 2400Gbps. These nodes are connected through RDMA over Converged Ethernet (RoCE), which can provide 100Gbps of inter-node bandwidth. The software environment includes CUDA-11.77 DeepSpeed-v0.10000 PyTorch-v1.922 Megatron-LM-v2..6 and NCCLLv2.14.3.
4.2 Experiment Settings
To evaluate the performance PaRO-DP, we compared them with current SOTA methods, including ZeRO, MiCS, ZeRO++, and FSDP-hsz, across different model sizes (denoted as V, such as 7B and 65B) and the number of trainable parameters (V'' categorized as full-parameters, partial-parameters, and PEFT). For each scenario, we select the recommended PaRO-DP strategies based on the through-
put indicator or TPS indicator log(1 IT), as outlined by the guideline in Section 3.1.3 The strategies
of PNN. (ZeRO-1 as PNNG. Pxni and DDP PNNN) were not considered due to its inability
7

---

50
1000
-10
Theoughpe
Memory
Indieator
Thanghout (Sangte
4.2
Single Morrary (G89
16
neeaghgg (Samllob
80.4
88,
SmilleGi Memo 6c
200
54
42.0
83
TPS Irdeam
68
-30
TPS DKState
2.6
76
12
48,8
14.0
24
20,8
100
0.0
ZdR03
100
MiGs
ZoRo2
ANG
ZoHo
ZaRO3
160
1kG
ZoR
MiCs
Zat023
FStP
(a) LLaMA 7B when W
(b) LLaMA 65B when W
Thocogippt
Menory
Thoogglgs
- Menay
Thraglpis (Smploshes)
Singl- SPC Menitey (GB))
Thropppe (Smphoree)
Singgt GPU Menuay (GB))
21,8
T19 hodicator
42
TPS Indioaeer
2.1
7.0
-0.7
0.5
00
-35
00
0.0
.60
To80-3
100
MBC
ZoR03
ZeRt
7:80-3
7aR0 FSDP-11
(c) LLaMA 7B when F
(d) LLaMA 65B when V 1/
Figure 3: The throughput and memory usage during LLaMA training with varying trainable parame- ters (V/)). The blue dashed line represents the trend of the throughput indicator, represented by the TPS Indicator using log (I)T), calculated based on the guideline. The cross indicates OOM.
Thruaghpul Indicaur
Memory
Single CPU Mieeory (30)
Thruggger (Somppantoo)
TP5 Indeoaorr
Time (ms)
14.
0.0
XXXX
XXXX
80
2:80-3
ING
ZoRO++
MICS
ZaRO-2
FSDP Az
16
32
64
96
128
Number GPUs
Figure 4: Throughput and Memory Usage of training LLaMA-65B in the PEFT(W' 3 11000 scenario.
Figure 5: Collective Communication Time (millisecond/ms) with the increasing number of GPUs.
run the large-scale model in our experiments. We use LLaMA-7B and LLaMA-65B 22 to evaluate the throughput and acceleration performance across 32 GPUs within 4 DGX nodes. For the LLaMA-65B model, we activate checkpointing ensure successful training. The C4 corpus in RedPajama was used as the training data. During training, we set the sequence length to 512, and the effective total batch size to 1280 (with 10 gradient accumulation steps) in a mixed precision training manner.
4.3 Efficiency of PaRO-DP
Incorporating group and refined partitioning, the PaRO-DP strategies adeptly distribute the volume of inter- and group communication to minimize overall communication costs. This results in sigmificant boost by up to 266% in training speed, accomplished with acceptable memory increment. The following section will describe the experimental results from various scenarios.
Full-parameters training: this scenario (V' VV), Figure 3(a) shows that the guideline recom- mends strategies such as PNII. PIGG and Pır (MiCS), which perform better throughput than others

---

ZeRO-2
ZeRO-3
NIG
12
Throughput (Samplesscco)
560.0
IGG
I1G
11
ZeRO NIG 1GG
11
10
280.0
Loss
140.0
70.0
35.0
16
Number of GPUs
32
64
128
300
Step
600
90
Figure 6: The throughput (samples/sec) when increasing the scale of GPUs.
Figure 7: Training convergence of PaRO against ZeRO using LLaMA-7B.
in LLaMA-7B. The throughput of PIIG is improved by 128% than ZeRO-3, and the throughput of PN11 improved by 67% than ZeRO-2. While in (LLMA-655B Figure F shows the throughput of PIG is improved by 86% than ZcRO-3. Since LLaMA-65B requires more fine- ane-grained partitioning, only ZeRO-3, ZeRO+t, PIGG and PnG perform training, while other solutions suffer from out-of-memory (OOM) issues.
Partial-parameters training: In experiments of partial- parameters training (V' 1/ 1/6)) we also reached a consistent conclusion. As shown in Figure 3(c) in LLaMA-7BB the throughput PNiG and PNTr is increased by 15% and 21% compared ZeRO-2. Figure 3(d) shows the throughput of PII (MiCS) and PIG is improved by 264% and 266% compared with ZeRO-3 LLaMA-65B, and PIıG uses less memory than MiCS.
PEFT: PEFT scenarios (V'' 3 / 11000) with LLaMA-65B shown as Figure the throughput of PING is improved by 44% compared to ZeRO-3.
Accuracy of the Guideline: The selection of high-throughput strategies across different scenarios determined by calculations based on the guideline. It chooses the strategy that owns a higher throughput indicator, log(1/T)), compared the baseline. The correlation between the guideline and actual performance in these scenarios is apparent, confirming that the throughput indicator reliably reflects the real world.
4.4 Efficiency of PaRO-CC
In the section, we performed experiments using up 16 DGX nodes with communication volume of 1GB. We measured the communication time of the all-gather operation with the traditional Ring (baseline) and PaRO-CC, shown Figure 5 Compared with the traditional Ring, there is pronounced performance enhancement by 25.2% %9.%% with different cluster scales (from 16 128 GPUs). An additional experiment with DGX nodes reveals that incorporating PaRO-CC optimization into end-to-end training tasks using Megatron- TPP results in a 17% increase in throughput, from 129 152 samples/second. Therefore, PaRO-CC can significantly improve communication efficiency by rearranging the communication topology.
4.5 Near-linear Scalability
To investigate the scalability of PaRO, we analyze the throughput of PaRO and ZeRO under varying counts of GPUs, as depicted in Figure 6 The experiments utilized the LLaMA-7B model. Using identical GPU resource configurations, PaRO consistently exceeded the speedup ratio of both ZeRO-2 and ZeRO-3 baselines. Notably, the speedup ratio for Piςg and PNiG from 16 to 128 GPUs are 0.92 and 0.84 respectively. These values are closer to 1, indicating near-linear scalability, compared to the speedup ratio of 0.63 for both ZeRO-2 and ZeRO-3.
4.6 Model Convergence
In this section, we demonstrate the consistent convergence of PaRO compared to ZeRO, which validates the correctness of our system. We used LLaMA-7B and C4 corpus in RedPajama to evaluate

---

the convergence of PaRO with 16 GPUs. During training, we set the sequence length to 128, the effective batch size to 1024, and the number of gradient accumulation steps 8. The loss validation process aims not to produce the same loss as ZeRO but to ensure identical convergence behaviours. As shown in Figure 7 PaRO provides the same convergence as ZeRO. The vertical axis represents the training step, but the duration per step varies with different strategies, affecting the training throughput.
Conclusion
In this paper, we dive into the trade-off between memory consumption and communication costs across various training scenarios in data parallelism training. Considering the performance gap between intra- and inter-group networks, introduce PaRO, refined and flexible partition strategy set. It contains the basic strategy, PaRO-DP, along with an effective collective communication topology. PaRO-CC. In heterogeneous networks, PaRO's advantages more prominently highlighted due to its effective adaption to the limited inter-group network, which outperforms the ZeRO-3 by up 266%. Furthermore, we conducted a comprehensive quantitative analysis of this trade-off and established a guideline to assist in selecting the suitable distributed training strategy in different scenarios.
Limitations and Societal Impacts
Limitations. Our method primarily focuses on data parallelism in training large language models. To train extremely large models, it is recommended to integrate our PaRO strategy with other parallelism techniques, such tensor and pipeline parallelism, though this may reduce the ease of use. Furthermore, performance improvements our method are largely due to the efficient utilization of heterogeneous networks, limiting its applicability in homogeneous networks.
Societal Impacts. By detailing our methodology and releasing our code, we aim to advance both research and industrial practices in the field of large language model training. By improving training efficiency, PaRO
potential to reduce the environmental impact of training I.LMs, which often require dozens to thousands of GPU devices. Additionally, our PaRO strategy is the first optimized distributed strategy for fine-tuning scenarios, especially PEFT scenarios. This enhances our social impact by benefiting large companies and democratizing the training of LLMs using limited number of GPUs.
Acknowledgements
We thank the anonymous reviewers for their useful comments and suggestions, and also thank Shangchun Zhao for his insightful advice.
10

---

References
Gibiansky Andrew. Bringing hpe techniques deep learning, 2017. httpa/////aadeeegiiiiaaaaa hwwp/bllogmmcchii--eeeeennnnnnno allreduee
[2] Zhengda Bian, Qifan Xu, Boxiang Wang, and Yang You. Doximmizing in distributed training for huge neural networks. CoRR, abs/2105, 14450, 2021. htt/:///s:/ar wwtp:/iii..rraaaaoe abs.abbb...... 14450
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, al. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, et al., editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL pptpsp:ss//p.ocesnngsssss negrssnnurrppp ips. cc/paerr/ 2020/hash/ 445770000bb0044977118186aaaaaaaa
Abstracctbbt httll
[4] Aakanksha Chowdhery, Sharan Narang. Jacob Devlin, Maarten Bosma, Gaurav Mishra, et Palm: Scaling language modeling with pathways. J. Much. Learn. Rex., 24:240:1-240:11 2023. URL
www.oo/gg/ppp///ppppee /vsprrs///2222// 1144.html
Together Computer. Redpajama: an open dataset for training large language models, October 2023. URL RedPaj ama- Data
[6]
Tri Dao. Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mohamed, et al.. editors. Advances in Neural Information Processing Systems 35: Annual Conference on Neu- ral Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http://papers wwwpaars..p................ www.:iipp.....ppeee cc/paper.ppppppcppcpc wttp:/cc.pppleeeeeeeaeepppspppee faarrffflleeppppeeee httpe//pppee/222//haaaee 2022//ash//
67d57c32e20f d0a7 www3330288166644600 www.eee6e000----aaaa----- Awstraat-tttt---e wwwnoreeeecc...mmmoommoo html html
[7] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. In Alice Oh, Tristan Naumann, al., editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http:: http::::/pppppra.aee nips. cc/ wwtp:per.ppppeeeeee leperrffill files/ paper/2023, 2023/ 1111/////////1363111000000022200
-eter/a009b--bbbrreeee Con.ssrrrac--ooooreneeeee Conference html
Zhengxiao Du Yujie Qian, Xiao Liu, Ming Ding. Jiezhong Qiu, al. GLM: general language model pretraining with autoregressive blank infilling, In Smaranda Muresan, Preslav Nakov, and Aline Villav- icencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Lin- guistics (Volume /: Long Pupers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pages 320-335. Association for Computational Linguistics, 2022. doi: 10.18653/V//0022ACCL 1202222L--L URL https:///oi. http:///toooorr/.../// wwwpg//...1666.//11 185/1/0111111/////2222 acl long..6.
[9] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
[10] Shiqing Fan, Yi Rong. Chen Meng. Zongyan Cao, Siyu Wang. et al. DAPPLE: a pipelined data parallel approach for training large models. Jacjin Lec and Frez Petrank, editors, PPoPP '21: 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Progaamming, Virtual Event, Republic of Korea, February 27- March 3, 2021, pages 431-445. ACM, 2021. doi: 10. 00.1145/3437801 0113433801.3441999 33411593, URL sttps:///oi. org/ 10. 1145/31378800 0437801..00 3441593
[11] Diandian Gu, Peng Sun, Qinghao Hu, Ting Huang, Xun Chen, Yingtong Xiong, Guoteng Wang, Qiaoling Chen, Shangchun Zhao, Jiarui Fang. et al. Loongtrain: Efficient training of long-sequence Ilms with head-context parallelism. arXiv preprint arXiv:2406. 18485, 2024.
[12] J.ward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, et al. Lora: Low-rank adaptation of large language models. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL nZeVKeePYEg
[13] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. 2019. htsL/tttpp//prrreedingg neurips ec/ Abstract.
[14] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang. Minjia Zhang, Shuaiwen Leon Song, et al. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309. 14509, 2023.

---

[15] Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. In Ameet Talwalkar, Virginia Smith, and Matei Zaharia, editors, Proceed- ingy of the SysML Conference 2019 (SysML 2019), Stanford. CA, USA, March 31 April 2019, mlsys.org, 2019. httphhtts:///proottt wwwpediiinnnssss.......mmmmmmmee mww.m. http://dna.aaa.ao.oooooppppeeeea 2t1/p//ee//////aaaaaa 120 Abstract.ttttteee
[16] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Reprexentations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Conference Truck Proceedings. 2015. URL bttp://arxiv org/ ab1/1412 1412
6980
[17]
Dmitry Lepikhin, HyoukJoong Lec, Yuanzhong Xu, Dehao Chen, Orhan Firat, et al. Gshard: Scaling giant models with conditional computation and automatic sharding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https://openreviett
[18] Shen Li. Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, al. Pytorch distributed: Experiences accelerating data parallel training. Proc. VLDB Endow., 13(12)33005-3018 2020. doi: 10. 00.14777//415478.4444 445.334 URI.tttp::/// //tp:/ vldb. wttp://ppw./oooo///////////////a 11.0000--1..p. pdf
[19] Shenggui Li, Hongxin Liu, Zhengda Bian, Jiarui Fang, Haichen Huang, et al. Colossal-ai: A unified deep learning system for large-scale parallel training. Proceedings of the 52nd International Conference on Parallel Processing, ICPP 2023, Salt Lake City, UT, USA, August 7-10, 2023, pages 766-775. ACM, 2023. doi: 10.1 145/3655573336666161
htl.https://ddd org/ 10. 10011111533888 360553366553333 http///665573333666611 3605613
[20] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01880 2023.
[21] Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, et al. Few-shot parameter-effiiccnt fine-tuning is better and cheaper than in-context learning. In Sanmi Koyejo, S. Mo- hamed, et al., editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans. LA, USA, November 28 December 9, 2022, 2022. URL http::/papers..a www..ccc.pppppeaaac Abb-aaattt-----tttattt-tt---- Conferencee html
wttppppp.alaeppppeeeepppppppppee httppapprr/////a///e 2ee///2222aaa///// /
[22] Xupeng Miao, Yujie Wang. Youhe Jiang, Chunan Shi, Xiaonan Nie, et al. Galvatron: Efficient transformer traiming over multiple gpus using automatic parallelism. Proc. VLDB Endow, 16(3):470-47,, 2022, doi: 10. 00.14778/37709903376677
/attps://ssst vldb. org/ www.b..rgg//....///ooll/...////. http://llllllollo//o/o/o miio.. pdf
[23] Deepak Narayanan, Mohammad Shoeybi. Jared Casper, Patrick LeGresley, Mostofa Patwary, et al. Efficient large-scale language model training on GPU clusters using megatron-In In Bronis R. de Supinski, Mary W. Hall, and Todd Gamblin, editors, International Conference for High Performance Computing, Networking. Storage and Analysis, SC 2021, St. Louis, Missouri, USA, November 14-19, 202 page 58. ACM, 2021. htRpsttts://// wrt. 1ee.l111/334555 34488177 3476209
doi: 10.1 145/3458881.34476009
[24] Long Ouyang. Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll L. Wainwright, et al. Training language models to follow instructions with human feedback. In Sanmi Koyejo, S. Mohamed, et al., editors, Advonces in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November
28 December 9, 2022, 2022. URL http::/ppppprsaappa/// http://ppprrs..pp..... www.:iiii...pp.eeeoee
wwtpppppr..ileeeppppeeeeeeee fiw.ee/apprr/////eeeeeepeaappe (aapper/2222 2022/hash/ (tt.:/bbb36641719194948888888880 www.aa001111111--taaceaa Awwtraatttttato- Conferencee html
[25] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In Christine Cuicchi, Irene Qualters, and William T. Kramer, editors. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9-19, 2020, page 20. IRFE/ACM, 2020. doi: 10.1
hntpa:/tt .ogg//.g.//.0 10111119///1111111 2020. .00024
[26] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, et al. Mesh- tensorflow: Deep learning for supercomputers. In Samy Bengio. Hanna M. Wallach, al.. editors. Advances in Neural Information Processing Systems 31: Annual Conference on Neu- ral Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 10435-10444, 2018. httttttttt///eeeeeeee wwwp:iiinnnnn..u........ wwwpipp..cppppprr.. 2018//a2h// www.:esssttbbbbbbbabaaeea Abstract.html
[27] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron- Im: Training multi-billion parameter language models using model parallelism. CoRR. abs/1909.08053, 2019. URL http::/tt//// www.aawwwww...raaaa
12

---

[28] Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, et al. Using deepspeed and megatron to train megatron--turing NLG 530b, A large-scale generative language model. CoRR, abs/2201. 19900, 2022. URL httppss//arriiv.oog/aaaaaa/////a
11990
[29] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, et al. Llama: Open and efficient foundation language models. CoRR. (ttpb//222..1.971... 2023. doi: 10.48550/ARXIV.2302..
13971. URL htttt:://dooo org/10. 48550/arXiv 2302.13977
[30] Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, et al. Unity: Accelerating DNN training through joint optimization of algebraic transformations and parallelization. In Marcos K. Aguilera and Hakim Weatherspoon, editors, 16th USENIX Symposiun on Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA. USA. July 11 11-13, 2022. pages 267 284. USENIX Association, 2022. URL httpa://www useniz. org/ conf erence/ caweencee/dd22.. wwtpppetetttnnnnee// unger
[31] Boxiang Wang, Qifan Xu, Zhengda Bian, and Yang You. Tesseract: Parallelize the tensor parallelism efficiently. Proceedings of the 51 International Conference on Parallel Processing, ICPP 2022, Bordeaux, France, 29 August 2022 September 2022, pages 12:1-12:11. ACM, 2022. doi: 10.1145/
35440087. URL httph: :ttp::::/doi... http://oooooogg///// 1145/3545008 3545087
[32] Guanhua Wang. Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, et al. Zero++: Extremely efficient collective communication for giant model training, 2023, httpttts:: //doi. org/
10. 0..44555///// arriv. 2306. 10209
[33] Youshao Xiao, Weichang Wu, Zhenglei Zhou, Fagui Mao, Shangchun Zhao, et al. An adantive placement and parallelism framework for accelerating RL.HF training. CoRR, abs/2312.1 1819, 2023. doi: 10.48550/
ARXIV.2312.11899 https: pettpps:/ddo..oo http:///oooo.r//1//// 10. 48550/arXiv. 2312. 11819
[34] Youshao Xiao, Shangehun Zhao, Zhenglei Zhou, Zhaoxin Huan, Lin Ju, et al. G-meta: Distributed meta learning in GPU clusters for large-scale recommender systems. In Ingo Frommholz. Frank Hopfgartner, et al., editors, Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, CIKM 2023, Birmingham, United Kingdom, October 21 25, 2023, pages 4365 4369. ACM, 2023, doi: 10.1 14455/3533780.6155223
org/ 10. 1145/ 3583780 3615208
[35] Qifan Xu and Yang You. An efficient 2d method for training super-large deep learning models. In JEEE International Parallel and Distributed Processing Symposium, IPDPS 2023, St. Petersburg, FL, USA. May 15-19, 2023, pages 222 232. IEEE, 2023. doi: 10.100//PDDS55459922223 00031. URL https://doi.org wtt.:o.........././000........
2023.00031
[36] Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake A. Hechtman, Yanping Huang, al. GSPMD: general and scalable parallclization for ML computation graphs. CoRR, abs/2105.04666 2021. URL
https://arxiv. https://arxiv.org/bbee abs/ 04663
[37] Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari. Xiaoxia Wu, et al. Deepspeed-chat: Easy, fast and affordable RL.HF training of chatgpt-like models at all scales. CoRR, abs/2308.00 01 01320, 2023. doi: A4A.550/ARXIV.22000122222 URL. httphttts://aad org/10. 0/e//10.485500 48550/arXix
2308. 01320
[38] Longteng Zhang, Lin Zhang, Shaohuai Shi, Xiaowen Chu, and Bo Li. Lora-fa: Memory-effiiient low-rank adaptation for large language models fine-tuning. CoRR, abs/2308.... 033003, 2023. doi: 10.48550/ARXIV. 2308.03303. URL httpssats////o.. otg/ 10. 48550/arX3s0
.03303
[39] Zhen Zhang, Shuai Zheng, Yida Wang, Justin Chiu, George Karypis, et al. Mies: Near-linear scaling for training gigantic model on public cloud. Proc. VLDB Endow. 16(1):37-50. 2022. doi: 10.14778/351261. 3561265. URL https://wws vldb. htgg/prrab/voll11//3---aaaa
pdf
[40] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, et al. Pytorch FSDP: experiences on scaling fully sharded data paraliel. Proc. VIDB Endow, 1 16(12)3384833866 2023. doi: 14.14778/
361 3611540.361 1569. URL htttp::// //ttps::/www.. vldb. org/ htts//vllbbvvollt 00066//3388--8 p3848-huang, pdf
[41] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, et al. Alpa: Automating inter- and intra-operator parallelism for distributed deep learning. In Marcos K. Aguilera and Hakim Weatherspoon. editors, 16th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA. USA. July 11-13, 2022, pages 559 578. USENIX Association, 2022. URL https://www: usenix. aesenii.or//eeee httpa//aw./oo.aaaaaa/////a/a//aa wtwp://ieeeeeeeeeettaeeeeeeeeeea whenggnn//zeggooeeee whw.ghhhann-mmiiin
13

---

A Appendix
A.1
Implementation of PaRO-DP
This section provides an overview the implementation strategies within the PaRO-DP. P. It details the partitioning of model states and elucidates the collective communications used during data- Mata-parallel training. The following subsections evaluate partitioning effectiveness and proceed with in-depth descriptions and pseudo-codes for each PaRO-DP strategy.
A.1.1 Effective Strategies Partitioning
Table presents all the combinations of model states partitioning Pp-gtoo- with ticks indicating those combinations that constitute effective solutions. The cross mark indicates the 13 strategies that have been eliminated by the insight in Section 3.1.1
Table 1: The 27 combinations of model states partitioning.
Strategy of Partition (pp+ Effective
Strategy of Partition (P (Pptgtos)
Effective
NNN (DDP)
IIG
NNI
IGN
NNG (ZeRO-1)
IGI
X
NIN
IGG
NII
GNN
NIG
GNI
NGN
X
GNG
NGI
GIN
X X
NGG (ZeRO-2) -2)
GII
INN
GIG
INI
GGN
ING
GGI
IIN
GGG (ZeRO-3)
III (MiCS)
A.1.2
Collective Communication Used in PaRO-DP
We divide the data evenly into n (the number of all GPUs) blocks, and the blocks held by GPU with different partitioning are shown as Table There are some additional notations:
à: The index of the current group in all groups, 0 S < ng
j: The index of the current GPU in the group, S j < m
G The j-th GPU in the i-th group
blocks: The array with the length is n of all blocks
Table 2: The partitioned blocks of data held by GPU
Partitioning Strategy
Data blocks held by GPU,,
No partitioning (N)
blocks 0 n
Intra-group partitioning (7) blocks[j x ng (j + 1) x ngl
Global partitioning (G)
blocks[j ng + 1]
The following introduces the collective communication operations used for synchronization between blocks of different partitioning.
14

---

Table 3: The Collective Communication used in PaRO-DP from the perspective of single GPU(GPUin.jo) jo. The values (N/I/G) of Inputs/Outputs blocks refer to Tabel
Input blocks
Output blocks
Participation ranks
Collective Communication
Description
(GPU)))
(0. i 9,
global all gather
G
N
optimized by PaRO-CC
0S < m]
{ 9.
global reeeueeeeeaaateeeeeeeeeeeeeeeeee _cctter
N
G
optimized by PaRO-CC
0 0 m) m)
{0<i<9
global all reduce
N
N
USj<m}
intra aara_group_att all gather
N
{t = io.0 1o:05ij<m)
hit_grrouu_rddueee scatter
N
{
i0,0
all all-gather with ju-th GPUs from different groups.
inter_gronn group all gather
G
{0<i<g)) jo}
reduce-scatter with jo-th GPUs from different groups.
inter geer_group_redtie__eee
G
(0<1< 9. = jo}
all-reduce with jo-th GPUs from different groups.
inter_grouy sver_group_au all_reduce
[05i<9. = jo}
A.1.3 Pseudo-code of PaRO-DP
We illustrate the PaRO-DP algorithms in this section. Each algorithm corresponds specific PaRO-DP strategy.
Algorithm 1: PsiG Algorithm
Input: model, dala
Output: model
for mini-batch epoch do
for micro-batee in mini-batch do
for layer in model.layers() do
4
intra group all gatherllayer.prraetters))
5
forwar::;::;
for layer in model. do
7
sara_group_ale La__alll_gttae
layer.backward(;
intra_rroup__ caailp_rdduee_eeeeee
scatter(layerggaaddennt))
layergrradentts.ccumulaee)
11
inter geeer_rroupp.
group_reducee reduce ceatter(model.gradieee
12
optimizer.stepO: // Apply gradient
13
inter ther_grouu__tt all gatheer(mmode..praaaeesss
15

---

Algorithm 2: PIge Algorithm
Input: model, dala
Output: model
for mini-batch epoch do
for micro-batch in mini-batch do for layer in model.layers/ do
intra group all gatherllayer.prrameter))
5
layer.forward):
for layer in model.reverse do
7
sara_group_at gattherlaayerr..rammeeeer
layer.backward))
9
global_reduce prredduce_ssaattrlllrrrrrrrradee
Luyergradients.ccumllaad;
11
optimizer.ste0O: // Apply gradient
12
inter group all gather(mmodel,praaeeeeer
inter_group
Algorithm 3: PNIG Algorithm
Input: model, data
Output: model
for mini-batch in epoch do
for micro-batel in mini-batch do for layer in model.layers() do layer.forwardO)
4
for layer in model. reverre do
aayer.backward);
intra intra_group crouppreduce, reduce aegluue_scatterllyyegraaadd
ayvergradientss.cccuukki
inter group_rouueee reduce scatter(model. graddennts
10
optimizer.step(: // Apply gradient
global all pacher/mmddelpprrmmeess)))
Algorithm 4: PNTI Algorithm
Input: model, dala
Output: model
for mini-batch epoch do
for micro-batch in mini-batch do for layer in model.layers() do
4
layer.forward(:
for layer in
do
wayer.backward);
7
Larr__grou__aa grrup_rrdduce cedtuce_ccatterllaeggraaadd
lavergrradentt..ccumuaeee
inter group all aaaaeee@@ooll...ddninttttt
inter_group..
10
optimizer.stepO:// // Apply gradient
11
intra Mara_group__tt gathher(moddll.aarmeeeeer
16

---

Algorithm 5: PING Algorithm
Input: model, dala
Output: model
for mini-batch epoch do
for micro-batch in mini-batch do for layer in model.layerst do
intra group all gatherllayer.prrameten))
5
layer.forward):
for layer in model. do
7
aura_group_as _a__alll___
layer.backwardO>
9
pneerggradients.accuullaa
10
Leabbal_reuue__gee_ peanterrmooell.rrdiinnttt optimizer.stepO: // Apply gradient
11
12 inter_ group all gather(model.
The implementations of other strategies of PaRO-DP can refer the pseudo-code provided for the noove-mentooned strategies, with the differences being in the details of how partitions are handled, which will not be reiterated here.
A.2 Detail of PaRO-CC
Algorithm and Algorithm are pseudocodes for PaRO-CC All-Gather and Reduce-Scattel Similar to traditional methods, in networks without NVIDIA SHARP configured, all-reduce operation consists of one Reduce-Scatter operation followed by one All-Gather operation.
Algorithm PaRO-CC All-Gather Algorithm
Input: Number of groups ng; Number of gpus in each group Current GPU index: i. ranges from 0, to ng 1, m
Function inter cicer_group______
2
for s=0 to ng do
do in parallel send block[ - 8 + i) * m%ang + + to GPU[(i + 1)%ng, receive (lo s 1) m/ong jl from
GPU[(1 - 1 jgl.
mark block[(ng- s
m/ng + j] as ready;
7 Function cutra_groupp_o do
select a ready data block or wait: send the block to the next GPU in the inner ring; receive the block from the previous CPU in the inner ring: mark the received block as ready; while not finished
11
split data into blocks as block[O...ng * m 1]: 16 mark blockçi m + +] as ready;
do in parallel
18 inter_group_part
intra_grouppppart;
17

---

Algorithm 7: PaRO-CC Reduce-Scatter Algorithm
Input:
Number of groups ng; Number of gpus in each group my Current
GPU index: i.j ranges from 0, 0 to ng - 1, 1,m
perform reduce-scatter using inner ring:
2
perform reduce-scatter using outer ring;
A.3 Detail of the Guideline
In this section, we present a detailed explanation of the time of communication cost comm. calculation formula, employing symbols that align with the definitions provided in the main text. And give an example the guideline based on the calculation formula.
A.3.1 Detail of the Calculation Formula
toom. (fparam + Egradient) tupdate
Table 4 and Table show the caleulation formula of param and Egradient for various partitioning states.
Table The calculation formula of 1param
Pp
time
Description
N
2 * 1/ m/mm * (m 1)/B'
intra-group all-gather parameters when forward and backward global all-gather parameters when forward and backward
2 V/ v/n= (n 1)/B B
Table 5: The calculation formula of Egradient
Pg
time
Description
N
0
T
V m * (m 1) B'
intra-group reduce-scatter gradients when backward global aduuee-scatter gradients when backward
G
V /n (n 1)/B
tupdate teyne, + tayne,
The fupdate includes two parts: the time of synchronizing gradients (teyney, shown as Table 6 required for the optimizer states step and the time synchronizing model parameters (fayne,, shown as Table after updated.
Table 6: The calculation formula of Usyne
Pg
Pos
time
Description
N
2 /n (n 1)/B
global all-reduce gradients inter-group reduce-scatter and all- reduce gradients global reduce-scatter gradients
N
I
V/nn+ (n 1)/B 3 v /////m/=> (g 1)/B
N
G
V n (n 1)/B
2 * /m/g 1)/B
inter-grous all-reduce gradients inter-group reduce-scate gradients ohters
G
V /m/ (g B
In light our experimental configuration and environment, the salient parameters informing our derived guidelines include: m = 32, 8, y n1/m 4, = 10, v 7 1e9 or 65 1e9, B 80 le9 (80Gbps), B' 2000 1e9 (2000 Gbps), as determined through testing within our experimental GPU cluster.
18

---

Table 7: The calculation formula of tagne,
Pp
Pos
time
Description
V /mm (m 1)/B'
intra-group all-gather updated parameters
N
G
V / (n 1)/B
global all-gather updated parameters inter-group all-gather updated parameters
V 1//m// g (g - 1)/B
0
ohters
Table 8: The Metric (1/7) of training throughput and GPU memory of LLaMA-7B. V and V/ /
y //16 mean the different ratios of trainable parameters model parameters.
V
V
V
Strategy (piglos)
1/T
Mem(GB)
1/T
Mem(GB)
NII (PaRO)
1.151
24.447
17.827
13.752
(IRRP)
0.489
17.113
7.716
13.293
NGG (ZeRO-2)
0.067
15,891
1.066
13.217
INI
0.386
24.447
0.841
3.056
ING (PaRO)
0.386
17.113
0.841
2.598
III (MiCS)
0.524
13.039
0.871
2.343
IIG (PaRO)
0.524
5.704
0.871
1.884
IGG (PaRO)
0.067
4.482
0.511
1.808
GNG
0.035
15.891
0.037
1.375
GIG
0.036
4.482
0.037
0.662
GGC (ZeRO-3)
0.024
3.26
0.036
0.586
Table 9: The Metric (1/T) of training throughput and GPU memory of LLaMA-65B. V' V T' 1 /66, and P /VV/1000) mean the different ratios of trainable parameters to model parameters. means out of GPU memory.
V
V 3 / /1000
Strategy (plglos)
1/T
Mem(GB)
1/T
Mem(GB)
1/T
Mem (GB)
NII
NIG (PaRO)
NGG (ZeRO-2)
INI
0.091
28.376
0.098
15.77
ING
0.091
24.12
0.098
15.565
III (MiCS)
0.094
21,755
0.098
15,452
IIG (PaRO)
0.057
52.969
0.094
17.499
0.098
15.247
IGG (PaRO)
0.007
41.618
0.055
16,789
0.095
15.213
GNG
0.004
12.769
0.004
4.215
GIG
0.004
41.618
0.004
6.148
0.004
3.897
GGG (ZeRO-3)
0.003
30.268
0.004
5,439
0.004
3.863
The outcomes of the guidelines presented in Table and Table indicate recommendations, distinguished by a higher throughput indicator 1/T.
Strategies highlighted in red
A.4 Detail of Experiments
A.4.1
Communication Performance of the Experimental Cluster
Table TO presents the measured inter- and intra-group GPU-to-GPU communication throughput in the experimental environment.
19

---

Table 10: The Throughput of GPU GPU communication.
GPU GPU
Size (Bytes)
Duration
Throughput (Gbps)
intra-node
31.74 GB
131.128
2079.138
inter-node
31.74 GB
3033.485 ms
89.875
A.4.2 Results of the Experiments
Tables 11 12 and 13 summarize the experimental results from Section 4.3
showcasing model
configurations, throughput metries, and GPU memory utilization for parl-par,meter, partaee,arameter, and PEFT training scenarios. These tables offer a clear and comprehensive evaluation of the effectiveness of PaRO-DP.
Table 11: Configuration and Result of Experiments when tull-parameters training (V'² W)
Model Size
Peak Memory
Strategy (p/glos)
Config
Throughput
GGG (ZeRO-3)
("stage":
1.93
37.13
IGG (PaRO)
("paro_strategy": ("paro_strategy"" "IIG") ("stage": 3, mnics_partition_size 8)
"IGG") |
2.84
44.76
IIG (PaRO)
4.40
46.04
III (MiCS)
4.41
49.20
NGG (ZeRO-2)
("stage": 2|
2.79
52.58
7B
NIG (PaRO)
("paro_strategy". "NIG")
4.44
44.06
NII (PaRO)
("paro_strategy" "NII")
4.65
51.98
ZeRO++
("stage": 3,
2.43
39.76
8)
FSDP-h2
HYBRID_SHARI ZER02
2.86
59.15
GGG (ZeRO-3)
("stage": 3)
0.21
65.56
IGG (PaRO)
"'par_strrateyy" "IGG")
0.29
76.12
65B
IIG (PaRO)
"yparosstrategy:: "IIG") ("stage": 3.
0.39
76.27
ZeRO++
0.25
77.54
"zero_hpp hwtpo//ppppaaitiionssssm
8)
Table 12: Configuration and Result of Experiments when partial-parameters training (V/ = V \/11).
Model Size
Peak
Strategy (p/glos)
Config
Throughput
Memory
GGG (ZeRO-3)
("stage":
4.65
10.63
IGG (PaRO)
("paro_strategy": "IGG"] | ("paro_strategy"" "IIG") ("stage": 3, imics_partiiion ("stage": 2
9,41
11.97
IIG (PaRO)
11.11
12.14
III (MiCS)
11.59
12.62
7B
8)
NGG (ZeRO-2)
11.94
18.86
NIG (PaRO)
("paro_stratogy" "NIG")
13.79
20.10
NII (PaRO)
("paro_stratogy": "NII"] ("stage": 3,
14.46
19.66
ZeRO++
4.79
12.34
pharo_hpz.partition. size": 8)
FSDP-hz
HYBRID_SHARD 2
12.50
21.23
GGG (ZeRO-3)
("stage": 3)
0.59
20.48
III (Mics)
("stage": 3,
2.15
44.23
65B
"alcs_partittin_sizee" 8)
IIG (PaRO)
aparrostrratey".. "IIG") ("stage": 3,
2.16
39.48
ZeRO++
0.60
34.72
zeroo_pp__pariiiion..zz
8)
20

---

Table 13: Configuration and Result of Experiments when PEFT (V' 3 //1000
Model Size
Peak Memory
Strategy (p/glos)
Config
Throughput
GGG (ZeRO-3)
("stage": 3/
1.20
29.54
65B
ING (PaRO)
"aroo_strateg": "ING") ("stage": 3, "zero_hpz_ pearo.bp..paritiion...... 8)
1.73
47.71
ZeRO++
1.28
44.68
A.4.3 Maximum Throughppt via Dynamic Effective Batch Size
In Section we standardized experimental conditions with batch seich_sizeeee 40 per GPU for a fair comparison. This section presents the impact of dynamic effective batch size training efficiency, specifically in the context of farl-parameter training scenario within the A100 cluster. As displayed in the Table [4] the throughput of PaRO-NNI is 48.7% higher than that of ZeRO-1 using the same effective batch size. Additionally, 0.62. surpassing the GGG strategy's 's 0.47 by 32%, as shown in Table 15 This demonstrates the significant effectiveness of the IIG strategy in this scenario, which balances communication costs and memory usage, reducing inter-group communication and consequently improving overall training throughput.
hat__batch_sii_ee of 4 and 10
accumulation__ steps, setting an effective Setectivv_bbtch____e
the PaRO-IIG strategy achieves a throughput
Table 14: Throughput Comparison of LLLMMA-BB on 32 GPUs when full palamearssserss training (V' V). The global batch size in one global step is fixed to 17280 or effective batch size per GPU is fixed to 540. Note: MBS micro batch size, AS accumulation steps, EBS effective batch size
Strategy (plglos)
Configuration (MBS, AS, EBS) of Single GPU
(180, 540)
(270, 2, 540)
(540, 1, 540)
NNG (ZeRO-1)
5.908
0OM
0OM
IIG
5.493
0OM
0OM
NII
5.416
8.785
00M
Table 15: Throughput Comparison of LLAMA-65B 64 GPUs when full-parameters training
(T' = W). Note: MBS = mnicro_batchh_aee AS = accumulation EBS effective_bbtcch _ize
Strategy (plglos)
Configuration (MBS, AS, EBS) Single GPU
(16, 10, 160)
(20, 8, 160)
(32, 5, 160)
(40, 4, 160)
(48, -,
IIG (PaRO)
0.62
0.57
0OM
0OM 0.47
0OM
GGG (ZeRO-3)
0. 0.23
0.28
0.41
0OM
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction sections of the paper clearly state the research objectives, methodologies, key findings, and conclusions, providing a coherent overview of the core contributions of the work. Moreover, the paper complements these claims with detailed experimental setups, results, and analyses that are consistent with the statements made in the abstract and introduction, ensuring that the readers have comprehensive understanding the work's ' scope and contributions.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Our method primarily focuses on data parallelism in training large language models. The PaRO strategy can be combined with tensor or pipeline parallelism techniques in the latest hybrid parallelism manner to train extremely large language models, albeit with the sacrifice in ease of use. Furthermore, performance improvements of our method are largely due to the efficient utilization of heterogeneous networks, limiting its applicability in homogeneous networks.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
22

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: Our paper focuses on the practical implementation and evaluation of a refined sharding strategy for large language model training, taking into account the impact of communication and machine resources on actual throughput. We provide detailed guideline for selecting optimal strategies various scenarios based on empirical data and performance metrics. Since the throughput results are inherently dependent on the specific environmental factors and hardware configurations, the paper does not engage in theoretical analysis or proofs for each throughput outcome.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper fully discloses all the information necessary to reproduce the main experimental results. Section 4 includes details on experimental setup. datasets, model configurations, and hyperparameter selections, all of which are essential for understanding the outcomes. Furthermore, Appendix providess specific configurations and results of the experiments, which support the claims and conclusions made in the paper, ensuring that other researchers can verify and replicate the findings.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
23

---

While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
The paper provides open access to the data and code, ensuring the reproducibil- ity of the main experimental results. We specifically provide GitHub link that contains detailed information about the codebase used for the experiments, which is essential for other researchers aiming to replicate the study. This level of transparency facilitates independent verification and promotes good scientific practice within the research community.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
peplicc/guides htte://uiidessoodeuummsssse
(CodeSulmissiionooloiooo for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.ccppubbl https/sccppulllcguuueeeeeec wttpi///guudee///Coeoeee
ssi ssionoolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper meticulously details all aspects of the training and testing procedures. Section and Appendix offfro comprehensive information that includes data partitioning strategies, hyperparameter configurations, the types of optimizers used, and the rationale
24

---

behind their selection. These details are indispensable for thorough understanding of the experimental results and allow for the replication and verification the findings by other researchers.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The paper does not report error bars or statistical significance measures because the experiments primarily focus throughput measurements, which are inherently stable and consistent when averaged over a significant number of steps (e.g., every 100 steps). Given the nature of the experiments conducted dedicated cluster, the results demonstrate high reliability, reducing the necessity for further statistical analysis to confirm the stability of the outcomes.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper provides detailed account of the computational resources utilized for each experiment. Section [.. meticulously describes the experimental environment setup. including the hardware used (such DGX nodes and Ampere A100 SXM4 80GB GPUs), the interconnections between these components (via NVLink/NVSwite. and RoCE), and the software environment (with versions of CUDA, DeepSpeed, PyTorch, and NCCL. specified). This level of detail equips readers with a clear understanding of the computational requirements necessary replicate the experimental results in similar setup.
Guidelines:
25

---

The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research presented in the paper strictly adheres to the NeurIPS Code of Ethics, ensuring that all aspects the work are conducted with integrity and responsibility.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Our paper aims to enhance research and industry practices in LLM training. By improving training efficiency, PaRO can reduce the environmental footprint of LLMs, which typically demand numerous GPUs. Furthermore, PaRO represents the first optimized approach for fine-tuning, particularly in PEFT scenarios, thereby increasing our social impact by aiding large companies and democratizing L.LM training with fewer GPUs.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and funetioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
26

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
The research does involve the release of models or data that could be potentially misused, as it is focused on theoretical contributions and experimental findings within controlled environment. There are no risks misuse associated with the presented work, and thus, no specific safeguards are required.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The paper meticulously acknowledges and adheres to the licensing terms of all existing assets utilized in the research, including datasets, code repositories, and models. The authors have ensured that all assets are properly credited, and the terms of use explicitly stated and respected, which is in line with the best practices for academic integrity and intellectual property rights.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwwppssiitchooeeoo thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not introduce any new assets such as datasets, or models that would require additional documentation or licensing information. The research is conducted using established assets and does not generate new ones that need to disseminated with the paper.
27

---

Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the datasct/codc/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The research conducted for this paper did not involve any form of crowd- sourcing experimentation with human subjects. Therefore, there is need to provide instructions to participants, screenshots, or details regarding compensation, as the work is based on computational experiments and analysis.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Since the research presented in the paper does not involve any human subjects, there was no requirement to obtain Institutional Review Board (IRB) approvals or conduct equivalent ethical review. The study is purely technical and does not engage with human participants, thus falling outside the scope of IRB considerations.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "7FSSifhKIc", "reviewer": "Reviewer_X1pV", "summary": "This paper recast basic known distributed training strategies (such as the ZeRo-1, 2, 3, MiCS, FSDP) in a unified framework that takes into account the trade-off between memory consumption and communication. By exhibiting natural levels of granularity for the partionning of different parts of the model and optimizer, this paper makes an exhaustive list of all 27 strategies possible for the distributed training of LLMs (that include previously known ones). Then, it proposes a simple model to predict the throughput and feasibility of all strategies based on physical measurements of the cluster used and task at hand, allowing to chose the fastest feasible distributed training strategy. This model is shown to predict effectively actual performances of LLMs training, allowing to bring significant speedup compared to previous \"go-to\" methods in some cases.", "questions": "* line 228: *“Assuming that communication and computation can be fully overlapped, $T$ can be approximately regarded as the maximum of $t_{comm.}$ and $t_{comp.}$”* Could you be more specific when this assumption is actually valid in practice?\n* line 237: *“the formula of $T$”* Do you mean $\\max \\\\{ t_{comp.}, t_{comm.} \\\\}$?\n* line 287: (throughput indicator) It not clear at first that this is the “TPS indicator” (otherwise not defined) in the subsequent figures, I would advice indicating it at this stage.\n* Fig 7: Small variations between the losses for the different methods are observed, although they should be mathematically identical, why is that?\n\n**Suggestion:**\n\nIn the paragraph line 232, I would advice to precise that the formulas (based on physical characteristics of the cluster/task) to estimate any $t_{\\times \\times}$ are provided in appendix A3.", "limitations": ".", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* **Exhaustive enumeration of possible distributed training strategies in a unified framework:** The 3 level of granularity for the partitioning of some “vector” are proposed, and seem natural in the context of GPU cluster (“no partition = copy on all worker”, partitioned inside a cluster node, or across all workers).  3 different model+optimizer parts are considered (model parameters, gradients, and optimizer’s state) as depending on the task and the use of mixed precision, these are treated differently. With the 3 level of granularity, and 3 type of vectors, it leads to $3^3=27$ possibilities for distributed strategies.\n\n* **Simple model to predict the throughput of a given strategy:** Based on simple physical characteristics of the GPU cluster used (e.g., intra and inter node communication bandwidth, GPU memory) and the task (pretraining, parameter-efficient finetuning), a very simple model is proposed to predict what would be the fastest feasible distributed training strategy.\n\n* **Model is shown to highly correlate to actual implementation throughput:** In part 4.2, through extensive experimentation, the model is shown to highly correlate to actual time measurements when implementing the different strategies.\n\n* **Significant speedup can be observed in practice by using the recommended strategy:** Thanks to this model, the best feasible strategy can be chosen in advance, which can be a completely new one compared to previously known ones (such as ZeRo). In the experiments performed, this can lead to significant speedup compared to standard methods.", "weaknesses": "* **Figures a bit hard to understand:** I found Fig.1 \\& 2 hard to read and to understand.\n* **Experiments in Part 4.2 do not seem exhaustive:** in line 287, 3 configurations out of the 14 considered are dismissed, leading to 11 remaining possible experiments. Yet, only 9 are reported in Figure 3, why is this the case?\n* **Comparison with standard collective communication strategies other than \"the ring topology\" is not done:** The only collective communication strategy investigated (other than the one proposed by the authors) is the one based on the ring topology. However, this is not the only standard implemented in the industry. For instance, [NCCL, in addition to \"Ring all-reduce\" also propose \"Tree all-reduce\", which can allow significant speedup at scale compared to the ring option]( https://developer.nvidia.com/blog/massively-scale-deep-learning-training-nccl-2-4/ ) (NCCL choses automatically which to use depending on the network/task, but it is possible to force a particular strategy). Moreover, other strategies such as Butterfly All-Reduce [[Patarasuk et al., 2009]](https://www.cs.fsu.edu/~xyuan/paper/09jpdc.pdf ) or Moshpit All-Reduce [[Ryabinin et al., 2021]]( https://arxiv.org/pdf/2103.03239 ) are not considered either.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive and constructive feedback. We will explain your concerns point by point.\n\n## Weaknesses:\n\n**Q1: Figure.1 & 2 a bit hard to understand.**\n\n**A1:** We have redrawn the figures, please refer to Figure 1 & 2 in the PDF.\n\n**Q2: Experiments in Part 4.2 do not seem exhaustive.**\n\n**A2:** We choose the best strategy for specific training scenarios based on the guidelines and verify it in our experiments.\n\n**Q3: Comparison with standard collective communication strategies other than \"the ring topology\" is not done.**\n\n**A3:** I appreciate you bringing this up.\n\nThe Butterfly topology sends and receives large data blocks each time, which makes it difficult to fully utilize the bandwidth and prone to delay jitter. Its performance is not as good as the Ring topology.\n\nCompared to the Ring topology, the Tree topology in NCCL speeds up communication primity locally. We focus on the huge difference in communication bandwidth between inter-group and intra-group interactions. Group communication is utilized to minimize inter-group communication and maximize intra-group communication. The communication can utilize either a ring or tree structure.\n\nThe Moshpit topology is a group communication strategy for allreduce that solves cluster instability and node failure problems through an iterative averaging protocol. Our algorithm combines multi-level sharding of large models and fully utilizes bandwidth between and within machines through a grouping strategy to improve communication.\n\n## Questions:\n\n**Q1: line 228: “Assuming that communication and computation can be fully overlapped, 𝑇 can be approximately regarded as the maximum of 𝑡𝑐𝑜𝑚𝑚. and 𝑡𝑐𝑜𝑚𝑝.” Could you be more specific when this assumption is actually valid in practice?**\n\n**A1:** When training transformer architecture models, the next layer can prefetch parameters while the previous layer is computing, allowing computation and communication to happen simultaneously. Here we simplify and assume perfect overlap, meaning the training time for a single step is determined by the longer of the two, computation or communication.\n\n**Q2: line 237: “the formula of 𝑇” Do you mean max{𝑡𝑐𝑜𝑚𝑝.,𝑡𝑐𝑜𝑚𝑚.}?**\n\n**A2:** Yes. We will make it a numbered formula in our manuscript for better clarification. \n\n**Q3: line 287: (throughput indicator) It not clear at first that this is the “TPS indicator” (otherwise not defined) in the subsequent figures, I would advice indicating it at this stage.**\n\n**A3:** Yes, throughput indicator(log(1/T)) is the “TPS indicator” in figures. We will declare this in the subsequent version.\n\n**Q4: Fig 7: Small variations between the losses for the different methods are observed, although they should be mathematically identical, why is that?**\n\n**A4:** Thanks for pointing this out. We have verified the values of the model parameter, gradient, and optimizer state before and after an update step, and the errors of different strategies remain within the normal machine precision range (1e-6). Small variations between the losses for different methods, despite their mathematical equivalence, can be attributed to floating-point representation errors in computers. These truncation errors accumulate during different steps of the calculations. As a result, the outputs may show slight discrepancies. However, these factors do not impact the statistical performance during training or the consistency in convergence behavior. The phenomenon is also consistent across ZeRO-1/2/3.\n\nThank you again for taking the time to review our paper. We hope we have addressed the reviewer's questions and the reviewer is willing to raise their score. Please let us know if we can provide any additional clarification or information.", "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": "TMjHpXAlNv", "reviewer": "Reviewer_63Xk", "summary": "This paper introduces the Partial Redundancy Optimizer (PaRO) to improve the efficiency of training large language models (LLMs) by optimizing the trade-off between memory and communication costs. PaRO includes two main strategies: PaRO Data Parallelism (PaRO-DP), which refines model state partitioning and training procedures, and PaRO Collective Communications (PaRO-CC), which rearranges the communication topology for faster collective operations. The proposed strategies demonstrate significant improvements in training speed, up to 266% over state-of-the-art (SOTA) data parallel strategies, and 17% when applied independently to model parallel strategies like Megatron.", "questions": "- How PaRO determines the group size (G)? Is it a hyper-parameter?\n- In the experiments, do other well-known computation optimizations (e.g., FlashAttention) are applied?\n- How can users enable PaRO? Could you provide an example of the programming interface?", "limitations": "No additional limitations exist.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- PaRO enables much more fine-grained parallelism strategy, compared to ZeRO. This facilitates much more optimized training performance in broad range of resource setup.\n- Not only partitioning strategy, but this paper also proposes PaRO-CC, intra-group-aware collective communication operation.\n- The paper provides extensive experiment results, including the training convergence.", "weaknesses": "- In full parameter training, 512 sequence length is too small. Do you have much longer sequence length, for example, 4K or 8K?\n- The latest Megatron supports more efficient training schemes. It would be helpful to show if PaRO can still achieve significantly higher training performance even when considering these methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your postitive and constructive feedback. We will explain your concerns point by point.\n\n## Weaknesses:\n\n**Q1: Do you have much longer sequence length, for example, 4K or 8K?**\n\n**A1:** We are currently conducting experiments with longer sequence lengths, but we have not had time to obtain the experiment results. Theoretically, longer sequence lengths should not significantly increase communication volume (only parameters and gradients are communicated in DP strategy) and are anticipated to result in similar performance improvements.\n\n**Q2: The latest Megatron supports more efficient training schemes. It would be helpful to show if PaRO can still achieve significantly higher training performance even when considering these methods.**\n\n**A2:**  PaRO could be an alternative to DP in the latest n-D parallel training and advance the training effiency. Please refer to global rebuttal (Author Rebuttal).\n\n## Questions:\n\n**Q1: How PaRO determines the group size (G)? Is it a hyper-parameter?**\n\n**A1:** The PaRO will automatically partition groups according to the nodes, making full use of the communication advantages within the nodes. It also supports other group sizes, such as using the lower-layer switch as the grouping basis in multi-layer switch networks, and of course, custom group sizes can also be defined.\n\n**Q2: In the experiments, do other well-known computation optimizations (e.g., FlashAttention) are applied?**\n\n**A2:** In order to avoid introducing more factors, we did not use optimization schemes such as flash-attn and quantization. But for the LLaMA-65B model, we activated checkpointing to ensure successful training.\n\n**Q3: How can users enable PaRO? Could you provide an example of the programming interface?**\n\n**A3:** We provide the open-source release of our code in paper. To enable PaRO, you just need to add \"paro_strategy\": \"NIG\" in the zero_optimization dict in the original ds_config.json, and start the training the same way as with deepspeed. Configuration examples can be found in the README.md file in the code repository.\n\nThank you again for taking the time to review our paper. We hope we have addressed the reviewer's questions and the reviewer is willing to raise their score. Please let us know if we can provide any additional clarification or information.", "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": "VuXYrrKQB7", "reviewer": "Reviewer_VsLY", "summary": "The paper proposed an unified space to enhance ZeRO-based partitioning strategies, providing a better trade-off between memory and communication. The paper also introduced a more efficient (than ring-based) collective communication method. The core motivation of this paper is to fully leverage the efficient intra-group communication, by introducing intra-group partitioning as an extra option (PaRO-DP) and making the collective communication group-based (PaRO-CC).\n\nThe paper is well written and the core insight is delivered clearly. However, the novelty is limitted. For example, \n- MiCS also improves the efficiency of ZeRO and collective communication based on a similar idea.\n- The partitioning space is similar to Alpa, where the devices are placed into a 2D mesh, with intra-node devices as inner dimension.\n- Megatron-LM suggests tensor parallelism should be within the same node generally.", "questions": "- What's the difference between partial-parameters training and PEFT?\n- If the scope of this paper is for LLM as in the title, could you please add more related work to elaborate the current state of LLM training?\n- Can you elaborate the difference between PaRO-CC and the Hierarchical Communication Strategy in MiCS (or other previous work if there is any)?", "limitations": "How to effectively integrate with current 3D parallelism (or other model parallel techniques) in LLM is a big concern, which may limit its application. In 3D parallelism,\n- ZeRO-1 is usually used to avoid per-microbatch communication, because PP requires a large number of microbatches to reduce pipeline bubbles.\n- tensor parallelism also leverages the high-speed intra-node communication, which will be in conflict with the proposed methods.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper provides a systematic view of ZeRO-based partitioning strategies.\n- The paper enlarges the option space of ZeRO strategies.\n- The proposed PoRA-CC can improve the communication efficiency in practice.\n- The problem is well formulated, easy to follow. \n- The paper considers various senarios, including full/partial parameters training, which helps to understand the motivations.\n- The experiments contain a large number of baselines, which makes it more convincible.", "weaknesses": "I have some concerns about the title and claims.\n- \"Rethinking\" is too ambitious for this paper, regarding it only focuses on ZeRO-based strategies and doesn't combine with any other stategies like TP and PP.\n- \"for Efficient Large Language Model Training\" seems unsuitable. It lacks evidence to prove it's efficient for LLM, as mentioned in Limitations, 3D parallelism + ZeRO-1 is usually used in LLM training. Actually, it's not necessary to couple with LLM, because the proposed methods is general for most of the models.\n- It's kind of overclaiming for \"266% that of SOTA basic DP strategies\", because it compares with ZeRO-3 in Figure 3(d), while MiCS has a similar performance.\n\nFor the experiments, \n- ZeRO-1 is missing, which makes it incomplete and less convincible. ZeRO-1 is a commonly used strategy, especially when there are many microbatches. It would be better if ZeRO-1 is also included for some senarios.\n- It would be better to compare PaRO-CC with the Hierarchical Communication Strategy in MiCS.\n\nFigure 1 is a bit complicated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable and constructive feedback. We will explain your concerns point by point.\n## Weaknesses:\n**Q1: concerns about the title and claims**\n\n**A1:** Thanks again for your insightful advice. We adjust the title to \"Rethinking  Memory and Communication Costs for Efficient Large Language Model Training in Data Parallelism\" to better fit the study of this paper. As stated in the global rebuttal (Author Rebuttal), PaRO can be used with other n-D parallel strategies to accelerate the training of large models. \n\nWe will update the description of performance improvements with more specific and accurate language, like \"PaRO improves the training speed of LLMs by up to 266% that of ZeRO-3 as basic DP strategy\".\n\n**Q2: For the experiments:**\n\na. ZeRO-1 is missing\n\nb. It would be better to compare PaRO-CC with the Hierarchical Communication Strategy in MiCS.\n\n**A2:**\n\na. We conducted a set of experiments comparing PaRO(NII) with ZeRO-1, using the 7B model on 32 GPUs. With the same effective batch size, the throughput of NII was found to be 48.7% higher than ZeRO-1. For detailed data, please refer to Table 1 in the PDF.\n\nb. The difference between PaRO-CC and MICS's Hierarchical Communication Strategy (MiCS-HCS) is that we overlap inter-group and intra-group communication, and we do not need to rearrange data blocks, avoiding unnecessary memory operations. Details can be compared between Figure 2 in the PDF and the MiCS paper. And the collective communication time with 128 GPUs of PaRO-CC, MiCS-HCS and Ring are 162, 183 and 288 ms separately.\n\n**Q3: Figure 1 is a bit complicated.**\n\n**A3:** We have redrawn the figure, please refer to Figure 1 in the PDF.\n\n## Questions:\n\n**Q1: What's the difference between partial-parameters training and PEFT?**\n\n**A1:** To tradeoff between training efficiency and statistical performance, model training may have diverse requirements of trainable parameters in different scenarios. The ratio (trainable-parameters/parameters) of partial-parameters and PEFT is 1/16 and 3/1000 separately. This difference will result in significant differences in memory usage and communication data volume, and the most suitable strategy is also different. Please refer to section 3.1.1 in our transcript and Table 9 in the Appendix for more information.\n\n**Q2: If the scope of this paper is for LLM as in the title, could you please add more related work to elaborate the current state of LLM training?**\n\n**A2:** As mentioned in Weaknesses A1 and global rebuttal (Author Rebuttal), we will adjust the title and give more description of PaRO usage in LLM training.\n\n**Q3: Can you elaborate the difference between PaRO-CC and the Hierarchical Communication Strategy in MiCS?**\n\n**A3:** Same as Weaknesses A2.b.\n\n## Limitations:\n**Q1: How to effectively integrate with current 3D parallelism (or other model parallel techniques) in LLM is a big concern, which may limit its application.**\n\n**A1:** As a future work, we currently working on further utilising PaRO in n-D parallel training for more large-scale LLM training. We find it is effective in two scenarios in the latest n-D parallel training:\n\n1. Scenario 1 is when there is a difference in bandwidth between intra-group and inter-group communication, which is common in GPU clusters. n-D hybrid parallelism usually performs intra-node TP and inter-node PP in a subgroup of nodes, and nesting DP or PaRO in the outer groups. However, the grouping criteria is no longer one machine per group, but rather based on the network topology of the cluster. For example, in a multi-layer switch scenario, grouping can be classified based on the lower-layer switches. For another example, when training with multiple Nvidia SuperPods or similar infrastructures, one SuperPod can be used as a group. \n\n2. Scenario 2. In a vast number of GPU scenario, n-D parallelism usually uses TP + PP + DP for efficient training of over ten thousand GPUs. Following the same principles as discussed earlier, PaRO could provide a more flexible strategy as an alternative to DP in such cases.\n\nPlease refer to global rebuttal (Author Rebuttal) for more detail.\n\nThank you again for taking the time to review our paper. We hope we have addressed the reviewer's questions and the reviewer is willing to raise their score. Please let us know if we can provide any additional clarification or information.", "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:
We thank the reviewers for these insightful comments and constructive advices. We offer a general response here and respond to each reviewer individually. 

The proposed PaRO can be used as a standalone DP strategy or combined with other parallel strategies for n-D parallel training. PaRO-CC communication optimization can also be widely applied to distributed training strategies with global collective communication operations. Next, we'll introduce the practical applications of PaRO.

Firstly, PaRO is not only effective in scenarios where there is a difference between intra-group and inter-group communication costs, but it is also particularly useful when there are a large number of nodes and GPUs (e.g. more than thousands of GPUs). This is because it becomes difficult to maintain a linear speedup ratio for collective communication operators when the number of GPUs is particularly large. The Pytorch community has recently proposed Hybrid Sharding Data Parallel (HSDP) as a 2D strategy for performing FSDP within a host and DDP across hosts to reduce the number of nodes involved in collective communication. While HSDP aims to solve the same problem as PaRO, it was not as mature at the time of drafting this article and therefore was not compared. Additionally, due to hardware limitations, scenarios with more than thousands of GPUs were not tested. Lastly, as demonstrated in our transcript, PaRO provides more flexibility to more complicated machine learning systems, such as distributed RLHF systems. 
 
As a future work, we currently working on further utilising PaRO in n-D parallel training for more large-scale LLM training. We find it is effective in two scenarios in the latest n-D parallel training:
* Scenario 1 is when there is a difference in bandwidth between intra-group and inter-group communication, which is common in GPU clusters. n-D hybrid parallelism usually performs intra-node TP and inter-node PP in a subgroup of nodes, and nesting DP or PaRO in the outer groups. However, the grouping criteria is no longer one machine per group, but rather based on the network topology of the cluster. For example, in a multi-layer switch scenario, grouping can be classified based on the lower-layer switches. For another example, when training with multiple Nvidia SuperPods or similar infrastructures, one SuperPod can be used as a group. 
* Scenario 2. In a vast number of GPU scenario, n-D parallelism usually uses TP + PP + DP for efficient training of over ten thousand GPUs. Following the same principles as discussed earlier, PaRO could provide a more flexible strategy as an alternative to DP in such cases.

It is worth noting that the reason why this article uses "intra-group" and "inter-group" instead of "intra-machine" and "inter-machine" to describe network topology is precisely because the grouping criteria for network topology can be diverse, as is the case when using n-D parallelism. 

It is also worth emphasizing again that PaRO is a non-intrusive distributed training solution for LLM training as ZeRO. We believe it is one of the few feasible  non-intrusive  acceleration solutions for more than thousands of GPUs.

Meta-review: The paper introduces the Partial Redundancy Optimizer (PaRO), a set of basic strategies designed to enhance the training speed of large language models (LLMs) by optimizing memory and communication costs. PaRO Data Parallelism (PaRO-DP) improves training speed through refined model state partitioning, while PaRO Collective Communications (PaRO-CC) enhances communication efficiency by rearranging topology. Experiments show that PaRO increases training speed significantly as compared with SOTA. Overall the results shown in the paper are interesting, and the reviewers all agree that this is an interesting work.
Decision: Accept (poster)