PaperlayerView exact revision on arXiv

Source-linked AI summary

MaskCoFT: Masked Co-Adaptive Fine-Tuning for Memory-Efficient MoE Inference

Junfeng Wu, Zehao Fan, Hadjer Benmeziane, Kaoutar El Maghraoui, Liu Liu, Yinan Wang

arXiv:2609.34077v2cs.LGcs.AI

TL;DR

Expert migrations limit decoding efficiency when MoE models exceed GPU memory, and caching or prefetching is constrained by routing. MaskCoFT jointly fine-tunes routers and experts under a learned mask; it reduces expert fetches and serving time while maintaining accuracy above the base model.

  • Problem

    During offloaded MoE decoding, selected experts absent from the GPU cache must be fetched from host memory, making expert migrations a major limit on decoding speed.

    Source: §1 · p. 1

  • Method

    MaskCoFT jointly fine-tunes routers and experts with cross-entropy alone, using a shared binary mask to concentrate training-time routing and a learned prior to re-rank experts at inference.

    Source: §1 · p. 2

  • Results

    MaskCoFT cuts expert fetches per token by 23.7% and 10.1%, and lowers MoE-Offloading time per output token by up to 16.4% and 5.5%, respectively, while average accuracy stays above the base model.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Fetch reductions hold under LRU, LFU, and FIFO at every tested cache budget, with average accuracy above the base model on both architectures.

    Source: §1 · p. 3

Abstract

from arXiv · show

Mixture-of-experts (MoE) language models often exceed the memory of a single GPU. Expert offloading keeps most experts in host memory and loads them on demand, so decoding speed depends on how many experts each token must fetch. Caching and prefetching reduce this cost only as far as the routing allows. Router-only fine-tuning can reshape the routing to reuse experts, but it keeps the experts frozen, so they cannot adapt to the tokens the new routing sends them. We propose MaskCoFT, a masked co-adaptive fine-tuning method that trains routers and experts together with the cross-entropy loss alone. During fine-tuning, a learnable binary mask restricts the Top-K routing of each layer to a subset of experts, and the experts adapt to the tokens redirected to them. At inference, the learned mask becomes a soft prior that re-ranks experts, so every expert remains selectable. We simulate a GPU cache of 4 experts per layer for Mixtral-8x7B and 12 for DeepSeek-V2-Lite. MaskCoFT cuts expert fetches per token by 23.7% and 10.1% relative to the base model. In real offloading system serving, it lowers the time per output token by up to 16.4% and 5.5%, respectively. Its average accuracy over nine benchmarks stays above the base model by 0.92 and 0.53 points.

1 INTRODUCTION

Expert migrations constrain decoding speed when MoE experts are offloaded. MaskCoFT concentrates routing during joint router–expert fine-tuning, then uses the learned mask as an inference-time prior.

  • Expert loading over PCIe takes 86% of Mixtral-8×7B decoding time in MoE-Offloading, while expert computation takes 3%.On memory-constrained devices, decoding speed therefore depends mainly on the number of expert migrations.Source: §1 · p. 1
  • Caching and prefetching are constrained by expert access patterns, GPU capacity, and transfer bandwidth; long transfers can outlast available computation for overlap.Source: §1 · pp. 1–2
  • MaskCoFT restricts each layer’s training-time Top-K routing to a shared concentration set, jointly adapting routers and experts to redirected tokens.At inference, the learned mask re-ranks experts as a soft prior, leaving every expert selectable.Source: §1 · p. 2
  • Unlike ReMoE’s temporal-locality approach, MaskCoFT’s shared expert set makes reuse independent of nearby tokens selecting the same experts.ReMoE keeps experts frozen, limiting their adaptation to tokens rerouted by updated routers.Source: §1 · p. 2
  • Experiments assess router–expert co-adaptation on coarse-grained Mixtral-8×7B and fine-grained DeepSeek-V2-Lite.Coarse-grained MoEs use fewer, larger experts; fine-grained MoEs use more, smaller experts.Source: §1 · p. 2
  • MaskCoFT uses cross-entropy alone, without auxiliary locality or load-balancing losses or their associated loss hyperparameters.Source: §1 · p. 2
  • MaskCoFT cuts expert fetches per token by 23.7% and 10.1%, and lowers MoE-Offloading time per output token by up to 16.4% and 5.5%, respectively.Average accuracy over nine benchmarks stays above the base model on both architectures.Source: Figure 1 · p. 2, §1 · p. 3

2 RELATED WORK

