OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

OOM: O FFLOAD -O PTIMIZED F ULL -PARAMETER M O E T RAINING VIA M EMORY-E XECUTION C O -D ESIGN Ruijia Yang 1 Shiyuan Lin 1 Chengxi Liao 1 Zeyi Wen 1

arXiv:2611.11364v1 [cs.DC] 8 Nov 2026

{ryang379,slin954,cliao118}@connect.hkust-gz.edu.cn wenzeyi@hkust-gz.edu.cn

A BSTRACT Offloading model state to host memory and streaming layer weights let a single multi-GPU node train large mixture-of-experts (MoE) models. GPU memory then becomes a performance budget, shared by resident expert state, retained activations, and execution worksets. The best allocation depends on the expert execution path, batch size, and sequence length. We present OOM (Offload-Optimized MoE), a full-parameter MoE training system that co-designs expert execution and memory allocation. For expert-parallel execution, OOM updates resident experts on their owning GPUs and jointly plans state residency and MoE activation retention according to pipeline dependencies. For local execution, OOM bounds GPU worksets through dependency-aware partitioning; in hybrid models, it also retains attention results on the host to avoid costly recomputation and selectively regenerates layer boundaries through cheaper layers to make room. Calibrated cost models guide joint configuration within GPU and host budgets to minimize step time; OOM selects among the configured paths through short executions. On four H100s, OOM trains with 1M-token contexts, whereas FSDP2, ZeRO-Offload, and MegaTrain exhaust GPU memory at 512K. It also trains a 122B-parameter MoE model, on which all three exhaust host memory. On four 24 GB RTX 4090s, it trains 512K-token contexts and 1,024 1K-token sequences per step without gradient accumulation. Across matched-global-batch comparisons, OOM delivers 1.58× and 4.08× geometric-mean (up to 2.14× and 7.77×) throughput over CPU-offloaded FSDP2 on the two platforms while using the least host memory among the evaluated systems—as little as 1/3.0 of FSDP2’s. 1

I NTRODUCTION

Mixture-of-experts (MoE) models expand model capacity by activating only a few experts per token (Fedus et al., 2022; Jiang et al., 2024). Qwen3.6-35B-A3B, for example, selects eight of 256 routed experts per token (Qwen Team, 2026c). Yet full-parameter training must maintain model state for all experts. This gap between sparse computation and total training state makes GPU memory (HBM) the bottleneck of full-parameter training. State offloading and layer streaming keep model state in host memory and materialize working weights on GPUs as layers execute (Ren et al., 2021; Rajbhandari et al., 2021; Yuan et al., 2026; Yang & Wen, 2026; Yang et al., 2026b). GPU memory then becomes a performance budget with competing uses: retaining state to reduce host service, preserving activations to avoid recomputation, and supporting larger execution worksets. With four-way expert parallelism 1

The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China. Correspondence to: Zeyi Wen <wenzeyi@hkust-gz.edu.cn>. Code: https://github.com/RegiaYoung/OOM-MoE.

on four H100s at global batch 32 and 4K context, retaining selected expert state on GPUs instead of streaming all experts cuts Qwen3.6’s step time from 8.48 to 5.02 s (1.69×) without changing model computation; Figure 1b splits a profiled step of each plan into compute and non-compute time. Host capacity is also shared by model state, saved layer inputs, and retained operator results. Expert execution and memory allocation constrain each other. Execution determines data movement, live worksets, and the dependencies among computation, transfers, and updates; memory allocation determines which execution configurations fit and changes which work delays step completion. The two must therefore be designed together and judged by the time of the whole training step, not by isolated savings in transfers or recomputation. ES-MoE adapts expert offloading and pinning, and MoP specializes execution by component, but neither jointly budgets GPU memory between selective state residency and activation retention (Kim et al., 2024; Nguyen et al., 2026). ProTrain jointly plans state and activations through chunkand block-level policies, while Mist co-optimizes parallelism and memory with an efficient tuner that assumes

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design (a) EP vs. Local EP S

S resident state

GPU 0

Shard prep

0

A retained activations

GPU 1

1

D2H

A2A

GEMM

GPU 0

Bank prep

A2A

GEMM

0

1

2

3

GPU 1

GEMM

0

1

2

3

Reduce Replay

A

GEMM

W

GEMM

A2A

E per GPU D2H

CPU Adam W

GEMM

Forward

Backward

(b) HBM spent on resident state non-compute

1

8.5 s 5.7 s

(c) Faster path flips

k tokens/s

R0 11 GiB R23 65 GiB

A2A

S

Backward

H2D

compute

CPU Adam

Replay

A2A

Forward

Local

E/P per GPU

3

2

H2D A

A2A

W live workset

40

R: resident state K: retained activations

20 EP4 Local

−60% non-compute

0 0

4 8 Profiled step (s)

32 64 128 Global batch (S = 4K)

Figure 1. Expert execution and GPU-memory value: (a) EP versus Local; (b) residency reduces step time; (c) configured-path crossover. (b,c) use four H100s at 4K context.

identical layers (Yang et al., 2026a; Zhu et al., 2025). These policies leave two host-backed MoE problems open. First, choosing expert-shard residency, MoE activation retention, and token-routing chunk sizes together requires weighing host-state service against replay with token exchange. Second, within each GPU, operator partitioning, result retention, and boundary regeneration must respect operator-specific dependencies and both memory budgets, dividing host capacity between layer inputs that are cheap to regenerate and attention results that are expensive to recompute. We present OOM (Offload-Optimized MoE), a single-node multi-GPU system that co-designs expert execution and memory allocation for full-parameter MoE training. It supports two complementary expert-execution paths: expert parallelism (EP) and local execution (Local) (Figure 1a,c). EP assigns disjoint expert shards to GPUs, so each GPU materializes fewer expert weights per layer but exchanges routed tokens (Lepikhin et al., 2021). Local instead materializes all expert weights for the current layer on each GPU to process local tokens, avoiding that exchange. C1: EP execution with state–activation budgeting. In EP, each expert shard is both an execution unit and an optimizer update unit. This ownership lets us retain and update selected shards on their GPUs, reducing host transfers and CPU optimizer work while streaming the rest through a bounded pipeline. Retaining MoE activations competes for this budget but eliminates corresponding replay and token

exchange. Residency changes which work remains exposed when computation, transfers, and updates overlap, and thus the step-time benefit of activation retention. A dependencybased EP planner allocates GPU memory jointly to resident state, retained activations, and routing worksets. C2: Dependency-aware hybrid activation execution. Local amortizes expert-bank materialization over more tokens, but larger token workloads expand activation and operator worksets. Recurrent mixing, full attention, and expert computation have different dependencies; we partition each along compatible axes and coordinate lifetimes to bound GPU worksets. Fitting the workload is not enough: in hybrid models that interleave full attention with cheaper recurrent or sliding-window layers, long-context attention results grow linearly in size but require quadratic computation to reproduce (Qwen Team, 2026c; Yang et al., 2025; Gemma Team, 2026). We retain these results on the host and, when host capacity is tight, regenerate selected layer boundaries through cheaper layers to make room. A Local planner jointly chooses partitions, result retention, and boundary regeneration within GPU and host budgets. At 512K on four H100s, the combined design delivers 1.26× the throughput of the partitioned baseline at nearly the same host footprint (Section 7.4). C3: Workload-aware executable configuration. Memory allocation can change EP and Local’s relative performance, so we configure each eligible path before comparing step times. We jointly configure state placement and activation execution under a common step-time objective and memory budgets for each GPU and host NUMA node. Calibrated cost models guide within-path searches; where both paths are eligible, short executions select the plan with lower measured step time. Within one runtime, the best configured path is 1.1–2.0× faster than the best unconfigured one. On four H100s at 4K context and global batch 96, EP4 is 1.17× faster before configuration, yet Local is 1.21× faster afterward (Section 7.5). Across 16 workloads on four models, short executions select the path that full-model runs measure as faster at every point. We implement OOM in PyTorch and evaluate five MoE models with 8–256 experts and up to 122B parameters on H100 and RTX 4090 nodes. On four H100s, OOM trains Qwen3.635B-A3B with 1M-token contexts, the model’s YaRNextended length, whereas FSDP2, ZeRO-Offload (Feng et al., 2025; Ren et al., 2021), and MegaTrain exhaust GPU memory at 512K. Beyond the four-GPU 1M-token workload, a single H100 completes a 2M-token training step with YaRN ×8. OOM also trains Qwen3.5-122B-A10B, whereas all three baselines exhaust the node’s host memory. On four 24 GB RTX 4090s, it completes 512K-token steps, a length at which the baselines exhaust 80 GB H100s, and trains 1,024 1K-token sequences per step without accumu-

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

lation; FSDP2 and ZeRO-Offload need it at that batch, and MegaTrain runs out of GPU memory.

how to divide GPU memory among resident state, retained activations, and worksets to reduce step time.

Across matched-global-batch sweeps, OOM delivers 1.58× and 4.08× geometric-mean (up to 2.14× and 7.77×) throughput over CPU-offloaded FSDP2 on H100 at 4K and RTX 4090 at 1K context (2.92× over ZeRO-Offload on RTX 4090) while using the least host memory among the evaluated systems—as little as 1/3.0 of FSDP2’s. Across four architectures at 1K context, it delivers 1.66–2.44× FSDP2 throughput at batch 32 and 1.23–1.38× at 256. On eight H100s at 1K context, it delivers 1.16–1.52× the throughput of offloaded FSDP2 and 2.38–3.00× that of Megatron-LM.

ES-MoE pipelines expert offloading and computation with token-load balancing; MoP combines component-specific parallelism with heterogeneous updates (Kim et al., 2024; Nguyen et al., 2026). LazyTrain jointly schedules checkpoints, recomputation, and transfers (Wu et al., 2026). KTransformers-FineTune (KT-FT) takes a different route, computing routed experts on the host CPU and replacing weight offload with CPU–GPU co-computation (Li et al., 2026a). ProTrain jointly tunes residency, buffers, checkpointing, and offloading, with GPU updates for persistent state; Mist co-optimizes memory and parallelism using overlap-aware modeling (Yang et al., 2026a; Zhu et al., 2025). ProTrain’s persistent chunks remain unpartitioned and Mist’s efficient tuner assumes identical layers; neither policy space jointly configures expert-owned residency, MoE activation retention, and hybrid result retention.

