feat: add training throughput, FLOPs and MoE balance metrics

[yulangz]

Description

Add training-side token throughput, model-aware estimated FLOPs/s, and per-layer MoE load diagnostics across FSDP, Megatron, and Archon. Training metrics report interval and cumulative rates using global work counts and maximum accumulated rank training time.

The FLOPs estimator registry supports custom model factories and includes Qwen3-30B-A3B and Qwen3.5-35B-A3B. Estimates account for sequence-length-dependent causal attention and Qwen3.5's linear attention, rather than assuming constant FLOPs per token.

CUDA/ROCm timing uses events on the captured training stream, with one device synchronization at export instead of two per training batch. CPU/NPU retain synchronized wall timing.

Dense model parts register no diagnostic hooks. TensorBoard reports scalar summaries without per-expert series.

MoE loads are accumulated across microbatches and relevant parallel ranks before normalization. Each layer reports maximum load divided by ideal load as a scalar; W&B receives expert counts and percentages in a separate moe_balance/expert_loads Table. Preserve all 10,240 Qwen3.5 expert rows in run-media previews by raising the SDK's default 10,000-row limit.

Include English and Chinese metric references and a standalone visualization guide with the tested heatmap JSON, historyTable field mappings, step selector instructions, and recovery guidance for older truncated previews.

Related Issue

No linked issue; implements the requested training observability improvements.

Type of Change

• 🐛 Bug fix
• ✨ New feature
• 💥 Breaking change
• 📝 Documentation update
• ♻️ Refactoring
• ⚡ Performance improvement
• ✅ Test coverage improvement

Checklist

• I have read the Contributing Guide
• Pre-commit hooks pass (pre-commit run --all-files)
• Relevant tests pass; new tests added for new functionality
• Documentation updated (if applicable; built with ./docs/build_all.sh)
• Branch is up to date with main
• Self-reviewed via /review-pr command
• This PR was created by a coding agent via /create-pr
• This PR is a breaking change

Additional Context

Validation:

• pre-commit run on all staged PR files: passed; Conventional Commit check passed.
• pytest -q tests/test_training_metrics.py -m 'not slow': 27 passed, 3 deselected.
• Earlier validation exercised real CUDA MoE routing with and without recomputation and two-GPU NCCL aggregation.
• A three-update Qwen3.5 validation used the configuration from feature/optimizer_offload with staged optimizer offload disabled and W&B enabled. All three updates finished; all three artifact Tables contained 40 layers × 256 experts with per-layer percentages summing to 100. A separate W&B upload verified that the fixed media preview retains all 10,240 rows. This integration validation used a separate checkout with the metrics overlaid on that branch.
• The remote pre-commit check passed for revision 929d79fb. Local checks use the PyPI index recorded in the lockfiles to avoid unrelated dependency re-resolution. Documentation formatting and embedded JSON were checked; the full documentation site was not built.
• Multi-node integration suites were not run; distributed validation used available GPUs on one node.

Limitations and review focus:

• FLOPs are analytical estimates, not hardware counters; they exclude recomputation, optimizer updates, elementwise operations, vision, and auxiliary MTP work.
• MoE counts measure executed routing assignments, including packing/alignment padding and Megatron's post-capacity/drop routing map. They can differ from logical throughput token counts. Shared experts and auxiliary MTP routers are excluded.
• CUDA event timing excludes side-stream work not joined before the end event; this is training-stream elapsed time rather than device-wide wall time.
• Review collective participation, parallel reductions, router hooks, and exclusion of backward recomputation. Unsupported routers do not emit MoE diagnostics.
• The preview fix affects future uploads. Existing truncated previews must be recovered from complete artifacts; server-side query or panel limits remain separate concerns.

[Copilot]

🟡 Changes recommended