Prior work addresses MoE serving through offloading systems, training methods, and expert pruning or masking. MaskCoFT differs by fine-tuning a pretrained model’s routers and experts to concentrate token activations and reduce expert transfers.

  • Offloading systems improve caching, prefetching, or heterogeneous execution, while cache-aware routing favors experts already resident in cache without training.MaskCoFT changes routing through fine-tuning and is compatible with different cache policies.Source: §2 · p. 3
  • Prior training methods either train from scratch, fine-tune routers while freezing experts, or train auxiliary cache routers; MaskCoFT fine-tunes pretrained routers and experts with cross-entropy alone.Its learned mask concentrates token activations on a subset of experts.Source: §2 · p. 3
  • Pruning and token-adaptive masks target model size or computational redundancy, which does not necessarily increase reuse across tokens or reduce migration latency.MaskCoFT’s binary routing masks target expert transfers during offloaded decoding.Source: §2 · p. 3

3 METHODOLOGY

MoE layers route tokens to selected experts, which can drive repeated host-to-GPU fetches when a small cache misses. MaskCoFT searches for shared concentration sets, co-adapts routers and experts within them, and uses a learned prior to re-rank experts at inference.

3.1 PRELIMINARIES

  • Each MoE layer uses a router to select Top-K experts for a token, while changing selections can miss a small GPU cache and trigger PCIe fetches.The selected experts' gated outputs are summed; in the studied offloading systems, every cache miss triggers a fetch.Source: §3.1 · pp. 3–4

3.2 MASKCOFT: MASKED CO-ADAPTIVE FINE-TUNING

  • MaskCoFT trains routers and experts jointly under a learnable mask, restricting each layer to a shared concentration set during fine-tuning.The learned sampling distribution later acts as a prior that re-ranks experts.Source: §3.2 · p. 4, Figure 2 · p. 4
  • Stage 1 samples E′ experts from a learned distribution and confines Top-K routing to that concentration set, resampling masks while accumulating mask gradients over T steps.The real-valued mask is updated once per T steps, averaging gradients across sampled sets.Source: §3.2 · p. 4
  • After the sampled sets stabilize, Stage 2 fixes them and fine-tunes routers and experts so the experts adapt to rerouted tokens.LoRA adapters learn to process tokens previously served by masked experts; this is the paper's knowledge transfer to the concentration set.Source: §3.2 · p. 5
  • Both stages use standard cross-entropy without an auxiliary load-balancing loss, while a straight-through estimator passes mask gradients through the nondifferentiable sampling step.The loss gradient sums contributions across tokens sharing the mask and is passed from the mask through the sampling distribution to the real-valued mask.Source: §3.2 · p. 5
  • At inference, the learned prior re-ranks experts rather than hard-excluding those outside the concentration set, and it changes active-expert selection without changing gate weights.An outside expert can still be active when its router probability multiplied by its prior ranks among the Top-K.Source: §3.2 · p. 5

4 EXPERIMENT

The experiments evaluate MaskCoFT’s accuracy, simulated expert-cache behavior, and end-to-end offloading performance on Mixtral-8×7B and DeepSeek-V2-Lite, then compare it with ReMoE on DeepSeek-V2-Lite.

4.1 EXPERIMENTAL SETUP

  • The evaluation covers nine accuracy benchmarks, trace-driven cache simulation under LRU, LFU, and FIFO, and real-system measurements of TTFT, TPOT, and throughput.It tests Mixtral-8×7B and DeepSeek-V2-Lite; the cache simulation uses 128 ShareGPT prompts, while system measurements use 32 prompts per model and 64- or 128-token outputs.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2 MODEL CAPACITY EVALUATION

  • MaskCoFT’s nine-benchmark average exceeds the base model by 0.53 points on DeepSeek-V2-Lite and 0.92 points on Mixtral-8×7B, while trailing CE-only by 0.46 and 1.22 points, respectively.The methods use the same cross-entropy objective and differ only in masking; MaskCoFT therefore largely preserves fine-tuning’s performance benefits.Source: §4.2 · pp. 6–7, Table 1 · p. 6

4.3 SIMULATION RESULTS

  • Under LFU, MaskCoFT’s hit-rate gain is +9.10 percentage points on Mixtral-8×7B versus +5.73 on DeepSeek-V2-Lite, and it has the highest hit rate and fewest fetches across tested policies and cache budgets.The simulation uses cache budgets of 4 experts per layer for Mixtral and 12 for DeepSeek; Figure 3 varies the budget.Source: §4.3 · p. 7, Table 2 · p. 7, Figure 3 · p. 8

