feat: track Bridge-preferred MFU in SFT v2 and single-controller

[terrykong]

Summary

• Add MFU to SFT v2 and shared single-controller training, including the SFT v2 L1 check.
• Prefer Bridge FLOPs from actual sequence lengths and image/video grids. Keep the backend-agnostic fallback and log which calculator was used.
• Let SFT v2 use a plain tokenizer for text models, with regression tests for the vision path.

Accounting check

Ten steps per run; MFU below averages steps 2–10. Text: Qwen3-0.6B, 8 H100, GBS 16, 607 tokens/sample. Omni: Qwen2.5-Omni-3B image/text, 2 H100, GBS 2.

Run	Model TFLOPs/step	Mean MFU	W&B
Text, agnostic baseline	36.7575	1.60%	run
Text, Bridge SFT v1	36.7575	1.56%	run
Text, Bridge SFT v2	36.7575	1.54%	run
Omni, Bridge SFT v2	13.275–13.630	1.19%	run

All three text runs count exactly 36,757,522,612,224 FLOPs on every step. Every uploaded MFU recomputes exactly as FLOPs / policy seconds / GPU peak capacity; the MFU differences come from timing. Independently calculated Omni image work also matches Bridge: 2,514,773,606,400 FLOPs/image.

These pre-rebase accounting fixtures are not a throughput comparison. Audio is not validated.

Tests

• Rebased onto 4aaa48fab: 394 affected tests passed on H100, including eight real Bridge/Megatron-LM parity cases, packed lengths, frozen-model fallback, text/vision loaders, and shared-controller MFU.
• Actual SFT v2 L1: three steps on 2 H100s; checkpointing and all MFU/source assertions passed.
• Pinned Ruff/format checks and project pyrefly passed. Four HF-gated tokenizer tests deselected; full suite and hosted CI not yet verified.

100-step text comparison: strict match not yet established

SFT v1 / SFT v2: Qwen3-0.6B, 1 H100, BF16, GBS8/MBS1, identical policy settings and 800 distinct examples. All 100 input batches match exactly. FLOPs/MFU recomputation matches within TensorBoard float32 precision.

Both completed, but loss fails the predefined rtol=atol=0.001 check: max difference 0.005913, mean 0.001433. A diagnostic-only run normalizing before backward matches v1 to 1.1e-16, with identical gradient norms. The normal split path scales gradients afterward; BF16 rounding explains this difference. No gradient behavior changed in this PR. Plot and diagnosis (NVIDIA-internal).

Plain-tokenizer support still requires Energon, not a regular HF dataset backend.

[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.

[terrykong]

The validated text and image cases match the pinned calculators; the runtime-freezing gap is called out below. The leader and six lanes verified the one-value-per-DP-shard sum, the absence of a vision double count, timer windows that match the docs, and the Bridge and Megatron-LM API usage at the pinned submodules; the 8 mcore parity cases were supplied on H100. Two choices worth keeping: the pure function treats NotImplementedError as "the caller may fall back" and lets everything else propagate, and MFU is omitted instead of logged as zero when no calculator applies.

Merge blockers

1. The PR is marked CONFLICTING against main. All five conflicts are import blocks: single_controller.py, tq_policy.py, megatron_policy_worker.py, test_sft_v2.py, test_split_api_wrappers.py.
2. No CI lane has run. The only label is the auto-applied Documentation. Please add CI:L1, because the PR edits the L1-only tests/functional/sft_v2_energon.sh and the Megatron train path, and post /ok to test <sha> on the rebased head.
3. Two items will be red on the first CI run: the new config test at test_sft_v2.py L40 fails from the tests/ working directory and aborts the Algorithms shard (see the inline comment), and the lint status below.

Lint status

The repo pre-commit hooks were run in a clone of the head over the 29 changed files. end-of-file, trailing-whitespace, ruff (F/D), ruff-format, taplo and minimize-check passed. The ruff import-order hook failed on two PR files: nemo_rl/utils/flops_tracker.py L15-L18 (stdlib order: import math and import warnings belong above from dataclasses import asdict) and tests/unit/models/megatron/test_flops.py L186-L187 (a blank line is needed before the first-party import). The author will fix these mechanically. pyrefly could not run locally; the new module is listed in pyrefly.toml. The GitHub Lint job has not run.

Rebase heads-ups (owned by the author; not defects at this head)

• Energon-owned packing (#4105) on main ships one input_lengths entry per pack and keeps the per-conversation boundaries in cu_seqlens (origin/main packing.py L176-L188). After the rebase, compute_bridge_batch_flops must take its lengths from cu_seqlens when present, the way Bridge's own _real_subseq_lengths does; otherwise it squares the pack length instead of each conversation length.
• main's build_energon_sft_loader reads processor.tokenizer unconditionally, which breaks the plain-tokenizer path this PR adds.
• The Bridge bump to 1f8873bb leaves flop_utils.py byte-identical, so the Bridge calls here are unaffected.
• The #4052 telemetry imports and the #4198/#4207 DTensor v1 removal are the other conflicts.

What the evidence shows

The W&B table demonstrates identical text FLOPs and an exact MFU recomputation. It does not show identical performance or convergence. Audio is not validated. SFT v2 still requires Energon. A 100-step SFT v1/v2 convergence comparison is planned on the rebased head and is not done.

What we ran, CPU only, with the mcore cases filtered:

File	Result
tests/unit/algorithms/test_metric_utils.py	22 passed
tests/unit/algorithms/test_sft_v2.py	22 passed, from the repo root only
tests/unit/models/megatron/test_flops.py	20 passed, 8 mcore cases filtered
tests/unit/models/policy/test_split_api_wrappers.py	10 passed, 4 pre-existing NVTX failures on CPU-only torch
tests/unit/single_controller/test_single_controller_actor.py	62 passed
tests/unit/single_controller/test_utils.py	26 passed
tests/unit/utils/test_flops_tracker.py	34 passed
tests/unit/data/test_energon_sft.py, test_energon_sft_v2.py, tests/unit/models/policy/test_megatron_split_state.py	skipped at importorskip

The review team's CPU checks did not rerun the earlier H100 tests, and the full suite has not passed anywhere.

Note. No run exercised the single-controller GRPO/PPO MFU path, TP/PP>1, or the fallback path end to end. The in-tree Nemotron Omni SFT v2 recipe (vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.yaml, freeze_vision_model: true, listed in tests/test_suites/disabled.txt) logs no MFU today because the provider-level flag trips the frozen-model guard. Under the scoped fix proposed inline, its text-only batches would count again.

Generated by Claude Code

[terrykong]

tests/unit/algorithms/test_sft_v2.py:40
1 action item.
TL;DR — This new test opens examples/configs/... relative to the current directory, so it fails in the CI Algorithms shard and -x then aborts the whole shard.

PR-introduced: the test is new here. How it fails in CI:

1. tests/run_unit.sh L8 runs cd $SCRIPT_DIR, so every unit lane runs pytest from tests/. The Algorithms lane calls it at L0_Unit_Tests_Algorithms.sh L22.
2. load_config opens the path as given. There is no repo-root lookup.
3. The relative path resolves under tests/, where no examples/configs directory exists, so both parametrizations raise FileNotFoundError.
4. pyproject.toml L670 sets addopts = "... -x", so the first failure stops the Algorithms shard.

The repo precedent for reading these YAMLs from a test is test_recipes_and_test_suites.py L26-L27 and test_config_v2.py L54: both build the path from __file__.

Reproduction from the tests/ directory

Reproduced twice, by two independent runs, with the source unchanged. From tests/:

pytest unit/algorithms/test_sft_v2.py -k tokenizer_mode -o addopts=
# 2 failed: FileNotFoundError for tests/examples/configs/sft.yaml and sft_vlm_3B.yaml

From the repo root the same two cases pass, which is why they passed in the PR's own record.

AI-1

Build the path from the test file. This needs two edits in this file: add the import to the block at L17-L22, and replace the anchored line. It is not a one-click block because the import lives outside this anchor; applying the L40 edit alone would raise NameError.

# 1. imports (tests/unit/algorithms/test_sft_v2.py, L17-L22): add
from pathlib import Path

# 2. test body: replace L40 with
    config = load_config(
        Path(__file__).resolve().parents[3] / "examples" / "configs" / config_name
    )

[terrykong]

nemo_rl/models/megatron/flops.py:48
1 action item.
TL;DR — This guard only reads provider fields, but megatron_cfg.freeze_config freezes at runtime without setting them, so Bridge still bills a frozen vision tower 3x although it runs forward only.

PR-introduced. How it shows up, with megatron_cfg.freeze_config: {freeze_language_model: false, freeze_vision_model: true, freeze_vision_projection: false} on a Qwen2.5-VL SFT v2 run (all three keys are required, because Bridge's freeze() has no defaults):

1. setup.py L2229 calls model_module.freeze(**freeze_config). Bridge's freeze() only sets requires_grad = False; the provider fields keep their defaults of False.
2. This guard passes, _batch_flops returns a number, and the run logs flops_from_bridge=1.
3. Bridge's vit_flops_from_patch_stats multiplies the vision tower and merger by 3 without checking whether they train, so the frozen tower is overcounted.
4. sft.md L252 promises the opposite: partially frozen models "warn and use NeMo-RL's backend-agnostic tracker".

AI-1

Read frozen flags from both sources, and raise only when the frozen part would actually run: always for a frozen language model, and for frozen vision flags only when the batch carries image or video grid rows. A frozen vision tower does zero work on a text-only batch, because Bridge runs the ViT only if pixel_values is not None. Multi-file (this file, the worker call at megatron_policy_worker.py L1152, and tests), so not a one-click block.

--- a/nemo_rl/models/megatron/flops.py
+++ b/nemo_rl/models/megatron/flops.py
-from typing import TYPE_CHECKING, Any
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
@@
-def compute_bridge_batch_flops(
-    config: "ConfigContainer", data: BatchedDataDict[Any]
-) -> float:
+def compute_bridge_batch_flops(
+    config: "ConfigContainer",
+    data: BatchedDataDict[Any],
+    *,
+    freeze_config: Mapping[str, bool] | None = None,
+) -> float:
@@
     model = config.model
-    if any(
-        getattr(model, name, False)
-        for name in (
-            "freeze_language_model",
-            "freeze_vision_model",
-            "freeze_vision_projection",
-        )
-    ):
-        raise NotImplementedError("Bridge FLOPs for partially frozen models")
+    # NeMo-RL freezes VL towers two ways: provider fields on config.model
+    # (Omni recipes) and megatron_cfg.freeze_config applied at runtime via
+    # model.freeze(**freeze_config); only the first is visible on config.model.
+    frozen = {
+        name
+        for name in (
+            "freeze_language_model",
+            "freeze_vision_model",
+            "freeze_vision_projection",
+        )
+        if getattr(model, name, False) or (freeze_config or {}).get(name, False)
+    }
+    if "freeze_language_model" in frozen:
+        raise NotImplementedError("Bridge FLOPs for a frozen language model")
+    has_media = any(
+        _grid_rows(data.get(key)) > 0 for key in ("image_grid_thw", "video_grid_thw")
+    )
+    if frozen and has_media:
+        # A frozen vision tower runs forward only; Bridge bills it 3x.
+        raise NotImplementedError("Bridge FLOPs for frozen vision towers with media")

Helper, worker call, and the four CPU tests

Helper next to the function:

def _grid_rows(grid: Any) -> int:
    if isinstance(grid, PackedTensor):
        grid = grid.as_tensor()
    return 0 if grid is None else int(grid.numel())

Worker call:

--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
     def _batch_flops(self, data: BatchedDataDict[Any]) -> float | None:
         """Use Bridge, explicitly marking unsupported cases for driver fallback."""
         try:
-            return compute_bridge_batch_flops(self.mcore_state.cfg, data)
+            return compute_bridge_batch_flops(
+                self.mcore_state.cfg,
+                data,
+                freeze_config=self.cfg["megatron_cfg"].get("freeze_config"),
+            )

Tests, CPU with the mocked calculator: (a) freeze_config={"freeze_vision_model": True} plus a text-only batch returns 120 and the calculator is called once; (b) the same plus image_grid_thw raises NotImplementedError matching "frozen vision"; (c) freeze_config={"freeze_language_model": True} raises NotImplementedError matching "frozen language"; (d) provider field freeze_vision_model=True plus a text-only batch returns 120. Case (d) replaces test_frozen_model_explicitly_requests_fallback, which today expects a raise for all three provider flags on a text-only batch. The docs sentence at sft.md L252 then reads "frozen language models, or frozen vision towers on batches with media"; see the comment on that line.

Why the guard misses freeze_config, and the in-tree users

NeMo-RL freezes Megatron VL models two ways. Nemotron Omni's legacy keys are mapped onto config.model with setattr in _apply_multimodal_config, so the guard at L40-L48 sees those. policy.megatron_cfg.freeze_config is read at setup.py L2210 and applied as a pre-wrap hook; nothing is written to config.model. The two in-tree freeze_config users today are text-only GRPO recipes (grpo-qwen3.5-35ba3b, grpo-qwen3.5-9b); no vision work runs there, so nothing is miscounted yet.

Ruled-out fixes: requires_grad scan, raise on any freeze_config

Context — no action. Do not fix this with a requires_grad scan over model.parameters(). freeze_moe_router: true also clears requires_grad on router weights and would push every MoE run onto the fallback for a tiny FLOPs difference.

Raising on any truthy freeze_config is also wrong: the fallback tracker has no entry for Qwen3.5 or Qwen2.5-VL, so FLOPTracker.from_config raises ValueError and the tracker is None. That would remove MFU entirely from the two in-tree Qwen3.5 GRPO recipes once the legacy packing gate is removed (see the comment on L32), replacing a correct number with none.

[terrykong]

nemo_rl/models/megatron/flops.py:32
1 action item.
TL;DR — This PR removes the reason for the legacy not packing_enabled MFU gate but leaves the gate in place, so GRPO, GRPO-sync and PPO Megatron runs with packing still log no MFU.

Adjacent, pre-existing code that this PR makes stale. The gate and the worker formula it protected landed together in #2790 (commit 6ab16882a):

1. Base worker: megatron_policy_worker.py L1138-L1141 at the merge base skipped FLOPs under packing because flops_per_sample * gbs overcounted by the packing factor. This PR deletes that block.
2. Head: this function counts real per-sample input_lengths on the unpacked batch, as this docstring line says, and the fallback tracker is fed per-sample lengths too (lm_policy.py L877-L881). Neither producer overcounts under packing any more.
3. Driver: print_performance_metrics still reads sequence_packing.enabled and skips the MFU block when it is true.
4. Symptom: a Megatron GRPO, GRPO-sync or PPO recipe with policy.sequence_packing.enabled: true gets a valid total_flops from Policy.train and then no train_fp_utilization. Packing is the normal Megatron configuration for these recipes.

Scope: only the callers of print_performance_metrics (grpo.py L3999, L5855, grpo_sync.py L1364, ppo.py L2082, L3078). SFT v1, DPO and distillation compute MFU inline with no gate (sft.py L758-L763, dpo.py L855-L860, distillation.py L1108-L1113).

AI-1

Delete packing_enabled and the and not packing_enabled clause in nemo_rl/algorithms/utils.py L864-L867. num_ranks is set whenever total_flops is, so removing the gate alone is safe. Not a one-click block: the target file is outside this PR's diff. This pairs with the scoped frozen-model fix on L48; together they make the legacy loops report MFU for packed runs when a supported calculator supplies a total.

--- a/nemo_rl/algorithms/utils.py
+++ b/nemo_rl/algorithms/utils.py
-    packing_enabled = master_config.policy.get("sequence_packing", {}).get(
-        "enabled", False
-    )
-    if "total_flops" in train_results and not packing_enabled:
+    if "total_flops" in train_results:

[terrykong]

nemo_rl/utils/flops_tracker.py:265
1 action item.
TL;DR — This elif is unreachable at head and would stamp flops_from_bridge=1.0 on a non-Bridge total; the helper it replaced, _aggregate_megatron_flops_metrics, is now dead code with three live tests.

PR-introduced: this PR removed the only producer and the only consumer of a worker-side total_flops but kept the leftovers.

1. Base: lm_policy.py L936-L941 at the merge base called _aggregate_megatron_flops_metrics when a worker returned total_flops. This PR replaced that call with resolve_flops_metrics and kept the helper at lm_policy.py L70-L89.
2. No worker under nemo_rl/models/policy/workers/ returns total_flops any more. The Megatron worker now returns local_flops. grep -rn '"total_flops"' nemo_rl/models/policy/workers/ finds nothing.
3. So this branch never runs in-tree. If an out-of-tree worker did hit it, its total would be labelled flops_from_bridge: 1.0, which the docs define as "Bridge's formulas with real input lengths".
4. Three tests at test_flops_tracker.py L276-L334 still pass against the dead helper, which tells the next reader that Megatron workers still return total_flops.

AI-1

Delete this elif (L265-L270), the helper at lm_policy.py L70-L89, and the three tests plus their import at test_flops_tracker.py L21. resolve_flops_metrics then has exactly two states: local_flops on every result, or fallback. Multi-file, so not a one-click block. The deletion in this file:

--- a/nemo_rl/utils/flops_tracker.py
+++ b/nemo_rl/utils/flops_tracker.py
-    elif results and "total_flops" in results[0]:
-        # Older Megatron workers report a replicated global total.
-        return {
-            "total_flops": float(results[0]["total_flops"]),
-            "flops_from_bridge": 1.0,
-        }
     if fallback_flops is not None:
         return {"total_flops": fallback_flops, "flops_from_bridge": 0.0}

resolve_flops_metrics keeps its own five tests in the same file.

[terrykong]

nemo_rl/models/policy/lm_policy.py:944
1 action item.
TL;DR — Attaching train_elapsed_seconds whenever total_flops exists changes the MFU time denominator for legacy GRPO, GRPO-sync and PPO Megatron runs on models the NeMo-RL tracker knows.

PR-introduced. At the merge base this key was attached only on the no-tracker branch; print_performance_metrics prefers it over the driver policy_training timer, so train_fp_utilization for the same GRPO or PPO Megatron recipe changes across this PR with no user-side change. The denominator changed.

AI-1

Keep the base timing: attach the key only when no NeMo-RL tracker exists. Tracker-known models keep the driver timer as before. Tracker-less models, now Bridge-counted, get the worker time, as the old elif did.

Suggested change

	if "train_elapsed_seconds" in results[0]:
	if self.flops_tracker is None and "train_elapsed_seconds" in results[0]:

Base vs head trace, and which loops are affected

1. Base: lm_policy.py L924-L941 attached train_elapsed_seconds only in the elif for models without a NeMo-RL tracker. Tracker-known models such as Qwen2, Qwen3, Llama, DeepSeek-V3 and GLM never got it.
2. Head: this line attaches it whenever total_flops exists.
3. print_performance_metrics prefers that key over the driver policy_training timer. The driver timer wraps a blocking policy.train(...) (grpo.py L3584, which waits on the gather at lm_policy.py L911), so the two windows differ.
4. The numerator also changed for tracker-less models, which is the PR's stated purpose.

Scope: only the callers of print_performance_metrics (grpo.py L3999, L5855, grpo_sync.py L1364, ppo.py L2082, L3078). SFT v1, DPO and distillation divide by the driver timer directly (sft.py L758-L763, dpo.py L855-L860, distillation.py L1108-L1113), so they are untouched. SFT v2 and the single controller ignore this key.

[terrykong]

tests/unit/reference_configs/sft.yaml:30
1 action item.
TL;DR — Same comment fix as on examples/configs/sft.yaml L35: only run_sft_v2.py reads this key.

PR-introduced comment. This file is the reference twin that test_reference_configs_up_to_date compares against the exemplar, so keep the wording in sync.

AI-1

Suggested change

	use_processor: false # Set true for image, video, or audio inputs.
	use_processor: false # Read by run_sft_v2.py only: true for image, video, or audio inputs.

[terrykong]

docs/guides/sft.md:248
1 action item.
TL;DR — "Plain Hugging Face tokenizer and text conversations" can be read as "a Hugging Face dataset loader now works"; SFT v2 still requires an Energon-prepared dataset.

PR-introduced paragraph. SFTMegatronPolicyWorker.setup_sft_dataloader always builds the Energon loader, and setup_sft_v2 requires data.backend=energon. The paragraph never states that constraint.

AI-1

Suggested change

For text-only models, set `policy.tokenizer.use_processor=false`. The existing `hf_multimodal` adapter also accepts a plain Hugging Face tokenizer and text conversations. Media samples still require a multimodal processor; existing vision recipes keep using one by default.

For text-only models, set `policy.tokenizer.use_processor=false`. The training source is still an Energon-prepared dataset; text-only means the conversations carry no media, not that a Hugging Face dataset loader is supported. The existing `hf_multimodal` adapter also accepts a plain Hugging Face tokenizer and text conversations. Media samples still require a multimodal processor; existing vision recipes keep using one by default.

[terrykong]

docs/guides/sft.md:252
1 action item.
TL;DR — The unsupported-cases list omits audio batches, which the code rejects, and "partially frozen models" does not match the scoped frozen-model rule proposed for flops.py L48.

PR-introduced sentence. flops.py L57-L58 raises NotImplementedError for audio batches and the PR body says audio is not validated, but this list never names audio. The frozen-model rewording below assumes the scoped fix on flops.py L48 is applied: frozen language models always fall back, and frozen vision towers fall back only on batches with media.

AI-1

Suggested change

Megatron workers use Megatron-Bridge's formulas with real input lengths, including prompt tokens. For Qwen vision models, they add vision encoder and merger FLOPs from image/video grids, preserving image and frame boundaries. `flops_from_bridge=1` identifies this estimate. Explicitly unsupported cases, including partially frozen models, PEFT, and media without the required metadata, warn and use NeMo-RL's backend-agnostic tracker when available (`flops_from_bridge=0`). If neither calculator supports the model, or GPU peak capacity is unknown, MFU is omitted. Unexpected Bridge errors are not silently replaced by fallback estimates. Treat a calculator change as a change in measurement when comparing MFU across runs.

Megatron workers use Megatron-Bridge's formulas with real input lengths, including prompt tokens. For Qwen vision models, they add vision encoder and merger FLOPs from image/video grids, preserving image and frame boundaries. `flops_from_bridge=1` identifies this estimate. Explicitly unsupported cases, including frozen language models, frozen vision towers on batches with media, PEFT, audio batches, and media without the required metadata, warn and use NeMo-RL's backend-agnostic tracker when available (`flops_from_bridge=0`). If neither calculator supports the model, or GPU peak capacity is unknown, MFU is omitted. Unexpected Bridge errors are not silently replaced by fallback estimates. Treat a calculator change as a change in measurement when comparing MFU across runs.

[terrykong]

nemo_rl/models/policy/lm_policy.py:926
1 action item.
TL;DR — Policy.train's use of resolve_flops_metrics has no CPU unit test: the per-DP-shard sum, the gpus_per_worker capacity math, and the fallback when one shard reports local_flops=None.

PR-introduced path. resolve_flops_metrics is tested on its own at test_flops_tracker.py L33-L66, and the TQ finish_train_step path is tested at test_split_api_wrappers.py L133-L183. Nothing exercises this call site. The only tests that touch total_flops on this path are GPU-marked worker tests.

AI-1

Add a CPU test file at tests/unit/models/policy/test_lm_policy_flops.py with the NVIDIA copyright header. It is unmarked, so it lands in the Models_1..4 lanes. Not a one-click block: different file than this anchor. Both tests pass against the head source.

from unittest.mock import MagicMock

import pytest
import torch

from nemo_rl.models.policy.lm_policy import Policy


def _make_policy(results):
    p = Policy.__new__(Policy)
    p.cfg = {"train_global_batch_size": 8, "train_micro_batch_size": 2}
    shard = {"input_lengths": torch.tensor([7, 11])}
    p._shard_for_train = MagicMock(return_value=[shard, shard])
    p._report_sharded_payload = MagicMock()
    p.flops_tracker = MagicMock(total_flops=999.0)
    p.worker_group = MagicMock()
    p.worker_group.cluster.world_size.return_value = 4
    p.worker_group.get_all_worker_results.return_value = results
    return p


def _result(**extra):
    return {
        "global_loss": 1.0,
        "grad_norm": 0.5,
        "all_mb_metrics": {"loss": [0.1]},
        "gpu_name": "NVIDIA H100 80GB HBM3",
        "model_dtype": torch.bfloat16,
        **extra,
    }


def test_train_sums_bridge_flops_from_each_dp_shard():
    p = _make_policy([_result(local_flops=11.0), _result(local_flops=17.0)])
    out = p.train(MagicMock(), loss_fn=MagicMock())
    assert out["total_flops"] == 28.0
    assert out["flops_from_bridge"] == 1.0
    assert out["num_ranks"] == 4
    # Two DP-shard results stand in for four GPUs: 2 GPUs per result.
    assert out["theoretical_tflops"] == pytest.approx(4 * 989.5)
    assert out["all_mb_metrics"]["loss"] == [0.1, 0.1]


def test_train_falls_back_to_tracker_when_a_shard_is_unsupported():
    p = _make_policy([_result(local_flops=11.0), _result(local_flops=None)])
    with pytest.warns(UserWarning, match="unsupported"):
        out = p.train(MagicMock(), loss_fn=MagicMock())
    assert out["total_flops"] == 999.0
    assert out["flops_from_bridge"] == 0.0
    assert out["theoretical_tflops"] == pytest.approx(4 * 989.5)

[terrykong]

tests/unit/single_controller/test_single_controller_actor.py:2110
1 action item.
TL;DR — Single-controller MFU is tested only on the PPO path; the GRPO path, which reaches compute_mfu_metrics through the aggregate_step_metrics passthrough, has no end-to-end test.

PR-introduced path. This test uses _ppo_train_pump_controller only. On the GRPO path, total_flops, theoretical_tflops and flops_from_bridge reach compute_mfu_metrics through the passthrough at single_controller_utils/utils.py L70-L73. That passthrough is tested alone at test_utils.py L94-L100 but not through _train_pump, and no test checks that flops_from_bridge is logged by the controller.

AI-1

Add this test next to the anchored one. Not a one-click block: it adds a new function, so a suggestion on the existing test header would replace that test. It passes against head. The helpers _train_pump_controller, _OneThenEmptySampler and _single_group_meta already exist in this file.

def test_train_pump_logs_mfu_on_a_grpo_run(monkeypatch) -> None:
    """Off the PPO path, total_flops comes straight from aggregate_step_metrics."""
    meta = _single_group_meta()
    ctrl = _train_pump_controller(sampler=_OneThenEmptySampler(meta))
    ctrl._master_config.grpo.num_prompts_per_step = 1
    ctrl._trainer = MagicMock()
    ctrl._trainer.finish_train_step.return_value = {
        "loss": 1.0,
        "total_flops": 1e15,
        "theoretical_tflops": 500.0,
        "flops_from_bridge": 1.0,
    }
    ctrl._advantage_stage = AsyncMock(return_value=(meta, True))
    ctrl._sync_weights = AsyncMock(return_value=1)
    ctrl._logger = MagicMock()
    monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {})
    monkeypatch.setattr(
        ctrl._timer,
        "get_timing_metrics",
        lambda **kwargs: {"policy_training": 4.0, "total_step_time": 100},
    )

    asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0))

    train_metrics = ctrl._logger.log_metrics.call_args_list[0].args[0]
    assert train_metrics["total_flops"] == 1e15
    assert train_metrics["flops_from_bridge"] == 1.0
    assert train_metrics["train_fp_utilization"] == pytest.approx(0.5)

[terrykong]

/ok to test 200b2eb

[terrykong]

Addressed the submitted review in 200b2eb after rebasing onto 4aaa48f: runtime freeze guards, packed source lengths, legacy timing, dead helper removal, plain-tokenizer loader compatibility, CI working-directory paths, and policy/controller regression tests. The affected H100 suite passed 394 tests; the actual vision L1 wrapper and project pyrefly also passed. The PR description separately records the strict 100-step loss mismatch and its normalization diagnosis. Hosted L1 CI is running.