Hierachical Balance Packing: Towards Efficient Supervised Fine-tuning for Long-Context LLM
Authors: Yongqiang Yao
Jingru Tan
Kaihuan Liang
Feizhao Zhang
Jiahao Hu
Shuo Wu
Yazhe Niu
Ruihao Gong
Dahua Lin
Ningyi Xu
Venue: NeurIPS 2025 poster ()
Keywords: Deep Learning
LLM
SFT
LongContext
Primary area: deep_learning
TL;DR: Hierachical Balance Packing, Efficient SFT

Abstract: Training Long-Context Large Language Models (LLMs) is challenging, as hybrid training with long-context and short-context data often leads to workload imbalances. Existing works mainly use data packing to alleviate this issue, but fail to consider imbalanced attention computation and wasted communication overhead. This paper proposes Hierarchical Balance Packing (HBP), which designs a novel batch-construction method and training recipe to address those inefficiencies. In particular, the HBP constructs multi-level data packing groups, each optimized with a distinct packing length. It assigns training samples to their optimal groups and configures each group with the most effective settings, including sequential parallelism degree and gradient checkpointing configuration. To effectively utilize multi-level groups of data, we design a dynamic training pipeline specifically tailored to HBP, including curriculum learning, adaptive sequential parallelism, and stable loss. Our extensive experiments demonstrate that our method significantly reduces training time over multiple datasets and open-source models while maintaining strong performance. For the largest DeepSeek-V2 (236B) MoE model, our method speeds up the training by 2.4$\times$ with competitive performance. Codes will be released at https://github.com/ModelTC/HBP.

Hierarchical Balance Packing: Towards Efficient Supervised Fine-tuning for Long-Context LLM
Yongqiang Yao
Jinru Tan Kaihuan Liang 2. !Shanghai Jiao Tong University 2SenseTime Research
Feizhao Zhang", Jiahao Hu" Shuo Wu", Yazhe Niu 4 Ruihao Gong Dahua Lin?, Ningyi Xu
CCentaal South University "The Chinese University of Hong Kong 'Beihang University
soundbupt@gmail.cm com, tanjingru@csu. edu. edu.cm cn, Liangkaihuannkenssiine com zhangfeizhao, hujiahaol, gongruihaotbuaa. enu dhlintsensetime. com, suningyi@sjtu. edu. cn
com, niuyazhe311Routtlookk com,
Abstract
Training Long-Context Large Language Models (LLMs) is challenging, as hybrid training with long-context and short-context data often leads workload imbal- ances. Existing works mainly use data packing to alleviate this issue, but fail to consider imbalanced attention computation and wasted communication overhead. This paper proposes Hierarchical Balance Packing (HBP), which designs a novel batch-construction method and training recipe address those inefficiencies. In particular, the HBP constructs multi-level data packing groups, each optimized with distinct packing length. It assigns training samples their optimal groups and configures each group with the most effective settings, including sequential parallelism degree and gradient checkpointing configuration. To effectively utilize multi-level groups of data, we design a dynamic training pipeline specifically tai- lored to HBP, including curriculum learning, adaptive sequential parallelism, and stable loss. Our extensive experiments demonstrate that our method significantly reduces training time over multiple datasets and open-source models while main- taining strong performance. For the largest DeepSeek-V2 (236B) MoE model, our method speeds up the training by 2.4x with competitive performance. Codes will be released at https:/gittubb.oo/Moeeeee
Introduction
Large Language Models (LLMs) 2 have achieved state-of-the-ar performance in tasks like machine translation, summarization, and code generation. However, many applications demand to process and understand long-context information legal documents, or retaining context in multi-turn conversations. This underscores the necessity for long-context LLMs that can efficiently process long input sequences.
6 such as summarizing books, analyzing
As mentioned in [] incorporating long context during the Supervised Fine-Tuning (SFT) stage is highly necessary. On the other hand, general short-context data is also erucial to maintain the model's general capabilities. However, this hybrid dataset composition introduces significant challenges, primarily in terms of speed and accuracy. For speed, long-context data intensifies training inefficiencies due to imbalanced workloads; For accuracy, long- long-context data can degrade performance on short-context tasks, affecting the model's general capabilities. These challenges hinder the efficiency and effectiveness of SFT for long-context LLMs.
The workload imbalance caused by the hybrid of long and short data arises from two main aspects: (1) within mini-batch imbalance 8), which caused by excessive padding from randomized mini-batch
"ooresponding authors: gongruihao huw..cdu.cm xuningyi@sjuu.cm
t: equal contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