4.4 SYSTEM PERFORMANCE MEASUREMENT

  • TPOT falls by 15.5% and 16.4% on Mixtral-8×7B, and by 3.7% and 5.5% on DeepSeek-V2-Lite, for 64- and 128-token outputs, respectively.Figure 4 reports end-to-end decoding in MoE-Offloading; TTFT changes by at most 4.4%.Source: §4.4 · pp. 7–8, Figure 4 · p. 8

4.5 COMPARISON WITH REMOE

  • On DeepSeek-V2-Lite, ReMoE gains more hit rate under LRU and FIFO, while both methods gain about the same under LFU; ReMoE is ahead overall on accuracy.The comparison uses each method’s own baseline because some ReMoE evaluation settings are unreported.Source: §4.5 · p. 9, Table 3 · p. 9

5 CONCLUSION

MaskCoFT co-adapts routers and experts during fine-tuning, then uses the learned mask as an inference-time prior. Its evaluations report fewer expert fetches, lower serving time, and average accuracy comparable to the base model.

5 CONCLUSION

  • MaskCoFT co-adapts routers and experts under a learnable binary mask, which becomes a soft prior that re-ranks experts at inference.The mask concentrates routing during fine-tuning; the inference prior re-ranks experts.Source: §5 · p. 9
  • With 4 and 12 cached experts per layer, MaskCoFT cuts expert fetches per token by up to 23.7% on Mixtral-8×7B and 10.1% on DeepSeek-V2-Lite.The reductions hold under three eviction policies and at every tested cache budget.Source: §5 · p. 9
  • Serving time per output token falls by up to 16.4% and 5.5%, while average accuracy stays comparable to the base model.Without masking, the same cross-entropy loss gains more accuracy but fetches more experts.Source: §5 · p. 10
  • The authors report using Claude to help implement experiment code, refine figures, polish text, and check cited-work support, while retaining responsibility for the paper.They state it was not used to generate research ideas, design methods or experiments, prove results, generate or clean datasets, or interpret results.Source: AI USE STATEMENT · p. 10
  • The paper describes its model, data, benchmark, cache-simulator, and training details, and plans to share anonymized code and checkpoints with reviewers before public release after publication.Section 3 covers the method; Section 4 covers evaluation, with specifications and hyperparameters in Appendices A and B.Source: REPRODUCIBILITY STATEMENT · p. 10

A MODEL DETAILS

Table 4 summarizes the specifications of the evaluated MoE models, with expert counts reported per layer.

A MODEL DETAILS

B HYPERPARAMETER SETTINGS

The experiments use shared LoRA settings and separate learning-rate schedules for experts and routers versus the mask; Stage 2 begins after mask convergence.

B HYPERPARAMETER SETTINGS

  • Both models use seed 42, LoRA rank 8, alpha 16, dropout 0.05, and a peak learning rate of 2e-4 for experts and routers.The mask uses a cosine schedule with peak learning rate 0.05 and floor ratio 0.2; its binary mask is resampled every step.Source: B HYPERPARAMETER SETTINGS · p. 12
  • The real-valued mask gradient accumulates over T=100 steps for DeepSeek-V2-Lite and T=1 for Mixtral-8×7B before each update.Stage 2 starts after mask convergence; training totals 5,000 steps for DeepSeek-V2-Lite and 3,000 for Mixtral-8×7B.Source: B HYPERPARAMETER SETTINGS · pp. 12–13

C CONVERGENCE OF THE CONCENTRATION SET

The concentration mask converges early in both models during Stage 1.

C CONVERGENCE OF THE CONCENTRATION SET

  • Mixtral-8×7B reaches 100% overlap with the final concentration set by step 1.7k, while DeepSeek-V2-Lite reaches 99% by step 4.5k.Figure 5 tracks overlap every 100 steps.Source: Figure 5 · p. 13

D DECODING LATENCY BREAKDOWN

Figures 6 and 7 present decoding-latency breakdowns across different output token lengths for Mixtral-8×7B and DeepSeek-V2-Lite.

D DECODING LATENCY BREAKDOWN

  • Figures 6 and 7 cover decoding-latency breakdowns across different output token lengths for Mixtral-8×7B and DeepSeek-V2-Lite.The figures provide model-specific views of latency as output length varies.Source: Figure 6 · p. 13, Figure 7 · p. 14
  • Figure 6 presents the decoding-latency breakdown for Mixtral-8×7B under different output token lengths.Source: Figure 6 · p. 13
  • Figure 7 presents the decoding-latency breakdown for DeepSeek-V2-Lite under different output token lengths.Source: Figure 7 · p. 14
Loading 2609.34077v2…