Source-linked AI summary
CIERA: Cross-Iteration Exponent Reuse for Lossless Allgather in Sharded MoE Training
Ali Zafar Sadiq, Haiying Shen, Masahiro Tanaka
TL;DR
Sharded MoE training suffers from an Allgather bottleneck, while existing compression approaches either lose numerical fidelity or impose lossless-compression overhead. CIERA reuses stable exponents across iterations and selectively overlaps lossless compression with communication and computation, achieving substantial speedups while preserving exact reconstruction.
Problem
Sharded MoE training makes Allgather a critical communication bottleneck, while prior methods either introduce numerical error or can slow training through lossless compression overhead.
Method
CIERA caches weight exponents across iterations, transmits sign and mantissa when exponents are unchanged, selectively compresses profitable shards, and overlaps compression with communication and computation.
Results
CIERA achieves 3.70× speedup over ZeRO-3 and 3.68× over ZeRO++ on OLMoE-1B-7B at 16 GPUs while preserving bitwise-exact parameter reconstruction in all evaluated runs.
Takeaways & Limitations
CIERA enables bit-exact, lossless compression with low decompression overhead and outperforms evaluated baselines on communication-bound MoE models.
Takeaways & Limitations
CIERA targets sparse MoE training, while results beyond 16 GPUs rely on trace-driven simulation and 64- and 128-GPU projections require more caution.
Abstract
from arXiv · showhide
In training Mixture-of-Experts (MoE) models, sharded data parallelism partitions each expert's parameters across GPUs, requiring an Allgather operation to reconstruct the full weight matrix before each layer executes. This communication often dominates iteration time. Prior work often reduces this overhead using lossy compression methods that sacrifice numerical fidelity, while existing lossless methods do not exploit cross-iteration exponent stability. In this paper, we propose Cross-Iteration Exponent Reuse Allgather (CIERA), a lossless, system-aware communication method for sharded MoE training. We observe that after a brief warmup phase, the exponent values of most weights remain unchanged across iterations. Based on this, we cache exponents locally and transmit only the sign and mantissa when exponents are unchanged. The receiver reconstructs the original weights exactly by combining the cached exponents with the received data. Since parameter matrices vary in shape across layers, compression is applied only when it yields net time savings. Moreover, the compression operation is overlapped with both Allgather communication and computation. Our real experiments and large-scale trace-driven simulator show that on OLMoE-1B-7B at 16 GPUs, CIERA achieves a 3.70x speedup over the lossless baseline and 3.68x over the lossy baseline, projected to reach 4.28x and 4.42x respectively at 128 GPUs, while preserving bitwise-exact parameter reconstruction in all evaluated runs.
1 Introduction
Sharded MoE training makes Allgather a major bottleneck, while prior compression methods trade numerical fidelity for communication savings or add lossless overhead. CIERA exploits stable exponents across iterations with selective, overlapped lossless compression.
- Sharded data parallelism requires Allgather to reassemble expert weights before forward computation, creating a critical bottleneck because MoE offers little computation to hide communication latency.
- Existing approaches either introduce numerical error through lossy compression or can slow training by inserting lossless compression and decompression into communication.
- Allgather accounts for 51.8–69.0% of communication time in sharded MoE training.
- After a brief warmup, at least 99% of weights retain unchanged exponents across iterations, although shard change rates vary substantially by type and model.
- CIERA caches exponents and transmits only sign and mantissa when unchanged, selectively compressing profitable shards while overlapping checking, compression, decompression, communication, and computation.
- 3.70× speedup over ZeRO-3 and 3.68× over ZeRO++ is achieved on OLMoE-1B-7B at 16 GPUs with bitwise-exact reconstruction in evaluated runs.
2 Motivation and Observations
Sharded MoE training makes Allgather a major communication bottleneck, while post-warmup exponent stability creates an opportunity for reuse. However, reuse varies by granularity and shard type, and naive exponent checking can impose substantial overhead.
- Allgather Bottleneck: Allgather accounts for 51.8–69.0% of total communication time across four MoE models.It reconstructs full parameters during forward and backward passes under sharded data parallelism.
- Exponent Stability: After 30 warmup iterations, at least 99% of exponent values remain unchanged between consecutive iterations across three MoE models.At learning rate 10−4, mean exponent change rates range from 0.30% to 0.51% over 1000 iterations.
- Reuse Granularity: Shard-level reuse is limited because 59–99% of shards change each iteration, whereas 84–99% of B=64 blocks remain reusable.At B=1024, the reusable-block range decreases to 25–86% across models.
- Shard-Type Heterogeneity: Block-change rates vary substantially by shard type, with layer norm consistently most stable and attention and routing-gate shards changing most frequently.Expert FFN shards are highly stable for OLMoE but less stable in some other models.
- Time Overhead: Checking exponents every iteration without overlap adds 25–30% iteration-time overhead.Checking every 10 iterations reduces the overhead to 3–4%.
3 System Design
CIERA combines exponent-reuse compression, benefit-driven shard selection, and computation–communication pipelining to reduce Allgather payloads without sacrificing exact reconstruction. Its design caches stable exponents, transmits changed information selectively, and overlaps compression work with communication and computation.
- Exponent-Reuse Compression: CIERA applies exponent-reuse compression after warmup, splitting weights into exponents and sign-plus-mantissa while reusing cached exponents when unchanged.The approach operates on selected shards at the Allgather path.
- Block-Level State: Each sender tracks per-block exponent changes with hashes and a bitmap, transmitting changed exponents plus sign-and-mantissa values for all blocks.The receiver uses the bitmap and cached exponents to identify and update only changed blocks.
- Reconstruction: CIERA reconstructs each received shard by merging decoded changed exponents with cached unchanged exponents and received sign-plus-mantissa values.Receiver-side reconstruction remains after Allgather, whereas several preparation and compression operations can overlap with communication.
- Selective Compression: Benefit-driven selective compression chooses shards whose expected communication savings exceed hashing, compression, decompression, and cache costs.Low-change shards are favored, while high-change attention and routing-gate shards are typically skipped.
- Computation–Communication Pipelining: The FX-graph schedule moves hashing, predecompression, and compression earlier, overlapping them with preceding computation and concurrent Allgather communication.Only receiver-side decompression remains on the critical path before the corresponding computation launches.
4 Performance Evaluation
CIERA improves iteration time across communication-bound MoE workloads by selectively applying lossless compression and overlapping its overhead with execution. Its benefits increase with communication scale, while exact reconstruction is preserved.
- Scaling: 1.16–4.28× iteration-time speedup over ZeRO-3 is reached from 8 to 128 GPUs as inter-node communication becomes more dominant.Measured and simulator-projected results are combined, with real measurements marked separately in the evaluation setup.
- Baseline comparison: 1.13–3.38× iteration-time speedup over ZeRO-3 and 1.21–3.80× over ZeRO++ is achieved on five communication-bound models at 4 GPUs.Scout remains at 1.00× because no shards are selected.
- Ablation study: 1.27–1.76× speedup over ZeRO-3 comes from exponent reuse on four models, but Mixtral slows to 0.96× when hashing and packing exceed byte savings.Exponent reuse is evaluated with the execution support required for compressed Allgather, rather than as an isolated codec.
- Ablation study: BSC raises Mixtral from 0.96× to 1.15× and increases speedups on OLMoE, DeepSeek, Qwen2, and MiniCPM by skipping unprofitable shards.Scout selects zero shards on 4-GPU NVLink and remains at 1.00×.
- Ablation study: CCP adds 0.04–0.09× on OLMoE, DeepSeek, and Qwen2, while improving MiniCPM and Mixtral by under 1%.It overlaps hash, compression, and decompression with preceding backward computation.
- Sensitivity analysis: 2.86× is DeepSeek’s speedup over ZeRO-3 at four layers, while increasing sequence length reduces but does not eliminate compression benefits.OLMoE falls from 2.89× to 2.61× and MiniCPM from 1.75× to 1.35× as sequence length increases; B=512 balances kernel-launch overhead and exponent reuse.
- Numerical fidelity: Zero maximum absolute loss difference across 20,000 iterations confirms bitwise-identical training dynamics to uncompressed ZeRO-3.Every gathered parameter is reconstructed bit-for-bit at each Allgather boundary.
5 Related Work
Prior communication-compression work includes lossy gradient and parameter methods plus system-aware scheduling approaches, whereas this paper evaluates CIERA in the sharded MoE setting.
- Gradient compression: Lossy gradient methods use quantization, sparsification, momentum correction, or error feedback to reduce backward communication.TAGC instead applies format-aware lossless compression to selected transformer-gradient layers.
- System-aware communication: System-aware methods adapt compression, reduce collective traffic, or reschedule communication with computation, but the cited adaptive gradient methods target ReduceScatter or AllReduce.The related systems include approaches such as Cupcake, OmniReduce, SqueezeNIC, and gZCCL.
6 Conclusion
CIERA reuses stable exponents to reduce Allgather overhead through lossless, selective compression and overlapping execution. Its supported scope is sparse MoE training, with larger-scale results relying increasingly on simulation.
- Conclusion: CIERA reuses cached exponents and selectively compresses shards whose expected benefits exceed costs.The method targets Allgather overhead while preserving lossless reconstruction and overlapping compression with computation and communication.
- Conclusion: CIERA outperforms the evaluated baselines on communication-bound MoE models.The conclusion states this result without extending it beyond the evaluated setting.
- Limitations and future work: CIERA targets sparse MoE training; on dense LLMs, the selection rule prunes most shards and speedup collapses.Beyond 16 GPUs, results use a trace-driven simulator; 64- and 128-GPU projections require more caution.
- Broader impacts: Bitwise-exact reconstruction avoids the accuracy risks introduced by lossy communication compression.The broader-impact discussion connects this property to reduced communication cost and potentially lower training cost.
A Block-level exponent reuse algorithm
The block-level algorithm detects changed exponents with sender-side hashes, transmits changed exponent payloads plus mantissas, and reconstructs each shard using receiver-side cached exponents.
- State and inputs: Each shard is divided into blocks, with persistent sender-side HashCache and receiver-side ExpCache state.The algorithm processes nb-value blocks after splitting each weight into exponent and mantissa components.
- Sender: The sender hashes each block’s exponents and sets its bitmap bit when the hash differs from the cached value.Changed exponent values are added to the payload for transmission.
- Transmission: The sender transmits the bitmap, mantissas, and changed-exponent payload only when at least one block has changed.Unchanged blocks reuse their cached exponent representation.
- Receiver: The receiver updates cached exponents from the payload and merges them with mantissas to reconstruct the shard.The merge operates over the receiver’s persistent exponent cache and the received mantissa data.