High Attention Complexity Varlance-
Complexity.
Pocking
3PK Packing
125K Packing Group
Shart
Shart
Wedium
Medium
Long
Medium
Shart
Shart
Short
Medium
Wecium
Comm
High oR
High
Comm
Lang
128K Group
Low
CR
Shert
Long
Long
High
Lang
Long
Comm
a) Naive Packing
b) Hierarchical Balance Packing
Figure 1: Difference between naive packing and hicrarchical balance packing. Short, medium, and long represent different length samples, and SP Comm refers to the additional communication overhead introduced by enabling sequence parallel (SP) training. ABR (Attention Balance Ratio) measures imbalanced attention computation, and CR (Communication Ratio) measures additional communication overhead, described in Section 3.1
construction, and (2) across mini-batch imbalance 18, which is caused by uneven computation distribution over data parallel replicas. Existing approaches mainly address this issue using data packing 16 (9]. which combines variable-Iength data into fixed-length mini-batches. While data packing helps mitigate workload imbalances, it can also introduce new challenges. First, data packing alters the data distribution, which might affect the models' performance. Second, the complexity of attention computation for short- and contect data differs significantly. The attention complexity of the packed samples the sum of the attention complexities of all samples within each packing group. As illustrated in Figure] 1 naive packing leads to high variance in attention complexity by directly mixing long and short samples, which causes imbalanced attention computation and workload imbalance. Third, handling long-context data requires sequential parallelism (SP) and collective communication for attention computation [10] TT, which short data do not need. When mixed, short-context data also requires sequence parallel communication, leading to wasted communication time. Larger models, such as DeepSeek- -2 (236B) [12) MoE (Mixture of Experts) models, introduce higher communication overhead due to the increased number of parameters.
To overcome the limitations of data packing, we propose Hierarchical Balance Packing (HBP), an innovative method that proposes multi-level data packing instead of conventional single-level data packing. HBP consists of three key components: (1) What are the optimal packing groups? (2) How to assign training samples to their optimal group? (3) How can a long-context L.MM be trained with that data? Firstly, we propose hierarchical group auto-selection to determine the optimal packing-length group set and corresponding configurations, including the packing length, Gradient Checkpointing configuration [33, and the SP degree (how many partitions the data is divided into). Secondly, we propose balanced packing to allocate each sample to the optimal group, aiming to minimize imbalanced attention computation and communication overhead. Thirdly, we adopt alternative training between different packing groups, along with curriculum learning and stable loss normalizer to stabilize the training process.
We validate the effectiveness of our method through extensive experiments in multiple settings. For example, on datasets Tulu3 [4] (32K) + Longeite [55 (128K), our approach speeds up 2.4x (57.1 to 23.8 GPU Days) on DeepSeek- V2 (236B) 12] On datasets OpenHerme [16] (4K) + Longeite (128K), our approach reduces training time from 2.95 to 2.04 GPU Days about 1.45 x speeds up on LLama3.1-8B []]. More importantly, method preserves performance on both short- and long-context datasets while achieving significant efficiency gains. Experiments on various models at different scales like LLama3. [I. Qwen2.5-320 (2]. Qwen2.5-72B (2], and DeepSeek-V2 (236B) demonstrate consistent improvements, showing the effectiveness and generalizability of HBP.
2 Related Works
2.1 Long-Context LLM
Long-Context Extension. Long-Conter Extensions aim to enhance L.LMs' capabilities in handling long contexts. Current research can be broadly categorized into approaches that require fine-tuning

---

and those that operate in a zero-shot manner. Zero-shot approaches often leverage techniques such as prompt compression [77] or specially designed attention mechanisms [8] 19. On the other hand, fine-tuning methods primarily focus on extending position encoding, such as RoPE-bbased
approaches [ or utilizing memory-augmentee architectures 20
Long-Context Supervised Fine- Tuning For Long-Context SFT, research mainly concentrates on generating long-context datasets [21] and establishing corresponding benchmarks [22 15 23)
LongAlign [ also focuses on workload balance and accuracy degradation issues. LongAlige mitigate these issues. However, they failed to recognize the imbalanced attention computation and wasted communication overhead due to the packing of short- and long-contex: data. FlexSP [24] dynamically organizes training samples into different data groups and uses flexible sequence parallelism to enable hybrid training. Its online data organization introduces significant overhead (5-15 seconds) each iteration and fails to handle attention computation complexity across data-parallel (DP) groups. In contrast, our method introduces only negligible overhead and achieves better computation balance of attention.
proposed using packing and loss- reweightng
2.2 Data Packing
Data packing [9] is a more practical approach compared to randomly organizing data batches in LLM training. It reduces padding within batches and minimizes idle time across different data-parallel groups. Common packing methods include Random Packing [25]. Sorted Batching [6]. First Fit
Shuffle (FFS), First Fit Decrease (FFD), Best Fit Shuffle (BFS), Shortest-Pack-Fins Histogram- Packing (SPFHP) (26). Iterative sampling and filtering (ISF) (8]. However, those packing methods operate on a fixed length. HBP operates data globally by introducing multiple packing groups of
varying lengths, enabling more flexible and efficient handling of hybrid training with both short- and long-context data.
Problem Analysis
In this section, we first define the notations in Table introduce performance metries in Section 3.1 We then conduct preliminary analysis of the commonly used packing methods in Section 3.2 and training strategies Section 3.3
Table 1: Notation and Definitions
Table 2: Results of ABR and CR in different sequence lengths using packing. Lower metrics indicate more efficient training.
Symbol
Delintion
token mumber one device number of devices
N
local batch size in one token number of i-th sample in B
Packing Len
SP
ABR
CR
DBR
PR
4
computation complexity of attention ~ O maximal token number actoss devices
4K
0.343
0
0.003
0.0
Toux
32K
0.456
1.0
0.001
0.0
Anse
maximal attention computation across N devices
Bee
total number of training iterations
128K
8
0.506
1.0
0.003
0.0
Tomm
token number for SP communication in one iteration
3.1 Measuring Metries
Dist Balance Ratio (DBR) (8] quantifies the computational balance inter-devices based on length. A lower DBR D indicates more balanced workload distribution across different devices.
ΣN (TTmxx T:) PR
ΣB 1max B
(1max - .t)
DBR
Tixx N
Padding Ratio (PR) [ measures the proportion of wasted computations resulting from intra-device padding based on input length. A lower PR I indicates fewer padding tokens.
Attention Balance Ratio (ABR) is proposed to quantify the imbalance in attention computation for different data inter-devices. Previous metries estimate the computational cost of attention based solely on the length of the input with full attention. However, using packing algorithms, attention computation becomes a significant factor in the overall cost. Consider the packing of a 4 sequence example, both {1K, 1K, 1K, LK3 and {2K, 2K} have the same total length. However, their

---

actual attention computation differs significantly, with complexities of 4K² and 8K2, respectively. The Attention Balance Ratio (ABR) and example are given by lower ABR indicates a more balanced attention computation between different DP groups, resulting in less idle time.
ABR Σ (Amax A,) N
8K2 - 4K2
ABR (example) = 8K2 2
0.25
(2)
Amax
Communication Ratio (CR) is proposed to measure the overhead of SP communication. Tomm,, in CR represents the total number of tokens that require SP communication in the current iteration, while T, denotes the total number of tokens in the current iteration. A lower CR indicaatss that fewer tokens require SP communication, resulting in reduced communication overhead.
Average Tokens (Ave-T) is defined as the average number of tokens processed per iteration, serving as measure of the model's workload. A larger Ave-T inndictes higher training efficiency.
Σ Ilefnaa
ΣN N T₁
Σx Toomm, Σ Ti
CR
Ave-T
max X N
3.2 Packing Analysis
Importance of Packing. In Table conduct three batching alternatives, random batching, ISF packing batching [8] (comparison between different packing methods is shown in Appendix A)) sorted batching [6] (ensures that the sequences within each batch have similar lengths) three sequence length 4K, 32K, 128K using Tulu3 [14] dataset. Both sorting and packing reduce DBR and PR, speeding up training. On longer sequences like 1 128K, packing achieves higher Ave-T than sorting, improving GPU utilization and benefiting mixed-length datasets.
Limitation of Packing. Although packing can partially address the efficiency issues associated with hybrid training, several limitations remain. Imbalance Problem of Attention Complexity: Packing samples with varying sequence lengths leads to differing attention computation complexities. Directly mixing these samples can cause imbalanced workloads and inefficient resource utilization. As illustrated in Table 2 a rise device idle time. More details are shown Appendix [ SP Communication Overhead: long sequences require communication for attention computation, whereas short sequences do not. Directly mixing them can lead to extra communication overhead. As shown in Table [ the CR reaches 1 when the sequence length increases, indicating that all short-context data are involved in unnecessary communication.
the ABR increases significantly with the sequence length growth, indicating
Table 3: Comparison of batching strategies different sequence lengths. The local batch size B is adjusted dynamically under the 32K and 128K settings in sorted batching.
Table 4: Results of SP, minimum GC layers, and memory cost across different sequence lengths. Iter Time represents the average time taken over ten iterations.
Sog Len
Batching
DBR
PR
Ave-T
GPU Days (speed up)
Seq Len
SP
GC Layer
Memory
Iter Time
32K
77G
4.45s
4K
random
0.540
0.416
2.4K
(1.00)
32K
78G
4.35s
4K
sorted
0.001
(2.4)
32K
8
78G
4.12s
4K
packing
0.003
0.0
4K
3.1 (2.6%)
64K
2
0OM
32K
random
0.64
0.0
0.8K
16.7 (1.0x)
64K
78G
6.3s
32K
sorted
0.01
0.02
30.2K
4.8 (3.5x)
64K
8
24
79G
6.2s
32K
packing
0.0007
0.0
32K
4.4 (3.8%)
128K
00M
128K
0.639
(1.x)
128K
78G
10.2s
128K
sorted
0.01
0.02
125K
5.5 (6.9%)
128K
16
23
79G
10.5s
packing
(7.3x)
3.3 Training Strategy Analysis
Given the GPU resources and the data to be trained, we can select from various training strategies, provided that the VRAM requirements are satisfied. The main factors to consider are the degree of SP and the configuration of Gradient Checkpointing (GC), which is the number of layers where GC is enabled. If the SP degree is small, the VRAM demand is high, which forces increase in

---

Adoptive 5P
Curriculam Learring
Iteviies-1
Data Rool
Preflling
Packing
Groups
Mode Treley
Duration skit
Tonal Treaning
Comman
Tite
lectioon
Pocking
Audiin
Node Traneng
Tonal Troiing
DiE Ralhya
Deermmin-000 with 3t
Tonal Theck
Data Prepare Pipeline: Hierorchical Packing Groups
Dynomic Training Pipeline
Figure 2: Hierarchical Balance Packing training framework.
the number of GC layers and leads to excessive additional computation. On the other hand, if the SP degree is large, although the VRAM demand is reduced, it introduces additional communication overhead. Therefore, there is trade-off between the SP degree and GC configuration. Moreover, the optimal strategy varies for different sequence lengths. 4 that the optimal strategies for 32K, 64K, and 128K are different.
Hierarchical Balance Packing
In this section, we first describe how determine optimal packing length groups (Section 4.1 then explain how each sample is assigned to a group (Section 4.2). Finally, we introduce dynamic training pipeline for HBP (Section The overall framework is shown in Figure
4.1 Hierarchical Groups Auto-Selection
To determine the optimal packing length groups, we design a profile-based auto-selection algorithm described in Algorithm T operates in stages: (1) finding the best training strategy for predefining possible sequence length set (e.g., 8K, 16K, 32K, 64K, 128K) based on naive packing. (2) Deriving the final packing groups by optimizing communication overhead.
Algorithm 1 Hierarchical Groups
Auto-
Algorithm 2 Balance Packing
Selection
1: Inputs: Dataset D {2), 22, 22. ......... ). hier- archical packing groups Lpp = (11, 12......4.. .......,, (n)
1: Inputs: Lengths L, Profile Time P, Strategy S
2: G {G1, G2, ................................ Gn J: packed data group
2: Stage-1: Find the best training strategy
3: = [B,, B2..... B.): final batched data splits D subsets D1. Dz...., Dm. by packing groups Lp
3: Initialize P + []. -
4: GroupData D, Lp):
4: for each 1 do
5: 8 s= =spp + FindBestSpCkp(()) Paddd(PoolleTime(s)) Sa
5: Packing G; = -D,, 1,): packing D, by 4
6:
6: GreedyFill G; = (G,,4, [D.-...... D.D:: D: fill
7: end for
data from smaller groups to reduce PR.
8: j argmin(P) P))
7: Balance Batching(Gs): construct balance batches by sorting packed data according to attention complexity in Section 3.1
9: Shear S(j), hest L[j
8max <- -[--1], 1,1,, + [[]]
11: Stage-2: Optimize packing groups for communication
8:
9: Initialize B +
12: 4 + /hees/ /hess//es..pp.
la +- /lmax flmxx//mxxxu sp_
[D..... Dn. +- GroupData(D, Lp)
if 1/2 > Thest then
11: for = n to do
14: Lp + [1, lbests lz, mmax]
G, <- Packing(D,.... 4)
else
13: G, t GreedyFill(G, 4. D,-1....,Dil
16:
Lp + [11, loest, Imax]
B, + Balance Batching(G.). B.add(B))
end if
15: end for
18:
return I.p
return Shuffle(B)
Stage I: Find the best training strategy. The algorithm begins by initializing two empty lists, and S, to store the profiling times and the corresponding strategies, respectively. For each possible
5

---

16K Pocking Group
16k Pocking Group
SP
16K
No
Naive 32k Packing Grouup
Shart
Shart
Shart
Short
Comm
Mediar
Mediun
Mesian
32K Pocking Group
SP
SP
Medium
Meskam
Mediun
Naive Packing Groups [32K]
Optimized Packing Groups [16K, 32K]
Figure 3: Example of optimizing packing group for communication.
input length we compute the optimal degree of SP and the configuration of the GC 8 = (sp, ckpl) with the FindBestSpCkpt function, which achieves the optimal trade-off between the degree of SP and the configuration of the GC. More details of FindBestSppCpt are shown in Appendix [.1
Specifically, given packing length I, we iterate all possible sp degrees and ckpt GC configurations and profile their iteration time. The best combination (sp, ckpt) then found by a greedy method. Once all input lengths are processed, the algorithm selects the best strategy by identifying the index that minimizes the profiling times in P. The optimal strategy and input length are:
j = argmin(P), P);
8best = Sli], luest L[j]
(4)
Additionally, the maximum input length lmxx and its corresponding strategy 8max are also retained.
Stage 2: Optimize packing groups for communication.
After calculating the optimal besst and the corresponding Sbeats we can optimize it further to reduce communication overhead and obtain the final packing group lengths. Taking 32K packing groups Figure 3 with SP=2 as an example, each SP process handles a split 16K tokens. For samples longer than 16K, additional communication is required anyway. For samples shorter than 16K, can group them into 16K packing groups, which is equivalent to training with 16K packing groups using SP=1, thereby reducing communication.
best
1max
1
and
l2
(5)
Sbostsp
Smax.sp
For best, the smallest packing group 1 is derived based on its optimal degree of SP Sbestsp- ensuring communication overhead during sequential parallelism. Similarly, the smallest packing group l2 for mmxx is also considered. If l2 is smaller than heem. it will be merged into the lhest range. Finally, the hierarchical packing groups Lp {11, Jbest 12, [maz] is obtained. More implementation details shown in Appendix
4.2 Balance Packing
After obtaining the optimal hierarchical pack groups, we distribute the dataset samples into different groups while ensuring that the metries outlined Section 3.1 (DBR, PR, ABR, and CR) are well-optimized, which conventional packing struggles with. We first divide the entire dataset into sub-datasets [D1, D P,) based on T.p. For each D,, the following steps are executed:
(1) Packing: We pack the data in D, to length 4, This ensures low PR and DBR. Note that arbitrary packing methods are feasible.
(2) GreedyFill. We use the remaining unpacked data for greedy fill, as larger groups struggle to fill the packed data using only their samples. A simple example illustration is shown in Appendix E
(3) Balance Batching: Constructing balance batches B, based on attention complexity using heuristic solution sorting support maintaining balanced attention computation across different packing groups. Heuristic solutions are an efficient, approximate solution for large-scale datasets. More evidences are shown in Appendix D

---

Sum Lows
Average Token Lown
1.4
CL
1208 50h
130c
1.5e-5
Orod Mattn
1.0
De-5
0.8 60
20
100
1000
420
Stop
Stop
Figure 4: Loss with and without curriculum learning (CL).
Figure 5: Grad Norm of Sum loss and Average Token loss.
The procedure of balanced packing is illustrated in Algorithm 2 Since our approach inherently involves multiple levels, i.e., hierarchical packing groups, it automatically separates short- and long- context data, avoiding wasted communication overhead and imbalanced attention computation and reducing CR and ABR significantly. We also achieve low PR and DBR at multiple levels, thanks to GreedyFill. More details about the implementation are shown in Appendix J
4.3 Dynamic Training Pipeline
Since HBP involves multi-leve inputs, it is essential design dynamic training pipeline, enabling hot switching of different packing groups. Adaptive SP (pre-initialized multi-SP achieving zero overhead) is proposed to ensure efficient training with multi-level packing. Curriculum Learning and a Stable Loss Normalizer are proposed to improve the performance in hybrid training.
Each packing group is assigned an optimal training config configuration (sp, ekpt), and we adopt an alternating scheme that selects the best SP degree and GC setting per packing group, as illustrated in Figure 2
Adaptive Sequential Parallel:
Curriculum Learning Strategy: Training on long-context tasks presents challenges because initi- ating training without instructional capabilities can result in significant fluctuations in the training loss, as illustrated Figure Thanks to our inherent hierarchical structure, is straightforward to adopt a curriculum learning strategy that begins with general short-context data in the early stages of training. As training progresses, we transition a hybrid approach that alternates training both short- and long-context data, as illustrated in Figure 2
Stable Loss Normalizer: The training stability introduced by data packing is an important problem, as it impacts the data distribution. Previous work [6] has analyzed loss calculation, identifying two primary loss normalizers Lioken (Token-Mean) and Sample (Sample-Mean):
Σ loss; Σ B: T
ΣB Bi loss;
ΣB' T B,
Σ loss;
Lroken
Sample
Tave
Litable
(6)
B
B, Tave
where Bj represents the local batch size within the DP group, and T; denotes the number of loss tokens in the batch i. They argued that instability lies in the inconsistent loss of normalization. [27) also proposes sum loss without normalization to mitigate the effect of norm discrepancies. However, the sum loss introduces a trade-off: as the sequence length increases, the gradient values escalate disproportionately (1e+5), as shown in the left part of Figure 5) To address the above issues, we empirically observe Tave (Average Token) of the global batch size By across all DP groups can serve as stable loss normalizer. The stable loss normalizer is derived from theoretical considerations, with the objective of guaranteeing that each token contributes equally the aggregate loss. This formulation mitigates biases introduced by heterogeneous sequence lengths, varying data-parallel group sizes, and gradient accumulation strategies. We illustrate the effect of different normalization strategies under simple setting: a data-paralle, (DP) group of size 2, where each rank processes local batch of samples.
Laken-Mean- In the token-mean approach, the loss for each DP rank is normalized by its local token:
loss + loss2
loss + loss
Lap,
Cupa =
Ti
T2
T3 Ta
7

---

The global loss is then averaged across ranks:
Lap, + Lap
l0ss + lossa
Josas
Lümal
25+
2
2
This formulation normalizes losses rank-wise without accounting for global token distribution, potentially introducing bias when sequence lengths vary across ranks.
Sample-Mean: In the sample-mean method, losses are first normalized per sample and then averaged:
Joss
lossq Ta
losss
losss
Eap
Ti
Capa
2
Ta
2
The global loss becomes:
Lap, + Ldpa
lussi Ti
loxss T2
+
losss
loss4
Cinal
T'
T4
2
4
This strategy balances individual samples but fails to reflect the total token count, giving undue weight to shorter sequences.
Stable Loss Normalizer: To eliminate such bias, we normalize by the global average token length:
T1
T2 + T3 T
loss;
loss
loss + loss
Lave
Lap.
Lap,
2. Tave
2 Tave
Each DP's loss then normalized by Tave. And the final loss is:
Lap, + Cap:
loss + loss + lossy + loss
loss; + loss + loss; + loss
Ginal
2
4 Lave
T. + T2 T3 T4
This normalization ensures that the contribution of each token is equal, regardless of which sample or rank it belongs to, thereby producing an unbiased global loss.
5
Experiments
5.1 Experimental Setup
Implementation Details. We use large-scale datasets: Tulu3 (32K) for general tasks and LongCite (128K) for long-context tasks. Both have shown strong performance across various benchmarks. They also enhance the model's ability to handle input lengths from O.1K to 128K tokens. We conduct experiments with the following models: LLaMA 3.1 [. Qwen-2.5 2 and DeepSeek-V2 (236B). Most models are trained on 32x H100 80GB GPUs using the DeepSpeed 2 while DeepSeek-Ve V2 (236B) is trained with the Megatron-LMM 29 with 256x H100 80G GPUs. In our experiments, we use DeepSpeed-Ulysseee
[10] sequence parallelism approach, and the ring-attention [ method is also applicable. We conducted ablation experiments using the LLaMA3. 1-BB model. For the Longsite dataset, approximately 2k samples uniformly sampled. The loss normalizer for baselines without special instructions Token Mean, while HBP uses Ave-Token. GPU days are the evaluation metric to estimate the total training time. We use learning rate of le-5, weight decay of 0.01, and adopt Adam as our optimizer.
Evaluation. We comprehensively evaluated the LLM's performance using OpenCompass 30 For general tasks, several benchmark datasets were assessed, including MMLUU 3I, MMLU PRO CMMLU 33 BBH (34 Math (35]. GPQA Diamond [36], GSM8K 37 HellaSwar [8) Math-
32
Bench (39) HvmanEval 40) MBPP 41 IFEval 42 and Drop 43 For long-context tasks,
evaluation including Ruler
44 NeccleBench [55. LongBench [46] and Longcite.

---

Table 5: Results of different models. The naive packing baseline ISF uses the Token-Mean loss normalizer. LongAlign uses their proposed lossrrewegg. AVE represents the average performance on general tasks. Deepseek- V22336BB is trained in 32K setting due to resource constraints.
Model
General Tasks
I Tosks
GPU Days
(lype)
Math
GSMSK
Humankval
Roler (32K)
LongBiench
LongCite
Ispeed
Lano3 J-80
LongAlign packing
(035736)
Long Aligosooottt
53.8
16x)
ISP
56.0
54.5
65.3
70.4
33.7
81.7
62.8
85.0
44.0
71.6
5.22 (1.0x)
HBP
58.2
63.00
67.2
67.7
33.0
81.9
63.4
85.6
70.8
43.1
71.5
377 (1.4%)
Quen2.5 320
LSF
73.5
74.8
X3.5
75.6
56.5
95.7
86.6
K8.2
59.3
51.0
602
21.3 (1.UX)
HBP
76.2
76.6
83.6
76.0
57,1
94.4
84.2
R8.3
5900
$1.9
61.7
16,0 (1.334)
LLannd.
ISF
934
(1.42x)
DeepSeek V2 (2360)
ISF
71.8
76.8
84.0
71.1
41.3
89.0
78.6
86.6
47.1
57.1 (1 0x)
HBP
72.9
76.5
83.1
72.6
41.4
89.9
78.1
87.3
50.3
23,8 12.48)
Table 6: Results of HBP Components. This experiment the Token-Mean Loss Normalizer. Hierarchical indicates the enabled hierarchical packing. Balance refers to balance batching.
Model
Hierarchical
Balance
ABR
CR
AVE
GPU Days (speed up)
LLaMA3.1 1LaMA3.1-8P
0.506
1.0
56.0
5.22 (1.0x)
0.288
0.173
56.4
4,51 (1.2x)
LLaMA3.1-8B
0.002
0.173
3,73 (1.40))
1LaMA3.1-70B
0.506
1.0
72.2
44.4 (1.0x)
LLaMA3. 1-70B
V
0.288
0.173
72.8
33.3 (1.25x)
LLaMA3.1-70B
0.002
0.173
72.2
31.1 (1.43))
5.2 Main Results
In Table we compare our method HBP with LongAlign and packing method ISF S. We also notice that Align improves the average general tasks to some extent; it sacrifices performance in long tasks e.g., Ruler- 128K. In contrast, our method maintains strong performance both in general short tasks and long tasks. The improvements from 8B to 236B parameters. Notably, for the largest MoE model, DeepSeck- V2 (236B), our method achieves impressive 2.4x training speed-up, reducing training time from 57.1 to 23.8 GPU days. Full results are shown in Appendix
consistent across wide range of model sizes,
5.3 Ablation Results
Importance of Hybrid Training. Table showw that both short- and long-context data are essential for maintaining general and long-text capabilities. "Longcite (8K)" includes only sequences up 8K. All settings use naive packing. Removing long-context data (row 1) harms long-text performance, while removing short-context data (row 2) weakens the general ability. Row 3, with partial short data, confirms these effects.
Components of HBP. In Table we show that the Attention Balance Ratio (ABR) and Communi- cation Ratio (CR) can be reduced significantly with hierarchical packing. In particular, ABR drops from 0.506 0.288, and the CR drops from 1.0 0.173. By balancing the batching data with similar complexity of attention computation, we have much more balanced mini-batches with low ABR, from 0.288 to 0.002. Overall, achieve 1.4x speedup. Similar results can be observed in the larger LLaMA-3.1 model.
Curriculum Learning. Table presents the impact of curriculum learning HBP. Starting with short tasks and gradually mixing in long-cone long-context tasks benefits training and convergence. We applied similar curriculum strategy to the naive ISF baseline using advanced sampling, which also showed improvements. This confirms the general effectiveness of curriculum learning for long-conted SFT. Notably, HBP naturally separates short and long contexts, making curriculum learning easier apply.
Iterations for Curriculum Learning training. Table g presents detailed ablation studies across different models. As shown, even simple curriculum learning setup, adding 100 steps of carly-stage

---

Table 7: Results of curriculum learning (CL).
Table Results of different CL iterations.
Model
CL.
AVE
LongBench
Model
CI-Bteratoon
AVE
LongBench
LLam33.114 1-88-HEP
56.6
41.6
LLamak 30 H1BP
56.5
41.6
LLam33.1-88 1-BB-BBPP
58.2
43.1
LLam33 4-
58.0
43.1
LLama.. I 708 HCP
51,5
1-RR-HRP
Lama3. 1-708-H0PP
74.2
52.2
1Lama3
L
44,0
1Lama3
73.8
3 80--SF
$7.4
434
H-BP-HHPP
74.2
Table 9: The importance of hybrid training.
Table 10: Results of different loss normalizers.
Tistanet
AVR
Langlouch
Raaees
Normalioer
LangSench
Rule--12KK
Longrite
32K
43.0
52.5
138K
18.5
16.7
68.5
Sanple
55.
42.9
Todu3 e Longatoo 583
32K
58.6
43.2
62.1
Token Mean
56.6
416
67.
20.6
Tuby Longait
176K
56.0
44.0
67.5
Token
57.6
4311
708
71.2
training with short samples, already yields noticeable improvements. To ensure training stability, we ultimately adopt a setting with 500 steps of short-sample-oniy training in our final experiments.
Stable Loss Normalizer. We compared several loss normalization methods by training models using different normalizers while keeping all other training configurations consistent. As shown in Table I0 the Ave-Token loss normalizer achieved the highest performance in both general tasks, Average (AVE) 57.6, and long context tasks LongBench 43.1, Ruler-128K 70.8, and LongSite 71.2.
Importance of Hierarchical Groups Auto-Selection. Table 11 compares manual (row 2) and automatic (row 3) packing groups. Automatic selection of groups (based on Appendix F) achieves greater speedup. reducing training time from 4.25 3.73 GPU days.
Table 11: Ablation results of optimizing packing groups for communication. Groups represent the
packing group lengths used during training, while SP refers to the Sequence Parallel degree.
Model
Groups
Selection
LungBench
Days (speed up)
LLamm'. 60-ISP
[128K]
manual
[8]
0.506
44.0
(1000
LLama3. 1-88 HBP
BEK, 128K]
manual
12. N1
0002
1.0
57.9
432
4.25 (1 22x)
E.Tama3 1-88-HBP
[16K
anto
[1.8]
0002
0.173
58.2
431
3,73 (1.4%)
Compared with FlexSP. As shown in Table [2] HBP achieves lower ABR, DBR, and PR than FlexSP, while maintaining a similar CR, contributing to faster training (4.35 to 3.73 GPU days), even considering only the model computation time. Furthermore, FlexSP introduces per-batch grouping overhead (5- 15 seconds), which further slows down training (see Appendix B). HBP performs global data-level operations in advance with negligible data overhead.
Negligible Overhead. The overhead of HBP is negligible: Auto-Selection takes ~ minutes and Packing 5 seconds, together contributing less than 2% of training time (see Appendix G
Table 12: Ablation results with FlexSP. All metrics are calculated based on the official code. For FlexSP, we only consider the model computation time and ignore the impact of data grouping time.
Model
CL
PR
DER
CR
ABR
AVE
LongHench
GPU Days (speed
LLamad. $D IST
0
0.001
1.0
0.506
55.0
44.0
(1
LLama3.1 FiesPP 24
11244.1
0U4
01064
0.171
0.36
56.1
(1
Hamms' 1-88-HRP
0.123
1.73(14xz)
H-R--HRP
(.73(142)
Dataset Generalization. verify the generalization our method, we conducted experiments on different datasets (OpenHermess[])) (16), NongWriter]e (3))) Our method also achieves effective results on these datasets. Results are shown in the Appendix
Conclusion and Limitations
In this paper, proposed Hierarchical Balance Packing (HBP), a novel strategy to address workload imbalances in long-context LLM training through multi-level data packing and a dynamic training pipeline. Due to limitations in computational resources, we have not conducted experiments on longer contexts, such as 256K or 512K. Additionally, we have not validated HBP on other post-training tasks, such RLHF or DPO. We leave these explorations to future work.
10

---

Acknowledgements
This work was supported by the National Key Research and Development Program of China (Grant No. 2023 102), the Postdoctoral Fellowship Program and the China Postdoctoral Science Foundation (Grant No. BX20250487), and in part by the Natural Science Foundation of Hunan Province (Grant No. 2024JJ6525).
References
[1] A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten,
A. Yang, A. Fan et al., "The Ilama 3 herd of models," arXiv preprint arXiv:2407.2178. 2024.
[2] A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei al., "Qwen2. 5 technical report," arXiv preprint arXiv:24/2. 15115, 2024.
[3] 7. Cai, M. Cao, H. Chen, K. Chen, K. Chen, X. Chen, X. Chen, Z. Chen, Z. Chen, P. Chu,
X. Dong, Duan, Q. Fan, Z. Fei, Y. Gao, J. Ge, C. Gu, Y. Gu, T. Gui, A. Guo, Q. Guo, C. He,
Y. Hu, T. Huang, T. Jiang, P. Jiao, Z. Jin, Z. Lei, J. Li, J. Li, L. Li, S. Li, W. Li, Y. Li, H. Liu,
J. Liu, J. Hong, K. Liu, K. Liu, X. Liu, C. Lv, H. Lv, K. Lv, L. Ma, R. Ma, Z. Ma, W. Ning,
L. Ouyang, J. Qiu, Y. Qu, F. Shang, Y. Shao, D. Song. Z. Song, Z. Sui, P. Sun, Y. Sun, H. Tang,
B. Wang, G. Wang. Wang. J. Wang. R. Wang, Y. Wang. Z. Wang. X. Wei, Q. Weng, F. Wu,
Y. Xiong, C. Xu, R. Xu, H. Yan, Y. Yan, X. Yang, H. Ye, H. Ying, Yu, J. Yu, Y. Zang, C. Z.Zhang,
L. Zhang, P. Zhang, P. Zhang. R. Zhang, S. Zhang. S. Zhang, W. Zhang, W. Zhang, X. Zhang,
X. Zhang, H. Zhao, Q. Zhao, X. Zhao, F. Zhou, Z. Zhou, J. Zhuo, Y. Zou, X. Qiu, Y. Qiao, and
D. Lin, "Intermm22 technical report," 2024.
[4] S. Chen, S. Wong, L. Chen, and Y. Tian, "Extending context window of large language models via positional interpolation," arXiv preprint arXiv:2306. 15595, 2023.
[5] B. Peng, J. Quesnelle, H. Fan, and E. Shippole, "Yarn: Efficient context window extension of large language models," arXiv preprint (XXiv:230900077/ 2023.
[6] Y. Bai, X. Lv, J. Zhang, Y. He, J. Qi. L. Hou, Tang, Y. Dong, and J. Li, "LongAlign: A recipe for long context alignment of large language models," in Findings of the Association for Computational Linguistics: EMNLP 2024. Miami, Florida, USA: Association for Computational Linguistics, Nov. 2024, pp. 1376-136 1395. [Online]. Available: https//aclanthologyyy orgey.org//22/4 2tt///222...nn findiigss emmlp..-4mmllp744...
[7] L. Ouyang. J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray et al., "Training language models to follow instructions with human feedback."
Advances in neural information processing systems, vol. 35, pp. 27 730-27 744, 2022.
[8] Y. Yao, J. Tan, J. Hu, F. Zhang, X. Jin, B. Li, R. Gong, and P. Liu, "Omnibal: Towards fast instruct-tuning for vision-language models via omniverse computation balance," arXiv preprint arXiv:2407.20761, 2024.
[9] S. Wang. G. Wang, Y. Wang, J. Li, E. Hovy, and C. Guo, "Packing analysis: Packing more appropriate for large models or datasets in supervised fine-tuning," arXiv preprint arXiv:2410.0800 2024.
[10] S. A. Jacobs, M. Tanaka, C. Zhang, M. Zhang, S. L.. Song, S. Rajbhandari, and Y. He, "Deep- speed ulysses: System optimizations for enabling training of extreme long sequence transformer models," " arXiv preprint Stti::2309.145009 2023.
[11] H. Liu and P. Abbeel, "Blockwise parallel transformer for large context models," Advances in neural information processing systems, 2023.
[12] DeepSeek-AI, "Deepseek-v2: A strong, economical, and efficient mixture-of-expers language model," 2024.
[13] M. Li, D. G. Andersen, J. W. Park, A. J. Smola, A. Ahmed, V. Josifovski, J. Long, E. J. Shekita, and B.-Y. Su, "Scaling distributed machine learning with the parameter server," in 11th USENIX Symposium on Operating Systems Design and Implementation (OSDI 14). Broomfield, CO: USENIX Association, Oct. 2014, pp. 583-598. [Online]. Available: mu

---

[14] N. Lambert, J. Morrison, V. Pyatkin, S. Huang, H. Ivison, F. Brahman, L. J. V. Miranda,
A. Liu, N. Dziri, S. Lyu. Y. Gu. S. Malik, V. Graf, J. D. Hwang. J. Yang, R. L. Bras, O. Tafjord, C. Wilhelm, I. Soldaini, N. A. Smith, Y. Wang, P. Dasigi, and H. Hajishirzi, "Tülu 3: Pushing frontiers in open language model post-training," 2024. [Online]. Available:
[15] J. Zhang. Y. Bai. X. Lv, W. Gu, D. Liu, M. Zou, S. Cao, L. Hou, Y. Dong, L. Feng, and J. Li, "Longcite: Enabling Ilms to generate fine-grained citations in long-context qa." arXiv preprint (rXiv2209.0289/ 2024.
[16] Teknium, "Openhermes 2.5: An open dataset of synthetic data for generalist Ilm assistants,"
2023. [Online]. Available:
[17] H. Jiang, Q. Wu, X. Luo, D. Li, C.-Y. Lin, Y. Yang, and L. Qiu, "Longllmlingua: Acceler- ating and enhancing 11ms in long context scenarios via prompt compression." arXiv preprint arXiv:2310.088 1..06999. 2023.
(18) Y. Sun, L.. Dong, B. Patra, S. Ma, S. Huang, A. Benhaim, V. Chaudhary, X. Song, and transformer," in Proceedings of the 6lst Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers). A. Rogers, Boyd Grabeer. and N. Okazaki, Eds. Toronto, Canada: Association for Computational Linguistics, Jul. 2023,
F. Wei. "A
pp. 14590-14 604. [Online]. Available: httpss//aclanthologgy.ot wtwp://ogg..rg//2.aaaooooooo
2023 816/
[19] C. Han, Q. Wang, H. Peng, W. Xiong, Y. Chen, H. Ji, and S. Wang. "LM-infinite: Zero-shot extreme length generalization for large language models," " in Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Pupers), K. Duh, H. Gomez, and Bethard, Eds. Mexico City, Mexico: Association for Computational Linguistics, Jun. 2024, pp. 3991-4008, [Online]. Available: httts://cclmmhooooy.org/// m00g:200020m..llloo22222222
[20] C. Packer, S. Wooders, K. Lin, V. Fang, S. G. Patil, I. Stoica, and J. E. Gonzalez, "MemGPT: Towards Ilms as operating systems," arXiv preprint arXiv:2310.08560, 2023.
[21] W. Xiong. J. Liu, Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz et al., "Effective long-context scaling of foundation models," arXiv preprint arXiv:2309, 16039, 2023.
[22] Y. Chen, S. Qian, H. Tang, X. Lai, Z. Liu, S. Han, and J. Jia, "Longlora: Efficient fine-tuning of long-context large language models," arXiv preprint arXiv:2309. 12307, 2023.
[23] Y. Bai, J. Zhang, X. Lv, L. Zheng, Zhu, L. Hou, Y. Dong, Tang, and J. Li, "Longwriter: Unleashing 10,000+ word generation from long context llms," arXiv preprint 4rXiv22408.07055 2024.
[24] Y. Wang. S. Wang. S. Zhu, F. Fu, X. Liu, X. Xiao, H. J. Li. F. Wu, and B. Cui, "Data-centric and heterogeneity-adappie sequence parallelism for efficient Ilm training," arXiv preprint arXiv:2412.01553, 2024.
[25] X. Contributors, "Xtuner: A toolkit for efficiently fine-tuning Ilm," 2023.
[26] M. M. Krell, M. Kosec, S. P. Perez, and A. Fitzgibbon, "Efficient sequence packing without eross-contamination: Accelerating large language models without impacting performance," arXiv preprint wrXiv:2107.00022 2021.
[27] N. Lambert, J. Morrison, V. Pyatkin, S. Huang. H. Ivison, F. Brahman, I., J. V. Miranda, A. Liu,
N. Dziri, S. Lyu et al., "T)"" ulu 3: Pushing frontiers in open language model post-training." arXiv preprint 2w.
2024.
[28] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. "Zero: Memory optimizations toward training trillion parameter models." " in SC20: International Conference for High Performance Computing. Networking, Storage and Analysis. IEEE, 2020, pp. 1-16.
[29] M. Shocybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, "Megatron-Im Training multi-billionn parameter language models using model parallelism," arXiv preprint 190i 2019.
[30] O. Contributors, "Opencompass: A universal evaluation platform for foundation models." 2023.
[31] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt, "Measuring massive multitask language understanding." Proceedings of the International Conference on Learning Representations (ICLR), 2021.
12

---

[32] Y. Wang, X. Ma, G. Zhang, Y. Ni. A. Chandra, S. Guo, W. Ren, A. Arulraj, X. He, Z. Jiang et al., "Mmlu-pro: A more robust and challenging multi-task language understanding benchmark," arXiv preprini arXiv:2406.01.54 2024.
[33] H. Li, Y. Zhang, F. Koto, Y. Yang, H. Zhao, Y. Gong, N. Duan, and T. Baldwin, "Cmmlu: Measuring massive multitask language understanding in chinese," 2023.
[34] M. Suzgun, N. Scales, N. Schärli, S. Gehrmann, Y. Tay, H. W. Chung, A. Chowdhery, Q. V. Le, E. H. Chi, D. Zhou, and J. Wei, "Challenging big-bench tasks and whether chain-of- of-thhuggtt solve them," arXiv preprint arXiv:2210.09261. 2022.
[35] H. Saxton, Grefenstette and Kohli, "Analysing mathematical reasoning abilities of neural models," " arXiv: arXiv: 995540050.5 2019.
[36] D. Rein, B. L. Hou, A. C. Stickland, Perty, R. Y. Pang. J. Dirani, J. Michael, and S. R. Bowman, "GPQA: A graduate-level googlc-proof q&a benchmark." in First Conference on Language Modeling, 2024. [Online]. Available: https://openreiiew.eee/oooo 100
[37] K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek,
J. Hilton, R. Nakano, C. Hesse, and J. Schulman, "Training verifiers to solve math word problems," arXiv preprint arXiv:2 10. 14168, 2021.
[38] R. 7ellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi, "Hellaswag: Can machine really finish your sentence?" in Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, 2019.
[39] H. Liu, Z. Zheng. Y. Qiao, H. Duan, Z. Fei, F. Zhou, W. Zhang, S. Zhang, D. Lin, and K. Chen, "Mathbench: Evaluating the theory and application proficiency of ilms with a hierarchical mathematics benchmark," 2024.
[40] M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. de Oliveira Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, A. Ray, R. Puri, G. Krueger, M. Petrov, H. Khlaaf, G. Sastry, P. Mishkin, B. Chan, S. Gray, N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, Tillet, F. P. Such, D. Cummings, M. Plappert, F. Chantzis, Barnes, A. Herbert-Voss, W. H. Guss, A. Nichol, A. Paino, N. Tezak, J. Tang, I. Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leike, Achiam, V. Misra, Morikawa, A. Radford. M. Knight, M. Brundage, M. Murati, K. Mayer, P. Welinder, B. McGrew, D. Amodei, S. McCandlish, I. Sutskever, and W. Zaremba, "Evaluating large language models trained on code," 2021.
[41] J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le et al., "Program synthesis with large language models," arXiv preprint arv: 108.02008..732.. 2021.
[42] J. Zhou, T. Lu, S. Mishra, S. Brahma, S. Basu, Y. Luan, D. Zhou, and L. Hou, "Instruction-followwe
evaluation for large language models," 2023. [Online]. Available: htss://axiivoor//abs3311111
[43] D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner, "DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs." in Proc. of NAACL, 2019.
[44] C.-P. Hsich, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg, "Ruler: What's the real context size of your long-context language models?" arXiv preprint 1rXiv:2404.06654 2024.
[45] M. Li, S. Zhang, Y. Liu, and K. Chen, "Needlebench: Can Ilms do retrieval and reasoning in 1 million context window?" 2024. [Online]. Available: https://anxi.orr/abb/24000
[46] Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng. L. Hou, Y. Dong, J. Tang, and J. Li, "Longbench: A bilingual, multitask benchmark for long context understand- ing," 2023.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the contributions and scope of the paper?
Answer: [Yes]
Justification: The main claims in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in 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.
It is fine to 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:
The paper provides brief discussion of the limitations in the Conclusion section.
Guidelines:
The answer NA means that the paper has 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 out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiitiinn 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 on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't '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?
14

---

Answer: [NA]
Justification: The paper is primarily based on empirical observations and proposes heuristic approaches, which not fully supported by formal mathematical proofs or complete set theoretical assumptions.
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 the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's 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 affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
The experimental results can be fully reproduced based on the information provided in the paper, and the authors have indicated that the code will be released in the future.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to 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 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.
While NeurlPS 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 a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way 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 describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [NA]
Justification: The code and data are not yet publicly available, and although the authors plan to release them later, they are currently not provided with sufficient instructions in the supplemental material to ensure faithful reproduction of the results.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips... cc/
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 rum to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.co tttpcc//pbbb. muublc//guidee wttggdddee/eeeeuubmmmmmeeee www./Sooieesisiiaaooaa wommissssoooiccy 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 as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code 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 full details provided in the main text and supplemental material.
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 level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or the standard error of the mean.
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 range (e.g. negative еπτог 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 text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper provides sufficient information about the computational resources used, including details in both the main text and the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///.pppapp...poooooooccctaa
Answer: [Yes]
Justification: The research presented in the paper fully conforms the NeurIPS Code of Ethics in all respects.
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 to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
17

---

If the authors answer NA or 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 or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess 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 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 out that an 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 to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: 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?
Answer: [Yes]
Justification; All resources used in the paper comply with their respective licenses and terms use.
Guidelines:
The answer NA means that the paper does 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 used and, possible, include a 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.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it 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 release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

A Packing Strategy Results
Table I3 illustrates the complexity of different packing strategies and their corresponding final training times. It also shows that the training indicators Data Balance Ratio (DBR), Padding Ratio (PR), and Attention Balance Ratio (ABR) are strongly correlated with the training time, emphasizing the effectiveness of these indicators. Based on the computational complexity of packing strategies and their training time, and accuracy performance, we ultimately selected ISF as our naive packing baseline.
Table 13: Results of different packing strategies in the training setting of 128K with hybrid data. N: Number samples; M: Number of samples per pack; S: Maximum pack length; C: Number of iterations.
Packing Strategy
Complexity
DBR
ABR
PR
Ave
LongBench
GPU (speed up)
No Random
0.63
0.72
0
56.5
41.8
38.3 (1.00x)
OIN)
56.2
42.0
5:57 (6.87x)
ISF
C*O(N+M)
0.506
0.01
56.6
41.6
5.22 (7.33x)
FES
OINM)
0.508
0.01
56.3
42.0
5.24 (7.30x)
FED
O(NM)
0.515
0.01
56.4
42.4
5,48 (6.98x)
RFS
O(NM)
0.508
0.01
56.8
41.7
5.24 (7.30x)
SPEHP
O(N+S*)
0.512
0.01
56.2
42.3
(7.100))
B Results of FlexSP data grouping
The experiments were conducted with a sequence length of 128K using 32 GPUs. The results for FPexSP were obtained by directly testing with the officially released code. The FlexSP data group method incurs a non-negligible overhead of 5-15 seconds per batch, which cannot always be hidden by the training time shown in Table 14 Rows and 4.
Table 14: Results of FlexSP data grouping. Data Time refers to the time taken to fetch a batch after overlapping with model training. Model Time refers to the time spent model training. Total Time represents the overall time, including both data and model processing. All times are averaged over 100 iterations.
Model
Dataset githuh CommonCrawl
Group method
Data Time
Model Time
Total Time
GPT-7b
FlexSP
0).55
15.4s
15.6s
GPT-7b
FlexSP
0.3s
13.0s
13.3s
GPT-7b
Wikipedia
FlexSP
4.41s
2.4s
6.81s
GPT-7b
Tulu3 Longeite
FlexSP
3,43s
3.02s
6.45s
GPT-7b
Tulu3 Longeite
HBP
0s
2.75s
2.75s
C Full Results
Tables 15 and 16 present our complete results for general tasks and long-context tasks, respectively.
These results collectively validate the effectiveness of our HBP method. The other dense models all use the GQA architecture, while Deepseek- V2 adopts the MLA structure, which involves more communication due to larger number of K and V heads, similar to MHA. HBP effectively reduces the communication overhead, making the speedup more significant.
D More Details of Attention Imbalance Problem
Imbalance of attention complexity: As deduced from Megatron-LLM paper, gives packed sequence containing sub-sequences with total lengths, each transformer layer computes cost can be approxi- mated as:24Bshh2 + 4B Σ of, where is the length of the i-th sub-sequence, is the batch size, and h is the hidden size. Therefore, variance in attention complexity Σ s diriectyy impacts compute cost.
Effectiveness of Balance Batching:
21

---

Table 15: Full results of general tasks. The naive packing baseline ISF uses the Token-Mean Mean loss normalizer. LongAlign uses their proposed loss-reweigtting AVE represents the average performance general tasks. Deepseek- V2(236B) is trained in a 32K training setting due resource constraints.
Model
BBI
Math
Iireandical
exade, pr
conala
GPgA
Dee
MEPP
bellorng authbeadha reatheadee
Bays ripeed api
packing
Langhighoets
623
65.4
67.
12.8
82.3
91.6
363
62.6
29
60-1
60.0
69.6
73.8
57.6
15.5-00 BkA
07.2
33.
384
43.4
66.7
(5.0
683
N64
56.2
21.8
015
156
56.5
91.
966
946
724
17.6
30.7
8611
915
66.1
63.5
25.5
01.5
631
16.3
M1
77.1
903
119
S1.1
75
940
624
181
760
839
55.0
19.3
92.4
16.6
830
9612
$6.7
43.4
94.6
84.5
933
76.1
159
20A
DivsSul-32 (263)
840
413
542
769
11.2
76.7
62.
923
71.
063
811
126
014
911
781
365
753
16.4
19.5
714
902
101
903
700
Table 16: Full results of Long tasks. The naive packing baseline ISF uses the Token-Mean loss normalizer. LongAlign uses their proposed loss-reweighting. AVE represents the average performance on general tasks. Deepseck-V2(236B) is trained in a 32K training setting due resource constraints.
Model
Raler (32K | 128K)
NeedleBench (32K (28K)
LongHencch
Langeite
GPU Days (speed up)
L.mme 1-88
Long LongAAlig-ppan lign-packing
(0.7%) 33.3 16x)
Long/llin-sornts
60.0
92.4 88.9
46.6
64.0
LSF
85.01 67.4
92.1 99.1
41.5
71.6
5.22 (L.0x)
HBP
85.61 70.8
91.8 90.0
43.2
71.5
3,73 (LA4))
Qwee2255.328
ISF
59.2
94.5 84.6
$1.0
60.2
21.33 (1.0x)
HDP
15.0
96.01 88.9
$1.9
61.7
16.00 (1.333)
LLaMA3,1-700
ISF
95.5
72.7
HBP
93.4
98.2 92.4
52.2
75.3
(1.4(1)42)))
Qween2.5-77 2-78
LSF
V2.66 58.5
94.5 90088
50.0
64.3
47.1((.0x))
HRP
58.22
95.22 91.4
64.7
(1.40)))
DeepSeek- (2368)
ISF
86 6
96.01
47.1
$7.1 (1.0x)
HBP
87.31
95.91
50.3
23.8 (2.43)
Sorting is indeed heuristic solution, and cannot be strietly proven mathematically guarantee perfect balance. However, based on the results across different datasets, achieves near-perfect balance (low ABR approaching 0).
Heuristic solutions are a common and efficient approximate solution for solving problems on large- scale datasets. Our method ensures that iterations maintain attention to balance across different datasets over 99.6% of training shown Table 17 demonstrating the strong generalization capability of our method.
Balance Evidence: The left side in Figure shows baseline packing with high attention cost variance across DP groups, while the right side shows our porting-basedd method, which significantly reduces this variance and improves efficiency.
Table 17: ABR comparison of different datasets using HBP.
Dataset
Packing method
ABR
Tulu3
HBP
0.001
OpenHermes
HBP
0.001
Tulu3 + Longcite
HBP
0.002
OpenHermes Longcite
HBP
0.004
Tulu3 + Longwriter
HBP
0.003
22

---

Forward Time sert)
Forward Time (wih exr))
1819
Attuttinn Comploetty twio yort)
Attirtion Compliaity (with sort)
Step
St0p
Figure 6: The plots show the forward time for different DP ranks (32 GPU DP-4, SP=8). The left column shows results without sorting, while the right column shows results after applying sort-based load balancing. The bottom plots illustrate the attention complexity (measured as the sum of sequence length squared) for each DP rank during forward passes. The top and bottom figures are used to verify the correlation between attention complexity and forward time. The left and right figures illustrate the imbalance in forward time across different DP ranks before and after applying sort-based balancing.
Lang
Long
Podt
High AVE-T
AVE-T
Long
128K Packing Group With Padding
128K Packing Group With GreedyFill
Figure 7: Example of GreedyFill. For longer samples, relying solely on data within the corresponding range (e.g.. [32K. 128K)) makes it difficult achieve full packing, resulting low Ave-T. It operates within local packing groups by selecting a small number of sequences from other shorter packing groups to fill residual space and optimize Ave-T efficiently.
E
More Details of GreedyFill.
GreedyFill aims to reduce intra-pack padding tokens and improve Ave-T. It operates within local packing groups by selecting small number of sequences from other shorter packing groups to efficiently fill residual space, as shown in Figure 7
F
More Details of Hierarchical Groups Auto-Selection.
The necessity of auto-selection: Table 18 presents example of 32K token-length training, showcasing various SP degrees GC configurations. The second row highlights the minimal
configuration = (2, 28) that satisfies memory constraints. While this configuration adheres to
23

