[faresobeid]

MFU in the trainer is computed the way torchtitan computes it, and the orchestrator packs micro-batches by the same FLOP count. Both config-parsing estimators (the trainer's get_active_mm_params and the orchestrator's Megatron-style _packing_cost_coeffs) are deleted.

One FLOP counter

FlopCounter(model) in trainer/perf.py is the single source of truth:

counter = FlopCounter(model)  # walks the model once
counter(seq_lens)             # Σ_docs  n · 6 · num_active_params  +  Σ_layers attention_flops(n)

• PreTrainedModelPrimeRL.num_active_params(): parameters a text token multiplies through, as torchtitan's get_nparams_and_active_nparams. The default walks the built model (count_active_params): routed experts (MoE.experts) count by top_k / num_experts (exact integer arithmetic, expert weights lead with num_experts), embedding tables count only when tied to the lm_head. Qwen3.5 overrides it to drop the vision encoder.
• AttentionFlops (runtime-checkable protocol in models/base.py): attention_flops(seq_len) = training FLOPs of a layer's sequence mixing for one document of seq_len tokens. The counter sums it per document, so a packed micro-batch is charged attention within each document, not over the packed length. On Python 3.12 the protocol check uses inspect.getattr_static, so attribute-forwarding wrappers (MultiLoRAModule, CheckpointWrapper, OptimizedModule) do not match and cannot double count.
• layers/attn.py has the two torchtitan rules as document totals: quadratic_attention_flops (6·heads·(qk+v)·seq_len·attended, no causal discount) and delta_rule_flops (18·v_heads·dk·dv·seq_len). Implementations:
  • FlashAttention (Llama, Qwen3, Qwen3-MoE, GLM-4.x, MiniMax, Laguna, Qwen3.5 full attention, Nemotron-H attention), AfmoeAttentionBase, GptOssAttention: quadratic, attended = min(seq_len, sliding_window) on sliding layers.
  • Qwen3.5 / Qwen3.8 GatedDeltaNet: delta rule. Nemotron-H Mamba2: 12·heads·head_dim·state_size·seq_len (state update x ⊗ B and readout state·C).
  • DeepSeek-V4 (as torchtitan): local window plus compressed entries; CSA layers add the indexer over the compressed length and read index_topk of them.
  • GLM-5.x DSA (MLA, qk = nope + rope, v = v_head_dim): attention reads min(seq_len, index_topk); the indexer term 2·index_heads·index_head_dim·seq_len² only on layers that own an indexer.
  • Qwen3.8 IndexedGatedAttention: indexer over seq_len // compression_ratio blocks, attention over min(seq_len, token_budget).
