feat(integrations): expert_offload — single-GPU CPU offload of frozen 4-bit MoE experts#3797
Conversation
📝 WalkthroughWalkthroughAdds a new "Expert Offload" integration package for single-GPU QLoRA MoE training that streams frozen 4-bit expert weights between pinned CPU RAM and GPU via a forward pre-hook with single-resident-slot staging/eviction. Includes config args, a validating plugin, core offload logic, tests, README, and example config. ChangesExpert Offload Integration
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/integrations/test_expert_offload.py (1)
293-338: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
pytest.importorskipfor the bitsandbytes dependency.The class is gated on
torch.cuda.is_available()only; if CUDA is present butbitsandbytesisn't installed,_buildraises a rawImportErrorinstead of a clean skip.♻️ Proposed fix
`@staticmethod` def _build(d=512, n_experts=16, n_layers=4): - import bitsandbytes as bnb + bnb = pytest.importorskip("bitsandbytes")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integrations/test_expert_offload.py` around lines 293 - 338, The TestCudaLinear4bit integration test only checks torch.cuda.is_available(), so _build() can still fail with an ImportError when bitsandbytes is missing. Update the TestCudaLinear4bit setup to use pytest.importorskip for bitsandbytes before importing it inside _build, so the test is cleanly skipped instead of erroring. Keep the existing CUDA skip condition and make the dependency guard explicit in the _build helper.src/axolotl/integrations/expert_offload/args.py (1)
14-30: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd
ConfigDict(extra="forbid")to reject unknown/typoed keys.
ExpertOffloadArgshas nomodel_config, so a typo likeexpert_offlaod_pin_memoryin YAML would silently be ignored rather than erroring.🛡️ Proposed fix
-from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict class ExpertOffloadArgs(BaseModel): """Input args for the expert_offload plugin. See the integration README. ... """ + model_config = ConfigDict(extra="forbid") expert_offload: bool = FalseBased on learnings, "enforce strict unknown-key rejection by setting
ConfigDict(extra=\"forbid\")on the nested integration-specific model" rather than the merged top-level config.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/integrations/expert_offload/args.py` around lines 14 - 30, ExpertOffloadArgs currently accepts unknown fields silently, so typoed YAML keys can be ignored without error. Add a model_config with ConfigDict(extra="forbid") on ExpertOffloadArgs to make this nested integration-specific model reject unknown keys, and keep the change localized to this class rather than the merged top-level config.Source: Learnings
src/axolotl/integrations/expert_offload/offload.py (1)
70-76: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNarrow the blind
except Exceptionto match_to_home's pattern.Static analysis flags this as overly broad; the codebase already uses
(RuntimeError, AssertionError)for the analogous pin-memory fallback in_to_home(lines 174-181).♻️ Proposed fix
try: return bool(t.is_pinned()) - except Exception: # pragma: no cover - platform dependent + except (RuntimeError, AssertionError): # pragma: no cover - platform dependent return False🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/integrations/expert_offload/offload.py` around lines 70 - 76, The _is_pinned helper currently catches Exception too broadly; narrow it to the same fallback used in _to_home by catching only RuntimeError and AssertionError around t.is_pinned(). Keep the existing False fallback and leave the rest of the _is_pinned logic unchanged.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/axolotl/integrations/expert_offload/offload.py`:
- Around line 243-250: install_expert_offload currently warns and returns an
empty list when find_moe_expert_blocks(model) finds no expert blocks, which
turns a required expert_offload configuration into a silent no-op. Update the
install_expert_offload path in offload.py to fail loudly instead of continuing:
replace the warning-only branch after blocks is empty with an explicit exception
that includes the expert_offload context and model-layout mismatch. Keep the
check near find_moe_expert_blocks(model) so misconfigured fused-expert or
non-MoE models are rejected before training starts.
---
Nitpick comments:
In `@src/axolotl/integrations/expert_offload/args.py`:
- Around line 14-30: ExpertOffloadArgs currently accepts unknown fields
silently, so typoed YAML keys can be ignored without error. Add a model_config
with ConfigDict(extra="forbid") on ExpertOffloadArgs to make this nested
integration-specific model reject unknown keys, and keep the change localized to
this class rather than the merged top-level config.
In `@src/axolotl/integrations/expert_offload/offload.py`:
- Around line 70-76: The _is_pinned helper currently catches Exception too
broadly; narrow it to the same fallback used in _to_home by catching only
RuntimeError and AssertionError around t.is_pinned(). Keep the existing False
fallback and leave the rest of the _is_pinned logic unchanged.
In `@tests/integrations/test_expert_offload.py`:
- Around line 293-338: The TestCudaLinear4bit integration test only checks
torch.cuda.is_available(), so _build() can still fail with an ImportError when
bitsandbytes is missing. Update the TestCudaLinear4bit setup to use
pytest.importorskip for bitsandbytes before importing it inside _build, so the
test is cleanly skipped instead of erroring. Keep the existing CUDA skip
condition and make the dependency guard explicit in the _build helper.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 25cc1cd3-67fa-45aa-a0a8-241397fc0a1f
📒 Files selected for processing (7)
src/axolotl/integrations/expert_offload/README.mdsrc/axolotl/integrations/expert_offload/__init__.pysrc/axolotl/integrations/expert_offload/args.pysrc/axolotl/integrations/expert_offload/examples/olmoe-1b-7b-qlora-expert-offload.yamlsrc/axolotl/integrations/expert_offload/offload.pysrc/axolotl/integrations/expert_offload/plugin.pytests/integrations/test_expert_offload.py
| if getattr(cfg, "fsdp_config", None) or getattr(cfg, "fsdp", None): | ||
| errors.append("is single-GPU only and incompatible with FSDP") | ||
| if getattr(cfg, "deepspeed", None): | ||
| errors.append("is single-GPU only and incompatible with DeepSpeed") | ||
| if (getattr(cfg, "expert_parallel_size", None) or 1) > 1: | ||
| errors.append( | ||
| "is incompatible with expert_parallel (both manage the expert weights)" | ||
| ) |
There was a problem hiding this comment.
these could be moved to args.py and use the native pydantic schema validation
There was a problem hiding this comment.
Done in fa42e43 — moved the whole cross-field check into a model_validator(mode="after") on ExpertOffloadArgs. Since merge_input_args folds plugin args into the config by inheritance, the validator sees the merged config (same pattern as liger's tensor_parallel_size check), so it now fires at config-parse time — before the model download — and the plugin's pre_model_load is gone entirely.
One tightening that fell out: use_reentrant must now be explicitly false. Parse-time validation runs before normalize_config defaults omitted gradient_checkpointing_kwargs to {"use_reentrant": True}, so an omitted value can't be trusted at that point — the validator demands the explicit setting, with the reason in the error message.
Verified end-to-end through the real merge_input_args (bare enable rejected with all requirements listed; omitted kwargs rejected; valid config accepted; plugin-off config unaffected), plus the unit tests now assert ValidationError through a merged-shape model instead of calling a plugin helper.
|
Review addressed in fa42e43:
Declining one nit — Suite after the changes: 25 tests (23 CPU + 2 CUDA against real bitsandbytes |
|
Ran the benchmarks through
Host: 2×RTX 3090 24 GB, fresh Running it through your trainer caught something the unit suite couldn't—current transformers stores OLMoE's experts as fused 3D stacks, so plain Numbers worth being straight about: loaded footprint drops 76%, but steady-state s/it costs +35% in this setup—the parametrized layout dequantizes the full 64-expert stack per access and the training path stages synchronously by design at seq 2048. The PR body's ~11% figure came from the per-expert Happy to run Qwen3-30B-A3B through the trainer on a bigger-RAM host as a follow-up if that's useful for the review. No rush—branch is stable as of |
… 4-bit MoE experts Adds an `axolotl.integrations.expert_offload` plugin that streams a MoE's frozen 4-bit experts from pinned CPU RAM to the GPU one block at a time, so a model whose experts exceed VRAM can QLoRA-train on a single small card. Only the experts move; attention, router, norms and the trainable LoRA adapters stay resident, so per-step PCIe traffic is limited to the experts (the bulk of a MoE's parameters). Mechanism: home each offloaded block's Linear4bit `weight.data` in pinned CPU RAM; a forward pre-hook stages a block's experts before it runs, and a single-resident-slot policy evicts the previously-staged block (no evict post-hook). Under use_reentrant=False gradient checkpointing the backward recompute re-runs the pre-hook, so exactly one block is GPU-resident in forward and backward alike — without depending on when PyTorch stops a recompute. gradient_checkpointing + use_reentrant=false is required (and enforced): it is what makes eviction correct and actually memory-freeing rather than pinning every staged weight alive as a matmul_4bit saved-for-backward tensor. Scope: single-GPU QLoRA only (refuses under FSDP/DeepSpeed/expert_parallel); per-expert Linear4bit layouts (Mixtral / Qwen2-3-MoE / OLMoE / DeepSeek). Complements layer_offloading (whole-layer) and is orthogonal to expert_parallel (shards experts across GPUs). Tests (tests/integrations/test_expert_offload.py): args + cross-field validation; expert discovery incl. PEFT base_layer unwrap; and the core correctness on CPU — offloaded gradients match the non-offloaded reference under gradient checkpointing (proves the recompute reads staged weights, not eviction placeholders), exactly one block resident, and state_dict home-substitution. A CUDA-gated pair repeats the grad-match against real bitsandbytes Linear4bit and checks install frees the experts off the GPU. Measured (OLMoE-1B-7B QLoRA, single 12GB RTX A2000, seed-matched A/B): peak VRAM 6.00 -> 2.60 GB (-57%), load footprint 4.70 -> 1.08 GB; convergence preserved (eval 1.6448 -> 1.2213 vs 1.6448 -> 1.2270). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Move all cross-field validation from the plugin's pre_model_load into a
pydantic model_validator on ExpertOffloadArgs (per review). merge_input_args
folds plugin args into the config by inheritance, so the validator sees the
merged config (same pattern as liger's tensor_parallel_size check) and fires
at config-parse time, before the model download. pre_model_load is gone.
One tightening this implies: use_reentrant must be EXPLICITLY false, since
parse-time validation runs before normalize_config defaults omitted kwargs
to {"use_reentrant": True}.
- install_expert_offload: raise RuntimeError instead of warn + silent no-op
when no 4-bit MoE expert blocks are found (per review), + regression test.
- tests: validate through a merged-shape subclass asserting ValidationError
(exercises the real mechanism); pytest.importorskip for bitsandbytes in the
CUDA class.
- _is_pinned: narrow the blind except to (RuntimeError, AssertionError).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
One GPU per replica: under DDP each rank homes its own pinned copy of the experts and stages to its own device. install_expert_offload now registers every offloaded expert weight on DDP's parameter ignore list (_ddp_params_and_buffers_to_ignore, read at DDP construction after post_model_load) so the initial module-state sync never touches the evicted 0-element placeholders; the frozen weights never enter gradient buckets. Validator keeps refusing FSDP/DeepSpeed/expert-parallel (they shard/move the same weights); docs note pinned CPU RAM scales with world size. Unit tests: ignore list covers exactly the offloaded weights, appends without duplicating, never captures trainable LoRA params. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…xperts) On current transformers, OLMoE/Qwen3-MoE store experts as fused 3D nn.Parameter stacks; plain load_in_4bit leaves those unquantized, so the shipped example OOMs a 24GB card resident and install_expert_offload finds no Linear4bit experts. Caught by running the example through axolotl train — the unit suite constructed Linear4bit experts directly and never exercised the real loader path. - discovery: find_parametrized_expert_stacks matches the Bnb4bitParametrization layout produced by quantize_moe_experts (bitsandbytes.nn.parametrize.replace_parameter_4bit); the packed parametrizations.<p>.original is homed/evicted exactly like a Linear4bit weight.data. The parametrization dequantizes under no_grad, so autograd never captures the packed tensor; checkpoint recompute re-stages it through the same pre-hook. - _BlockOffload now works on (param, owner, state-dict keys) slots shared by both layouts; the state-dict hook also covers bitsandbytes' renamed clean key. - DDP ignore list registers the parametrized packed params by identity, same as Linear4bit weights. - example yaml sets quantize_moe_experts: true (required for OLMoE); README documents both layouts; install error points at quantize_moe_experts instead of calling OLMoE a ModuleList model. - tests: CPU structural coverage (discovery, install/evict, recompute gradient parity, state-dict homes, DDP ignore) on a fake parametrized grouped model; CUDA end-to-end against real replace_parameter_4bit stacks. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…arly-stop leaks the global cache and retains every dequantized expert (pool x4 bytes) bitsandbytes' replace_parameter_4bit registers a forward_pre/forward hook pair that gates torch's GLOBAL parametrize cache by counter. use_reentrant=False checkpointing aborts the backward recompute mid-forward by design (early stop), skipping the post-hook of the module holding the last recomputed save: the counter leaks +1 per checkpointed region per step, the cache is never cleared again, and every dequantized bf16 expert weight is retained for the rest of training (~13 GiB OLMoE, ~53 GiB Qwen3-30B — 30B can never fit a 24 GB card). Expert params are single-access per forward, so the cache buys them nothing: install now removes the pair from offloaded expert modules (escape hatch AXOLOTL_EXPERT_OFFLOAD_KEEP_BNB_CACHE_HOOKS=1) and zeroes the global state. The state-dict quantization-state hook is preserved. Regression tests reproduce the torch-level hook skip with bnb-shaped stubs (no bnb import) and pin the stripper. Upstream bnb fix is register_forward_hook(..., always_call=True); to be filed separately.
|
Rebased onto current main (230fe76) and folded in the bitsandbytes parametrize-cache workaround: Re-validated post-rebase on GPU (RTX A5000): the full |
fd12f92 to
93757e6
Compare
3 participants
For #3787 (as requested there — a demonstrating PR).
Adds an
axolotl.integrations.expert_offloadplugin that streams a MoE's frozen 4-bit experts from pinned CPU RAM to the GPU one block at a time, so a model whose experts exceed VRAM can QLoRA-train on a single small card. Only the experts move — attention, router/gate, norms and the trainable LoRA adapters stay GPU-resident, so per-step PCIe traffic is limited to the experts (the bulk of a MoE's parameters, and the reason it doesn't fit).Measured (seed-matched A/B, OLMoE-1B-7B QLoRA, single 12 GB RTX A2000)
Peak VRAM −57%, load footprint −77%; both arms share a seed (identical data order), and the loss curves overlay — convergence is preserved. Throughput cost is the per-block H2D copy, ~+11% s/step uncontended on this card. Full telemetry (per-step JSONL,
wandb sync-able offline runs, charts, the A/B runner + methodology) is inab-telemetry/.How it works
Each expert is a bitsandbytes
Linear4bitwhose big tensor is the packedweight.data(thequant_statescales are ~1/32 the size and stay resident). Every offloaded block'sweight.datais homed in pinned CPU RAM; a forward pre-hook on the MoE block stages its experts just before it runs, and eviction is a single-resident-slot policy — staging a block evicts the previously-staged one; there is deliberately no evict post-hook. Underuse_reentrant=Falsegradient checkpointing, the backward recompute re-runs the pre-hook, so exactly one block's experts are GPU-resident in forward and backward — without depending on when PyTorch stops a recompute.Gradient checkpointing is required and enforced (not just recommended):
matmul_4bit's backward re-reads the packed weight viasave_for_backward, and eviction repointsweight.dataat a 0-element placeholder. The recompute is what makes eviction both correct and actually memory-freeing (instead of pinning every staged weight alive as a saved tensor). While evicted, modules hold only the 0-element placeholder, so a stray.to()or checkpoint-save never drags the experts back to the GPU; astate_dictpost-hook substitutes the CPU homes so full-model saves stay correct.Scope and guards
pre_model_loadrefuses under FSDP / DeepSpeed /expert_parallel, and requiresload_in_4bit+adapter: qlora+gradient_checkpointing: truewithuse_reentrant: false, each with an actionable error message.Linear4bitlayouts (expertsModuleList): Mixtral, Qwen2/3-MoE, OLMoE, DeepSeek-MoE, Jamba, … PEFTbase_layerwrappers are unwrapped. Fused-expert layouts (GPT-OSS, DBRX) aren't 4-bit-quantized per-expert and are left untouched.layer_offloading(whole-layer) — this moves only the experts, so it streams far fewer bytes per step for nearly the same VRAM win on a MoE. Orthogonal toexpert_parallel(multi-GPU sharding — the opposite regime). Chose the standalone-plugin shape to keep it self-contained; happy to refold it into thelayer_offloadingmixin instead if you'd prefer.Tests
tests/integrations/test_expert_offload.py— 23 tests (21 CPU + 2 CUDA-gated):base_layerunwrap, dedup, and non-4bit/singleton skipstate_dicthome-substitutionLinear4bit): install frees the experts' bytes off the GPU; offloaded grads match the non-offloaded reference through the realmatmul_4bitbackwardRun on a 12 GB RTX A2000: 23/23 pass. Example config in
src/axolotl/integrations/expert_offload/examples/.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
Developed with Claude Code (Anthropic) — implementation, tests, and the validation runs described in this PR were produced with AI assistance under the author's direction and review.