From 45733af4c887745b696ebce7cdd1ff412d3bca11 Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:56:20 +0900 Subject: [PATCH] fix(moe): unregister offloaded layers in OffloadReport.release() (#173) 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 --- opt_prime/opt_prime/accel/moe/api.py | 15 ++++++++++----- opt_prime/opt_prime/accel/moe/ops.py | 4 ++++ opt_prime/tests/test_moe_offload.py | 11 +++++++++++ 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/opt_prime/opt_prime/accel/moe/api.py b/opt_prime/opt_prime/accel/moe/api.py index 3e6bdebe..d612df5e 100644 --- a/opt_prime/opt_prime/accel/moe/api.py +++ b/opt_prime/opt_prime/accel/moe/api.py @@ -12,7 +12,7 @@ from opt_prime.accel.spec.adapters import adapter_for, moe_layers from opt_prime.accel.moe.config import MoEOffloadConfig from opt_prime.accel.moe.layer import OffloadedLayer -from opt_prime.accel.moe.ops import OffloadedExperts, register_layer +from opt_prime.accel.moe.ops import OffloadedExperts, register_layer, unregister_layer from opt_prime.accel.moe.planner import device_fraction, plan GiB = 1 << 30 @@ -27,9 +27,13 @@ class OffloadReport: expected_accelerator_fraction: Optional[float] = None #: opt_prime.accel.platforms leases of the accelerator tiers (role "experts"); held until release() leases: list = field(default_factory=list, repr=False, compare=False) + #: ids of the layers registered for the custom op (they hold the tier handles); unregistered by release() + layer_ids: list = field(default_factory=list, repr=False, compare=False) def release(self): - """release the device leases (after the offloaded model is no longer used).""" + """release the device leases and the registered layers (after the offloaded model is no longer used).""" + for lid in self.layer_ids: + unregister_layer(lid) for lease in self.leases: lease.release() @@ -38,7 +42,7 @@ def hot_experts(self) -> int: return sum(t["experts"] for t in self.tiers[:-1]) def as_dict(self): - d = asdict(dataclasses.replace(self, leases=[])) + d = asdict(dataclasses.replace(self, leases=[], layer_ids=[])) d["leases"] = [f"{l.device} gen {l.generation}" for l in self.leases] return d @@ -79,7 +83,7 @@ def apply_moe_offload(model: nn.Module, cfg: MoEOffloadConfig) -> OffloadReport: adapters = [adapter_for(i.module) for i in infos] specs = [a.spec(i.module) for a, i in zip(adapters, infos)] E = specs[0].num_experts - leases = [] + leases, lids = [], [] tiers, names = build_tiers(cfg, leases) try: if cfg.placement is not None: @@ -98,6 +102,7 @@ def apply_moe_offload(model: nn.Module, cfg: MoEOffloadConfig) -> OffloadReport: handles = [tier.load(a.weights(info.module, ids)) if ids else None for tier, ids in zip(tiers[:-1], pl[:-1])] handles.append(tiers[-1].load(a.weights(info.module, list(range(E))))) lid = register_layer(OffloadedLayer(E, tiers, handles, pl, names, cfg, adapter=a, source=info.module)) + lids.append(lid) parent, _, child = info.name.rpartition(".") parent_mod = model.get_submodule(parent) if parent else model setattr(parent_mod, child, OffloadedExperts(lid, E, [e for r in pl[:-1] for e in r])) @@ -108,7 +113,7 @@ def apply_moe_offload(model: nn.Module, cfg: MoEOffloadConfig) -> OffloadReport: for n, c in zip(names, per_tier)]) if cfg.profile is not None: rep.expected_accelerator_fraction = round(device_fraction(cfg.profile, placement, range(len(tiers) - 1)), 4) - rep.leases = leases + rep.leases, rep.layer_ids = leases, lids except BaseException: for lease in leases: lease.release() diff --git a/opt_prime/opt_prime/accel/moe/ops.py b/opt_prime/opt_prime/accel/moe/ops.py index f4d9e56f..4656570e 100644 --- a/opt_prime/opt_prime/accel/moe/ops.py +++ b/opt_prime/opt_prime/accel/moe/ops.py @@ -29,6 +29,10 @@ def register_layer(layer: OffloadedLayer) -> int: return lid +def unregister_layer(layer_id: int): + _LAYERS.pop(layer_id, None) + + def get_layer(layer_id: int) -> OffloadedLayer: return _LAYERS[layer_id] diff --git a/opt_prime/tests/test_moe_offload.py b/opt_prime/tests/test_moe_offload.py index 2650a071..eccb2700 100644 --- a/opt_prime/tests/test_moe_offload.py +++ b/opt_prime/tests/test_moe_offload.py @@ -368,3 +368,14 @@ def test_backward_through_offloaded_experts(family, frac): assert grads[1].keys() == {n for n in grads[0] if "experts." not in n} and grads[1] for n, g in grads[1].items(): torch.testing.assert_close(g, grads[0][n], atol=1e-5, rtol=1e-3, msg=n) + + +def test_release_unregisters_layers(): + """release() drops the model's layers (and their tier handles) from the op registry (#173)""" + from opt_prime.accel.moe import ops + m = build_tiny("qwen3_moe") + rep = apply_moe_offload(m, MoEOffloadConfig(hot=[[0, 1]] * len(find_moe_experts(m)))) + lids = [mod.layer_id for mod in m.modules() if hasattr(mod, "layer_id")] + assert lids and all(lid in ops._LAYERS for lid in lids) + rep.release() + assert not any(lid in ops._LAYERS for lid in lids)