• LoRA: 4·(active − adapters) + 6·adapters per token (main froze every non-LoRA parameter in refactor(lora): single-adapter LoRA layers, drop modules_to_save #3855). Adapters go through count_active_params, so MoE expert adapters count by top_k / num_experts like the base experts.

MFU

The RL and SFT trainers sum perf_counter.flop_counter(seq_lens) over their micro-batches, all-reduce it with the token count over the dp mesh, and pass it to get_step_mfu. Log keys are unchanged.

Example of the per-document attention: GLM-4.5-Air at seq_len = 131072 (2x8 H200, cp4 ulysses, EP8, packed SWE samples, ~15 documents per micro-batch) logged MFU 127.7% / 152.3% on main; the same steps give 24.4% / 29.0%.

FLOPs/token for one document (main vs this PR; torchtitan = its get_nparams_and_flops on the matching flavor, built on meta):

model	seq_len	main	this PR	Δ	torchtitan
Qwen3-8B	8k	5.9904e+10	5.9906e+10	+0.0%	5.9906e+10 (exact)
Qwen3-8B	64k	1.6137e+11	1.6137e+11	+0.0%	1.6137e+11 (exact)
Qwen3-30B-A3B	8k	3.7577e+10	3.7579e+10	+0.0%	3.7579e+10 (exact)
Qwen3-30B-A3B	64k	1.7287e+11	1.7287e+11	+0.0%	1.7287e+11 (exact)
GLM-5.3-BF16	8k	3.4925e+11	2.7461e+11	-21.4%	n/a
GLM-5.3-BF16	64k	1.2286e+12	2.8447e+11	-76.8%	n/a
Qwen3.5-9B	8k	5.6032e+10	5.1068e+10	-8.9%	5.1068e+10 (exact)
Qwen3.5-9B	64k	1.4623e+11	7.3616e+10	-49.7%	7.3616e+10 (exact)
Qwen3.5-35B-A3B	8k	2.9853e+10	2.1988e+10	-26.3%	2.1988e+10 (exact)
Qwen3.5-35B-A3B	64k	1.4260e+11	5.0174e+10	-64.8%	5.0174e+10 (exact)
DeepSeek-V4-Flash	8k	2.5126e+11	8.7221e+10	-65.3%	8.8630e+10 (indexer at 6x)
DeepSeek-V4-Flash	64k	1.2208e+12	9.5676e+10	-92.2%	1.0695e+11 (indexer at 6x)
gpt-oss-20b	8k	3.1307e+10	2.6557e+10	-15.2%	n/a
gpt-oss-20b	64k	9.8952e+10	6.0380e+10	-39.0%	n/a
NVIDIA-Nemotron-3-Nano-30B-A3B-BF16	8k	5.8477e+10	2.1927e+10	-62.5%	n/a
NVIDIA-Nemotron-3-Nano-30B-A3B-BF16	64k	2.0504e+11	3.8839e+10	-81.1%	n/a
Trinity-Mini	8k	2.9891e+10	2.4026e+10	-19.6%	n/a
Trinity-Mini	64k	1.2009e+11	4.6574e+10	-61.2%	n/a

• Dense and plain-MoE models (Qwen3) are unchanged and match torchtitan exactly (the tiny difference is norm weights, which torchtitan also counts).
• Hybrid models go down: linear-attention, Mamba and sliding-window layers are no longer counted as full attention.
• GLM-5.3: sparse attention reads 2048 tokens and the indexer runs on the 21 IndexShare full layers. Active params 40.3B (main: 37.3B, it missed the shared expert and the indexer weights).
• DeepSeek-V4: main treated every layer as full attention and counted 18.8B active params (13.3B here).

Packing

BatchPacker builds the model once on the meta device with the trainer's own get_model(..., device=meta), so it gets the same modules as the trainer (IndexShare, pad token), and uses its FlopCounter as the bin cost: ranks are balanced against the cost the trainer reports. Measured: build 0.1–0.4 s and counter construction ≤ 0.02 s for GLM-5.3, DeepSeek-V4-Flash/Pro, Qwen3.5-35B-A3B, GLM-4.5-Air and Qwen3-30B-A3B configs; one bin-cost call over 16 documents takes 0.2–2.6 ms. The orchestrator already imported prime_rl.trainer.model, so import time is unchanged.

The orchestrator runs in the trainer's environment, so if the model cannot be built (no config, or an architecture without a PrimeRL implementation, which the trainer rejects too), it fails at startup instead of falling back to token-count balancing.

Tests

tests/unit/train/test_perf.py builds three architectures on meta from small configs (dense Qwen3, GLM-MoE-DSA with IndexShare, Qwen3.5 GDN) and checks the attention term against a hand-computed value, and checks the dense Qwen3 active-parameter count (attention and MLP projections, norms, lm_head) against a hand count.

End to end

1x8 H200, Qwen3-30B-A3B-Instruct-2507 RL on reverse-text (6 trainer GPUs, ep=2, 2 inference GPUs, seq_len 4096, batch 256x~300 tokens): the orchestrator built the meta model and packed with the FlopCounter; all 15 steps completed (reward 0.28 -> 0.45). Steps 2-15: 12-19 s, 3.8-7.3k tokens/s, MFU 1.2-2.4%. The MFU matches throughput x 1.83e10 FLOPs/token over 6 x 989 TFLOPS (e.g. step 15: 6040 tok/s -> 1.9%); it is low because this batch is tiny for a 30B model on 6 GPUs.

---

Note

Medium Risk
Changes MFU reporting and micro-batch load balancing for hybrid/sparse models; orchestrator now requires a buildable PrimeRL model at startup instead of degrading to token balancing.

Overview
Replaces separate HuggingFace-config FLOP heuristics in the orchestrator and trainer with a single FlopCounter built from the actual PrimeRL model on meta.

Trainer MFU now sums per-document FLOPs from packed seq_lens (matmul term from num_active_params / MoE-aware count_active_params, plus each layer’s attention_flops) and feeds that total into get_step_mfu in RL and SFT instead of a fixed flops-per-token derived from config.

Model hooks: AttentionFlops protocol and attention_flops(seq_len) on flash, sliding-window, sparse/MLA (GLM DSA, DeepSeek V4, Qwen3.8 indexed), delta-rule, and Mamba layers; shared helpers quadratic_attention_flops / delta_rule_flops in layers/attn.py.

Orchestrator packing: BatchPacker drops AutoConfig + build_bin_cost and uses the same FlopCounter(get_model(...)) as bin cost so DP balancing matches trainer compute (no silent fallback to token count).

Cleanup: Large PerfCounter config parsers removed; unit tests use injectable bin_cost stubs and meta-model FlopCounter checks.

^{Reviewed by Cursor Bugbot for commit a306149. Bugbot is set up for automated code reviews on this repo. Configure here.}

[garrett361]

Thanks for this, definitely good to clean up all the MFU computations. I have a proposal for a
slightly different design, lmk what you think? Also left some small comments throughout.

Proposed design: one FLOP counter

In this PR there are two FLOP models for the same computation. perf/mfu uses the trainer's counter, 6 * active_params plus per-module attention hooks:

prime-rl/src/prime_rl/trainer/perf.py

Lines 134 to 145 in 0cf50b1

	def _get_matmul_flops_per_token(self) -> int:
	"""6 FLOPs per active matmul parameter, as torchtitan's `get_nparams_and_flops`."""
	active_params = get_active_params(self.model)
	if not has_lora_layers(self.model):
	return 6 * active_params
	# LoRA: frozen base matmuls still cost forward (2x) and dX (2x); fully trainable non-LoRA
	# params (modules_to_save) add their weight gradient (2x); adapters cost 6x.
	return (
	4 * active_params
	+ 2 * self._count_fully_trainable_params_excluding_lora()
	+ 6 * self._count_lora_adapter_params()
	)

while packing uses 24 * L * d^2 plus a full-attention pair term:

prime-rl/src/prime_rl/orchestrator/batch.py

Lines 41 to 42 in 0cf50b1

	token_cost = 24 * config.num_hidden_layers * config.hidden_size**2
	pair_cost = num_attention_layers * config.num_attention_heads * (qk_head_dim + v_head_dim)

Example: for GLM-5.x they disagree in a way that plausibly matters. MFU charges DSA layers only min(n, index_topk) plus the indexer, while packing charges all 78 layers a full n^2 (the config has no layer_types). At long context that gap could be significant, and ranks get balanced against a cost the trainer itself says is wrong.

Proposal: a single source of truth that both MFU and packing use.

class PreTrainedModelPrimeRL:
    def num_active_params(self) -> int:
        """Matmul params used per token; the default is the module walk in `get_active_params`."""


@runtime_checkable
class AttentionFlops(Protocol):
    def attention_flops(self, seq_len: int) -> int:
        """Training FLOPs of this layer's sequence mixing for one document of `seq_len` tokens."""


class FlopCounter:
    def __init__(self, model: PreTrainedModelPrimeRL):
        self.matmul_flops_per_token = 6 * model.num_active_params()
        self.attention_modules = [m for m in model.modules() if isinstance(m, AttentionFlops)]

    def __call__(self, seq_lens: Sequence[int]) -> int:
        return sum(
            n * self.matmul_flops_per_token + sum(m.attention_flops(n) for m in self.attention_modules)
            for n in seq_lens
        )

• Models override num_active_params for their own quirks (vision encoder, tied embeddings) instead of special cases in perf.py.

• FlopCounter walks the model once at construction; packing calls the bin cost many times per step, so the walk must not repeat per call.

• Each sequence-mixing layer (full, sliding, or linear attention, Mamba, DSA) describes one document, and the counter sums over documents. This is the PR's attention_flops_per_token hook, returning the document total instead (so no factor of n to divide out and multiply back), as a typed protocol. On Python 3.12, isinstance against a runtime protocol uses inspect.getattr_static, so attribute-forwarding wrappers (MultiLoRAModule, CheckpointWrapper) don't match and can't double count.

• Packing: the orchestrator builds the model on the meta device once, as you did for the FLOPs table, and uses a FlopCounter as its bin cost. That removes the config parsing entirely. The price is importing the model code into the orchestrator plus the meta-build time; do you think that's acceptable?

  # orchestrator/batch.py
  from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL
  from prime_rl.trainer.perf import FlopCounter
  
  
  class BatchPacker:
      def __init__(self, config: OrchestratorConfig):
          ...
          model_config = AutoConfig.from_pretrained(config.model.name, trust_remote_code=config.tokenizer.trust_remote_code)
          with torch.device("meta"):
              model = AutoModelForCausalLMPrimeRL.from_config(model_config)
          self.bin_cost = FlopCounter(model)

  A side benefit: importing prime_rl.trainer.models runs its AutoConfig.register calls, so models whose configs aren't in transformers (e.g. DeepSeek-V4, Laguna) would no longer fall back to token-count balancing.

• LoRA is left out of the sketch; its matmul term would stay as in the PR (trainer side only). Note the PR's LoRA term currently overcounts MoE expert adapters: get_lora_param_counts includes every expert's slice, while the base experts are scaled by top_k / num_experts:

  prime-rl/src/prime_rl/trainer/perf.py

  Lines 139 to 145 in 0cf50b1

  	# LoRA: frozen base matmuls still cost forward (2x) and dX (2x); fully trainable non-LoRA
  	# params (modules_to_save) add their weight gradient (2x); adapters cost 6x.
  	return (
  	4 * active_params
  	+ 2 * self._count_fully_trainable_params_excluding_lora()
  	+ 6 * self._count_lora_adapter_params()
  	)

• Worth checking with tests: every registered model built on meta has a positive attention term that matches a hand-computed value, and no hooked module is nested inside another. Up to you whether those get committed.

[garrett361]

nit: I find the token_cost and pair_cost naming somewhat confusing.

Also, when do we expect model_config is None? Know that is pre-existing.

[garrett361]

This seems like a worse estimate than what we had previously? 24 * n_layers * hidden_size ^2 rather than 6 * n_active_params?

Also disagrees w/ the usage of 6 * n_active_params in the actual reported MFU computation. Should unify the bin cost and MFU computations w/ a single source of truth.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit b30c5dd. Configure here.}