---

memory requirements, it fails to achieve optimal performance. In contrast, the fourth row illustrates more balanced and effective configuration with (8, (), achieving the fastest speed. These experiments demonstrate that, given specific packing length, there are significant performance differences among various SP degrees and GC configurations.
The final selection groups result: Table 19 presents the training speed of the optimal SP degrees and GC configurations for various sequence lengths L. The second row of the table shows that our optimal configuration is groups = 16K, with the corresponding training strategy being = (1, 28). This indicates that the optimal training strategy is distinct for different packing groups. The evidence above emphasizes the importance of searching for the best groups and their corresponding training strategies, i.e, Auto-Group Selection.
Table 18: Results of different SP degrees and Table 19: Results of different packing groups GC configurations. Iter Time is the average time using hybrid data. Iter Time is the average time over ten iterations.
over ten iterations.
sr
GC Layer
AVE
Mensory
Iee Time
GPL Days
Parking Groupe
SP
GC Layer
Mimory
Her Time
27
2.69
1861
26
57.6
77G
78
36.6
760
3,60
399
G
Overhead Analysis
Balance Packing and Auto-Selection overhead relative to training time. Both steps together remain consistently negligible: Auto-Selection requires at most 15 minutes and Packing only 5 seconds, while training spans from 168 1400 minutes. Overall, the combined overhead below 2% of training time, confirming that preprocessing costs are insignificant compared to model training.
Table 20: Balance Packing and Auto-Selection overhead relative to training time.
Model
Dataset
Auto-Selection Time
Balance Packing Time
Training Time
Overhead Ratio
LLaMA3,1-88
Tulu3 Longeite
3,25 min min
5 sec
168 min
2%
11aMA3.1-70B
Longcite
min
G.1 Overhead of Profiling
Memory Profiling.
The memory profiling cost depends on the sequence length set (Seq) processed
by the device:
Cost, len(Seq) x profile_ite iter x tteration_time.
Time Profiling.
The total number of search strategies (S) is determined by the packing length set (I.) and the Sequence Parallel (SP) settings. profile_iter can be adjusted based on actual iteration times, with typical values ranging from 10:
S = len(L) len(SP), Cost, len(S) (S) x profile iter iteration time.
Total Cost (Example). For instance, with profile_ite jter = 5, T = (16K, 32K, 128K) and SP (2, 82), Seq {4K, 8K, 16K), for L.IMAA.1-8BB
Cost, 3 x 3 x 5 45 x iteration time,
Cost, = x 5 = 15 x iteration time,
Costwall Cost, Cost,m 60
Considering the 3000+ training iterations, totaling 60 iterations (2%) is negligible compared to the
overall training time.
G.2 Overhead of Balance Packing
For our training dataset of million samples, the overhead introduced by Balance Packing consistently within seconds.
24

---

Complexity. The computational complexity of Balance Packing can be expressed as:
Cost L. (C O + M) M log log.og M)),
where N is the number of samples, M is the number of packing groups, C is the number of ISF iterations, and denotes the selected packing set.
Since C < N. M & N. and L < N, the overall complexity reduces to sublinear class relative to N. Hence, scaling substantially larger datasets incurs only minimal additional overhead.
G.3 Scalability of the Proposed Method
Large-scale dataset. Our profiling is performed for only a small number of iterations (< 100). As the dataset size increases, the overall overhead ratio (< 1%) becomes even lower. Meanwhile, the complexity of Balance Packing is ONN), so the additional cost introduced negligible as the dataset scales.
H Dataset Generalization
H.1
OpenHermes
We also provide some of experimental results on OpenHermes, For example, Table 21 shows that the results are consistent with Tulu3 under different packing strategies. The label 22 shows that our HBP is also effective in 128K training setting.
Table 21: Results of different packing strategies in the training setting of 4K the OpenHermes dataset.
Packing Strategy
Complexity
DBR
ABR
PR
Ave
LongBench
GPU Days (speed up)
46.4
(1.0x)
random
O(N)
0.648
0.089
51.44
47.8
3.64 (3.1x)
C*O(N+M)
(3.5%)
FFS
O(NM)
0.638
0.022
52.4
48.4
3.29
(3.3x)
BES
OINM)
0.64
0.022
52.3
47.7
3.33 (3.4%)
OIN+S")
0.71
52.0
47.8
3.47 (3.3x)
Table 22: Results of HBP Components on OpenHermes Longcite dataset. This experiment uses the Token-Mean Loss Normalizer. Balance refers enabling Attention Balance Sort, while Hierarchical indicates the activation of hierarchical packing.
Hierarchical
Balance
ABR
AVE
LongBench
Ruler-32K
Longsite
GPU Days (speed
0.513
52.1
47.2
87.7
72.1
2.95 (1.0x)
0.269
53
47.1
89.1
72.0
2.33 (1.27x)
0.004
52.4
47.1
88.3
72.3
2.04 (1.4x))
H.2 LongWriter
Figure @ and 22 present the results of our HBP hybrid training on Tulu3 and LongWriter. The x-axis represents the user instruction required length, and the y-axis represents the model output length based on the Long-Writer paper. A stronger linear correlation indicates better instruction- following capability of the model. The left side of the Figure shows the original baseline results without using the Long- Writer dataset. The middle shows the results of naive packing, and the right shows the results of HBP. The experiment demonstrates that our method is equally effective, achieving consistent acceleration across both models.
Table 23: Model Evaluation Results of Tulu3 and LongWriter Hybrid Training.
Model
Dataset
AVe
LongBench
LongWriter
GPU Days (spoed up)
Llama3.1-88 1-8B-1SF
Tulu3 LongWriter
57.3
40.9
67.0
4.4 (1.0x)
Llama3.1-8 1-cba33.1-BBBB1 1-8-HHBP
Tulu3 LongWrite Writer
42.7
66.4
(1.26x)
25

