Skip to content

fix: release native MoE layer cycles after collection callbacks - #598

Closed
simone-chen wants to merge 1 commit into
simonec/vllm031-moe-situ-qualificationfrom
simonec/vllm031-moe-callback-gc-review
Closed

simone-chen wants to merge 1 commit into
simonec/vllm031-moe-situ-qualificationfrom
simonec/vllm031-moe-callback-gc-review

Conversation

@simone-chen

Copy link
Copy Markdown
Contributor
  • Native vLLM MoE layers can retain expert weights through bound-loader reference cycles after a collection callback returns. A paired B200 experiment reproduced accumulation through three original callbacks and an OOM on the fourth; post-return collection released each layer and completed all four full token sweeps.
  • Collect unreachable cycles once after a successful run_moe_torch return. Keep the native callback body, token sweep, timing, dispatch, case keys, exception propagation and multiprocessing callback identity unchanged.
  • Validate with 83 focused CPU tests covering cycle release, callback boundaries, failure propagation, argument forwarding, pickle identity, import compatibility and getter/EP contracts; Ruff and diff checks pass. Production GPU smoke remains pending on the separate pinned campaign source.
  • Stacked draft dependency: fix: preserve native MoE activation and routing contracts #589 at d67e3ce06fc27c1e7cb6c308c9ee5f2f103fc501. This PR adds only the cleanup, its tests and collector documentation. Retarget/rebase onto main after the dependency lands; do not merge into the component branch as campaign publication.
  • The source closure changes. Existing checkpoint loading alone is not permission to reuse old measurements; preserve their provenance and qualify the fresh four-callback smoke before continuation. Llama4 EP2 failures remain separate.

Signed-off-by: Simone Chen <simonec@nvidia.com>
@coderabbitai

coderabbitai Bot commented Oct 26, 2026

Copy link
Copy Markdown

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

@simone-chen

Copy link
Copy Markdown
Contributor Author

/ok to test e505429

@simone-chen

Copy link
Copy Markdown
Contributor Author

Superseded by consolidated recurrent/MoE compatibility PR #605, stacked on attention #604. The MoE cleanup is preserved once, including the duplicate in #603. Communication remains #591. Original branches and campaign commit references are preserved; closing without merging.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant