refactor: MFU as torchtitan computes it; relative packing cost

src/prime_rl/orchestrator/batch.py

@@ -2,151 +2,22 @@
 import heapq
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
-from typing import Any
 
 import numpy as np
-from transformers import AutoConfig
+import torch
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
+from prime_rl.configs.trainer import ModelConfig
+from prime_rl.trainer.model import get_model
+from prime_rl.trainer.perf import FlopCounter
 from prime_rl.transports.batch.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, SamplingMask, TrainingSample
-from prime_rl.utils.logger import get_logger
 
 # Backfill value per component weight stream when a packed sample doesn't
 # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl
 # means no component (weight 0.0).
 STREAM_FILL = {"rl_weights": 1.0, "ce_weights": 0.0, "ref_kl_weights": 0.0}
 
 
-def _text_config(model_config: Any) -> Any:
-    """Unwrap a multimodal model's text config, else return the config unchanged."""
-    return getattr(model_config, "text_config", model_config)
-
-
-def _is_mla(config: Any) -> bool:
-    return bool(getattr(config, "multi_latent_attention", False) or hasattr(config, "q_lora_rank"))
-
-
-def _kv_channels(config: Any) -> int:
-    return getattr(config, "kv_channels", getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
-
-
-def _qkv_projection_flops_per_token(config: Any) -> int:
-    """Linear-in-seqlen FLOPs of the Q/K/V projections (per token)."""
-    hidden_size = config.hidden_size
-    num_attention_heads = config.num_attention_heads
-    kv_channels = _kv_channels(config)
-    is_mla = _is_mla(config)
-    qk_head_dim = getattr(config, "qk_head_dim", 0)
-    qk_pos_emb_head_dim = getattr(config, "qk_pos_emb_head_dim", 0)
-
-    if is_mla and getattr(config, "q_lora_rank", None) is not None:
-        q_flops = 2 * config.q_lora_rank * (hidden_size + num_attention_heads * (qk_head_dim + qk_pos_emb_head_dim))
-    else:
-        q_head_dim = (qk_head_dim + qk_pos_emb_head_dim) if is_mla else kv_channels
-        q_flops = 2 * hidden_size * num_attention_heads * q_head_dim
-
-    if is_mla and getattr(config, "kv_lora_rank", None) is not None:
-        v_head_dim = getattr(config, "v_head_dim", 0)
-        kv_flops = 2 * (
-            config.kv_lora_rank * (hidden_size + num_attention_heads * (qk_head_dim + v_head_dim))
-            + hidden_size * qk_pos_emb_head_dim
-        )
-    else:
-        num_query_groups = getattr(
-            config, "num_query_groups", getattr(config, "num_key_value_heads", num_attention_heads)
-        )
-        kv_flops = 4 * hidden_size * num_query_groups * kv_channels
-    return q_flops + kv_flops
-
-
-def _attention_flops_per_token_squared(config: Any) -> int:
-    """Quadratic-in-seqlen FLOPs of the attention scores and values (per token^2)."""
-    num_attention_heads = config.num_attention_heads
-    kv_channels = _kv_channels(config)
-    if _is_mla(config):
-        qk_head_dim = getattr(config, "qk_head_dim", 0)
-        qk_pos_emb_head_dim = getattr(config, "qk_pos_emb_head_dim", 0)
-        v_head_dim = getattr(config, "v_head_dim", kv_channels)
-        return num_attention_heads * (qk_head_dim + qk_pos_emb_head_dim) + num_attention_heads * v_head_dim
-    return 2 * num_attention_heads * kv_channels
-
-
-def _ffn_flops_per_token(hidden_size: int, ffn_hidden_size: int) -> int:
-    return 6 * hidden_size * ffn_hidden_size
-
-
-def _dense_ffn_hidden_size(config: Any) -> int:
-    return getattr(
-        config, "ffn_hidden_size", getattr(config, "intermediate_size", getattr(config, "moe_intermediate_size", 0))
-    )
-
-
-def _moe_ffn_hidden_size(config: Any) -> int:
-    """Effective FFN width of an MoE layer: routed experts (topk) plus shared experts."""
-    dense_ffn = _dense_ffn_hidden_size(config)
-    routed_topk = getattr(config, "moe_router_topk", getattr(config, "num_experts_per_tok", 1))
-    moe_ffn = getattr(config, "moe_ffn_hidden_size", getattr(config, "moe_intermediate_size", dense_ffn))
-    return moe_ffn * routed_topk + (getattr(config, "moe_shared_expert_intermediate_size", None) or 0)
-
-
-def _count_dense_and_moe_layers(config: Any) -> tuple[int, int]:
-    """Split the model's layers into (dense, moe) counts."""
-    num_experts = getattr(config, "num_experts", getattr(config, "n_routed_experts", None))
-    if num_experts is None:
-        return config.num_hidden_layers, 0
-
-    moe_layer_freq = getattr(config, "moe_layer_freq", None)
-    if isinstance(moe_layer_freq, list):
-        num_dense = sum(1 for freq in moe_layer_freq if freq == 0)
-        num_moe = sum(1 for freq in moe_layer_freq if freq > 0)
-    elif isinstance(moe_layer_freq, int):
-        num_dense = sum(1 for i in range(config.num_hidden_layers) if i % moe_layer_freq != 0)
-        num_moe = config.num_hidden_layers - num_dense
-    elif getattr(config, "first_k_dense_replace", None) is not None:
-        num_dense = config.first_k_dense_replace
-        num_moe = config.num_hidden_layers - num_dense
-    else:
-        num_dense = 0
-        num_moe = config.num_hidden_layers
-    return num_dense, num_moe
-
-
-def _packing_cost_coeffs(config: Any) -> tuple[int, int]:
-    """Return the (linear, quadratic) coefficients of the per-sequence forward FLOPs."""
-    hidden_size = config.hidden_size
-    num_dense_layers, num_moe_layers = _count_dense_and_moe_layers(config)
-    qkv_per_token = _qkv_projection_flops_per_token(config)
-    attn_per_token_squared = _attention_flops_per_token_squared(config)
-
-    def layer_linear(ffn_hidden_size: int) -> int:
-        return qkv_per_token + 2 * hidden_size * hidden_size + _ffn_flops_per_token(hidden_size, ffn_hidden_size)
-
-    linear = (
-        num_dense_layers * layer_linear(_dense_ffn_hidden_size(config))
-        + num_moe_layers * layer_linear(_moe_ffn_hidden_size(config))
-        + 2 * hidden_size * config.vocab_size
-    )
-    quadratic = (num_dense_layers + num_moe_layers) * attn_per_token_squared
-    return linear, quadratic
-
-
-def build_bin_cost(model_config: Any | None) -> Callable[[Sequence[int]], int]:
-    """Build a closure scoring a packed bin by estimated forward compute.
-
-    With ``model_config=None`` the linear/quadratic coefficients are ``(1, 0)``, so
-    the cost reduces to the token count and balancing falls back to sequence length.
-    """
-    if model_config is None:
-        linear, quadratic = 1, 0
-    else:
-        linear, quadratic = _packing_cost_coeffs(_text_config(model_config))
-
-    def bin_cost(seqlens: Sequence[int]) -> int:
-        return linear * sum(seqlens) + quadratic * sum(n * n for n in seqlens)
-
-    return bin_cost
-
-
 @dataclass
 class _WeightedSet:
     total: int = 0
@@ -869,17 +740,11 @@ def __init__(self, config: OrchestratorConfig):
         self.seq_len = config.seq_len
         self.num_train_workers = config.num_train_workers
         self.pad_to_multiple_of = config.pad_to_multiple_of
-        try:
-            model_config = AutoConfig.from_pretrained(
-                config.model.name, trust_remote_code=config.tokenizer.trust_remote_code
-            )
-        except Exception as e:
-            get_logger().warning(
-                f"Could not load model config for {config.model.name} ({e}) - "
-                "packing balances by token count instead of estimated FLOPs"
-            )
-            model_config = None
-        self.bin_cost = build_bin_cost(model_config)
+        # Built as the trainer builds it. The flash-attention version selects kernels, not modules.
+        model_config = ModelConfig(
+            name=config.model.name, attn="flash_attention_3", trust_remote_code=config.tokenizer.trust_remote_code
+        )
+        self.bin_cost = FlopCounter(get_model(model_config, device=torch.device("meta")))
 
     def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
         return prepare_batch(

src/prime_rl/trainer/models/afmoe/modeling_afmoe.py

@@ -18,6 +18,7 @@
     flash_attn_2_varlen_op,
     flash_attn_3_varlen_op,
     flash_attn_4_varlen_op,
+    quadratic_attention_flops,
 )
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput
 from prime_rl.trainer.models.layers.mlp import FeedForward
@@ -82,6 +83,15 @@ def __init__(self, config: AfmoeAttentionConfig):
         self.q_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
         self.k_norm = RMSNorm(RMSNormConfig(hidden_size=self.head_dim, eps=config.rms_norm_eps))
 
+    def attention_flops(self, seq_len: int) -> int:
+        return quadratic_attention_flops(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            attended_tokens=min(seq_len, self.sliding_window or seq_len),
+        )
+
 
 class AfmoeFlashAttention(AfmoeAttentionBase):
     """AFMoE attention using Flash Attention varlen functions."""

src/prime_rl/trainer/models/base.py

@@ -1,9 +1,47 @@
-from torch import Tensor
+from collections.abc import Iterable
+from typing import Protocol, runtime_checkable
+
+from torch import Tensor, nn
 from transformers.modeling_utils import PreTrainedModel
 
+from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.utils.cp import CPContext
 
 
+@runtime_checkable
+class AttentionFlops(Protocol):
+    """A sequence-mixing layer (softmax, sliding-window, sparse or linear attention, Mamba)."""
+
+    def attention_flops(self, seq_len: int) -> int:
+        """Training FLOPs of this layer's sequence mixing for one document of `seq_len` tokens."""
+        ...
+
+
+def count_active_params(model: nn.Module, params: Iterable[nn.Parameter] | None = None) -> int:
+    """Elements of `params` (default: all of the model's) that a token multiplies through.
+
+    Every parameter counts (attention and MLP projections, norms, routers, the lm_head) except the
+    routed experts a token skips (`1 - top_k / num_experts` of them) and embedding tables, which
+    count only when tied to the lm_head.
+    """
+    params = {id(param): param for param in (model.parameters() if params is None else params)}
+    lm_head_params = {id(param) for param in model.lm_head.parameters()}
+    inactive = 0
+    for module in model.modules():
+        if isinstance(module, nn.Embedding):
+            inactive += sum(
+                param.numel()
+                for param in module.parameters(recurse=False)
+                if id(param) in params and id(param) not in lm_head_params
+            )
+        elif isinstance(module, MoE):
+            # Expert weights lead with `num_experts`, so the division is exact.
+            num_experts, top_k = module.experts.num_experts, module.router.top_k
+            expert_params = sum(param.numel() for param in module.experts.parameters() if id(param) in params)
+            inactive += expert_params * (num_experts - top_k) // num_experts
+    return sum(param.numel() for param in params.values()) - inactive
+
+
 class PreTrainedModelPrimeRL(PreTrainedModel):
     """
     Base class for all PrimeRL models that extends HuggingFace PreTrainedModel.
@@ -66,6 +104,10 @@ def get_correct_experts_implementation(self, requested_experts: str | None) -> s
         """PrimeRL models always use eager experts implementation."""
         return "eager"
 
+    def num_active_params(self) -> int:
+        """Parameters a text token multiplies through, LoRA adapters included."""
+        return count_active_params(self)
+
     @classmethod
     def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool:
         """

src/prime_rl/trainer/models/deepseek_v4/attention.py

@@ -138,6 +138,7 @@
 )
 from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_rope import dsv4_q_norm_rope, dsv4_rope
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops
 from prime_rl.trainer.models.layers.grouped_linear import GroupedLinear
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
 from prime_rl.utils.cp import CPContext, gather_for_cp
@@ -717,6 +718,26 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
 
         self.cp_context = CPContext()
 
+    def attention_flops(self, seq_len: int) -> int:
+        # Each query attends the local window plus the compressed entries (CSA: the indexer scores
+        # all of them forward only, under no_grad; attention reads `index_topk`).
+        attended_tokens = min(seq_len, self.config.sliding_window)
+        indexer_flops = 0
+        if self.compressor is not None:
+            compressed_len = seq_len // self.compressor.compress_rate
+            indexer = getattr(self.compressor, "indexer", None)
+            if indexer is not None:
+                indexer_flops = 2 * indexer.num_heads * indexer.head_dim * seq_len * compressed_len
+                compressed_len = min(compressed_len, indexer.index_topk)
+            attended_tokens += compressed_len
+        return indexer_flops + quadratic_attention_flops(
+            num_heads=self.num_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            attended_tokens=attended_tokens,
+        )
+
     def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[torch.Tensor, None]:
         """`packed` carries the document boundaries every pathway below is clipped at."""
         # Shape keys in the comments below:

src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py

@@ -6,6 +6,7 @@
 
 from prime_rl.configs.trainer import DSABackend
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
+from prime_rl.trainer.models.layers.attn import quadratic_attention_flops
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
 from prime_rl.utils.cp import CPContext, gather_for_cp