A few concrete issues remain (unnecessary MoE hook attachment overhead on non-MoE models, overly verbose TensorBoard logging of per-expert scalars, and a likely broken docs link) that should be addressed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds training-side observability across AReaL engines (FSDP, Megatron, Archon) by tracking (1) token throughput, (2) model-aware estimated FLOPs/s, and (3) MoE routing balance diagnostics, and documents how to interpret and visualize the new MoE outputs in W&B.

Changes:

• Introduce a FLOPs estimator registry and per-sequence FLOPs estimators for Qwen3/Qwen3.5 MoE variants, plus training-time aggregation into train_perf/* metrics.
• Add MoE routing load collection/normalization and integrate it into engine training + stats export, including W&B Table logging for expert matrices.
• Add English/Chinese documentation for metric semantics and a W&B Custom Chart (Vega) walkthrough, and extend docs TOCs.

File summaries

File	Description
areal/utils/flops.py	FLOPs estimator registry + Qwen3/Qwen3.5 estimators.
areal/utils/training_metrics.py	Train-time token/time/FLOPs aggregation and export helpers.
areal/utils/moe_metrics.py	MoE routing count collection, reduction, normalization, and table-splitting helper.
areal/utils/stats_logger.py	Split per-expert metrics into a W&B Table and raise W&B run-media row limit.
areal/engine/fsdp_engine.py	Integrate training + MoE metrics collection/export for FSDP engine.
areal/engine/megatron_engine.py	Integrate training + MoE metrics collection/export for Megatron engine.
areal/experimental/engine/archon_engine.py	Integrate training + MoE metrics collection/export for Archon engine.
areal/experimental/models/archon/moe/moe.py	Add non-persistent routing_counts buffer for diagnostics.
tests/test_training_metrics.py	CPU + distributed tests for FLOPs, throughput, MoE aggregation, and W&B table serialization.
docs/en/reference/metrics_tracking.md	Document new train_perf/* and moe_balance/* metrics and conventions.
docs/zh/reference/metrics_tracking.md	Chinese version of metric documentation.
docs/en/reference/moe_visualization.md	W&B visualization guide and Vega spec for MoE expert loads.
docs/zh/reference/moe_visualization.md	Chinese visualization guide and Vega spec.
docs/en/_toc.yml	Add MoE visualization doc to English TOC.
docs/zh/_toc.yml	Add MoE visualization doc to Chinese TOC.

Review details

• Files reviewed: 15/15 changed files
• Comments generated: 5
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[yulangz]

Addressed the failures from both CI variants in 86bc9278:

• Isolate the metrics test's fake clock to the training-metrics module. Patching the shared time.perf_counter allowed other components to exhaust the finite timeline in the SGLang suite.
• Update the Megatron loss-scaling fixture to initialize _weight_residency and isolate unrelated metrics collection.
• Initialize mopd_execution_plan in the PPO evaluation-order fixture that bypasses __init__.
• Validate remotized shard contents against the local store where remotize() wrote them, instead of fetching from a separate RPC process with no corresponding shards.

The latter fixture/storage mismatches also exist in this branch's base. These changes only touch tests; production timing and training behavior are unchanged.

Local regression: 50 passed, 3 deselected, covering every failure from both CI variants. Full pre-commit was run, with formatting changes applied and the affected hooks rerun successfully. A new SGLang + vLLM CI run was requested for the updated branch; its result is pending.

[Le8r0nJames]

What guarantees that all ranks in reduce_group have the same ordered layer IDs and per-layer expert counts at every export?

A rank with no counts skips the collective, while equal-length tensors with different layer IDs could silently aggregate unrelated layers. Could we document the invariant and add coverage for mismatched rank-local state, with coordinated handling where needed?

[Le8r0nJames]

Is this callback reached inside a compiled region in any supported FSDP/Megatron configuration? If so, @torch.compiler.disable can introduce graph breaks or fail with fullgraph=True. The current fullgraph test covers Archon's attach_buffers path rather than this router-hook path.

Could we clarify the supported compilation modes and add coverage for this path where applicable?