Repository navigation
fix: unregister offloaded MoE layers on release (#173) - #182
Merged
Merged
Conversation
The custom-op registry (ops._LAYERS) only grew: every OffloadedLayer, with its host expert tensors and accelerator slots, stayed reachable after release(), so reloading a model in a server kept the old experts in memory. apply_moe_offload now records the registered layer ids in the report, and release() pops them before releasing the leases. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
opt_prime.accel.moe.ops._LAYERSonly grows. EachOffloadedLayerholds the tier handles (host expert tensors, accelerator slots), andOffloadReport.release()freed the device leases but never the layers, so reloading a model in a server kept the old experts in host and GPU memory (#173: apply -> release -> del -> gc three times gave 2 -> 4 -> 6 registered layers).Changes
opt_prime/opt_prime/accel/moe/ops.py:unregister_layer(layer_id)pops a layer from the registry.opt_prime/opt_prime/accel/moe/api.py:apply_moe_offloadrecords the ids it registers in a newOffloadReport.layer_idsfield (repr=False, compare=False, left out ofas_dict()likeleases);release()unregisters them before releasing the leases.opt_prime/tests/test_moe_offload.py:test_release_unregisters_layers, which checks that the model's layer ids are registered afterapply_moe_offloadand gone afterrelease().Verification
CPU plus one A100 (
CUDA_VISIBLE_DEVICES=1), Python 3.10, torch 2.10.0+cu128, transformers 5.1.0, base optimus-rm @253be20.TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests/test_moe_offload.py -k release-> 1 failed (assert not any(lid in ops._LAYERS for lid in lids): the ids are still registered afterrelease()).pytest -q tests/test_moe_offload.py tests/test_moe_review.py tests/test_moe_train.py-> 116 passed, 2 skipped.opt_prime/:TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=1 PYTHONPATH=$PWD python3 -m pytest -q tests: before (@253be20, without the new test): 400 passed, 23 skipped, 135 subtests passed; after: 401 passed, 23 skipped, 135 subtests passed.build_tiny("qwen3_moe")with two cpu tiers) now prints 0 registered layers after each round (was 2 -> 4 -> 6).Compatibility and limits
OffloadReportgains one field with a default;as_dict()output is unchanged. Callers that never callrelease()behave as before.release(), the offloaded model can no longer run (itsOffloadedExpertsmodules refer to removed ids);release()was already documented as "after the offloaded model is no longer used".apply_moe_offloadare still left in the registry (the except path releases only the leases). Not part of this issue.Related Issue
Closes #173
🤖 Generated with Claude Code