Layerwise Repacking: Variable-Length Hybrid-Attention Training without an Extra Redistribution Pass

1 Introduction

• Full attention needs SP16/32 groups for the longest samples, while sparse-layer work balances within one eight-GPU node in the profiled model.Evaluated longest samples are 9.8–175× the shortest (median 23×); full-attention computation is quadratic, while sparse-layer work is approximately length-linear.Source: §1 · p. 1
• Using wide groups for sparse layers lowers observed effective Ulysses rate by 7.0× versus node-local SP8 groups.Wide-group all-to-all occupies 21.4% of a step on wide ranks, versus 4.3% on SP8 ranks.Source: §1 · p. 1
• Layerwise Repacking keeps wide parent groups for full attention and repacks samples into node-local sparse children without a separate redistribution pass.The existing full-layer Ulysses exchanges perform both head exchange and sample repacking.Source: §1 · p. 2
• 5.95% lower mean end-to-end time with repacking alone; the complete system lowers it by 7.61% and raises MFU by 4.83 percentage points on 22 wide-parent steps.The evaluation covers steps with SP16/32 parents among the first 200 training steps; gains do not apply to the full 200-step window.Source: §1 · p. 2
• Repacking improves 20 of the 22 evaluated steps but slows two where an extremely long sample limits balance within node-local children.The evaluation isolates repacking and the auxiliary parent solver’s contributions.Source: §1 · pp. 2–3
• An auxiliary parent solver accounts for full- and sparse-layer costs when assigning samples, improving the load balance achievable by node-local children.Child balance depends on which samples share each parent.Source: §1 · p. 2

2.1 Sample and Layer Heterogeneity

• Sparse-layer work is dominated by near-linear projections and MoE in the profiled model, while the attention-cost gap widens with sample length and is carried by long samples.Across 22 wide-parent steps, the longest sample is 9.8–175× the shortest (median 23×).Source: §2.1 · p. 3

2.2 Layerwise Topology Mismatch

• For most evaluated workloads, sparse-layer work balances within node-local SP8, while SP16/32 Ulysses crosses node boundaries.The observed effective Ulysses rates are 275.6 GiB/s on node-local SP8 ranks and 39.5 GiB/s on cross-node SP16/32 ranks; the wide-rank all-to-all time is 7.97 s versus 1.60 s on SP8 ranks.Source: §2.2 · p. 3, Figure 2 · p. 3

2.3 Switching Layouts Moves Token Ownership

• Changing from an SP16 parent to node-local SP8 children changes which ranks own sample tokens, so a separate layout-conversion phase would add cross-node activation traffic.A dedicated pass would cost about 0.9–1.9 s per step, roughly 30–60% of repacking’s exposed-all-to-all reduction.Source: §2.3 · pp. 3–4

2.4 Parent Composition Limits Child Balance

• Child balance is limited by parent membership, and the baseline’s full-attention-calibrated weights can miscompose samples for the distinct full- and sparse-layer costs.The wide-parent region’s forward rank maximum exceeds that of ordinary SP8 ranks by 2.418 s with baseline membership and repacking disabled.Source: §2.4 · p. 4

2.5 Why Pure Ulysses rather than USP

• Pure Ulysses is chosen because sparse attention offers too little computation for effective ring overlap and subsequent MoE work falls outside that overlap window.This is a design choice, not a measured comparison with USP; ring-based overlap for the remaining full-layer exchange is described as orthogonal and left for future work.Source: §2.5 · p. 4

3.1 Problem Formulation

• The step objective is critical-path time—the slowest rank’s computation and exposed communication—subject to adding no redistribution pass.Parent rank boundaries remain fixed; child partitioning is restricted to samples within a parent, and per-rank memory is checked empirically rather than modeled explicitly.Source: §3.1 · p. 4

3.2 Layer-Specific Layout Formulation

• Full attention uses parent layout F, while sparse layers use node-local child layout S; hidden states remain in S except inside full attention.The full layout exists between the full layer’s two Ulysses exchanges; child padding is local, and the S shards must partition the parent’s real packed tokens with an inverse permutation restoring the original packing order.Source: §3.2 · pp. 4–5

3.3 Communication-Fused Repacking

• Variable-size splits let the existing Q/K/V and output Ulysses sites perform the S-to-F gather and inverse return, including sample repacking.Every tensor traverses each mandatory exchange once; split sizes and local permutation change, but no payload is transferred twice.Source: §3.3 · p. 5

3.4 What the Fused Exchange Must Preserve

• The fused exchange must preserve inverse layout transitions, identical deterministic permutations across ranks, child-local padding bookkeeping, and consistent backward and recomputation splits.The loss path must also reconstruct global packed order without adding a phase.Source: §3.4 · p. 5

4.3 Communication-Fused Execution

• At each full layer, the existing Ulysses exchanges temporarily realize F and return to S, combining head exchange and layout conversion without a new redistribution phase.After the trunk, an uneven all-gather reconstructs global packed order for the loss path, replacing the baseline uniform all-gather.Source: §4 · p. 5, §4.3 · p. 6

4.1 Child Planning and Repair

• Before child planning, repair ensures every rank can join a non-empty node-local child of at most eight ranks; SP16 parents need two samples and SP32 parents four.Repair moved no samples in the evaluated workloads because every wide parent already met its minimum.Source: §4.1 · p. 6

4.2 Sparse-Trunk Execution

• The sparse trunk retains hidden states and related state in S, where sparse layers run with child-specific packed lengths under node-local groups of at most eight ranks.SP1 children need no attention Ulysses exchange; larger children use the usual Ulysses primitive within the child group.Source: §4.2 · p. 6
• A 1F3S repacked parent has 75% fewer cross-node Ulysses call sites in its forward graph, falling from 8N_F to 2N_F.Backward and recomputation preserve the same layer-type ratio.Source: §4.2 · p. 6

5 Improved Parent Solver

• The auxiliary solver changes parent membership, not repacking execution; it improves the children repacking can produce.Round 1B preserves parent topology, while repair and round 2 remain the same across the baseline and improved configurations.Source: §5 · p. 6, §5.1 · p. 7

5.1 Why Membership Matters

• Child planning cannot move samples across parents, so a parent containing several expensive samples can limit achievable sparse-layer balance.The improved solver targets the parent-level ideal-average term, while the heaviest-sample bound remains dependent on which parent holds that sample.Source: §5.1 · p. 6

5.2 Parent Reallocation

• Round 1B uses modality- and SP-degree-calibrated work predictions, then scores rebuilt assignments with an aggregate predictor of step-level critical-path time.The evaluator scores membership before repair and child planning; it does not model those stages or score the executed child plan.Source: §5.2 · p. 7

5.3 Candidate Portfolio and Selection

• The solver generates a portfolio of candidate memberships and selects the tolerance-tied minimum predicted step time; measured solve time is at most 4 ms per step.Routes include Compute-LPT, Memory-LPT, and optional Anchor-Fill; the round-1A assignment remains a fixed fallback if no candidate is valid.Source: §5.3 · p. 7

6.1 Experimental Setup

• The evaluation covers 22 wide-parent workloads among the first 200 training steps; the other 178 steps use unchanged SP8 plans and are untouched by either lever.On control steps, configuration means differ by at most 0.346 s, about one eighth of the 2.83 s headline gain; loss agreement is a necessary check, not proof of correctness.Source: §6.1 · pp. 7–8, §6.2 · p. 8

6.2 Overall End-to-End Performance

• The reported 7.61% saving is a per-triggered-step figure; wide-parent steps are 37.221 s, and window-level savings vary with how often workloads trigger wide parents.Wide parents occurred on 22 of the 200 collected steps.Source: §6.2 · p. 8
• 2.214 s (5.95%) and 1.497 s (4.17%) mean E2E reductions accompany repacking under baseline and improved solvers; gains span all 20 one-SP16 workloads but reverse on the two-SP16 and SP32 workloads.The corresponding MFU gains are 3.75 and 2.78 pp; the SP8 child cap leaves one long sample dominating the parent in those two exceptional workloads.Source: Figure 3 · p. 8, §6.2 · p. 8

6.3 Communication Reduction

• Figure 4 shows exposed A2A time falling from 7.97/7.59 s to 4.76/4.60 s under baseline/improved solvers, while total payload changes by less than 5%.More than 70% of repacked payload takes the node-local child path; the fused full-layer exchanges perform the S↔F conversion without a second payload traversal.Source: Figure 4 · p. 9, §6.3 · p. 9

6.4 Parent-Solver Gains

• The improved solver repairs parent-composition imbalance that predates repacking, lowering parent-region forward rank max gaps over ordinary SP8 from 2.418/0.696 s to 1.692/0.452 s with repacking OFF/ON.Its solver gains overlap with repacking because the child planner absorbs part of the mis-membership.Source: §6.4 · p. 9
• 1.334 s with repacking OFF versus 0.617 s with it ON are the improved solver’s mean E2E reductions, and it does not lower mean E2E on 7 of 22 workloads in either state.Solver membership reduces parent-region forward rank max by 0.414 s versus −0.071 s in ordinary-SP8 regions; backward reductions are similar at 0.367 and 0.389 s.Source: §6.4 · pp. 9–10, Figure 5 · p. 9

6.3 Communication Reduction

• A separate redistribution pass would return roughly 30–60% of repacking’s 3.20/2.99 s exposed-A2A reduction, depending on which activations it moves.Table 3 estimates about 1.89/1.78 s for attention tensors or 0.94/0.89 s for hidden states at 39.5 GiB/s, excluding extra launch and synchronization overheads.Source: §6.3 · p. 9, Table 3 · p. 10

The SP8 child cap

• The SP8 child cap leaves one long sample dominating its parent on the two-SP16 and SP32 workloads, erasing gains and leaving the SP32 step unimproved.The heaviest sample’s eight-rank load is 1.47–1.70× its parent’s average per-rank load on those regressing workloads.Source: §7 · p. 10

Scope of the evaluation

• Results cover one 1F3S model on one cluster and 22 wide-parent steps among the first 200 training steps.The length distribution, and thus how often wide parents arise, may change later in training.Source: §7 · p. 10

Applicability

• The method may apply to other near-linear layers, such as linear attention, only when their state, communication, and operator structure permit the same split; this remains unvalidated.A larger high-bandwidth domain postpones the mismatch, but it reappears when full attention requires a group wider than that domain.Source: §7 · p. 10

Memory and calibration

• Repacking can raise per-rank activation residency by up to R_b/r_bj, while the improved solver requires cost-predictor recalibration when the model or cluster changes.Evaluated steps show little change in whole-device peaks, but larger samples may behave differently.Source: §7 · p. 10

Future directions

• Future work could gate repacking by predicted gain, make the child cap a per-sample choice, and score executed plans after repair, child planning, and feasibility checks.The proposed plan evaluator would account for critical-path time, communication, permutation and padding overhead, and peak memory; thresholds inferred from 22 workloads need validation elsewhere.Source: §7 · p. 10

8.1 Sequence Parallelism

• Layerwise Repacking changes a sample’s participating ranks and group width between full and sparse layers, using variable splits and an explicit sample permutation within the same exchange.DSP also moves tokens between stages with all-to-all, but keeps each sample’s rank set and group width fixed.Source: §8.1 · p. 10

8.2 Sparse Attention Parallelism

• Unlike sparse-attention methods that shrink attention work or expert-parallel systems that balance token routing, this work targets layer-level mismatch and leaves expert-routing code unchanged.The profiled sparse layers are MoE-dominated, with per-token expert work treated as approximately length-linear when placing samples.Source: §8.2 · p. 11

8.3 Layer-Wise Parallelism Planning

• Unlike systems that select mappings per layer or stage before training, Layerwise Repacking changes which ranks hold each sample as execution moves between layer types in every step.HotSPa varies strategies across sequence groups within a step, but each group keeps one strategy across all layers.Source: §8.3 · p. 11

8.4 Variable-Length Workload Placement

• Existing variable-length placement systems use one group for every layer of a sample, whereas Layerwise Repacking changes that group between layer types within a step.KnapFormer redistributes tokens at both ends of the transformer trunk; the per-boundary dedicated pass estimated in Table 3 is an alternative design, not a KnapFormer measurement.Source: §8.4 · p. 11

9 Conclusion

• Layerwise Repacking retains wide parent groups for full attention and repacks samples into node-local sparse children, fusing layout conversion into the two mandatory full-layer Ulysses exchanges.An auxiliary parent solver reassigns samples across fixed parent groups to improve the load balance achievable by the children.Source: §9 · p. 11
• Across all 22 steps with SP16/32 parents in the first 200 training steps, the complete system reduces E2E by 7.61% and raises MFU by 4.83 pp over the baseline with repacking disabled.Repacking alone reduces E2E by 5.95% under the baseline solver and 4.17% under the improved solver.Source: §9 · p. 11
• The SP8 child cap leaves one long sample dominating its parent on the two-SP16 and SP32 workloads, erasing the gain; future steps include benefit-based gating, adaptive child sizing, and broader evaluation.The proposed extensions also include scoring the executed plan.Source: §9 · p. 11

Algorithm 1 Resource-Priority Parent Allocation

• The algorithm accepts samples, fixed parents, a keep set, and a priority, and returns a valid parent assignment or failure.After rebuilding the assignment, it computes b_T(F).Source: Algorithm 1 Resource-Priority Parent… · p. 12
• Kept samples are placed in descending (a_i, −i) order among distinct empty SP>8 parents, with failure if no candidate survives.Each placement uses Equation (8).Source: Algorithm 1 Resource-Priority Parent… · p. 12
• Remaining samples are ordered by Q_i^ν and allocated using candidate parents, with failure when no candidate remains.When remaining samples equal the number of empty parents, the algorithm assigns them to those parents; it then verifies coverage, non-empty parents, anchors, and length caps.Source: Algorithm 1 Resource-Priority Parent… · p. 12

• Keep set K contains heavy samples whose predicted work exceeds a fixed multiple of mean per-rank work; each must anchor a distinct wide parent.If |K| exceeds the number of wide parents, the solve falls back to round-1A membership.Source: A Solver Route Constructions · p. 12
• Anchor-Fill initializes parent membership from Equation (5), disables the length cap, and assigns each keep sample to a distinct empty SP>8 parent.It fails if no distinct empty SP>8 parent remains.Source: A Solver Route Constructions · p. 13
• Anchor-Fill fills anchored parents, may swap an accepted filler for an unassigned sample while allowing one budget overflow, then closes each parent.The route reserves one sample for each unanchored empty parent.Source: A Solver Route Constructions · p. 13
• Residual samples are allocated by Equations (7)–(8), followed by checks for coverage, non-empty parents, and anchors.A successful assignment rebuilds F, computes bT(F) by Equation (6), and returns F.Source: A Solver Route Constructions · p. 13

A.1 Resource-Priority Allocation

• Compute-LPT and Memory-LPT pre-place required heavy samples, then choose among least-loaded feasible parent representatives by projected global makespan.Compute-LPT prioritizes expensive samples for scarce wide parents; Memory-LPT prioritizes long samples to reduce fragmentation by short samples.Source: §A.1 · p. 13
• Near-ties prefer smaller projected parent load, then smaller SP class, then fixed rank order, making assignments deterministic for identical inputs.Source: §A.1 · p. 13

A.2 Anchor-Fill Allocation

• Anchor-Fill anchors threshold-selected keep samples in distinct eligible large parents, fills them with cheap work within a budget, and may use one swap before closing a parent.Remaining work goes through the shared compute-LPT constructor; unlike Zellige, this route preserves the round-1A parent count, sizes, and rank intervals.Source: §A.2 · p. 13

A.3 Candidate Selection Rule

• Among enabled routes returning valid candidates, selection minimizes predicted E2E; candidates tied within numerical tolerance use the fixed order Compute-LPT, Memory-LPT, Anchor-Fill.Source: §A.3 · pp. 13–14