2

BACKGROUND AND M OTIVATION

2.1

MoE Models and Training Memory

An MoE layer routes each token to k of E experts and combines their outputs (Shazeer et al., 2017; Jiang et al., 2024). For equal-sized routed experts, the full bank contains E/k times as many parameters as the experts selected for one token. Full-parameter training maintains state for every expert: FP32 master weights and two FP32 Adam moments alone require 12 bytes per parameter (Kingma & Ba, 2015; Ren et al., 2021), plus working weights and gradients. We call the masters and moments canonical state; their updates produce the next working weights. Activations and operator/communication worksets also consume memory, and their peak footprint depends on which objects are simultaneously live. The model’s sequencemixing operators also affect activation storage and recomputation costs: Mixtral and Qwen3 use attention throughout, Gemma4 interleaves sliding-window and global attention, and Qwen3.6 combines full attention with recurrent Gated DeltaNet (GDN) layers (Jiang et al., 2024; Qwen Team, 2025; Gemma Team, 2026; Qwen Team, 2026b). 2.2

Training Beyond GPU Memory

ZeRO and FSDP shard state across GPUs; ZeRO-Offload and ZeRO-Infinity extend capacity through CPU and hierarchical memory (Rajbhandari et al., 2020; Zhao et al., 2023; Ren et al., 2021; Rajbhandari et al., 2021). SlideFormer, SlideDP, and MegaTrain bound active parameter storage through layer streaming (Yang & Wen, 2026; Yang et al., 2026b; Yuan et al., 2026). For streamed state, CPU masters supply working weights; backward uses the same weight version as forward, and aggregated gradients return for CPU updates. Transfers and updates overlap computation on other layers. Within this streaming framework, expert execution still changes weight delivery, token exchange, activation worksets, and replay costs. Bounding active parameter storage and overlapping host service therefore leave open which expert-execution path to use and

2.3

Execution Costs in Host-Backed MoE Training

EP partitions experts across P GPUs, assigning E/P per GPU under balanced placement, and exchanges tokens and results through all-to-all (A2A) communication (Lepikhin et al., 2021). Local execution materializes the working weights of all E experts in the current layer on each GPU to process local tokens, avoiding expert-token exchange while still requiring weight delivery and gradient aggregation. Optimizer state need not follow the placement of these working weights (Skiadopoulos et al., 2026). Retaining MoE activations in EP avoids the corresponding backward replay and its A2A; the retained activations compete with resident state and execution worksets for GPU memory. At fixed sequence length, larger local batches amortize expert materialization over more tokens while increasing activation storage. Sequence length also changes the balance of operator arithmetic. For B local sequences of length S, fixed architecture and routing degree give leading terms F (B, S) ≈ aBS + bBS 2 ,

(1)

where projections, experts, and recurrent mixing contribute to a, and full attention to b. At equal token count T = BS, the full-attention term bT S still depends on S, so batch size and sequence length must be modeled separately. Within each GPU, partitioning each operator along the axes its dependencies permit (Table 1) bounds worksets: GDN sequence chunks pass recurrent state (Yang et al., 2025), full attention admits head groups over complete sequences, and expert FFNs admit token chunks after routing. Each partition must preserve backward state and gradient flow. Checkpointing trades storage for recomputation (Chen et al., 2016; Jain et al., 2020); FlashAttention avoids quadratic

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design Table 1. Leading arithmetic costs and operator dependencies at fixed widths and routing degree. B counts local sequences and S their length. Operator

Compute

Dependency

GDN mixing

O(BS)

Full attention

O(BS 2 )

Expert FFN

O(BS)

Recurrent state across sequence segments Causal token interactions within each head Token-wise evaluation after routing

score storage, and MoEBlaze reduces routing buffers with selective recomputation (Dao et al., 2022; Zhang et al., 2026). For dense LLMs, OOMB reaches million-token contexts by recomputing activations in chunks and offloading the KV cache (Li et al., 2026b). Long sequences expose an asymmetry: full attention’s computation grows quadratically with length, while its output and softmax statistics grow only linearly, so keeping them can beat recomputation. Host memory bounds retention: decoder inputs (layer boundaries) and retained results compete with model state.

3

S YSTEM OVERVIEW

OOM converts a workload and its memory budgets into a plan for a shared training runtime (Figure 2). Host memory supplies capacity for model state and saved activations; bounded parameter windows feed GPU computation. Routed experts use either disjoint EP shards with token A2A or a complete bank serving local tokens. Attention, routers, and shared experts use the common parameter stream. The backward pipeline overlaps state delivery, GPU computation, activation transfers, and host updates. Completion events preserve weight versions and release storage after its last consumer. GPU memory serves three purposes: resident state Mstate , retained activations Mkeep , and the active workset Mwork (labeled S, A, and W in the figure). EP plans state residency and activation retention (Section 4); Local plans partitioning, result retention, and boundary regeneration (Section 5). Under a common step-time objective and GPU/NUMA budgets, Configure-then-Compare uses calibrated costs to configure each path and short executions to compare them (Section 6). Our PyTorch runtime (Paszke et al., 2019) installs the plan before training.

replay, including its token routing. We budget memory by the work these choices remove from the training step. 4.1

Expert Ownership and the Shared Pipeline

Within an EP group, each GPU owns disjoint routed experts. A streamed shard keeps its FP32 master and Adam moments on the host: a worker prepares BF16 weights, and backward returns gradients for CPU update. Promoting a shard replaces host preparation, parameter H2D, gradient D2H, and CPU update with local weight generation and GPU AdamW. Each shard still has a single updater; resident and streamed experts both participate in token dispatch/return A2A. Attention, normalization, routers, and shared experts continue through the common stream. Figure 2 shows how this service overlaps computation within bounded parameter windows and gradient slots. Three dependencies preserve the execution: computation waits for its required weight version, including the original version during replay; updates wait for all required gradients; storage reuse waits for the last consumer. The GPU can therefore execute layer ℓ − 1 while the CPU updates layer ℓ. Interleaving resident layers gives host updates additional GPU work to overlap, making both the number and positions of resident layers useful planning choices. 4.2

Spend Memory on the Completion Path

Activation checkpointing supplies the complementary choice (Chen et al., 2016; Korthikanti et al., 2023). Fulldecoder checkpointing offloads the input and reconstructs internal tensors during backward. Splitting at the attention/MoE boundary allows either sublayer’s backward tensors to remain on GPU; retaining both removes decoderlevel replay. In EP, MoE retention saves expert recomputation and the token exchanges required for replay. The value of retention depends on what delays completion. Let Th be the duration of a submitted host service and Tg the independent GPU work before its consumer. If that service is the final prerequisite, the exposed wait is [Th − Tg ]+ . Residency shortens host service and thereby changes which work is exposed, so retention must be valued against the resulting critical path. Joint budgeting scores complete candidates rather than adding independent gains. For candidate z, the per-GPU peak is peak MGPU (z) = M0 + Mstate (z)

4

E XPERT-PARALLEL E XECUTION AND M EMORY B UDGETING

An EP shard is both an execution unit and an update unit. OOM uses this ownership to keep selected experts’ master weights, optimizer moments, and updates on their GPUs. The same memory can instead retain activations and remove

+ max{Mkeep (t, z) + Mwork (t, z)} + ϵ,

(2)

t

where M0 is fixed occupancy and ϵ is headroom. The workset includes working weights, gradients, routing buffers, and operator temporaries. CPU checkpoint inputs consume host capacity. State cost depends on expert parameters and EP partitioning, whereas activation and routing costs grow with

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design (a)

Configure–then–Compare Model, workload, budgets

Cost-guided EP / Local planning

Compare step time & select

Install selected plan

EP: State–Activation Budgeting

Local: Split–Keep–Replay

GPU 0

Token A2A

GPU P − 1

Owned expert shard

Owner-GPU updates

Owned expert shard

Per-GPU HBM

W

S weights

Host DRAM

A grads

weights

save/load

GPU

··· ℓ + 1

State I/O

ℓ: O+E

Activation pipeline

H2D

save/load

Saved boundaries Attention results

A Retained activations time

Replay + A2A D2H

ℓ+1

O: ordinary E: experts Prefetch ℓ

A

grads

Model state CPU updates

Host DRAM

(b) EP backward pipeline

CPU update

W

W Worksets

S Resident state

Full-layer expert weights Per-GPU HBM

Saved layer inputs

Streamed state CPU updates

Each GPU

ℓ · streamed experts A2A

Expert backward

A2A

Other backward

H2D O, ℓ − 1

D2H E, ℓ

Expert backward

D2H O, ℓ

A2A

Update

Other backward

H2D

ℓ − 2 ···

D2H O, ℓ − 1

ℓ−2

O, ℓ − 1

ℓ: O+E

ℓ+1

Use ℓ → release

ℓ − 1 · resident S + retained A A2A

Prefetch ℓ−2

Use Aℓ−1

Prefetch ℓ−3

Figure 2. Overview of OOM. (a) Cost-guided planning and EP/Local resource placement. EP and Local show alternative placements on 1 the same node. (b) EP backward pipeline; O/E denote ordinary/expert components. Widths are schematic.

token load. Smaller routing chunks free memory for residency or retention but require more communication rounds. Consequently, batch size and sequence length determine how much memory each purpose deserves. 4.3

Produce an Executable EP Plan

For a fixed EP layout, the planner selects resident layers R, checkpoint policies, and routing chunk c. It explores contiguous and interleaved resident sets; the calibrated checkpoint family retains MoE backward tensors in a set K and fully checkpoints other layers. Reference occupancy and incremental costs filter candidates using Equation 2. The time model combines calibrated forward work, a dependency-simulated backward region, and the final update drain. For a stage v with service time dv and predecessor set Pv , completion follows fv = dv + max fu . u∈Pv

(3)

Edges encode data readiness, stream order, and buffer reuse. Thus replay, routing, transfers, and updates overlap only where execution permits; waiting emerges from dependencies instead of being counted again as service. The fastest predicted feasible candidate determines the plan. As activation demand grows with batch, the configured EP4 plan shifts memory from state to activations: 23 resident layers at global batch 32, three resident and four MoEretained layers at 64, and one of each at 128 (Figure 1c).

5

L OCAL E XECUTION WITH S PLIT–K EEP –R EPLAY