---

10
10a
10*
Output Langth
Output Length
Output Length
10a
10a
10a
10a
10a
102
10a
10a
10"
10a
10a
10
10a
10a
10
Required Length
Required Length
Required Length
Figure 8: Results of Long- Writer Dataset for evaluating the impact of our method on accuracy. The x- axis represents the user instruction required length, and the y-axis represents the model output length based on the Long-Writer Writer paper. A stronger linear correlation indicates better instruction-folloviin capability of the model. The left side of the Figure shows the original baseline results without using the Long- Writer dataset. The middle shows the results of naive packing, and the right shows the results of HBP.
I Implementation Details of Hierarchical Groups Auto-Selection
L1 FindBestSpCkpt
Algorithm 3 FindBestSpCKpt Function
Algorithm 4 GreedyProfileCkpt
1: Initialize + [. + +
Inputs: 1, sp, Cmin Cmax
2: for each sp SP do
2: f- (sp, 1.Cin))
3: ckpt + (roeeyyProfllCkpp(()
3: 82 f- (sp, L, Cmaz..
4: Padd((ProfileTimee) sp, ckpt)),
4: mi <- ProfileMemory(st) (1)
O.add(spp, ckpt)
mz < ProfileMemory (8))
5: end for
6: Mave +- (m² mi) (Cmaz Cmin)
6: j + argmin(P) P))
7: Co f- Cmar - mm//Maue
7: return Olj]
8: return
Algorithm determines the optimal gradient checkpointing strategy by evaluating all possible sp strategies.
Initialization:
Start with empty lists /' and O for profiling times and configurations, respectively.
Iterate Over Strategies:
For each strategy sp SP:
Compute the best gradient checkpointing configuration (ckpt) using the GreedyPro- fileCkpt function.
Use ProfileTime to profile the execution time for the model with the given configuration and append it to
Find Optimal Strategy: Identify the index j of the minimum profiling time in P using argmin(P).
Return Best Configuration: Return the SP degree and corresponding gradient checkpoint- ing configuration Olj).
I.2 GreedyProfileCCpt
The algorithm essiinates the number of gradient checkpoint layers required for a given strategy and sp. Here, provide more detailed explanation. We use the "pynvml" library to monitor GPU memory usage for memory analysis.
Memory Profiling:
26