@@ -149,6 +150,20 @@ def __init__(self, args: SparseMlaAttentionArgs):
 
         self.cp_context = CPContext()
 
+    def attention_flops(self, seq_len: int) -> int:
+        # MLA counted un-absorbed (qk = nope + rope); sparse attention attends `index_topk` tokens, and
+        # the indexer scores every token with its own heads, forward only (it runs under no_grad).
+        flops = quadratic_attention_flops(
+            num_heads=self.num_heads,
+            qk_head_dim=self.qk_head_dim,
+            v_head_dim=self.v_head_dim,
+            seq_len=seq_len,
+            attended_tokens=min(seq_len, self.args.index_topk),
+        )
+        if self.indexer is not None:
+            flops += 2 * self.args.index_n_heads * self.args.index_head_dim * seq_len * seq_len
+        return flops
+
     def mla_latents(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         q_latent = self.q_a_layernorm(self.q_a_proj(hidden_states))
         compressed_kv = self.kv_a_proj_with_mqa(hidden_states)

src/prime_rl/trainer/models/gpt_oss/attention.py

@@ -10,7 +10,7 @@
 
 from prime_rl.trainer.models.fusions import fuse_qkv_projections
 from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig
-from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func
+from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func, quadratic_attention_flops
 from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb
 
 
@@ -51,6 +51,15 @@ def __init__(self, config: GptOssConfig, layer_idx: int) -> None:
         nn.init.normal_(self.sinks, mean=0.0, std=config.initializer_range)
         self.flash_attn = flash_attn_4_varlen_func
 
+    def attention_flops(self, seq_len: int) -> int:
+        return quadratic_attention_flops(
+            num_heads=self.num_attention_heads,
+            qk_head_dim=self.head_dim,
+            v_head_dim=self.head_dim,
+            seq_len=seq_len,
+            attended_tokens=min(seq_len, self.sliding_window or seq_len),
+        )
+
     def compute_attention(
         self,
         query: torch.Tensor,