Local execution avoids token exchange and amortizes a full expert bank over local tokens. Its advantage grows with token load, making activation capacity central. OOM manages hybrid activations with Split–Keep–Replay: split each operator along axes its dependencies permit, keep expensive results, and regenerate selected boundaries by replaying cheaper decoders. A Local planner coordinates these choices within GPU and host budgets. 5.1

Deliver the Expert Bank

As in layer-streaming systems (Rajbhandari et al., 2021; Yuan et al., 2026; Yang & Wen, 2026; Yang et al., 2026b), Local keeps canonical state on the host and materializes each layer’s full expert bank in a double-buffered parameter window. Full delivery copies a shared BF16 bank to each rank; Shard delivers disjoint slices and reconstructs the bank through GPU AllGather, cutting each rank’s parameter H2D to about 1/P for P GPUs. Gradients are reduced on the GPUs or in host shared memory, accumulating BF16 partials in FP32 before the CPU update. 5.2

Partition along Dependency-Compatible Axes

Large batches admit independent sub-batches of complete sequences. Long sequences require different partitions for recurrent mixing, attention, and expert computation (Fig-

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design (a)

Bsub , projection and gate workspace scales with Bsub Sgd. Head grouping complements FlashAttention’s score-matrix tiling (Dao, 2024) by bounding projection workspace.

Partition by dependency Seq 0

Batch

Seq 1

Seq 2

Independent complete sequences

GDN

Chunk 0

Attention

q group 0

MoE

Tokens 0

Chunk 1

Chunk 2

state → ← state gradient

q group 1

MoE: token chunks. Normalization, routing, routed experts, and the shared branch process token slices; backward reconstructs and releases one slice at a time. Chunking bounds temporaries, while input, output, and gradient buffers still scale with BS.

q group 2

Shared K/V · full-sequence attention

Tokens 1

Tokens 2

One token chunk at a time

(b)

Retain / regenerate

backward: regenerate x1 –x3 from x0

GDN + MoE

GDN + MoE

GDN + MoE

x1 Stored on host

x2 x3 Regenerated on GPU

(c)

Compute and host-memory tradeoff 10 1

GDN mix MoE fwd 30×

0.1 0.01

0.105–0.117 s

32K

128K 256K 512K Sequence length

Host payload (GiB)

Time (s, log)

Measured on H100 6

5.3

reuse

Full attn + MoE

x0

Attn fwd Reuse O

Attn results on host: O, LSE save

x4

Payload per 3+1 block 6.0

Freed boundaries Retained O/LSE 4.03

4 2 0

32K

128K 256K 512K Sequence length

Figure 3. Hybrid activation execution: (a) dependency-aware parti1 tioning; (b) retention and regeneration; (c) Qwen3.6 operator times on H100 and calculated host payloads per sequence.

ure 3(a)). Each partition bounds temporary storage, preserves forward/backward dependencies, and accumulates parameter-gradient partials in FP32. GDN: sequence chunks with state handoff. Chunks pass an FP32 recurrent state and causal-convolution left context (Yang et al., 2025; Yang & Zhang, 2024). Backward rematerializes chunks in reverse, passing the final-state derivative into each chunk and propagating its initial-state derivative to the preceding one. Convolution-context gradients join the preceding slice. For chunk length C, Hv value heads, and dk × dv state matrices, one GDN layer’s chunk-boundary states occupy MGDN states = 4B(⌈S/C⌉ − 1)Hv dk dv

bytes.

(4)

Smaller chunks reduce token workspace but increase state storage and calls. Complete-sequence sub-batches provide an additional control for large B. Full attention: head groups over complete sequences. OOM splits the query heads of each GQA group, which share one K/V head, into groups of g. Each group retains the full causal sequence, materializes its projections and intermediates, and releases them after use. Backward rebuilds groups and accumulates shared K/V and input gradients in FP32. For head dimension d and sub-batch size

Retain Expensive Results, Regenerate Boundaries

Hybrid models make recomputation costs sharply asymmetric: attention arithmetic grows quadratically with sequence length, while its output grows only linearly. At 512K, attention forward takes 3.180 s, versus 0.117 s for GDN mixing and 0.105 s for MoE, and costs 30× its replay with output reuse, up from 2.6× at 32K (Figure 3(c), left). OOM therefore saves attention output O and softmax log-sum-exp statistics (LSE) on the host. Checkpoint replay reuses O, reducing attention-sublayer replay to 0.105 s; backward consumes O/LSE while rebuilding Q/K/V. Reuse across both phases removes redundant attention forwards from nested rematerialization inside layer checkpointing (Chen et al., 2016; Jain et al., 2020; Korthikanti et al., 2023). Selective recomputation discards large yet cheapto-rebuild attention activations at moderate lengths (Korthikanti et al., 2023). Long context reverses that trade: OOM retains attention’s linear-size results and rebuilds only cheap projections during replay. Decoder inputs compete for the same host capacity: L BF16 boundaries of hidden width h occupy 2LBSh bytes. OOM keeps selected boundaries and regenerates the others with no-gradient forwards from the nearest kept input. In Figure 3(b), keeping x0 and x4 around three GDN decoders and a full-attention layer recovers x1 –x3 through the three complete GDN decoders, including normalization and MoE, without repeating attention. At 512K, these three boundaries free 6.0 GiB of host memory per sequence, more than the 4.03 GiB of O/LSE (Figure 3(c), right). Sliding-window hybrids show the same asymmetry: Gemma4-26B-A4B interleaves five 1K-window layers with each global-attention layer; at 256K, replay with O reused cuts a global-attention forward from 17.2 to 0.04 s, and one block’s five boundaries free 6.9 GiB per sequence against 4.0 GiB of O/LSE. Regenerated boundaries stay on GPU for reverse-order consumption; a cap on their number bounds the footprint, with spare GPU memory permitting longer runs. The planner weighs saved attention replay and freed host capacity against whole-decoder regeneration, temporary GPU storage, and changes in exposed transfer time. Ample host capacity allows both boundaries and results to remain stored; tighter budgets favor regenerating cheap boundaries to retain expen-

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

5.4

Coordinate Lifetimes and Plan the Workset

Operator workspaces coexist with recurrent states, regenerated boundaries, and prefetched inputs. OOM schedules them around consumption. The original GDN forward rolls its state; backward boundary states are created during the current layer’s replay and released after use. Activationprefetch buffers are allocated on demand. Grouped attention defers input prefetch until query-head workspace is released, then overlaps transfer with remaining backward computation. Regeneration uses original-version weights from the shared parameter window. Together, these choices turn local partitioning (Zhao et al., 2026; Pandit et al., 2026) into control over the step’s simultaneous peak. The Local planner searches these coupled choices jointly: partitions set each phase’s workspace, which competes with held boundaries for GPU memory, while O/LSE competes with kept boundaries for host memory. It enumerates attention head groups, GDN sequence/sub-batch chunks, and MoE token chunks. For each combination, a layerwise dynamic program chooses boundary policies and attentionresult placement. Retaining a boundary charges its storage; dropping it extends a regeneration run and charges the preceding decoder’s forward. Retaining attention results replaces replay work with storage and transfer costs. DP states preserve open regeneration runs, GPU objects carried across layers, and pending transfer dependencies. Labels record accumulated time, aggregate host storage, and the running GPU peak. Within the same DP state, a label is discarded if another is no slower and uses no more of any resource. GPU peak is the largest live footprint across execution phases, including held boundaries and the output head, rather than the sum of separate workspace maxima. Surviving plans are checked against each NUMA node’s host budget, including model state and pinned pools, and the fastest predicted feasible plan is selected. The resulting plan specifies partitions, retained boundaries, result placement, and regeneration limits; FlashAttention, FLA, and Liger adapters execute it (Dao, 2024; Yang & Zhang, 2024; Hsu et al., 2024).

6

C ONFIGURE - THEN -C OMPARE

EP and Local convert memory into throughput differently, so OOM applies Configure-then-Compare: it first allocates resources within each path, then selects among them. Path choice thus reflects the execution each memory budget enables. Configuration raises each path’s throughput and can change which path is faster: the upper envelope in Figure 4 marks the faster configured path at each batch size.

60

Throughput (k tokens/s)

sive attention results. The regeneration interval is a planned quantity set by sequence length and both budgets.

EP4 Configured

Local Base

45

30

15

0 16

32

64

96

128

256

Global batch (4×H100, S = 4K)

Figure 4. EP4/Local throughput before and after configuration. Gray traces the upper measured Configured curve.

6.1

Model the Complete Configuration

The input supplies model structure, local batch B, sequence length S, precision, GPU topology, and resource budgets. A configuration z specifies expert execution and state delivery, resident layers, checkpoint and boundary policies, operator partitions, and result placement. For the supported candidates C(w) at workload w, the objective is z ∗ = arg min Tstep (z), z∈C(w)

s.t.

peak MGPU,i (z) ≤ Gi , peak Mhost,n (z) ≤ Hn ,

(5)

where Gi and Hn are the memory budgets of GPU i and NUMA node n. Time includes forward, replay, backward, gradient processing, optimizer updates, and weight publication. Memory follows simultaneous allocations, including transfer staging. This step-level accounting therefore values residency, retention, and partitioning consistently. Calibrate service and storage costs. Short reference runs on a few repeating decoder units characterize each path’s compute, replay, communication, and host services, as well as reference occupancy. Matching component profiles can be reused. EP composes these costs through execution dependencies (Section 4.3); Local sums operator costs, regeneration, and exposed transfers, and checks memory against its phase-wise live set (Section 5.4). Keeping B and S separate captures both token-scaled work and the sequencedependent cost of attention. Search within each path. State-delivery mode, transfer chunk sizes, and buffer settings are supplied runtime settings; their service and storage costs enter calibration. Within these settings, EP planning chooses residency, checkpointing, and routing chunks by its dependency model or, where none is calibrated, from measured per-layer costs;

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design ZeRO-Offload

MegaTrain

KT-FT

OOM

Split–Keep–Replay

20

5

10

0

0

(c) 4×H100 1 seq/GPU

0.72

0.75

0.83

20

0.80

10

0.78

Out of GPU memory

(b) 4×RTX 4090, S = 1K

40

0 1K

GAS > 1

0.71

(a) 4×H100, S = 4K

0.66

Host (GiB) Throughput (k tokens/s)

FSDP2

0 16

32