---

Step 1: Given an input length L, enable Gradient Checkpointing for all layers (32 for 7B), and record the remaining GPU memory as ml.
Step 2: Given the same input length I., enable Gradient Checkpointing for only a subset of layers (e.g., 20 layers), and record the remaining GPU memory m2.
Step 3: From the above, we can estimate the average memory saved per layer with Gradient Checkpointing: ave_m- = (ml m2) (32 20)
By running only 10 iterations, we can obtain reliable estimate of the memory saved per layer when using Gradient Checkpointing.
Memory And Time Consumption Report:
For given input length L., we set different sequence parallelism (SP) configurations. Based on the remaining GPU memory and the previously estimated ave_m, we determine the number of layers to apply Gradient Checkpointing. Then, we run 10 iterations to record the corresponding memory usage and iteration time.
Full workflow:
Initialization: Obtain the configurations using the minimum and maximum number of gradient checkpointing layers (based on empirical observations):
81 = (sp, Cmin),
82 (sp, L, Cmax)
Memory Profiling: Profile the remaining memory for 81 (m²j) and (mz).
Memory Slope Calculation: Compute the average memory slope (avem) as:
m mi
Mave
Cmax
Cman
Checkpointing Layer Estimation: Estimate the required number of checkpoints (c,) as:
m
Co Cyaa
Mave
Implementation Details of Balance Packing
J.1
GroupData
Given data set D and predefined hierarchical lengths Ly, we evaluate the length of each data set to determine the interval (4:1..4) 4;) to which it belongs, assigning it to the corresponding group D;. The detailed procedure is outlined in Algorithm
J.2 GreedyFill
For a given packing group G3 with the corresponding length 2, we iterate through the smaller dataset partitions Di-,,, D2) and greedily fill the group g within the current G;. The detailed procedure is illustrated in Algorithm 6
J.3 Attention Balance Sort Function
First, calculate the attention complexity for all data within the given packing group G. Then, sort the elements in G; based on their attention complexity and construct mini-batches according to the global token number requirements shown in Algorithm 7
27

---