[garrett361]

Was this change just for local testing?

[garrett361]

nit: maybe remove the direct torchtitan reference. Will go stale at some point and one more thing to maintain.

[garrett361]

don't think we need to be adding more checks here? Seems unrelated

[garrett361]

why did this logic change?

[garrett361]

same nit: drop the torchtitan reference please

[garrett361]

drop torchtitan ref

[garrett361]

drop torchtitan ref

[garrett361]

Is this again just for local testing? Is this protecting a scenario we'd expect to hit in prod? If not, I suggest dropping this.

[garrett361]

Just to reduce maintenance burden

[garrett361]

why do we need this for the flops counter?

[garrett361]

Are we missing the attention params? Thought we'd need another if case for nn.Linear, e.g.

[garrett361]

Could also fold those dense attn matmul costs into what attention_flops reports, but IMO more natural to include in active param count

[garrett361]

model signature is redundant? Can get params from model.paramters()? Single source of truth

[garrett361]

Oh, looks like this is for lora support and special casing? Would suggest making params optional and if None defaults to model.parameters()

[garrett361]

Think this is looking good and almost there, can you please address the new comments questions and verify this does work e2e?

[garrett361]

LGTM, just one small edit to indexer accounting, please, then can merge

[garrett361]

Bit a nit, but a detail that can seem to matter at long ctx: all indexers are run w/ no_grad, so the 6x from fwd + bwd is an overestimate which can significantly inflate MFU.