64

128

256

Global batch (sequences)

64

128

256

512

1024

64K

Global batch (sequences)

128K

256K

512K

1M

Sequence length (tokens)

Figure 5. Batch-size and context-length scaling. Labels give speedup over the fastest baseline; labels on KT-FT bars give k tokens/s.

Local enumerates operator partitions and searches boundary/result placement. Each planner filters capacity violations, ranks the remaining candidates by predicted step time Tbstep , and emits the settings required to execute its preferred plan on the shared runtime. 6.2

Compare Configured Paths

For N = BS tokens per GPU, top-k routing, width h, and b-byte elements, balanced EP routing transmits VEP = 2kN hb

P −1 P

Experimental Setup

We train the 40-layer Qwen3.6-35B-A3B language backbone (256 top-8 experts, 30 GDN and 10 full-attention layers) with AdamW on packed OpenThoughts3 inputs (Guha et al., 2026), updating all language-model parameters. We also evaluate Mixtral-8x7B, Qwen3-30B-A3B, Gemma426B-A4B, and Qwen3.5-122B-A10B. We write Rn and Kn for n resident expert layers and n MoE-retained layers.

(6)

bytes per GPU for dispatch and return in one MoE pass. Local’s expert-bank delivery is independent of N , so more tokens amortize its materialization cost. Calibrated models price these transfers with overlap, replay, and live worksets. Long-context workloads use Local’s Split–Keep–Replay planner, built for their activation worksets. Where both paths are candidates, each planner proposes its best feasible plan, and short executions at the same workload and budgets select the faster one: prediction narrows the space, and execution resolves the cross-path comparison. Because decoder layers repeat, each path first runs one and two repeating units, which separates its fixed cost from its per-layer cost and reads its GPU occupancy; EP then runs its plan’s resident and retained shares once more. Scaling these tests to the full depth predicts each configured path’s step time at the target workload. The runtime then installs the chosen path, its parameter windows, retained objects, and completion dependencies before training begins.

7

7.1

E VALUATION

We measure throughput, host memory, and capacity, then isolate the mechanisms and configuration choices behind them and test training quality and model coverage.

Baselines and resources. Our baselines span sharded offload (FSDP2 and DeepSpeed ZeRO-3 with CPU parameter/optimizer offload, ZeRO-Offload), layer streaming (MegaTrain), CPU co-computation (KT-FT), and model parallelism with CPU optimizer offload (Megatron-LM) (Feng et al., 2025; Ren et al., 2021; Yuan et al., 2026; Li et al., 2026a; Shoeybi et al., 2019; NVIDIA, 2026). FSDP2 and OOM share FA3 attention (Shah et al., 2024) and Liger MoE kernels in the main four-H100 sweeps, retaining their own memory policies. All use BF16 compute and FP32 master weights and moments. FSDP2 and Megatron reduce gradients in FP32; ZeRO-Offload communicates BF16 gradients. KT-FT runs its native AMX BF16 CPU experts on the H100 host; we add the standard FP32 masters and moments of the other systems in place of its native BF16 optimizer state. We found no public implementation of ProTrain, and released Mist and ES-MoE code lacks native support for the recent MoE models evaluated here. The H100 server has eight 80 GB GPUs on NVLink and two Xeon 8468 sockets with 1,008 GiB each; the RTX 4090 server has four 24 GB GPUs on PCIe and two Xeon 6338N sockets with 503 GiB each. Each comparison runs all systems on the same GPUs and host. OOM uses no gradient accumulation. Baselines accumulate when needed to match the global batch, keeping long sequences intact.

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

Measurement. Main sweeps use two independent processes, each with two warmup and five timed training steps. We average their medians of slowest-rank step times; whiskers span the two medians. Split–Keep–Replay studies use one process with one warmup and three measured steps. Step time covers computation, replay, transfers, reductions, optimizer updates, and weight publication. Throughput counts input tokens once; sweep summaries use geometric means of paired ratios. Resource runs sample process-tree PSS and peak GPU allocation; 512K/1M and Split–Keep– Replay host footprints sum rank-level PSS estimates. 7.2

Throughput Across Batch Sizes

On four H100s at 4K context, OOM delivers 1.29–2.14× FSDP2 throughput across global batches 16–256 (1.58× geometric mean; Figure 5a), 1.45–3.16× that of ZeROOffload, and 2.46–6.63× that of MegaTrain, reaching 47.0K tokens/s. EP4 serves batches 16–64 (R27, R23, and R3/K4) with smaller banks and resident state; Local serves 128–256, amortizing expert banks without token A2A. At batch 32, OOM runs 4.8× faster than Megatron-LM, which needs eight accumulation rounds. KT-FT reaches 659–803 tokens/s across batches 16–128, so OOM is 26–60× faster; at batch 256 KT-FT needs two accumulation rounds (751 tokens/s, 63×). Routed experts are dense matrix computation: even on a 48-core AMX Xeon, KT-FT’s CPU expert path takes 70% of its exposed step time at batch 32, against 4% for the FP32 optimizer update. On 24 GB RTX 4090s, the gains grow to 2.16–7.77× over FSDP2 and 1.91–4.50× over ZeRO-Offload across batches 64–1024, with geometric means of 4.08× and 2.92× (Figure 5b). At batch 1024, each GPU processes 256 complete 1K sequences without accumulation (11.6K tokens/s), whereas both baselines need eight accumulation rounds. Split–Keep–Replay’s partitioning enables this amortization: default operators fit only 64 sequences per GPU, and the Local planner shifts GDN/MoE chunks from 32K/16K tokens at 128 sequences per GPU to 16K/32K at 256, where keeping 32K/16K exhausts GPU memory.

Memory as a performance budget. Residency spends GPU capacity to reduce host service and storage; Local leaves more GPU space for the workset. Across H100 batches, OOM uses 197–571 GiB of host memory versus FSDP2’s 591–751 GiB and KT-FT’s 619–876 GiB, with resident experts reducing batches 16/32 to 197/230 GiB (Figure 5a). On RTX 4090, OOM uses 0.70–0.93× FSDP2’s host memory and 0.48–0.66× ZeRO-Offload’s; at 64K– 256K on H100, it uses 0.76–0.87× FSDP2’s (Figure 5b,c).

7.3

Long-Sequence Efficiency and Capacity

With one sequence per GPU, OOM trains at 22.4K, 22.2K, and 16.7K tokens/s for 64K, 128K, and 256K contexts, 1.19– 1.26× the throughput of FSDP2 (Figure 5c), 1.22–1.58× that of ZeRO-Offload, and 1.79–2.57× that of MegaTrain. KT-FT reaches 834 and 714 tokens/s at 64K and 128K, 27–31× slower, and exhausts GPU memory at 256K. Our default operators suffice up to 128K; from 256K, Split– Keep–Replay is faster and saves GPU memory. Split–Keep–Replay extends this range to 512K and 1M: with grouped attention, operator chunks, O/LSE retention, and boundary regeneration, OOM completes a step in 205.7 and 712.6 s (10.2K and 5.9K tokens/s at 733 and 894 GiB of host memory; YaRN (Peng et al., 2024) ×2 and ×4), whereas FSDP2, ZeRO-Offload, and MegaTrain exhaust GPU memory at both lengths. One H100 also completes 2M-token steps at 805 tokens/s (YaRN ×8). 7.4

State and Activation Mechanisms

Resident experts remove host service. At batch 32 and 4K context, the planner’s R23 plan with a 64K routing chunk cuts step time from 8.48 to 5.02 s (1.69×; predicted 5.14 s) without manual tuning, raising GPU allocation from 11.4 to 64.8 GiB for resident state. Profiled steps (8.5 and 5.7 s; Figure 1b) show why: computation with replay stays at 3.93 s, while non-compute time falls by 60%, from 4.55 to 1.80 s, as expert-owned updates remove host preparation, transfers, and CPU optimization. Placement and capacity shape the gain. Position also shortens the critical path. With the same 12 resident layers at batch 32 and 4K, interleaving them cuts step time from 7.53 to 6.69 s (1.13×; predicted 1.14×) at the same GPU and host memory. The two plans do the same GPU and CPU work and differ in exposed wait (Section 4.2): consecutive resident layers run after the streamed layers’ updates have already stalled the GPU, whereas interleaved ones give those updates GPU work to overlap. At batch 32 and 1K, step time falls from 8.11 s without resident layers to 3.54 s at the planned R27 (2.29×), and the gain lasts until capacity runs out: on Gemma4 at batch 128, the planned R18 is within 0.3% of the fastest point tested, whereas R22 leaves no GPU headroom and runs 1.35× slower. Retention’s value depends on residency. Retaining one decoder layer’s MoE activations removes its MoE replay and token exchange, shortening the step by 5.2% with 17 resident layers but by under 1% with 21. Residency changes the exposed work and retention’s benefit, so the EP planner budgets the two jointly (Section 4.3). At about 64 GiB, its R23/K0 plan (5.02 s) beats R21/K1 (5.33 s).

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design Step time (s)

Host memory (GiB)

Default Out of GPU memory 260

+ Split + Keep O/LSE

200

+ Replay

206 0

100

GPU (GiB) —

Out of GPU memory

200

731

43 893

733 300

0

250

500

750

43 45

1000

(a) Cumulative mechanisms at 512K, four H100s GPU

Context, 1 seq/GPU

Batch per GPU

Default

Default

S–K–R

S–K–R

H100 256K 1M 64 256 RTX 4090 64K 512K 64 256 (b) Largest completed workload, four GPUs Figure 6. Split–Keep–Replay (S–K–R) gains and capacity. Batch contexts: 4K on H100; 1K on RTX 4090.

A: all boundaries

B: + O/LSE

3+1 + O/LSE

(b) 512K Step time (s)

Speedup over A

(a) Length 1.2 1.0 0.8

1+1 + O/LSE

A

260

7.5

64K 128K 256K 512K

Sequence length

Workload-Aware Configuration

230 3+1

1+1

200 B

32K

