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
+import torch
 from transformers import AutoConfig
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
+from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL, set_pad_token_id
+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,16 @@ 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)
+        # The flash-attention version selects kernels, not modules, so it does not change the FLOPs.
+        model_config = AutoConfig.from_pretrained(
+            config.model.name,
+            attn_implementation="flash_attention_3",
+            trust_remote_code=config.tokenizer.trust_remote_code,
+        )
+        set_pad_token_id(model_config)
+        with torch.device("meta"):
+            model = AutoModelForCausalLMPrimeRL.from_config(model_config)
+        self.bin_cost = FlopCounter(model)
 
     def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
         return prepare_batch(

src/prime_rl/trainer/distributed/collectives.py

@@ -1,6 +1,5 @@
 """Autograd and compile-friendly distributed collectives."""
 
-import prime_kernels
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
@@ -100,6 +99,9 @@ def _mxfp8_all_to_all(
     group_name: str,
     quantized: bool,
 ) -> torch.Tensor:
+    # Imported here because prime-kernels is an optional linux-only extra.
+    import prime_kernels
+
     kernel = prime_kernels.load("mxfp8_moe")
     operation = kernel.all_to_all_dispatch if quantized else kernel.all_to_all_combine
     return operation(

src/prime_rl/trainer/model.py

@@ -23,7 +23,7 @@
 from torch.distributed.fsdp._fully_shard._fsdp_init import _get_mesh_info
 from torch.distributed.tensor import Shard
 from torch.distributed.tensor.parallel import parallelize_module
-from transformers import AutoConfig, AutoTokenizer, GenerationConfig, PretrainedConfig
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig
 from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.utils.import_utils import is_flash_attn_3_available
 
@@ -45,6 +45,7 @@
     cast_float_and_contiguous,
     get_custom_causal_lm_cls,
     get_custom_vlm_cls,
+    set_pad_token_id,
     supports_custom_impl,
 )
 from prime_rl.trainer.models.deepseek_v4.attention import DeepseekV4Indexer
@@ -356,40 +357,8 @@ def get_model(
         if subconfig is not None and hasattr(subconfig, "use_cache"):
             subconfig.use_cache = False
     model_config.dsa_backend = config.dsa_backend
-    # Auto-enable IndexShare from the model's own indexer schedule (e.g. GLM-5.2). The model
-    # reads `indexer_types` directly: shared layers reuse cached indices and carry no indexer weights.
-    indexer_types = getattr(model_config, "indexer_types", None)
-    if indexer_types and any(t == "shared" for t in indexer_types):
-        model_config.use_index_cache = True
-        logger.info(
-            f"Auto-enabled IndexShare from indexer_types schedule "
-            f"({sum(t == 'full' for t in indexer_types)}/{len(indexer_types)} full layers)"
-        )
 
-    # Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
-    # In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
-    if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
-        gen_config = GenerationConfig.from_model_config(model_config)
-        # Use `is not None` instead of truthiness: token ID 0 is valid.
-        pad_token_id = next(
-            (
-                v
-                for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)]
-                if v is not None
-            ),
-            None,
-        )
-        # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which
-        # crashes both huggingface_hub's strict setter and transformers'
-        # GenerationConfig.validate(). Unwrap before assigning.
-        if isinstance(pad_token_id, list):
-            pad_token_id = pad_token_id[0]
-        model_config.pad_token_id = pad_token_id
-
-    # Handle list pad_token_id that was already set on the config (not from our
-    # fallback above, but directly in the model's config.json).
-    if isinstance(getattr(model_config, "pad_token_id", None), list):
-        model_config.pad_token_id = model_config.pad_token_id[0]
+    set_pad_token_id(model_config)
 
     # NOTE: For VLM models, we do NOT propagate dtype to sub_configs.
     # The model should load in its default dtype (bf16) to match vLLM inference.

src/prime_rl/trainer/models/__init__.py

@@ -2,7 +2,7 @@
 
 from collections import OrderedDict
 
-from transformers import AutoConfig
+from transformers import AutoConfig, GenerationConfig
 from transformers.configuration_utils import PretrainedConfig
 from transformers.models.auto.auto_factory import _BaseAutoModelClass, _LazyAutoMapping, auto_class_update
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
@@ -97,6 +97,35 @@ def get_custom_causal_lm_cls(model_config: PretrainedConfig) -> type[PreTrainedM
     return _CUSTOM_CAUSAL_LM_BY_MODEL_TYPE[model_config.model_type]
 
 
+def set_pad_token_id(model_config: PretrainedConfig) -> None:
+    """Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
+
+    In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
+    """
+    if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
+        gen_config = GenerationConfig.from_model_config(model_config)
+        # Use `is not None` instead of truthiness: token ID 0 is valid.
+        pad_token_id = next(
+            (
+                v
+                for v in [gen_config.pad_token_id, gen_config.eos_token_id, getattr(model_config, "eos_token_id", None)]
+                if v is not None
+            ),
+            None,
+        )
+        # Some HF configs (e.g. Llama 3.2) set pad_token_id to a list, which
+        # crashes both huggingface_hub's strict setter and transformers'
+        # GenerationConfig.validate(). Unwrap before assigning.
+        if isinstance(pad_token_id, list):
+            pad_token_id = pad_token_id[0]
+        model_config.pad_token_id = pad_token_id
+
+    # Handle list pad_token_id that was already set on the config (not from our
+    # fallback above, but directly in the model's config.json).
+    if isinstance(getattr(model_config, "pad_token_id", None), list):
+        model_config.pad_token_id = model_config.pad_token_id[0]
+
+
 def supports_custom_impl(model_config: PretrainedConfig) -> bool:
     """Check if the model configuration supports the custom PrimeRL implementation.
 

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,15 +1,49 @@
+from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Protocol, runtime_checkable
 
-from torch import Tensor
+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
 
 CPStyle = Literal["ring", "ulysses"]
 ALL_CP_STYLES: frozenset[CPStyle] = frozenset({"ring", "ulysses"})
 
 
+@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]) -> int:
+    """Elements of `params` that a token multiplies through, as torchtitan's `get_nparams_and_active_nparams`.
+
+    Routed experts count by `top_k / num_experts`; embedding tables count only when tied to the lm_head.
+    """
+    params = {id(param): param for param in 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
+
+
 @dataclass(frozen=True)
 class CPSupport:
     """Which context-parallel styles an architecture can train under, and why not the rest."""
@@ -77,6 +111,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, self.parameters())
+
     @classmethod
     def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool:
         """