feat(moe): shared MoE expert-routing stats (capture/reduce/scalars + router hook + CLI)#15
Open
ilsa001 wants to merge 5 commits into
Open
feat(moe): shared MoE expert-routing stats (capture/reduce/scalars + router hook + CLI)#15ilsa001 wants to merge 5 commits into
ilsa001 wants to merge 5 commits into
Conversation
hscos
reviewed
Jun 26, 2026
| tracker["tokens_per_expert"][idx] += routing_map.float().sum(dim=0) # [E] | ||
| tracker["weight_per_expert"][idx] += probs.detach().sum(dim=0) # [E] | ||
|
|
||
| for fn in _EXTRA_OBSERVERS: |
| if "tokens_per_expert" not in tracker: | ||
| device = routing_map.device | ||
| tracker["tokens_per_expert"] = torch.zeros( | ||
| num_layers, num_experts, device=device, dtype=torch.float32 |
|
|
||
| with torch.no_grad(): | ||
| idx = layer_number - 1 | ||
| tracker["tokens_per_expert"][idx] += routing_map.float().sum(dim=0) # [E] |
| for key in ("tokens_per_expert", "weight_per_expert"): | ||
| values = tracker[key] | ||
| # PP: different stages hold different layers (idle stages contribute 0). | ||
| dist.all_reduce(values, group=parallel_state.get_pipeline_model_parallel_group()) |
There was a problem hiding this comment.
这里如果有 pp 只有 dense的情况下，会有 pp 不掉用 save_expert_stats 的，这里会不会 hang 住？
…router hook + CLI) Adds a framework-agnostic expert-routing stats module consumed by BOTH mcore-trainer (SFT/CPT/DPO) and slime-trainer (RL): - megatron/core/transformer/moe/expert_stats.py: configure/is_enabled, save_expert_stats (1-based layer guard), increment_step_count, reduce_snapshot_and_clear (PP -> TP+CP -> DP-without-CP; CP counted once), collect_scalar_metrics, render_and_dump (lazy matplotlib, no hard dep), resolve_output_dir, register_extra_observer. - TopKRouter.routing(): flag-gated capture hook (stock-router callers get it automatically; full-override subclasses call save_expert_stats themselves). - arguments.py: _add_expert_stats_args (--expert-stats-*), discoverable by slime submit's structured megatron.* passthrough. Disabled by default (--expert-stats-log-interval 0). Verified: py_compile + 15 CPU unit tests (validation, 1-based guard, reduce group selection incl. with_context_parallel=False, cadence, scalars, path resolution, clear-on-all-ranks). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
In-tree coverage for the shared expert_stats module (addresses the review note that the 15 CPU tests were not committed): - TestExpertStatsLogic (CPU): configure negative-raise + env kill-switch, 1-based layer guard, save accumulate/noop, heatmap cadence, output-dir resolution, collect_scalar_metrics (CV/std/max/min/entropy). 9 tests, verified passing via the megatron pytest harness on CPU. - TestExpertStatsDistributed (CUDA-gated, model-parallel via Utils): reduce_snapshot_and_clear clears + normalizes; real TopKRouter.routing() hook fires only on the grad-enabled training forward, excluded under no_grad/eval and when disabled. The strong multi-topology single-count / CP-no-double-count assertion stays in the external GPU smoke (needs a controlled world size). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Mirrors the cosmos/slime feishu workflow so new comments on Megatron-LM fork issues are pushed to Feishu. Untrusted event fields pass via env (no shell injection), jq builds the JSON payload, and it skips gracefully until the FEISHU_WEBHOOK_URL secret is set. Runner: [self-hosted, hisys-megatron-runner]. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Synchronous helper returning the rendered MoE-layer token-count heatmap (RGB array or PNG-path fallback) so trainers can log it into their own sink (e.g. swanlab.Image) — megatron owns rendering, the trainer owns the sink. The async local-PNG path (render_and_dump) is unchanged. Unit test added. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
slime/sglang runtimes ship a stubbed `IPython` module that lacks the attributes matplotlib's interactive-framework probe reads (get_ipython / version_info, varying by matplotlib version), raising AttributeError mid-render so the expert-stats heatmap silently fell back to 'skipped'. Force the non-interactive Agg backend and hide IPython from sys.modules for the duration of the render (restore afterward) so matplotlib never probes it. No-op on trainers with a real IPython (mcore SFT already rendered fine). Verified: 32-GPU UBio-A3B-160E RL smoke now produces heatmap PNGs + swanlab.Image (slime) in addition to the JSONL. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
ilsa001
force-pushed
the
feat/expert-stats-moe
branch
from
July 2, 2026 23:15
10602f1 to
cd698ee
Compare
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
2 participants
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.
Summary
Adds a framework-agnostic MoE expert-routing stats module to the shared vendor, consumed by BOTH
mcore-trainer(SFT/CPT/DPO) andslime-trainer(RL) — single source of truth for capture/reduce/scalars; each engine owns its own logging sink.Changes (+516, 3 files)
megatron/core/transformer/moe/expert_stats.py(new):configure/is_enabled,save_expert_stats(1-based layer guard, skips MTP/out-of-range rows),increment_step_count,reduce_snapshot_and_clear(SUM over PP → TP+CP → DP-without-CP, so CP is counted exactly once),collect_scalar_metrics,render_and_dump(lazy matplotlib — not a hard megatron dep),resolve_output_dir,register_extra_observer(v2 per-source hook).TopKRouter.routing(): flag-gated capture hook. Stock-router callers (mcore) get it automatically; full-override subclasses (slimeSiRLTopKRouter) callsave_expert_statsthemselves.arguments.py:_add_expert_stats_args(--expert-stats-{log-interval,heatmap-interval,output-dir,per-layer-logging}), named so structured submit passthrough discovers theexpert_statsgroup.Safety
--expert-stats-log-interval 0): the hook is a single bool short-circuit; zero behavior change when off. Dense models never build a router → never fires.MEGATRON_DISABLE_EXPERT_STATS=1.Tests
py_compile(3 files) + 15 CPU unit tests: negative-interval raise, env kill-switch, 1-based layer guard (rejects 0 and >num_layers, writes last layer), reduce group selection (assertswith_context_parallel=False, no CP double-count), cadence, scalars (CV/std/max/min/entropy/per-layer), output-dir resolution, clear-on-all-ranks, non-logging-rank → None.GPU/distributed smoke (real collective, multi-rank) to follow.
Downstream / review gate
Foundational change. Engine PRs depend on this landing + a vendor stable tag:
scitix/slime-trainer#…(RL glue)scitix/mcore-trainer#…(SFT/CPT/DPO glue)Per the vendor-stable-branch contract: merge here → tag
scitix-megatron-sirl-stable-YYYY.MM.DD.N→ engines bump gitlink to the tag.Reviewers: @ilsa001 (integration lead).
🤖 Generated with Claude Code