The best policy depends on length and budget. Policy A keeps every decoder boundary on the host and recomputes attention; B also keeps O/LSE; 1+1 and 3+1 add regeneration, keeping one boundary in every two and four layers. Within partitioned execution, retaining O/LSE adds 9% to step time at 32K but delivers 1.14× throughput at 128K and 1.30× at 512K (Figure 7a), as attention computation outgrows its retained results; from 128K, relative regeneration overhead likewise shrinks with length. At 512K, B is fastest at 893 GiB; 1+1 saves 160 GiB for under 3% more time, and 3+1 saves 240 GiB for 3.8% (Figure 7b). Compared with A, 3+1 is both 1.25× faster and 79 GiB smaller: inexpensive boundaries make room for valuable attention results. Predictions frozen before execution rank these policies in measured order, within 0.6% of step time and 0.2 GiB of host memory, so at every host budget the predicted-fastest feasible policy is the measured-fastest one. Under a 750 GiB budget, the planner holds the last 17 boundaries on GPU and runs in 202.1 s, faster than 1+1 (205.5 s), the fastest fixed policy that fits this budget.

600

700

800

900

Host memory (GiB)

Figure 7. Split–Keep–Replay policies on four H100s. Dotted line: fastest measured policy within each host-memory budget.

Partition, retain, and regenerate. Figure 6 shows how Split–Keep–Replay turns GPU and host capacity into throughput. At 512K with one sequence per GPU, the default path exhausts GPU memory in its first MoE backward. Dependency-compatible partitioning and lifetime coordination make the step fit in 43.4 GiB at 260.0 s. Retaining attention O/LSE on the host avoids recomputing the attention core and cuts step time to 200.5 s, a 1.30× gain, at the cost of 161 GiB more host memory. Regenerating half of the decoder boundaries recovers 160 GiB of it for 2.8% more step time and 2 GiB of GPU memory. The complete design thus runs 1.26× faster than the partitioned baseline at nearly the same host footprint (733 versus 731 GiB). At 256K, where the default path fits, partitioning with O/LSE retention still runs 1.06× faster (62.8 versus 66.9 s) and lowers peak GPU allocation by 27%, from 53.8 to 39.4 GiB. Split–Keep–Replay thus improves feasible workloads as well as extending capacity: on both platforms it raises the largest completed context by 4–8× and the per-GPU batch by 4× (Figure 6b). Default execution exhausts GPU memory at the next doubling.

Configured EP4 shifts memory from state to activations as the batch grows (Section 4.3), yet at batch 128 Local is 1.54× faster (12.18 versus 18.71 s), a gain fixed EP forgoes. Across batches 16–256, the best configured path is 1.1–2.0× faster than the best Base path, the same runtime without residency, retention, or tuned delivery (Figure 4). Configuration changes the chosen path. At global batch 96, configuration reverses the ranking: EP4 is 1.17× faster before configuration (14.9 versus 17.4 s), but Local is 1.21× faster afterward (11.4 versus 13.8 s; Figure 4). Selecting before configuration would pick EP4 and run 1.21× slower. Short-test path selection. Because decoder layers repeat, tests of one and two repeating units price each path and read its GPU occupancy, and one more test executes the plan’s own resident mix, all within 260 GiB of host memory. Across 16 workloads on four models at 1K and 4K context, the selected path matches full-model runs every time. On six of them, with decisions recorded before full-model execution (Table 10), predicted EP step time is within 8.3% (5.6% on average), and selection follows the crossover: at global batch 192, Local for Gemma4 (predicted 1.16×, measured 1.29×) and EP4 for Qwen3 (1.17×, 1.08×). Searching residency over these measured costs also yields the EP plans for models beyond the calibrated profile (Figure 8a). Candidate ranking and independent execution. At global batch 384 (96 per GPU) and 1K context on RTX 4090, two Local candidates with 16K or 32K GDN/MoE chunks fit a 23 GiB budget. Predictions recorded before mea-

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

Speedup over FSDP2

OOM (EP)

OOM (Local)

FSDP2

ZeRO-Offload

(a) Models, 4×H100 2.34

2.19

2

1.66

1.38

1.23

1.29

1.26

1 0

32

256

Mixtral

32

256

Gemma4

32

256

15

FSDP2, ZeRO, MegaTrain: host memory exhausted 10.0

Qwen3

(b) 8×H100 k tokens/s

Megatron-LM 2.44

50

10 5 128

Global batch

256

Qwen3.6

(c) 122B, 4K

1.16× 1.52×

25 0

32

256

0

10.9

5.7

64

128

Global batch

256

Figure 8. Model and node scaling. (a,b) 1K context; (c) Qwen3.5122B-A10B. (b,c) Thousands of tokens/s; labels in (b): speedup over the fastest baseline.

surement favor 32K (38.66 versus 39.85 s), as do measured times (40.34 versus 42.13 s). The selected plan executes in independent processes at 40.38 s within 21.34 GiB. Prediction and planning cost. Predictions made before execution are within 5.5% of measured step time over Local’s calibrated 128K–512K range (median 0.1%) and 0.3 GiB of host memory; the batch-32 and batch-96 EP4 plans are within 2.3%. Given each batch’s GPU headroom, a complete EP search takes under 9 CPU seconds and ranks the executed plan first or within 0.5% of its fastest prediction across global batches 16–256; Local searches its full partition menu offline on one core in under an hour. 7.6

Training Quality

OOM changes where state lives and when operators run, not the training computation: it uses no quantization, pruning, or approximate attention, and checkpointing recomputes the original operators. We train OOM and FSDP2 independently from the same checkpoint and data order. EP4 with 21 resident layers and Local at 1K context (500 steps, 16.4M tokens), and partitioned Local at 32K (128 steps, 16.8M tokens), stay within 0.11% of FSDP2’s held-out NLL at every checkpoint; final NLL differs by at most 0.08%. 7.7

Model Coverage and Scale

At 1K context, across four architectures with 8–256 experts, different routing degrees, and full or hybrid attention, OOM delivers 1.66–2.44× FSDP2 throughput at batch 32 and 1.23–1.38× at batch 256 on four H100s (Figure 8a). All four use EP4 at batch 32; at 256, Gemma4 and Qwen3 use Local, while Mixtral and Qwen3.6 retain EP4 at this lighter token load. Configuration thus adapts to model structure as well as workload size. OOM also trains Qwen3.5-122B-

A10B (Qwen Team, 2026a) on four H100s at 4K context, reaching 10.0K tokens/s at global batch 128 and 10.9K at 256 (Figure 8c); at 256, regenerating every other decoder boundary keeps host memory at 1,612 GiB. Offloaded FSDP2, ZeRO-Offload, and MegaTrain exhaust the node’s host memory already at batch 64, before completing a step. With the same FP32 masters and moments, KT-FT exhausts host memory even at global batch 32 and 1K context, with rank 0 alone reaching 1.9 TiB. Residency also pays on 122B: at 1K context and global batch 16, six resident expert layers make EP4 1.34× faster and use 162 GiB less host memory. On eight H100s at 1K context, OOM with EP8 reaches 40.5K and 47.2K tokens/s at global batch 128 and 256: 1.16– 1.52× the throughput of offloaded FSDP2, 1.57–2.31× that of ZeRO-Offload, and 2.38–3.00× that of Megatron-LM (Figure 8b). EP8 serves disjoint expert shards, so per-GPU expert state shrinks with node size.

8

C ONCLUSION

When host memory supplies model-state capacity, workload and expert execution determine the value of retaining state and activations, while memory allocation can change which path is faster. Through expert ownership, operator dependencies, and workload-aware configuration, OOM turns GPU and host capacity into throughput. On four H100s, it trains a 122B-parameter model, and a 35B model at 1Mtoken contexts; on four RTX 4090s, it processes 1,024sequence batches without accumulation. Across the main batch sweeps, it delivers 1.58× and 4.08× geometric-mean speedups over offloaded FSDP2 on H100 and RTX 4090 with the least host memory among the evaluated systems. Short executions select the faster configured path across 16 workloads on four models. These results make efficient full-parameter MoE training practical on a single node.

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

R EFERENCES Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost, 2016. URL https://arxiv.org/abs/1604.06174. Dao, T. FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations, pp. 35549–35562, 2024. URL https://proceedings.iclr.cc/pa per files/paper/2024/hash/98ed250b203d1ac6b24bbcf 263e3d4a7-Abstract-Conference.html. Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, volume 35, pp. 16344–16359. Curran Associates, Inc., 2022. doi: 10.52202/068431-1189. Fedus, W., Zoph, B., and Shazeer, N. Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. URL https://jmlr.org/papers/v23/21-0998.html. Feng, W., Constable, W., and Mao, Y. Getting started with fully sharded data parallel (FSDP2). PyTorch Tutorials, 2025. URL https://docs.pytorch.org/tutorials/intermediat e/FSDP tutorial.html. Gemma Team. Gemma 4 technical report, 2026. URL https://arxiv.org/abs/2607.02770. Guha, E., Marten, R., Keh, S., Raoof, N., Smyrnis, G., Bansal, H., Nezhurina, M., Mercat, J., Vu, T., Sprague, Z., Suvarna, A., Feuer, B., Chen, L. L., Khan, Z., Frankel, E., Grover, S., Choi, C., Muennighoff, N., Su, S., Zhao, W., Yang, J., Pimpalgaonkar, S., Sharma, K., Ji, C., Deng, Y., Pratt, S., Ramanujan, V., Saad-Falcon, J., Acharya, S., Li, J., Dave, A., Albalak, A., Arora, K., Wulfe, B., Hegde, C., Durrett, G., Oh, S., Bansal, M., Gabriel, S., Grover, A., Chang, K.-W., Shankar, V., Gokaslan, A., Merrill, M., Hashimoto, T., Choi, Y., Jitsev, J., Heckel, R., Sathiamoorthy, M., Dimakis, A., and Schmidt, L. OpenThoughts: Data recipes for reasoning models. In International Conference on Learning Representations, pp. 108059–108130, 2026. URL https: //proceedings.iclr.cc/paper files/paper/2026/file/b01024 1b9f1cdfc7d4c392db899cef86-Paper-Conference.pdf. Hsu, P.-L., Dai, Y., Kothapalli, V., Song, Q., Tang, S., Zhu, S., Shimizu, S., Sahni, S., Ning, H., and Chen, Y. Liger Kernel: Efficient Triton kernels for LLM training. arXiv preprint arXiv:2410.10989, 2024. doi: 10.48550/arXiv.2410.10989. Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Gonzalez, J., Keutzer, K., and Stoica, I. Checkmate: Breaking the memory wall with optimal tensor rematerialization. In Dhillon, I., Papailiopoulos, D., and

