[Dev] Refactor MoE loss logging

[yanring]

What does this PR do ?

main pr #3431

PR Design / Implementation Doc

TL;DR

Refactors MoE metric logging from ad-hoc global state into MoEMetricsTracker -- a structured, per-config tracker that centralizes metric collection, distributed reduction, aggregation, and logging. The singleton default preserves existing behavior; multi-model workloads can opt into isolation by assigning separate tracker instances to each TransformerConfig.

---

Problem Statement

MoE logging logic was spread across mutable global dictionaries and standalone utility functions (save_to_aux_losses_tracker, reduce_aux_losses_tracker_across_ranks, track_moe_metrics). The reduction semantics (which groups to reduce over, whether to DP-average) were implicit and easy to get wrong when adding new metrics. The global dict also made multi-model isolation impossible.

---

Goals

• Centralize the full MoE metric lifecycle (record, sync, aggregate, log, clear) in one component.
• Make distributed reduction semantics explicit and configurable per metric via needs_dp_avg.
• Enable per-model metric isolation through TransformerConfig without changing training-level call sites.
• Preserve legacy entry points (save_to_aux_losses_tracker, track_moe_metrics) as deprecated wrappers.

Non-Goals

• Change MoE loss formulas or optimizer/training semantics.
• Introduce new user-facing CLI flags.

---

High-Level Design

Forward pass (per layer)         Step end (once)
─────────────────────────        ──────────────────────────────────────
router.record(name, value,  ──►  tracker.report(loss_scale, iteration, ...)
  layer_number, num_layers,        ├─ _sync_metrics()     PP → reduce_group → avg_group → DP
  reduce_group, needs_dp_avg)      ├─ _aggregate()        per-layer values → scalar means
                                   ├─ _log_scalars()      TensorBoard / W&B
                                   ├─ _log_per_layer()    optional per-layer detail
                                   ├─ _format()           console log string
                                   └─ clear()             zero buffers for next step

Ownership model:

Layer	How it accesses the tracker	Isolation?
megatron/core/ (router)	self.config.moe_metrics_tracker	Per-config
megatron/training/ (training loop)	MoEMetricsTracker.get_instance()	Singleton

• Single model (default): TransformerConfig.__post_init__ sets moe_metrics_tracker to the global singleton. All paths point to the same instance -- behavior is identical to before.
• Multi model: User creates separate MoEMetricsTracker() instances and assigns them to each config.

---

Per-Metric Metadata (MetricEntry)

Each metric name maps to a MetricEntry dataclass:

Field	Type	Purpose
values	torch.Tensor	Per-layer accumulator (size = num_layers)
reduce_group	ProcessGroup?	Sum all-reduce group (e.g. TP+CP)
avg_group	ProcessGroup?	Average all-reduce group
needs_dp_avg	bool	Whether to average across DP after the above reductions

This replaces the old reduce_group_has_dp flag (inverted semantics: needs_dp_avg = not reduce_group_has_dp).

---

Reduction Semantics

For each metric, _sync_metrics applies reductions in this fixed order:

1. PP all-reduce (sum) -- collect across pipeline stages
2. reduce_group all-reduce (sum, optional) -- e.g. TP+CP group
3. avg_group all-reduce (avg, optional)
4. DP all-reduce (avg, if needs_dp_avg=True) -- average across data parallel ranks

Example: global_load_balancing_loss uses reduce_group=tp_dp_cp_group which already includes DP, so needs_dp_avg=False to avoid double-averaging.

---

Key File Changes

File	Change
moe_logging.py (new)	MoEMetricsTracker class with record(), report(), clear(), and private sync/aggregate/log helpers
transformer_config.py	Added moe_metrics_tracker field (defaults to singleton via lazy init in __post_init__)
router.py	Calls self.config.moe_metrics_tracker.record(...) instead of save_to_aux_losses_tracker(...)
training.py	Calls MoEMetricsTracker.get_instance().report(...) instead of track_moe_metrics(...)
cuda_graphs.py	Calls MoEMetricsTracker.get_instance().clear() instead of clear_aux_losses_tracker()
moe_utils.py	save_to_aux_losses_tracker, track_moe_metrics retained as deprecated wrappers forwarding to the tracker

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yanring]

/ok to test 9a8e019

[yanring]

/ok to test 1a18019

[yanring]

/ok to test 26a5025

[yanring]

/ok to test 5bf244f

[yanring]

/ok to test aa234a3

[Victarry]

/ok to test 25fb25d

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22707774878