Claude analysis below. Applies to all models which use indexers.

---

GLM-5.x MFU overcounts the DSA indexer by 3x

The trainer's FLOP counter charges the indexer's scoring at the training rate of 6 FLOPs per multiply-add:

prime-rl/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py

Line 164 in a306149

flops += 6 * self.args.index_n_heads * self.args.index_head_dim * seq_len * seq_len

But the indexer runs under @torch.no_grad(), so it has no backward pass and the factor should be 2:

prime-rl/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py

Lines 62 to 63 in a306149

	@torch.no_grad()
	def compute_sparse_indices(

Size

For a document of $N$ tokens, GLM-5.3's FLOPs per token are $F(N) = F_\text{rest} + F_\text{idx}(N)$, with $F_\text{idx}(N) = 6 \cdot 32 \cdot 128 \cdot N \cdot 21$ (32 indexer heads of dim 128, on the 21 of 78 layers that own an indexer). The PR's FLOPs table gives $F_\text{rest} = 2.73 \times 10^{11}$. The overcount is $\tfrac{2}{3} F_\text{idx}$, so true MFU is reported MFU times $1 - \tfrac{2}{3} F_\text{idx} / F$:

Document length	Indexer share of counted FLOPs	True MFU / reported MFU
8k	1.5%	0.990
32k	5.8%	0.961
64k	11.0%	0.927
128k	19.8%	0.868

To correct a reported MFU, multiply by the last column: at 128k, a reported 50% MFU is really $50% \times 0.868 \approx 43.4%$. The error grows with document length, not seq_len, so it matters for long documents, not for micro-batches packed with short samples.

Fix

Use 2 instead of 6 at L164. The DeepSeek-V4 and Qwen3.8 indexers also run under no_grad and need the same change.