Sze, V. (eds.), Proceedings of Machine Learning and Systems, volume 2, pp. 497–511, 2020. URL https://proceedings.mlsys.org/paper files/paper/2020/fil e/0b816ae8f06f8dd3543dc3d9ef196cab-Paper.pdf. Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.-A., Stock, P., Subramanian, S., Yang, S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts, 2024. URL https://arxiv.org/abs/2401.04088. Kim, Y., Lim, H., and Han, D. Scaling beyond the GPU memory limit for large mixture-of-experts model training. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 24342–24353. PMLR, 2024. URL https://proceedings.mlr.press/v235/kim24w.html. Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL https://arxiv.org/abs/1412.6980. Korthikanti, V. A., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., and Catanzaro, B. Reducing activation recomputation in large transformer models. In Song, D., Carbin, M., and Chen, T. (eds.), Proceedings of Machine Learning and Systems, volume 5, pp. 341–353. Curan, 2023. URL https://proceedings.mlsys.org/paper f iles/paper/2023/file/80083951326cf5b35e5100260d64 ed81-Paper-mlsys2023.pdf. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. GShard: Scaling giant models with conditional computation and automatic sharding. In 9th International Conference on Learning Representations (ICLR 2021). OpenReview.net, 2021. URL https://openreview.net/forum?id=qrwe7XHTmYb. Li, P., Hao, X., Chen, H., Xie, W., Zheng, Y., Wu, B., Yang, Y., Shen, H., Ou, Q., Zhang, B., et al. KTransformers-FineTune: Beyond offload with layout-aware and schedule-optimized heterogeneous MoE fine-tuning. Author preprint, version paper-ktft-v1.1, 2026a. URL https://github.com/kvcache -ai/ktransformers/tree/paper-ktft-v1.1/papers/kt-ft. Li, W., Yu, D., et al. Out of the memory barrier: A highly memory efficient training system for LLMs with million-token contexts. arXiv:2602.02108, 2026b. URL https://arxiv.org/abs/2602.02108. Nguyen, X.-P., Pandit, S., Zhao, Y., Yavuz, S., Savarese, S., and Joty, S. Mixture-of-parallelisms: Towards

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

memory-efficient training stack for mixture-of-experts models, 2026. URL https://arxiv.org/abs/2607.01844. NVIDIA. Megatron Bridge, 2026. URL https://github.com/NVIDIA-NeMo/Megatron-Bridge. Pandit, S., Nguyen, X.-P., Zhao, Y., and Joty, S. Flattening every memory peak in long-context mixture-of-experts training, 2026. URL https://arxiv.org/abs/2609.14306v1. arXiv:2609.14306v1. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32, 2019. Peng, B., Quesnelle, J., Fan, H., and Shippole, E. YaRN: Efficient context window extension of large language models. In International Conference on Learning Representations, pp. 31932–31951, 2024. URL https: //proceedings.iclr.cc/paper files/paper/2024/file/874a4d 89f2d04b4bcf9a2c19545cf040-Paper-Conference.pdf. Qwen Team. Qwen3-30B-A3B: Model configuration, 2025. URL https://huggingface.co/Qwen/Qwen3-30B-A3B/r aw/main/config.json. Qwen Team. Qwen3.5-122B-A10B: Model configuration, 2026a. URL https://huggingface.co/Qwen/Qwen3.5-122 B-A10B/raw/main/config.json. Qwen Team. Qwen3.6-35B-A3B: Model card, 2026b. URL https://huggingface.co/Qwen/Qwen3.6-35B-A3B. Qwen Team. Qwen3.6-35B-A3B: Model configuration, 2026c. URL https://huggingface.co/Qwen/Qwen3.6-35B -A3B/raw/main/config.json. Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. ZeRO: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–16. IEEE, 2020. doi: 10.1109/SC41405.2020.00024. Rajbhandari, S., Ruwase, O., Rasley, J., Smith, S., and He, Y. ZeRO-Infinity: Breaking the GPU memory wall for extreme scale deep learning. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’21, pp. 1–14. ACM, 2021. doi: 10.1145/3458817.3476205. Ren, J., Rajbhandari, S., Aminabadi, R. Y., Ruwase, O., Yang, S., Zhang, M., Li, D., and He, Y. ZeRO-Offload: Democratizing billion-scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551–564. USENIX Association, 2021. URL https://www.usenix.org/conference/atc21/presentation/r en-jie.

Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. In Advances in Neural Information Processing Systems, volume 37, pp. 68658–68685, 2024. doi: 10.52202/079017-2193. URL https://proceedings.neurip s.cc/paper files/paper/2024/hash/7ede97c3e082c6df10a 8d6103a2eebd2-Abstract-Conference.html. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. V., Hinton, G. E., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations (ICLR 2017). OpenReview.net, 2017. URL https://openreview.net/forum?id=B1ckMDqlg. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LM: Training multi-billion parameter language models using model parallelism, 2019. URL https://arxiv.org/abs/1909.08053. Skiadopoulos, A., Zhao, M., Gandhi, S., Norrie, T., Mukherjee, S., and Kozyrakis, C. SYMI: Efficient mixture-of-experts training via model and optimizer state decoupling. In 23rd USENIX Symposium on Networked Systems Design and Implementation (NSDI 26), pp. 75–92, Renton, WA, 2026. USENIX Association. URL https://www.usenix.org/conference/nsdi26/presentation/ skiadopoulos. Wu, X., Yang, C., Liu, H., Lin, X., Jiang, X., Xu, C., Li, J., and Guo, J. LazyTrain: Limited-resource allocation toward zero-waste yield optimization in large language model training, 2026. URL https://arxiv.org/abs/2608.11919v1. arXiv:2608.11919v1. Yang, H., Zhou, J., Fu, Y., Wang, X., Roane, R., Guan, H., and Liu, T. ProTrain: Efficient LLM training via automatic memory management. In Proceedings of Machine Learning and Systems, volume 8, pp. 1757–1770, 2026a. URL https://proceedings.mlsys.org/ paper files/paper/2026/file/56280ad5fb53967fd55e4ba1 b1cfc418-Paper-Conference.pdf. Yang, R. and Wen, Z. An efficient heterogeneous co-design for fine-tuning on a single GPU. In Proceedings of the 63rd ACM/IEEE Design Automation Conference, DAC ’26, New York, NY, USA, 2026. Association for Computing Machinery. doi: 10.1145/3770743.3804125. Yang, R., Lin, S., Ao, Y., Li, Z., Zhao, Y., Li, X., Lin, Y., and Wen, Z. SlideDP: Scaling host-resident LLM fine-tuning across multiple GPUs. arXiv:2609.34162, 2026b. URL https://arxiv.org/abs/2609.34162. Yang, S. and Zhang, Y. FLA: A Triton-based library for hardware-efficient implementations of linear attention mechanism.

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

https://github.com/fla-org/flash-linear-attention, 2024. Gated Delta Rule implementation; evaluated version 0.5.1. Yang, S., Kautz, J., and Hatamizadeh, A. Gated Delta Networks: Improving Mamba2 with delta rule. In The Thirteenth International Conference on Learning Representations, pp. 29687–29707, 2025. URL https://proceedings.iclr.cc/paper files/paper/2025/hash /4904fad153f6434a7bcf04465d4be2cc-Abstract-Confe rence.html. Yuan, Z., Sun, H., Sun, L., and Ye, Y. MegaTrain: Full precision training of 100B+ parameter large language models on a single GPU, 2026. URL https://arxiv.org/abs/2604.05091. Zhang, J., Liu, Y., Yan, S., Deng, L., Cao, J., Yang, S., Xue, B., Ni, M., and Li, S. MoEBlaze: Breaking the memory wall for efficient MoE training on modern GPUs. In Chowdhery, A. and Jia, Z. (eds.), Proceedings of Machine Learning and Systems, volume 8, pp. 1856–1869. MLSys, 2026. URL https://proceedings.ml sys.org/paper files/paper/2026/file/9032e5c9ec394ce7 68a2fa9bdc56af6c-Paper-Conference.pdf. Zhao, L., Shi, R., Zhang, S., Chen, Y., He, B., Sun, H., Yin, Z., Su, S., Cui, Z., Dong, L., Li, X., Wang, L., He, J., Ma, J., Huang, W., Tong, J., Gao, D., Zhang, J., Tian, H., Shen, H., Luo, Z., Sun, Z., Niu, H., and Sun, Y. MemFine: Memory-aware fine-grained scheduling for MoE training, 2026. URL https://arxiv.org/abs/2511.21431v4. arXiv:2511.21431v4. Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C., Xu, M., Wright, L., Shojanazeri, H., Ott, M., Shleifer, S., Desmaison, A., Balioglu, C., Damania, P., Nguyen, B., Chauhan, G., Hao, Y., Mathews, A., and Li, S. PyTorch FSDP: Experiences on scaling fully sharded data parallel. Proc. VLDB Endow., 16(12):3848–3860, 2023. doi: 10.14778/3611540.3611569. Zhu, Z., Giannoula, C., Andoorveedu, M., Su, Q., Mangalam, K., Zheng, B., and Pekhimenko, G. Mist: Efficient distributed training of large language models via memory-parallelism co-optimization. In Proceedings of the Twentieth European Conference on Computer Systems, EuroSys ’25, 2025. doi: 10.1145/3689031.3717461.

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

A

E XPERIMENTAL S ETUP

Unless stated otherwise, runs use four GPUs and the Qwen3.6-35B-A3B language backbone, and batch sizes are global. AdamW uses learning rate 10−5 , weight decay 0.01, betas (0.9, 0.999), and epsilon 10−8 . R and K count resident expert layers and MoE-retained layers; GAS denotes gradient-accumulation rounds. Platforms and CPU placement. Table 2 summarizes the evaluated systems, and Table 3 the routed-expert geometry of the evaluated models. The H100 host has eight NVLinkconnected 80 GB GPUs and two 48-core Xeon Platinum 8468 sockets with 1,008 GiB of memory each. The RTX 4090 host has four 24 GB PCIe GPUs and two 32-core Xeon Gold 6338N sockets with 503 GiB each. Each comparison runs all systems on the same GPUs and host; on H100, every system runs its CPU work on the socket local to its GPUs. Measurement. Each configuration runs as independent processes. A measured step covers forward, replay, backward, communication, optimizer updates, and weight publication, and ends after pending transfers finish, CUDA synchronizes, and all ranks reach a barrier. We take the slowest rank at each step and the median within each process. Main sweeps report the mean of two process medians, each over two warmup and five measured steps; mechanism studies (Figures 6a and 7) use one process with one warmup and three measured steps. Host memory is the peak PSS of the training process tree; for 512K/1M and the mechanism studies, it is the runtime’s rank-level PSS sum. GPU memory is the largest allocated peak across ranks.