Algorithm 5 GroupData Function
Algorithm 6 GreedyFill Function
1: Inputs: Dataset D (2 NN ) lengths
1: for y G, do
hierarchical
packing
2:
for j = 1i--1 1 + do
L = (4.12....,44 bn3
3:
for Dj do
2: Initialize:
(Di, D2. D2....,, Dn)
4:
if L,,, lcn(s) + len(x) s then g.add(z)
(I.O....:))
5:
3: for x E D do
Remove from Dj 7: end if
4: if len E (4-1,44 then
5:
(z.ddd(z)
end for
6: end if
9: end for
7: end for
10: end for
8: return (D1, D2....; Da)
11: return G;
Algorithm 7 Attention Balance Sort Function
Initialize A+1
2: for do
3: a = (Ien((e()))))
4: (aadd(a)
5: end for
6: Sort G based on A
7: return G
28


Reviews:
[{"review_id": "jNvd0CFcfO", "reviewer": "Reviewer_Nqmr", "summary": "This paper focuses on training Long-Context Large Language Models (LLMs), with particular emphasis on data packing when mixing long-context and short-context data. The paper points out that existing data packing methods suffer from inefficiencies and imbalances. To address these issues, the paper introduces Hierarchical Balance Packing (HBP), which tackles the challenges of optimal packing group selection, data assignment, and training with packed data through three components: Hierarchical Groups Auto-Selection, Balance Packing, and Dynamic Training Pipeline. Experimental results demonstrate that HBP effectively improves training efficiency while maintaining strong performance.", "questions": "In Section 4.3, the paper proposes the Stable Loss Normalizer and empirically selects the Average Token method, but lacks an in-depth analysis of the underlying reasons. Furthermore, it remains unclear whether the stability advantages of the Average Token normalizer apply exclusively to HBP, or if they also generalize to other methods. Could the authors clarify this?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper introduces a variety of measuring metrics for data packing , providing valuable insights for future analyses of data packing strategies.\n2. The paper constructs a comprehensive pipeline to address different challenges in data packing. Furthermore, the automatic packing capability of HBP reduces the difficulty of parameter tuning.\n3. HBP significantly improves the training efficiency of Long-Context LLMs while maintaining strong performance across different benchmark datasets, thoroughly validating its effectiveness.\nWeaknesses:\n1. Section 3.1 introduces multiple measuring metrics and explains their relevance to LLM’s training, but the paper lacks experiments demonstrating how variations in individual metrics concretely affect LLM’s training performance and efficiency.\n2. The paper claims that the HBP algorithm introduces only negligible overhead, but lacks detailed quantitative data showing how much time is cost by the Hierarchical Groups Auto-Selection and Balance Packing components, respectively.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful comment.\n\n**W1** \n > Section 3.1 introduces multiple measuring metrics and explains their relevance to LLM’s training, but the paper lacks experiments demonstrating how variations in individual metrics concretely affect LLM’s training performance and efficiency.\n\n**A1** \nWe agree that a deeper analysis of how individual metrics affect LLM training would be valuable.\n\nSince the metrics are interdependent, we did not isolate them in ablation studies. However, most related results are already shown in **Tables 3 and 6**, and we have **summarized** the experimental results from the paper into the following table to illustrate their impact more clearly.\n\n\n|      Model      | Seq Len |        Batching Method  |  DBR  |  PR  |  AVT  |  CR  |  ABR  | GPU days |\n|:---------------:|:-------:|:---------------------:|:-----:|:----:|:-----:|:----:|:-----:|:--------:|\n| LLama3.1-8B     |  128K   |        random         | 0.639 | 0.0  | 0.9K  |  1   |   0.61   |   38.0   |\n| LLama3.1-8B     |  128K   |        sorted         | 0.02  | 0.0  | 125K  |  1   |   0.51   |   5.5    |\n| LLama3.1-8B     |  128K   |     lsf packing       | 0.001 |  0   | 128K  |  1   | 0.506 |   5.2    |\n| LLama3.1-8B     |  128K   |         HBP （Hierarchical）           | 0.001 |  0   | 128K  | 0.17 | 0.288 |   4.2    |\n| LLama3.1-8B     |  128K   |         HBP （Hierarchical + balance）          | 0.001 |  0   | 128K  | 0.17 | 0.002 |   3.7    |\n\n**Results Analysis**\n\n- Comparing the first and second rows, **DBR drops** significantly and **AVT increases**, reducing training time from **38 to 5.5 GPU days**.\n\n- From the second to third row, AVT reaches the maximum 128K, slightly improving speed (**5.5 → 5.2 GPU days**).\n\n- The third to fourth row shows a sharp **drop in CR and ABR**, leading to further reduction in training time (**5.2 → 4.2 GPU days**).\n\n- Finally, the fifth row shows that **reducing ABR** even more translates into the lowest training cost of **3.7 GPU days**.\n\nThese results demonstrate how improvements in **individual metrics (e.g., DBR, AVT, ABR) correlate with more efficient LLM training**.\n\n**W2**\n>The paper claims that the HBP algorithm introduces only negligible overhead, but lacks detailed quantitative data showing how much time is cost by the Hierarchical Groups Auto-Selection and Balance Packing components, respectively.\n\n**A2:**\n\nThank you for your question. We have included some overhead details in **Appendix I**. To clarify this, we now provide a full **quantitative data analysis** and will incorporate these results into the main text in future versions.\n\n**The overhead introduced by the HBP algorithm is negligible.**\n\nUnder the setting of LLama3.1-8B model and Tulu3 dataset, the overhead brought by HBP includes：\n* Hierarchical Groups Auto-Selection (**3.25min**)\n* Balance Packing (**5 sec**)\n\n|    Model    |      Dataset      | Auto-Selection time | Balance Packing time | Training time | Overhead ratio |\n|:-------------:|:----------------:|:------------------:|:-------------------:|:-------------:|:--------------:|\n| LLama3.1-8B   | Tulu3 + longcite |      **3.25 min**        |        **5 sec**          |    168 min      |     **2 %**        |\n| LLama3.1-70B  | Tulu3 + longcite |      **15 min**          |        **5 sec**          |   1400 min      |     **1 %**        |\n\n\nCompared to the total training time (168 min),  the overall overhead of **3.25 min** is **negligible**. Note that HBP **ONLY** operates **before** the training process, so it introduces **NO** overhead during the training process.  \n\n**Q3:**\n>In Section 4.3, the paper proposes the Stable Loss Normalizer and empirically selects the Average Token method, but lacks an in-depth analysis of the underlying reasons. Furthermore, it remains unclear whether the stability advantages of the Average Token normalizer apply exclusively to HBP, or if they also generalize to other methods. Could the authors clarify this?\n\n**A3**\n\nThank you for your valuable feedback. We will clarify this in detail.\n\n* **Stable Loss Normalizer is not empirical**\n\nThe Stable Loss Normalizer is designed based on theoretical principles. Its core purpose is to **ensure each token contributes equally to the overall loss**, removing biases from sequence length, data-parallel size, and gradient accumulation. Further analysis supporting this design is provided below.\n\n\n* **Can generalize to other methods**\n\nWe conducted experiments on the loss normalizer under the base packing setting without HBP. The results show its effectiveness across different methods.\n\n|      Model        |      Dataset      | Loss Normalizer |  AVE  | LongBench |\n|:-----------------:|:----------------:|:---------------:|:-----:|:---------:|\n| LLama3.1-8B-ISF   | Tulu3 + longcite |   token-mean    | 56.0  |   44.0    |\n| LLama3.1-8B-ISF   | Tulu3 + longcite |  sample-mean    | 55.4  |   43.0    |\n| LLama3.1-8B-ISF   | Tulu3 + longcite |   **Ave-Token (Stable Loss)**     | **57.0**  |   **43.8**    |\n\n---\n\n**In-depth analysis**\n\nTo clarify, let’s consider a data-parallel (dp) group with a size of 2, where each rank runs with a local batch size of 2. \n\nLet $loss$ denote the sum loss for a sample, and $T$ represent the number of tokens involved in the loss calculation.\n\n\n  -  **Token-Mean**：\n\nIn the token-mean approach, we compute the loss for each DP rank independently as follows:\n\n   $$\n   L_{dp_{1}} = \\frac{loss_{1} + loss_{2}}{T_{1} + T_{2}}, \\quad L_{dp_{2}} = \\frac{loss_{3} + loss_{4}}{T_{3} + T_{4}}\n$$\n\nThe final loss for the global batch is calculated by averaging the individual DP rank losses:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1}+loss_{2}}{2·(T_{1}+T_{2})} + \\frac{loss_{3}+loss_{4}}{2·(T_{3}+T_{4})}$$\n\nThis method may introduce bias if the data-parallel ranks have different sequence length distributions, as it normalizes the losses independently for each rank **without considering the global batch’s overall token length distribution**.\n\n---\n - **Sample-Mean:**\n\nIn the sample-mean method, losses for each DP are first normalized by their respective token lengths and then averaged across the samples:\n$$L_{dp_{1}}=\\frac{\\frac{loss_{1}}{T_{1}}+\\frac{loss_{2}}{T_{2}}}{2} , \\quad                   L_{dp_{2}}=\\frac{\\frac{loss_{3}}{T_{3}}+\\frac{loss_{4}}{T_{4}}}{2}$$\n\nThe final loss is then:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1}}{4·T_{1}} + \\frac{loss_{2}}{4·T_{2}}+\\frac{loss_{3}}{4·T_{3}} + \\frac{loss_{4}}{4·T_{4}} $$\n\nThis method balances the contribution of each individual sample. However, it does not fully account for the total number of tokens in the global batch, resulting in **samples with fewer tokens having a disproportionately larger influence on the overall loss**.\n\n---\n\n-  **Stable Loss Normalizer:**\n\nThe Average Token normalizer ensures that the total number of tokens in a global batch is considered. Here, we define the average token length across all samples:\n\n$$T_{ave}   = \\frac{T_{1} + T_{2} + T_{3} + T_{4}}{4}$$\n\nEach DP's loss is then normalized by  $T_{ave}$\n\n$$L_{dp_{1}}=\\frac{loss_{1}+loss_{2}}{2·T_{ave}}, \\quad       L_{dp_{2}}=\\frac{loss_{3}+loss_{4}}{2·T_{ave}}$$\n\nThe final loss is:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1} + loss_{2} +loss_{3} + loss_{4}}{4·T_{ave}} = \\frac{loss_{1} + loss_{2} + loss_{3} + loss_{4}}{T_{1} + T_{2} + T_{3} + T_{4}}$$\n\nThis approach effectively normalizes the loss by the total number of tokens across all samples in a global batch, **ensuring that each token, regardless of which sample it belongs to, contributes equally to the final loss**.\n\nWe hope our response addresses your concerns, and we welcome any further questions or feedback."}, {"review_id": "EK1sJK3tUR", "reviewer": "Reviewer_no8c", "summary": "This paper proposed a new data packaging method with automatic training configuration optimization (wrt gradient checkpointing and sequence parallelism). \nThe author designed several useful metrics to effectively measure the training efficiency in terms of waste of padding, waste of communication, balance of attention FLPOs and communications, etc.. These metrics could serve as helpful optimization objectives for future works as well.\nThe author designed automatic optimization method to decide the optimal packing length, assign data to these groups, and design the optimal training configurations for these groups. \nExperiments demonstrate that the proposed method could speed up training time to 1.33-2.4x across different model sizes, while maintaining high performance on both short- and long-context tasks.", "questions": "1. Could the author provide more details about the overhead brought by the optimization process and discuss how scalable it is if someone considers applying it to pretraining-scale datasets?\n2. Do you think clustering the data by lengths would impact the training data distribution, thus leading to very different bahaviors when different domains have very different length biases? How do you mitigate this impact?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The set of metrics and optimization objectives is general and useful for future works to design better algorithms.\n2. The proposed method improves the training efficiency without harming the performance.\n3. The author identified an important but have been overlooked issue with existing data packaging approaches, and proposed a method that not only keeps the natural distribution of varied-length data but also minimizes the waste of compute and load imbalance in current training strategies.\n4. Extensive experiments show that the proposed method is effective and performant.\n\nWeaknesses:\n1. The writing could be further improved--the paper used a large space to repeatedly highlight the key contributions and the limitation of the current packaging methods. Limited space was left for explaining the method and showing the results. In particular, a brief hint on what the notations stand for would be helpful than just assuming everyone to repeatedly look back at the notation table. In addition, some details might be missing, e.g., there is no explanation on what is FindBestSpCkpt function in 4.1 while both pseudocode and the main text uses it. \n2. The proposed packaging method focuses only on the length distribution. However, it ignores the content distribution shift after clustering the data purely based on lengths. For example, if one domain only has short sentences, it will only appear at the early training stage, leading to catastrophic forgetting for this domain.\n3. The paper uses profile-based automatic optimization, however, it did not discuss the overhead brought by this optimization process. In addition, it applies the method to Tulu3, which is a mid-scale post-training dataset. However, sequence packaging is often used in pretraining stage, where the paper lacks discussion on the scalability of the proposed method to large-scale pretraining datasets.\n4. The search space only contains sequence parallelism and gradient checkpointing, while more dimensions, such as tensor parallelism, fsdp, etc., are not considered. In addition, sequence parallelism also has different categories, such as DeepSpeed version (using all-to-all communication and split over attention head) and DistFlashAttention version (using p2p communication and splitting over sequence length), while the method's optimization considers a very limited optimization space.", "was_revised": true, "final_justification": "The authors supplemented the additional prefilling overhead analysis as I requested and the results seem good. My concerns about the efficiency and randomness have been fully addressed with potential mitigation approaches. I have raised my score based this discussion.", "rebuttal": "Thank you for your helpful comment.\n\n**W1**\n> The writing could be further improved--the paper used a large space to repeatedly highlight the key contributions and the limitation of the current packaging methods. Limited space was left for explaining the method and showing the results. In particular, a brief hint on what the notations stand for would be helpful than just assuming everyone to repeatedly look back at the notation table. In addition, some details might be missing, e.g., there is no explanation on what is FindBestSpCkpt function in 4.1 while both pseudocode and the main text uses it.\n\n**A1:**\n\nThank you for your valuable feedback. In future versions, we will simplify the discussion of our contributions and the limitations of existing packing methods to allocate more space to the methodology details and experimental results (**part of which are currently in the Appendix**). \n\nWe will also ensure that important notations are reintroduced where necessary to improve clarity in the future version. \n\nAdditionally, the \"FindBestSpCkpt\" function, currently described in **Appendix I**, will be relocated to the main text for improved accessibility.\n\n---\n\n**W2 & Q2**\n> The proposed packaging method focuses only on the length distribution. However, it ignores the content distribution shift after clustering the data purely based on lengths. For example, if one domain only has short sentences, it will only appear at the \nearly training stage, leading to catastrophic forgetting for this domain.\n\n>  Do you think clustering the data by lengths would impact the training data distribution, thus leading to very different bahaviors when different domains have very different length biases? How do you mitigate this impact?\n\n**A2:**\n\nThank you for your question. \n\n**Short sentences not only appear at the early training stage**\n\nAs described in **Section 4.3** and illustrated in **Figure 2**, short samples are distributed throughout the entire training stage.\nWe employ an **alternating training strategy** between different packing groups with varying packing lengths (after a few iterations of Curriculum Learning)\n\n---\n\n**How to mitigate this impact**\n\nWhile clustering by length can change the data distribution, the randomness and stable loss normalizer in our method greatly mitigates this effect. Our experiments (**Tables 5, 7, 10, and Appendix H Dataset Generalization**) also confirm the robustness of our approach.\n\n * **randomness**\n    - The base packing method is inherently **stochastic**.\n    - The **alternating training strategy**—switching between short- and long-context data—introduces additional **randomness**, which helps prevent length bias.\n  * **Stable loss normalizer:** We incorporate a stable loss normalizer that ensures balanced contributions from different sentence lengths to the overall training objective.\n\n---\n\n**W3 and Q1**\n> The paper uses profile-based automatic optimization, however, it did not discuss the overhead brought by this optimization process. In addition, it applies the method to Tulu3, which is a mid-scale post-training dataset. However, sequence packaging is often used in pretraining stage, where the paper lacks discussion on the scalability of the proposed method to large-scale pretraining datasets.  \n\n>  Could the author provide more details about the overhead brought by the optimization process and discuss how scalable it is if someone considers applying it to pretraining-scale datasets?\n\n**A3:**\n\nThank you for your question. We have included some overhead details in **Appendix I**. To clarify this, we now provide a **full quantitative data analysis** and will incorporate these results into the main text in future versions.\n\n**Overhead of optimization process:**\n\nCompared to the total training time (168 min),  the overall overhead of **3.25 min** is negligible for  LLaMA3.1-8B.\n|        Setting        |       Dataset       | Profiling time | Training time | Overhead ratio |\n|:---------------------:|:-------------------:|:--------------:|:-------------:|:--------------:|\n|   LLaMA3.1-8B         | Tulu3 + longcite    |     **3.25 min**     |    168 min      |      **2 %**       |\n|   LLaMA3.1-70B        | Tulu3 + longcite    |     **15 min**       |   1400 min      |      **1 %**       |\n\n**Details of profile-based automatic optimization overhead**\n\n * **Memory Profiling:**\n\nThe memory profiling cost depends on the sequence length set (**Seq**) processed by the device. \n\n      Cost_m = len(Seq) × profile_iter × iteration_time\n\n* **Time Profiling:**\n\nThe total number of search strategies (S) is determined by the packing length set (L) and the Sequence Parallel (SP) settings. Profile_iter can be adjusted based on actual iteration times, with typical values ranging from 3 to 10.\nSpecifically:\n\n      S = len(L) x len(SP)\n      Cost_t = len(S) x profile_iter x iteration_time\n\n* **Total Cost** (Example)\n\nFor instance, with a  profile_iter 5,  L =  [16K,32K,128K] and SP ([2, 4, 8]), Seq = [4K, 8K, 16K], LLama3.1-8B\n\n     Cost_t = 3 × 3 × 5 = 45 x iteration_time\n     Cost_m = 3 x 5 = 15 x iteration_time\n     Cost_total = Cost_t + Cost_m = 60 × iteration_time\n\nConsidering the 3000+ training iterations,  totaling **60 iterations (2%**) is negligible compared to the overall training time.\n\n**Scalability of the proposed method**\n\n*  **Large-scale dataset**\n\n     Our profiling is performed only for a small number of iterations (**<100**), so as the dataset size increases, the overall overhead ratio (**<1%**) becomes even lower. Meanwhile, the complexity of Balance Packing is an **O(N)** level, so the additional **cost introduced is negligible** as the dataset scales.\n\n* **Pretraining setting**\n\n     The pretrain setting can be used **directly**.  However, in the pre-training setup, training is typically conducted in stages:\n    The model is first trained on shorter sequences (e.g., 8K), and only in the final stage is trained on longer sequences (e.g., 128K). Both Qwen2.5 and Llama3.1 follow this phased approach, **so the issue of mixed sequence lengths during pre-training is not significant.**\n\n---\n\n**W4**\n> The search space only contains sequence parallelism and gradient checkpointing, while more dimensions, such as tensor parallelism, FSDP, etc., are not considered. In addition, sequence parallelism also has different categories, such as DeepSpeed version (using all-to-all communication and split over attention head) and DistFlashAttention version (using p2p communication and splitting over sequence length), while the method's optimization considers a very limited optimization space.\n\n**A4:**\n\nThank you for your valuable suggestions. We agree that the search space could be further extended.  \nHowever, including too many dimensions would result in **high profiling overhead**. In this work, we therefore focus on the two most influential factors for long-context training—sequence parallelism and gradient checkpointing, and leave further extensions to future work.\n\n**different categories of sequence parallelism**\n\nOur method is agnostic to the specific implementation of sequence parallelism and works effectively with both commonly used SP approaches, including DeepSpeed-Ulysses and Ring-Attention. We have provided detailed implementation descriptions in Appendix A (*In our experiments, we use DeepSpeed-Ulysses’s  sequence parallelism approach, and the ring-attention method is also applicable*)\n\n|      Model            |      Dataset          |   Sequence Parallelism            |   GPU days         |\n|:--------------------:|:--------------------:|:---------------------:|:------------------:|\n| LLama3.1-8B-ISF      | Tulu3 + longcite     |     DeepSpeed-Ulysses                |   5.22 (1.0)       |\n| LLama3.1-8B-HBP      | Tulu3 + longcite     |    DeepSpeed-Ulysses  |   **3.73 (1.4x)**      |\n| LLama3.1-8B-ISF      | Tulu3 + longcite     |    Ring-Attention                |   5.3 (1.0)       |\n| LLama3.1-8B-HBP      | Tulu3 + longcite     |    Ring-Attention             |   **3.8 (1.39x)**     |\n\n---\n\nWe hope our response addresses your concerns, and we welcome any further questions or feedback."}, {"review_id": "1VP8LRso2X", "reviewer": "Reviewer_UkHx", "summary": "In summary, the authors propose an efficient hierarchical packing algorithm for long-context SFT for LLM. The traditional packing algorithm only has one packing group and faces two problems: 1) imbalanced attention complexity, 2) Inefficient Sequence-Parallelism (SP) communications. To tackle this challenge, the proposed method divides the training dataset into multiple packing groups with different sequence lengths to save the SP communications. Within each group, Balanced Packing is performed to maintain balanced attention computation among different batches. To support this multi-length input packing groups, a curriculum learning strategy is adopted. The proposed training strategy is tested on multiple models on an H100 cluster and is shown to improve the training efficiency with improved model performance.", "questions": "1. Table 2 seems important to motivate the proposed method but it lacks details. How are the numbers calculated/measured? Is it a general phenomenon? \n\n2. The Curriculum Learning Strategy seems important for the proposed training strategy but it is not introduced in detail. For example, the proposed method is applied in two different clusters, which results in a different number of packing groups with different sequence lengths. How would the curriculum learning strategy vary for these two cases? \n\n3. It would be great to know whether the proposed method is compatible with various parallelism strategies such as DP/PP/TP/FSDP.", "limitations": "Yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. This paper studies an important problem in large-scale training and is well-motivated\n2. The proposed solution achieves faster training speed on multiple models with improved model performance\n\nWeakness:\n1. The proposed method is purely heuristic and lacks in-depth analysis \n2. Some descriptions in the paper are not clear and confusing (Please refer to Questions)\n3. The proposed packing method seems to be much more complicated than the traditional packing, but such overhead is not measured in the paper\n4. The proposed training framework has to come with a curriculum learning strategy due to its multi-level nature. It is unclear whether this learning strategy could lead to worse accuracy for certain models/datasets\n5. Seems the proposed method is only validated on an H100 cluster. However, to show the generalizability of the proposed method, it would be better to test it on multiple cluster setups where the hardware resource is different (e.g., A100 clusters, different CPUs)\n6. The proposed method is limited only to SFT and does not apply to other tasks such as PPO/GRPO.", "was_revised": true, "final_justification": "I appreciate the authors' responses which have solved most of my concerns. I have raised my scores for originality and quality. However, I am still concerned about the clarity of the paper writing. So I am leaned to keep my initial rating of this paper", "rebuttal": "**W1**\n> The proposed method is purely heuristic and lacks in-depth analysis\n\n**A1**\n\nThank you for your valuable feedback. \nThe packing problem is a well-known **NP-hard problem**, and finding exact solutions is difficult for large-scale training. As a result, **heuristic or approximate methods** are commonly used in practice. Experimental results demonstrate the efficiency and effectiveness of our heuristic solutions, HBP.\n\nWe have conducted an **Extensive** Depth Analysis\n\n* **Packing Problem Analysis**\nIn Section 3, we provide a thorough problem analysis and utilize various metrics to quantify the limitations of existing methods (see **Tables 2, 3, and 4**).\n\n* **Stable Loss Normalizer** We provide a brief theoretical introduction and a comparison with existing methods in **Section 4.3 (see Figure 5)**. A more comprehensive and **in-depth analysis** is included at the end of the response, and we will also add it in future versions of the paper.\n\n---\n\n**W3**\n> The proposed packing method seems to be much more complicated than the traditional packing, but such overhead is not measured in the paper\n\n**A3**:\n\nThank you for your question. We have included some overhead details in **Appendix I**. To clarify this, we now provide a **full quantitative data analysis** and will incorporate these results into the main text in future versions.\n\n**The overhead of HBP  is negligible.**\n\nUnder the setting of LLama3.1-8B model and Tulu3 dataset, the overhead  includes：\n* Hierarchical Groups Auto-Selection (**3.25 min**)\n* Balance Packing (**5 sec**)\n\n|    Setting    |      Dataset      | Auto-Selection time | Balance Packing time | Training time | Overhead ratio |\n|:-------------:|:----------------:|:------------------:|:-------------------:|:-------------:|:--------------:|\n| LLama3.1-8B   | Tulu3 + longcite |      **3.25 min**        |       **5 sec**         |    168 m      |     **2 %**        |\n| LLama3.1-70B  | Tulu3 + longcite |      **15 min**         |        **5 sec**          |   1400 m      |     **1 %**        |\n\nCompared to the total training time (168 min),  the overall overhead of **3.25 min** is negligible. Note that HBP **ONLY** operates **before** the training process, so it introduces **NO** overhead during the training process.  \n\n---\n\n**W4** \n> The proposed training framework has to come with a curriculum learning strategy due to its multi-level nature. It is unclear whether this learning strategy could lead to worse accuracy for certain models/datasets\n\n**A4**\n\nWe appreciate your concern and acknowledge that this is a possibility in theory.\n\nHowever,  in our experiments, we did not observe **accuracy degradation** when applying the curriculum learning strategy across multiple datasets and model scales (see **Tables 5, 7, 8, and Appendix H**). We will continue to verify its effectiveness on future models and datasets.\n\n---\n\n**W5**\n> Seems the proposed method is only validated on an H100 cluster. However, to show the generalizability of the proposed method, it would be better to test it on multiple cluster setups where the hardware resource is different (e.g., A100 clusters, different CPUs)\n\n**A5**\n\nOur method is **hardware-agnostic**. To address this concern, we rented an A100 cluster and replicated the experiments. The method remained effective, delivering a **1.38× speedup**, comparable to the gains observed on H100. We will further clarify this in the revised version.\n\n|      Model       |      Dataset      |  GPU   |AVE |   GPU days    |\n|:---------------:|:----------------:|:-----:|:--------------:|:--------------:|\n| LLama3.1-8B-ISF | Tulu3 + longcite | A100  |     56.2 | 9.4       ||\n| LLama3.1-8B-HBP | Tulu3 + longcite | A100  |   58.1  |**6.79 (1.38x)**   |\n\n---\n\n**W6** \n>The proposed method is limited only to SFT and does not apply to other tasks such as PPO/GRPO.\n\n**A6**\n\nThank you for highlighting this point.\n* Our method is **not limited to SFT**,  it can be applied to the Actor training in  PPO or GRPO. \n\n* HBP results with GRPO \n\n|       Model        |   Dataset    | Len | Actor-Train Iter times |\n|:-----------------:|:------------:|:-------:|:----------------:|\n| Qwen2.5-7B   | deepscaler   |  32K    |      51 s        |\n| Qwen2.5-7B-HBP    | deepscaler   |  32K    |      **39 s**       |\n\nDue to time and resource constraints, we have not explored these extensions in this paper (see **Conclusion and Limitations**).\n\n---\n\n**Q1**\n> Table 2 seems important to motivate the proposed method but it lacks details. How are the numbers calculated/measured? Is it a general phenomenon?\n\n**A7**\n\n**All metric calculations based on Section 3:**\n\nIn Table 2, we use the Tulu3 dataset and apply the ISF packing algorithm to group sequences into 4K, 32K, and 128K. \n\n- **ABR:** As shown in **Equation (2)**,  ABR is based on attention computation complexity, which scales quadratically with sequence length. A concrete example is given in Section 3.1, and its impact on computational efficiency is discussed in **Appendix E**.\n\n- **CR**, defined in **Equation (3)**, measures the proportion of tokens involved in communication. When Sequence Parallelism (SP) > 1, all tokens are involved, so CR is 1 for the packed sequences in Table 2.\n\n- **DBR and PR**: Both are computed by **Equation (1)**. After ISF packing, sequence lengths are uniform (similar length), so DBR and PR are nearly zero.\n\n**General phenomenon across different packing algorithms and datasets.**\n\nIn **Appendix B**, we show that ABR, DBR, and PR exhibit similar trends across different packing algorithms. We also demonstrate these patterns in the OpenHermes dataset in **Appendix H.1**.\n\n---\n\n**Q2**\n> The Curriculum Learning Strategy seems important for the proposed training strategy but it is not introduced in detail. For example, the proposed method is applied in two different clusters, which results in a different number of packing groups with different sequence lengths. How would the curriculum learning strategy vary for these two cases?\n\n**A8**\n\n**Curriculum learning strategy is similar to the different case.**\n\n* **Cluster Results:**\n  Results with similar CL settings on both **A100 (see A5)** and H100 show that HBP works robustly across hardware platforms.\n\n|   CL-iter |GPU   |AVE |   GPU days    |\n|:-----:|:-----:|:--------------:|:--------------:|\n|  100 |A100  |     56.2 | 9.4       |\n |  100 |A100  |   58.1  |**6.79 (1.38x)**   |\n\n* **Detail of Curriculum Learning:**\n  Described in **Section 4.3**. We start with short sequences, then alternate short and long sequences in later training stages.\n\n* **Hyperparameter Robustness:**\n  As shown in **Table 8**, our strategy is robust to different hyperparameter settings, packing groups, and sequence lengths.\n\n---\n\n**Q3**\n> It would be great to know whether the proposed method is compatible with various parallelism strategies such as DP/PP/TP/FSDP.\n\n**A9**\n\n**Our method is fully compatible with DP, PP, TP, and FSDP**\n\nOur largest model, Deepseek-V2 (236B), uses **Megatron-LM 3D parallelism** (**Implementation Details in Appendix A**), and our method is **fully compatible with DP/PP/TP**. For other models, we use DeepSpeed ZeRO-3, which operates on the same principle as FSDP, so our approach is also **applicable to FSDP**.\n\n---\n\n**Stable Loss In-depth analysis**\n\nTo clarify, let’s consider a data-parallel (dp) group with a size of 2, where each rank runs with a local batch size of 2. \n\nLet $loss$ denote the sum loss for a sample, and $T$ represent the number of tokens involved in the loss calculation.\n\n   **Token-Mean**：\n\nIn the token-mean approach, we compute the loss for each DP rank independently as follows:\n\n$$\nL_{dp_{1}} = \\frac{loss_{1} + loss_{2}}{T_{1} + T_{2}}, \\quad L_{dp_{2}} = \\frac{loss_{3} + loss_{4}}{T_{3} + T_{4}}\n$$\n\n\nThe final loss for the global batch is calculated by averaging the individual DP rank losses:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1}+loss_{2}}{2·(T_{1}+T_{2})} + \\frac{loss_{3}+loss_{4}}{2·(T_{3}+T_{4})}$$\n\nThis method may introduce bias if the data-parallel ranks have different sequence length distributions, as it normalizes the losses independently for each rank **without considering the global batch’s overall token length distribution**.\n\n---\n\n **Sample-Mean:**\n\nIn the sample-mean method, losses for each DP are first normalized by their respective token lengths and then averaged across the samples:\n$$L_{dp_{1}}=\\frac{\\frac{loss_{1}}{T_{1}}+\\frac{loss_{2}}{T_{2}}}{2} , \\quad                   L_{dp_{2}}=\\frac{\\frac{loss_{3}}{T_{3}}+\\frac{loss_{4}}{T_{4}}}{2}$$\n\nThe final loss is then:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1}}{4·T_{1}} + \\frac{loss_{2}}{4·T_{2}}+\\frac{loss_{3}}{4·T_{3}} + \\frac{loss_{4}}{4·T_{4}} $$\n\nThis method balances the contribution of each individual sample. However, it does not fully account for the total number of tokens in the global batch, resulting in **samples with fewer tokens having a disproportionately larger influence on the overall loss**.\n\n---\n\n **Stable Loss Normalizer**\n\nThe Average Token normalizer ensures that the total number of tokens in a global batch is considered. Here, we define the average token length across all samples:\n\n$$T_{ave}   = \\frac{T_{1} + T_{2} + T_{3} + T_{4}}{4}$$\n\nEach DP's loss is then normalized by  $T_{ave}$\n\n$$L_{dp_{1}}=\\frac{loss_{1}+loss_{2}}{2·T_{ave}}, \\quad       L_{dp_{2}}=\\frac{loss_{3}+loss_{4}}{2·T_{ave}}$$\n\nThe final loss is:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1} + loss_{2} +loss_{3} + loss_{4}}{4·T_{ave}} = \\frac{loss_{1} + loss_{2} + loss_{3} + loss_{4}}{T_{1} + T_{2} + T_{3} + T_{4}}$$\n\nThis approach effectively normalizes the loss by the total number of tokens across all samples in a global batch, **ensuring that each token, regardless of which sample it belongs to, contributes equally to the final loss**.\n\n---\n\nWe hope our response addresses your concerns, and we welcome any further questions or feedback."}, {"review_id": "VAhP5ZFdFf", "reviewer": "Reviewer_ow9B", "summary": "This paper proposes Hierarchical Balance Packing (HBP), a method to enhance the efficiency of supervised fine-tuning for long-context large language models. HBP mitigates workload imbalances inherent in hybrid long/short-context data training through a novel approach that combines multi-level packing groups, strategic sample allocation, and a dynamic training pipeline. Experiments show that HBP significantly accelerates training speed across various models and datasets without compromising performance.", "questions": "1. In Section 4.1, Algorithm 1 leverages a profile-based approach for hierarchical group selection. However, the paper does not clarify how the profiling was conducted, such as the number of iterations or evaluation settings. Could the authors provide details on the profiling procedure, and quantify the computational overhead of this process?\n2. The GreedyFill step uses shorter samples to fill the packs of longer samples. Does it reintroduce the communication overhead and attention imbalance problems that HBP aims to solve? \n3. Table 9 shows the importance of both short- and long-context data. Does the mixture ratio of short- and long-context data in the training set affect the performance of HBP?", "limitations": "Please see the comments in Weakness and Questions.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**:\n\n1. **The paper is well motivated**. The paper clearly analyzes the limitations caused by naive packing strategies in hybrid long/short context training, and proposed Hierarchical Balance Packing as an effective solution.\n2. **The experiments are sufficient**. The paper validates HBP's effectiveness across multiple models of varying scales and several datasets. The ablation studies clearly demonstrate the contribution of each component of HBP to performance and efficiency, which supports the paper's core claims. \n3. The paper is well-written and easy to follow.\n\n**Weakness**\n1. **The hierarchical groups auto-selection is inefficient**. It requires an additional profiling stage to determine the optimal configuration, yet the manuscript lacks a thorough discussion of the associated computational overhead (e.g., time cost) and the robustness of this process across different hardware environments. \n2. The empirical findings would be strengthened by more rigorous theoretical analysis. Specifically, while the Stable Loss Normalizer proposed in Section 4.3 shows promising results, the paper does not provide a clear theoretical justification for why the Average Token strategy effectively overcomes the limitations of existing normalization methods.", "was_revised": true, "final_justification": "I appreciate the authors’ response and keep my initial rating as borderline accept.", "rebuttal": "Thank you for your helpful comment.\n\n**W1 and Q1**\n> The hierarchical groups auto-selection is inefficient. It requires an additional profiling stage to determine the optimal configuration, yet the manuscript lacks a thorough discussion of the associated computational overhead (e.g., time cost) and the robustness of this process across different hardware environments. \n\n> In Section 4.1, Algorithm 1 leverages a profile-based approach for hierarchical group selection. However, the paper does not clarify how the profiling was conducted, such as the number of iterations or evaluation settings. Could the authors provide details on the profiling procedure, and quantify the computational overhead of this process?\n\n**A1**\n\nThank you for your valuable feedback. \n\nAlthough our method introduces an additional profiling stage, this stage is **highly efficient**, accounting for **only 1–2%** of the total training time. Therefore, the overall profiling overhead is very low, and the framework remains **efficient**.\n\n\n**Profiling Process**\n\nThe profiling process is currently provided in **Appendix I**. To make the profiling procedure clearer, we will include it in the main text in future versions.\n\n**Overhead of profiling process:**\n\nCompared to the total training time (168 min),  the overall overhead of **3.25 min** is negligible for  LLaMA3.1-8B.\n|        Setting        |       Dataset       | Profiling time | Training time | Overhead ratio |\n|:---------------------:|:-------------------:|:--------------:|:-------------:|:--------------:|\n|   LLaMA3.1-8B         | Tulu3 + longcite    |     **3.25 min**     |    168 min      |      **2 %**       |\n|   LLaMA3.1-70B        | Tulu3 + longcite    |     **15 min**       |   1400 min     |      **1 %**       |\n\n**Details of profile-based automatic optimization overhead**\n\n * **Memory Profiling:**\n\nThe memory profiling cost depends on the sequence length set (**Seq**) processed by the device. \n\n      Cost_m = len(Seq) × profile_iter × iteration_time\n\n* **Time Profiling:**\n\nThe total number of search strategies (S) is determined by the packing length set (L) and the Sequence Parallel (SP) settings. \n\nProfile_iter can be adjusted based on actual iteration times, with typical values ranging from 3 to 10. \n\nSpecifically:\n\n      S = len(L) x len(SP)\n      Cost_t = len(S) x profile_iter x iteration_time\n\n* **Total Cost** (Example)\n\nFor instance, with a  profile_iter 5,  L =  [16K,32K,128K] and SP ([2, 4, 8]), Seq = [4K, 8K, 16K], LLama3.1-8B\n\n     Cost_t = 3 × 3 × 5 = 45 x iteration_time\n     Cost_m = 3 x 5 = 15 x iteration_time\n     Cost_total = Cost_t + Cost_m = 60 × iteration_time\n\nConsidering the 3000+ training iterations,  totaling 60 iterations (**2%**) is negligible compared to the overall training time.\n\n**Robustness of this process across different hardware environments.**\n\nOur method is **hardware-agnostic** and consistently delivers good results across different hardware environments. To demonstrate this, we have included experimental results on A100 GPUs as an example.\n\n|      Model       |      Dataset      |  GPU   |AVE |   GPU days    |\n|:---------------:|:----------------:|:-----:|:--------------:|:--------------:|\n| LLama3.1-8B-ISF | Tulu3 + longcite | A100  |     56.2 | 9.4       ||\n| LLama3.1-8B-HBP | Tulu3 + longcite | A100  |   58.1  |**6.79 (1.38x)**   |\n\n---\n\n**Q2**\n> The GreedyFill step uses shorter samples to fill the packs of longer samples. Does it reintroduce the communication overhead and attention imbalance problems that HBP aims to solve?\n\n**A2**\n\nThank you for highlighting this point.\n\n**GreedyFill does not introduce attention imbalance problems.**\n\nAs described in **Section 4.2**, the workflow is: **Packing → GreedyFill → Balance Batching**. The attention imbalance problems are primarily addressed by **Balance Batching**. Since Balance Batching is applied after GreedyFill, it does not reintroduce attention imbalance.\n\n**GreedyFill does introduce some additional communication overhead; however, it enables *faster training*.**\n\nWithout the GreedyFill step, shorter samples would simply be replaced with padding,  leading to **wasted computation**. GreedyFill is specifically designed to minimize this waste, enabling **faster training overall**.\n\n\n| **Model**         | **Dataset**           | **GreedyFill** | **GPU days** |\n|:-----------------:|:---------------------:|:--------------:|:------------:|\n| LLaMA3.1-8B-HBP   | Tulu3 + longcite      | No           | 3.92         |\n| LLaMA3.1-8B-HBP   | Tulu3 + longcite      | ✔              | **3.73**         |\n\n---\n\n**Q3**\n> Table 9 shows the importance of both short- and long-context data. Does the mixture ratio of short- and long-context data in the training set affect the performance of HBP?\n\n**A3:**\n\nThis is an interesting question.  In our current experiments, we follow a typical SFT setting where short-context samples dominate. For example, LLaMA3 adopts a similar distribution. In our setup, the token-level ratio of short-to-long-context samples is approximately 3:2.  We have not conducted an extensive study on different short/long-context ratios in the paper. However, we performed a **small-scale experiment** by reducing the proportion of long-context samples. \n\n**Impact of Short/Long-Context Ratio**\n\nAs shown in the results, the model’s performance remains stable (similar to Table 9), suggesting that HBP is robust to variations in the mixture ratio.\n\n| **Model**          | **Dataset**             | **Short_context / long_context ratio** | **AVE** | **LongBench** | **Ruler-128K** |\n|:------------------:|:-----------------------:|:-----------------------------:|:-------:|:-------------:|:--------------:|\n| LLaMA3.1-8B-HBP     | Tulu3                   | 3:0                           | 57.9    | 43.2          | 53.5           |\n| LLaMA3.1-8B-HBP     | Tulu3 + longcite-half     | 3:1                           | 58.0    | 43.5          | 70.5           |\n| LLaMA3.1-8B-HBP     | Tulu3 + longcite-full     | 3:2                           | 58.2    | 43.1          | 70.8           |\n\n---\n\n**W2**\n> The empirical findings would be strengthened by more rigorous theoretical analysis. Specifically, while the Stable Loss Normalizer proposed in Section 4.3 shows promising results, the paper does not provide a clear theoretical justification for why the Average Token strategy effectively overcomes the limitations of existing normalization methods.\n\n\n**A4**\n\nThank you for your valuable feedback. We will clarify this in detail and will add it in the future version.\n\n**Stable Loss Normalizer is not empirical**\n\nThe Stable Loss Normalizer is designed based on theoretical principles. Its core purpose is to **ensure each token contributes equally to the overall loss, removing biases from sequence length, data-parallel size, and gradient accumulation**. Further analysis supporting this design is provided below. \n\n**In-depth analysis**\n\nTo clarify, let’s consider a data-parallel (dp) group with a size of 2, where each rank runs with a local batch size of 2. \n\nLet $loss$ denote the sum loss for a sample, and $T$ represent the number of tokens involved in the loss calculation.\n\n   **Token-Mean**：\n\nIn the token-mean approach, we compute the loss for each DP rank independently as follows:\n\n$$\nL_{dp_{1}} = \\frac{loss_{1} + loss_{2}}{T_{1} + T_{2}}, \\quad L_{dp_{2}} = \\frac{loss_{3} + loss_{4}}{T_{3} + T_{4}}\n$$\n\n\nThe final loss for the global batch is calculated by averaging the individual DP rank losses:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1}+loss_{2}}{2·(T_{1}+T_{2})} + \\frac{loss_{3}+loss_{4}}{2·(T_{3}+T_{4})}$$\n\nThis method may introduce bias if the data-parallel ranks have different sequence length distributions, as it normalizes the losses independently for each rank **without considering the global batch’s overall token length distribution**.\n\n **Sample-Mean:**\n\nIn the sample-mean method, losses for each DP are first normalized by their respective token lengths and then averaged across the samples:\n$$L_{dp_{1}}=\\frac{\\frac{loss_{1}}{T_{1}}+\\frac{loss_{2}}{T_{2}}}{2} , \\quad                   L_{dp_{2}}=\\frac{\\frac{loss_{3}}{T_{3}}+\\frac{loss_{4}}{T_{4}}}{2}$$\n\nThe final loss is then:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1}}{4·T_{1}} + \\frac{loss_{2}}{4·T_{2}}+\\frac{loss_{3}}{4·T_{3}} + \\frac{loss_{4}}{4·T_{4}} $$\n\nThis method balances the contribution of each individual sample. However, it does not fully account for the total number of tokens in the global batch, resulting in **samples with fewer tokens having a disproportionately larger influence on the overall loss**.\n\n **Stable Loss Normalizer**\n\nThe Average Token normalizer ensures that the total number of tokens in a global batch is considered. Here, we define the average token length across all samples:\n\n$$T_{ave}   = \\frac{T_{1} + T_{2} + T_{3} + T_{4}}{4}$$\n\nEach DP's loss is then normalized by  $T_{ave}$\n\n$$L_{dp_{1}}=\\frac{loss_{1}+loss_{2}}{2·T_{ave}}, \\quad       L_{dp_{2}}=\\frac{loss_{3}+loss_{4}}{2·T_{ave}}$$\n\nThe final loss is:\n$$L_{final}=\\frac{L_{dp_{1}}+L_{dp_{2}}}{2} =\\frac{loss_{1} + loss_{2} +loss_{3} + loss_{4}}{4·T_{ave}} = \\frac{loss_{1} + loss_{2} + loss_{3} + loss_{4}}{T_{1} + T_{2} + T_{3} + T_{4}}$$\n\nThis approach effectively normalizes the loss by the total number of tokens across all samples in a global batch, **ensuring that each token, regardless of which sample it belongs to, contributes equally to the final loss**.\n\n---\n\nWe hope our response addresses your concerns, and we welcome any further questions or feedback."}, {"review_id": "NuTTFyyOHI", "reviewer": "Reviewer_4FuH", "summary": "This paper addresses the problem of workload imbalance when fine-tuning large language models on mixed short- and long-context data. The authors propose Hierarchical Balance Packing (HBP), which creates multiple packing-length groups, assigns each example to its optimal configuration group via \"balanced packing\", and employs a dynamic training pipeline (adaptive sequence-parallelism, curriculum learning, and a stable loss normalizer). Experiments on models from 8B to 236 B parameters (e.g., LLaMA3.1-8B, Qwen2.5-32B, DeepSeek-V2) and datasets (Tulu3, LongCite, OpenHermes, etc.) show that HBP significantly improves attention balance ratio and communication ratio, yielding up to 2.4$\\times$ speedup with negligible or positive impact on downstream task performance.", "questions": "* The term \"hierarchical balancing\" is somewhat misleading. While the algorithm does employ multiple packing groups, it doesn’t follow a clear hierarchical (tree-like) structure. Typically, \"hierarchical\" implies successive merging or nesting—e.g., combining shorter-length groups into longer ones. It would help if the authors clarified why they chose this terminology and described any implicit hierarchy in their grouping process.\n\n* Introducing curriculum learning and adaptive sequence parallelism departs from the usual random-sampling assumptions in stochastic optimization, which could, in extreme cases, impair convergence. Are there practical downsides—such as degraded performance on certain tasks, increased risk of loss spikes, or sensitivity to hyperparameters—that practitioners should be aware of?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n1. **Well motivated and Important Problem**: Data packing plays an important role in improving both the efficiency and effectiveness of LLM training. This paper presents a practical method to tackle that challenge.\n\n2. **Novelty in Multi-Level Packing** Unlike previous single-length packing, HBP's two-stage group auto-selection (Alg. 1) and balanced packing (Alg. 2) jointly optimize packing lengths, sequence-parallel degree, and gradient-checkpoint configuration. This approach effectively takes into account Attention Balance Ratio and Communication Ratio, significantly improving both of them.\n\n3. **Comprehensive Empirical Validation.** Experiments span small to very large models and standard general-task (MMLU, BBH, HumanEval, etc.) and long-context benchmarks (Ruler, LongBench, LongCite). HBP consistently yields 1.3–2.4× speedups while maintaining or improving average task accuracy (Table 5, 6).\n\n### Weakness\n1. **Ablation Study**: There is an inherent trade-off among Padding Ratio, Attention Balance Ratio, Communication Ratio, and Average Tokens. I strongly encourage a systematic, pairwise analysis of these metrics—ideally visualized as a Pareto front—to clarify how different packing configurations balance efficiency and performance across these objectives.\n\n2. **Clarity**: Figures 2 and 3 would benefit from more descriptive captions that walk the reader through each step of the main algorithm. Additionally, the manuscript uses numerous abbreviations (e.g., ABR, CR, DBR, PR, SP for sequence parallelism, GC for gradient checkpointing), which should be better illustratd in the context, especially for captions of figures and tables.", "was_revised": true, "final_justification": "The rebuttal has addressed most of my concerns, so I would like to keep my current rating.", "rebuttal": "Thank you for your helpful comment.\n\n**W1**\n> Ablation Study: There is an inherent trade-off among Padding Ratio, Attention Balance Ratio, Communication Ratio, and Average Tokens. I strongly encourage a systematic, pairwise analysis of these metrics—ideally visualized as a Pareto front—to clarify how different packing configurations balance efficiency and performance across these objectives.\n\n\n**A1**\n\nWe agree that a  **systematic, pairwise analysis of these metrics** of how individual metrics affect LLM training would be valuable.\n\nSince the metrics are interdependent, we did not isolate them in ablation studies. However, most related results are already shown in **Tables 3 and 6**, and we have **summarized** the experimental results from the paper into the following table to illustrate their impact more clearly. We will add this ablation study in the future version:\n\n\n|      Model      | Seq Len |        Batching        |  DBR  |  PR  |  AVT  |  CR  |  ABR  | GPU days |\n|:---------------:|:-------:|:---------------------:|:-----:|:----:|:-----:|:----:|:-----:|:--------:|\n| LLama3.1-8B     |  128K   |        random         | 0.639 | 0.0  | 0.9K  |  1   |   0.61   |   38.0   |\n| LLama3.1-8B     |  128K   |        sorted         | 0.02  | 0.0  | 125K  |  1   |   0.51   |   5.5    |\n| LLama3.1-8B     |  128K   |     lsf packing       | 0.001 |  0   | 128K  |  1   | 0.506 |   5.2    |\n| LLama3.1-8B     |  128K   |         HBP （Hierarchical）           | 0.001 |  0   | 128K  | 0.17 | 0.288 |   4.2    |\n| LLama3.1-8B     |  128K   |         HBP （Hierarchical + balance）          | 0.001 |  0   | 128K  | 0.17 | 0.002 |   3.7    |\n\n**Results Analysis**\n\n- Comparing the first and second rows, DBR drops significantly and AVT increases, reducing training time from **38 to 5.5** GPU days.\n\n- From the second to third row, AVT reaches the maximum 128K, slightly improving speed (**5.5 → 5.2 GPU days**).\n\n- The third to fourth row shows a sharp drop in CR and ABR, leading to further reduction in training time (**5.2 → 4.2 GPU days**).\n\n- Finally, the fifth row shows that reducing ABR even more translates into the lowest training cost of 3.7 GPU days.\n\nThese results demonstrate how improvements in **individual metrics (e.g., DBR, AVT, ABR) correlate with more efficient LLM training**.\n\n---\n\n**W2**\n> Clarity: Figures 2 and 3 would benefit from more descriptive captions that walk the reader through each step of the main algorithm. Additionally, the manuscript uses numerous abbreviations (e.g., ABR, CR, DBR, PR, SP for sequence parallelism, GC for gradient checkpointing), which should be better illustratd in the context, especially for captions of figures and tables.\n\n**A2**\n\nThank you for your valuable feedback. We will revise the captions of Figures 2 and 3 to provide more detailed explanations of each step in the main algorithm. Additionally, we will improve the clarity of all abbreviations by defining them more clearly in the text and in the captions of figures and tables.\n\n---\n\n**Q1**\n> The term \"hierarchical balancing\" is somewhat misleading. While the algorithm does employ multiple packing groups, it doesn’t follow a clear hierarchical (tree-like) structure. Typically, \"hierarchical\" implies successive merging or nesting—e.g., combining shorter-length groups into longer ones. It would help if the authors clarified why they chose this terminology and described any implicit hierarchy in their grouping process.\n\n**A3**\n\nThank you for pointing out the potential ambiguity regarding the term “hierarchical balancing.”.  \n\nOur use of the term “hierarchical” is intended to highlight how our method organizes and balances packing across multiple levels of granularity.\n\nSpecifically, our approach first partitions the data at a coarse level (such as by length intervals) and then further refines the packing within each of these groups. While our method does not construct a conventional hierarchical (tree-like) structure or employ successive merging, it does involve multi-level grouping and coordinated resource allocation. In this sense, “hierarchical” reflects the layered nature of our organizational strategy rather than a strict tree-based hierarchy.\n\nWe will clarify this in the revised manuscript to prevent any potential confusion.\n\n---\n\n**Q2**\n> Introducing curriculum learning and adaptive sequence parallelism departs from the usual random-sampling assumptions in stochastic optimization, which could, in extreme cases, impair convergence. Are there practical downsides—such as degraded performance on certain tasks, increased risk of loss spikes, or sensitivity to hyperparameters—that practitioners should be aware of?\n\n**A4**\n\nThank you for your question. In our approach, curriculum learning is used only for **a few early iterations** to stabilize training, so its impact on randomness and convergence is **minimal**.\n\n**Usual random-sampling assumptions**\n\nMoreover, HBP itself incorporates several sources of randomness:\n\n- The base packing method is **inherently stochastic**.\n- The alternating training strategy—switching between short- and long-context data—**introduces additional randomness**, which helps prevent length bias.\n\n**Robustness to Hyperparameter Settings**\n\n As shown in Table 8, our curriculum learning strategy is **not sensitive** to hyperparameter settings and maintains consistent effectiveness.\n\n|        Model        | CL-Iterations |  AVE  | LongBench |\n|:-------------------:|:-------------:|:-----:|:---------:|\n| LLama3.1-8B-HBP     |      0        | 56.6  |   41.6    |\n| LLama3.1-8B-HBP     |     100       | 58.0  |   43.1    |\n| **LLama3.1-8B-HBP** |    500        | **58.2** | **43.1** |\n\n In practice, we have **not observed significant degradation** in performance. Additionally, we conduct extensive experiments and evaluate the model's performance on both long and short sentence benchmarks to validate the robustness of our approach.\n\n---\n\nWe hope our response addresses your concerns, and we welcome any further questions or feedback."}]

Author rebuttal:

Meta-review: This paper proposes a data packing method to address training efficiency issues for long-context large language models (LLMs).The authors propose Hierarchical Balance Packing (HBP), which introduces hierarchical group auto-selection,  balance packing, and  a dynamic training pipeline with curriculum learning and a stable loss normalizer. Experiments on models ranging from 8B to 236B parameters demonstrate consistent 1.3–2.4× training speedups with negligible or positive impact on downstream performance. The rebuttal effectively addressed efficiency and generalization concerns, strengthening the case for acceptance. Remaining limitations include heuristic design choices and clarity of presentation. Overall, given the importance of the problem, the empirical gains, and the authors’ clarifications, the paper leans toward borderline accept.
Decision: Accept (poster)