B

E XECUTION S ETTINGS

Expert-owned state execution. EP4 uses two managed GPU buffers and two CPU gradient slots each for managed parameters and expert shards, with a dedicated stream for expert-gradient D2H. Multithreaded CPU AdamW updates host state; GPU AdamW updates resident expert state and refreshes its BF16 weights. Managed parameters and expert shards share one update executor per rank. MoE-retained layers keep their backward tensors on GPU and replay only attention; other layers checkpoint the full decoder with hoststored inputs. Split–Keep–Replay storage. Retained attention outputs O, log-sum-exp statistics (LSE), and decoder inputs live in pinned host pools whose NUMA placement is part of the configuration. A retained object is written once during the original forward, read by replay and backward, and released after its last reader completes; reuse waits for that completion. GDN recurrent states connect the chunks of one layer and remain on GPU through that layer’s backward.

Regeneration reuses the retained input’s prefetched copy and keeps the last traversed decoder’s parameters loaded for its backward.

C

E XECUTED C ONFIGURATIONS

H100 batch sweep. Table 4 lists the executed H100 4K configurations. EP4 uses full parameter transfer and two CPU gradient slots; Local uses full-layer checkpointing with host-stored inputs, shard transfer with 256 MiB chunks, and one CPU gradient slot. All use GAS 1; batch 32 runs the planner’s proposal, with 23 interleaved resident layers and 64K-token routing chunks. RTX 4090 batch sweep. At global batches 64–1024 (1K context), OOM uses per-GPU batches 16–256 with GAS 1; FSDP2 and ZeRO-Offload use microbatches 16/32/32/32/32 with GAS 1/1/2/4/8; MegaTrain, whose multi-GPU path does not accumulate gradients, exhausts GPU memory from 32 sequences per GPU. Up to 64 sequences per GPU, OOM uses the default operators. At 128 per GPU, GDN and MoE run in 32K- and 16K-token partitions; at 256 per GPU, the Local planner’s plan uses 16K-token GDN, 32K-token MoE, and 64K-token attention batch partitions. Long sequences. With one sequence per GPU, 64K and 128K use the default operators. At 256K, Split–Keep– Replay uses query group 2, 128K-token GDN and 256Ktoken MoE partitions, and retains every boundary and O/LSE; it takes 62.80 s versus 66.86 s for the default path, with allocated peak 39.4 versus 53.8 GiB. The 512K and 1M configurations use YaRN scaling factors two and four, 1+1 and 3+1 regeneration, and take 205.66 and 712.64 s per step. At 512K, FSDP2 exhausts GPU memory during GDN checkpoint replay, MegaTrain in GDN forward with checkpoint intervals of four and six, and ZeRO-Offload in its first step; all three do so again at 1M. Megatron-LM at 4K context. Megatron-LM uses TP2 with sequence parallelism, EP4/ETP1, PP1, and dense DP2, with full recomputation, vocabulary-parallel cross-entropy, and a CPU optimizer holding FP32 masters, gradients, and moments. Among the tested microbatch/GAS pairs, 16/1 exhausts memory in the first output projection and 4/4 in cross-entropy. The completing 2/8 setting takes 23.93 s per step (5.48K tokens/s) with 68.0 GiB allocated; our EP4 plan’s 5.020 s is 4.77× faster. KT-FT. KT-FT runs its official Accelerate and LlamaFactory full-parameter path with native AMX BF16 CPU experts on one 48-core socket (44 expert workers, the fastest of 40, 44, and 48; 24 Torch/BLAS threads; recompute). Our changes add compatibility fixes, packed inputs, timing, and FP32 master weights and Adam moments. Step times at

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design Table 2. Evaluated systems. All use BF16 compute and FP32 canonical masters and moments.

System

Parallelism / canonical state

Activation / transfer

Gradient processing

OOM

EP4: selected expert state on GPU; Local: host

FSDP2

CPU-sharded state

GPU reduction, or BF16 partials with FP32 host accumulation FP32 reduce-scatter

ZeRO-Offload

ZeRO-3 with CPU parameter and optimizer offload TP2/SP, EP4/ETP1, DP2, PP1; CPU optimizer Layer streaming with CPU offload Routed experts on CPU (AMX); others through FSDP2

Workload-specific checkpointing, retention, and partitioning; EP4 full and Local shard transfer Layer checkpointing and activation offload Layer checkpointing and activation offload Full recomputation Its streaming path with FA3 Recompute; official full-parameter path

Host optimizer BF16 expert gradients, FP32 update

Megatron-LM MegaTrain KT-FT

Table 3. Routed-expert geometry of the evaluated language backbones (Jiang et al., 2024; Qwen Team, 2025; 2026c;a; Gemma Team, 2026). Model Experts E Active k E/k Mixtral-8×7B Gemma4-26B-A4B Qwen3-30B-A3B Qwen3.6-35B-A3B Qwen3.5-122B-A10B

8 128 128 256 256

2 8 8 8 8

4 16 16 32 32

Table 4. Executed configurations in the H100 4K sweep. Batch

Path

State / activation

Step (s)

16 32 64 128 256

EP4 EP4 EP4 Local Local

R27/K0 R23/K0 R3/K4 Host state, checkpoint offload Host state, checkpoint offload

3.776 5.020 9.302 12.177 22.305

global batch 16, 32, 64, and 128 are 99.49, 168.38, 326.55, and 726.60 s; batch 256 exhausts GPU memory at GAS 1 and takes 1,396.19 s at GAS 2. With one sequence per GPU, 64K and 128K take 314.29 and 734.61 s; 256K exhausts GPU memory with both the default and expandable allocators. Sampled host peaks are 619, 656, 732, 876, and 876 GiB at batches 16–256. At batch 32, the CPU expert forward, recompute, and backward paths take 70.2% of rank 0’s exposed wall time and the FP32 update 4.2%. For Qwen3.5122B-A10B, FP32 masters and moments add 864 GiB over KT-FT’s native BF16 optimizer state; at global batch 32 and 1K context, the run exhausts the node’s host memory before its first update, with 1,946 GiB of anonymous memory on rank 0. Base and configured paths. Figure 4 measures each layout before and after configuration. Base uses full parameter

BF16 communication FP32 main gradients

transfer, 64 MiB H2D/D2H chunks, full checkpoint offload, and one CPU gradient slot; EP4 Base has no resident layers and 32K-token routing chunks. Configured EP4 sets resident layers, MoE retention, and 32K–256K-token routing chunks for each batch, with two CPU gradient slots and the dedicated gradient stream; configured Local uses shard transfer with 256–512 MiB chunks. At batch 96, configured EP4 uses R1/K1. Path comparison in Figure 1c. EP4/Local step times are 5.020/10.035, 9.302/10.245, and 18.714/12.177 s at global batch 32, 64, and 128. EP4 uses R23/K0, R3/K4, and R1/K1 with full transfer and two gradient slots; Local uses shard transfer with one slot.

D

M ECHANISM M EASUREMENTS

Residency and retention. The batch-32/4K residency pair in Section 7.4 runs the planned EP4 configuration with no resident layers (8.48 s) and with the planned 23 (5.02 s), raising GPU allocation from 11.4 to 64.8 GiB. The gain persists at 1K context, where 27 resident layers reduce step time from 8.11 to 3.54 s (2.29×). Interleaved sets keep decoders 0 and 1 resident, which finish backward last and are needed first by the next forward, and space the rest evenly: the 12-layer pair (64K routing chunk) compares decoders 0–11 with decoders 0, 1, 3, 7, 11, 15, 19, 22, 26, 30, 34, and 38. Table 5 gives the residency–retention pairs in Section 7.4 at global batch 32 and 4K context, two processes each. R21 adds four resident layers (decoders 33, 35, 37, and 39) to R17; K1 retains decoder 0’s MoE activations while its attention still replays. Cumulative 512K ablation. All arms use Local execution with shard transfer in 256 MiB chunks, GPU gradient reduction, full checkpoint offload, one CPU gradient slot, and YaRN×2. The default arm uses the unpartitioned oper-

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design Table 5. EP4 residency and retention at batch 32, 4K: step time (s) and allocated peak (GiB). K0

K1

Resident

Time

Peak

Time

Peak

R17 R21

6.048 5.376

51.3 60.3

5.733 5.331

55.0 64.1

Table 6. Cumulative Split–Keep–Replay arms at 512K: step time (s), GPU allocated peak and host memory (GiB). Arm Default + Split + Keep + Replay

Table 7. Policies by length: step time (s) / host memory (GiB).

Time

GPU

Host

Out of GPU memory 260.02 200.46 206.00

— 43.4 43.4 45.4

— 731.3 892.5 732.5

ators and exhausts GPU memory in the first MoE backward. +Split groups eight query heads, partitions GDN and MoE into 128K- and 256K-token chunks, releases GDN intermediates at last use, skips unused terminal MoE outputs during replay, and defers backward input prefetch until attention workspace is released. +Keep retains forward O/LSE on the host. +Replay keeps decoder inputs 0, 2, . . . , 38 and regenerates the rest, holding at most one regenerated boundary. Table 6 gives the values in Figure 6a. Policies by length. Table 7 gives the values in Figure 7. A retains every decoder boundary and recomputes attention; B also retains O/LSE; 1+1 and 3+1 retain O/LSE and keep one boundary in every two and four layers. Policies share the 512K operator partitions, clamped to the sequence length; at 512K, A and B are the +Split and +Keep arms above. The length-dependent benefit of retention and the host savings of regeneration make the best policy workload-specific. Under a 750 GiB budget at 512K, the Local planner stores boundaries 0–21 on the host, regenerates boundary 22, and holds boundaries 23–39 on GPU; this plan takes 202.06 s (predicted 201.66 s) at 63.8 GiB of GPU allocation. Sliding-window hybrid. Gemma4-26B-A4B interleaves five sliding-window layers (1K window, head dimension 256) with each global-attention layer (head dimension 512, two shared key/value heads). Table 8 gives per-layer forward times for one sequence on one H100 and the host payload of each 5+1 block. Global attention uses PyTorch’s memoryefficient kernel, since FlashAttention does not support head dimension 512. Largest completed workloads. For Figure 6b, the H100 context endpoint is the 1M configuration above, and the H100 batch endpoint runs Local at 256 sequences of 4K per GPU with 64K-token GDN, MoE, and attention batch parti-

Length

A

B

1+1

3+1

32K 64K 128K 256K 512K

10.69/431 12.32/451 26.48/491 77.83/571 260.0/731

11.70/441 11.62/471 23.27/531 62.55/651 200.5/893

11.75/431 12.94/451 26.46/491 66.85/569 205.5/732

14.62/426 15.17/441 26.29/471 66.88/531 208.2/652

Table 8. Gemma4-26B-A4B per-layer forward time (s) on one H100 and host payload per 5+1 block (GiB). Length Global Reuse O Sliding MoE O/LSE Freed 32K 64K 128K 256K

0.160 0.821 3.904 17.164

0.004 0.008 0.021 0.039

0.005 0.013 0.021 0.042

0.015 0.030 0.071 0.147

0.50 1.00 2.01 4.02

0.86 1.72 3.44 6.88

tions, at 66.3 GiB peak allocation. The RTX 4090 endpoints are the 256-per-GPU plan above and one 512K sequence per GPU with grouped attention and GDN/MoE partitions, at 21.95 GiB peak allocation. The default path completes 256K and 64 sequences per GPU on H100, and 64K and 64 per GPU on RTX 4090; each exhausts GPU memory at the next doubling. A single H100 with the same mechanisms and YaRN factor eight completes a 2M-token step in 2,606 s, with 67.0 GiB of GPU allocation and 915 GiB of host memory.

E

C ONFIGURE - THEN -C OMPARE D ETAILS

Cost and memory accounting. EP calibration scales GPU work with token count, treats parameter work as approximately token-independent, and adds a per-chunk scheduling charge. The Local time model sums operator time, boundary regeneration, exposed transfer chains, and a fixed overhead. GPU feasibility sums the allocations live in each execution phase and takes the largest phase total; views and their backing storage are counted once. Host estimates include model state, pinned pools, checkpoints, retained results, and staging, charged to the NUMA node of each allocation.

Selection rules and plan output. Local prices each candidate in under a millisecond and searches its full partition menu offline on one core in under an hour. Among candidates within 0.5% of the minimum predicted time, the EP search prefers the one with the most memory headroom. A saved plan records the layout, resident and retained layer indices, operator chunks, parameter-transfer mode and chunk sizes, buffer settings, the boundary keep set, host placement, and resource estimates; the runtime executes it directly.

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design EP4/R21, 1K

Table 9. Predicted versus measured step time (s), with predictions recorded before execution. Error is predicted/measured −1. Plan

Pred.

Meas.

Error

H100, B32, 4K H100, B96, 4K 4090, B384, 1K H100, 128K H100, 128K H100, 256K H100, 512K H100, 512K

EP4 R23 5.135 5.020 EP4 R1/K1 13.73 13.78 Local, 32K chunks 38.66 40.38 B 23.21 23.27 3+1 24.85 26.29 B 62.73 62.76 1+1 205.70 205.49 3+1 207.89 208.18

+2.3% −0.4% −4.3% −0.2% −5.5% −0.1% +0.1% −0.1%

Table 10. Predicted/measured step time (s) of both configured paths; the selected path is in bold. Model

S

Batch

Plan

EP4

Local

Qwen3.6 Qwen3.6 Qwen3.6 Qwen3.6† Qwen3.6 Mixtral Mixtral† Mixtral Gemma4 Gemma4† Gemma4† Gemma4 Qwen3 Qwen3† Qwen3† Qwen3

4K 4K 1K 1K 1K 1K 1K 1K 1K 1K 1K 1K 1K 1K 1K 1K

48 192 32 128 256 32 128 256 32 128 192 256 32 128 192 256

R15/K1 R1 R27 R23 R19 R15 R11 R7 R27 R18 R8 R1 R36 R30 R26 R18

7.08/7.29 25.91/26.72 4.03/3.54 5.54/5.12 8.54/9.29 6.26/6.31 9.23/10.06 16.85/18.06 2.78/2.89 6.98/7.07 9.91/10.17 13.78/13.32 3.19/2.95 5.32/5.63 7.11/7.72 10.04/10.55

9.27/10.20 17.09/17.51 10.20/10.35 10.27/10.26 9.55/10.08 11.62/12.21 14.30/14.68 21.78/21.67 7.87/7.79 8.24/8.34 8.51/7.90 9.79/9.45 7.77/8.12 8.05/8.87 8.30/8.34 8.93/8.90

Prediction accuracy. Table 9 lists plans whose predictions were recorded before execution; the long-sequence rows cover 128K–512K, the range on which the Local time model is calibrated. With operator partitions and O/LSE retention fixed at 512K, the model predicts 205.70 and 207.89 s for 1+1 and 3+1 regeneration, respectively, matching their measured order at 205.49 and 208.18 s. The five EP4 budget policies at batch 32/4K (64 and 72 GiB budgets) are predicted within 3.3–6.6% of measured time and 0.24 GiB of allocated memory. Short-test path selection. Table 10 lists the 16 workloads in Section 7.5, on four H100s. Each test runs one warmup and five measured steps on whole repeating units: four and eight layers, or six and twelve for Gemma4’s 5+1 block. Decisions for the workloads marked † were recorded before their full-model runs.

F

T RAINING Q UALITY

OOM and FSDP2 train independently from the same checkpoint and data order; held-out NLL is measured on token-

Local, 32K

0.10

Held-out NLL vs. FSDP2 (%)

Workload

Local, 1K

0.05 0.00 FSDP2 reference

−0.05 −0.10 0

2

4

6

8

10

12

14

16

Training input tokens (millions)

Figure 9. Held-out NLL gap to independent FSDP2 runs from the same checkpoint and data order for three configurations of our system.

and block-disjoint data. Local and EP4/R21 at eight 1K sequences per GPU run 500 steps (16.384M input tokens) with evaluation at step 0 and every 50 steps; Local with partitioned operators at one 32K sequence per GPU (query group 2, GDN 8K, MoE 4K) runs 128 steps (16.777M tokens) with evaluation every 32 steps. Evaluation runs between steps and restores the training state before training resumes. Heldout NLL decreases at every checkpoint for both systems, and the relative gap stays within 0.11%. With the same seed and inputs, policies B, 3+1, and 1+1 (Table 7) match policy A’s first-step loss to four decimals at 32K and 128K and stay within 0.0006 of it over the next three steps. Final OOM/FSDP2 NLL is 1.376510/1.375475 for Local at 1K, 1.375354/1.375475 for EP4, and 1.159174/1.158994 at 32K (Figure 9).

G

M ODELS AND E IGHT GPU S

Additional architectures. Mixtral-8x7B (32 layers, eight experts, top-2, 4K sliding window), Gemma4-26B-A4B (30 layers, 128 experts, top-8, with a dense MLP beside the experts), and Qwen3-30B-A3B (48 layers, 128 experts, top8) run on four H100s at 1K context. At batch 32, OOM uses EP4 with R15 (Mixtral), R27 (Gemma4), and R36 (Qwen3-30B) resident layers; step times versus FSDP2 are 6.308/10.503, 2.887/6.318, and 2.954/6.918 s. At batch 256, OOM uses EP4 for Mixtral and Local for Gemma4 and Qwen3-30B; step times versus FSDP2 are 18.057/22.129, 9.454/13.076, and 8.897/11.247 s. FSDP2 uses the same expert kernel as OOM, or its own faster Liger kernel. Qwen3.5-122B-A10B. The 48-layer model (256 top-8 experts, hidden size 3072) runs on four H100s. Local execution at 4K context, without resident layers, takes 46.0, 52.3, and 96.3 s per step at global batch 64, 128, and 256 (5.7K, 10.0K, and 10.9K tokens/s), with 44.9–65.4 GiB of peak GPU allocation and 1,516–1,612 GiB of host memory. Batch 256 uses 64K-token GDN, MoE, and attention batch

OOM: Offload-Optimized Full-Parameter MoE Training via Memory-Execution Co-Design

partitions and regenerates every other decoder boundary, since the pinned host memory for keeping every boundary exceeds what our test machine can allocate. With EP4 at 1K context and global batch 16, six resident expert layers reduce step time from 55.7 to 41.7 s (1.34×) and host memory from 1,420 to 1,258 GiB. CPU-offloaded FSDP2 (batch 64, microbatch 4, GAS 4) uses over 1.8 TiB of host memory and exhausts the node before completing its first step; ZeROOffload with the same settings exhausts it during setup at 1.82 TiB, and MegaTrain at batch 64 exhausts it during its first step. FSDP2 and ZeRO-Offload keep FP32 gradients, master weights, and Adam moments in host memory, at least 16 bytes per parameter (about 1.8 TiB for this model). Eight GPUs. The eight-H100 comparison uses 1K context and global batches 128 and 256, with each rank on its local NUMA node. OOM uses EP8/R21 with GAS 1. FSDP2 uses the CPU-offloaded recipe of the four-GPU model comparison and ZeRO-Offload its CPU offload, both with microbatches 16 and 32; Megatron-LM keeps its CPU optimizer and uses microbatch 32 with GAS 2 and 4, since microbatch 64 exhausts GPU memory. All use FA2 attention. Step times at batch 128 and 256 are 3.238/5.548 s for OOM, 4.932/6.445 s for FSDP2, 7.482/8.732 s for ZeROOffload, and 9.713/13.232 s for Megatron-LM.

Related documents

• ISER-MoE: Dynamic In-Switch Replication of Hot Experts for Load Balancing in Expert-Parallel#1598837
• Efficient CPU-GPU collaborative MoE inference via Adaptive Expert Assignment#1387295
• Accelerating Offloading-Based MoE Inference by Interleaving Concurrent Requests#1527889
• SlimWise: Decoupling Expert Pruning Across Prefill and Decode for Efficient MoE Serving#1119190
• MaskCoFT: Masked Co-Adaptive Fine-Tuning for Memory-Efficient MoE Inference#1119150