Don't compute the logprobs and their gradients for tokens no loss com…

docs/scaling.md

@@ -143,15 +143,26 @@ State-only optimizer offload remains enabled by default with `model.optim_cpu_of
 
 ### LM Head Chunking
 
-The vanilla LM head materializes a `[batch * seq, vocab]` logits tensor on every step — a major memory tax when the vocabulary is large (often >100K). `fused_lm_head_token_chunk_size` swaps in a custom fused linear + logprob/entropy kernel that streams through `chunk_size` tokens at a time, avoiding the materialization. It defaults to `1024` for RL training:
+The vanilla LM head materializes a `[batch * seq, vocab]` logits tensor on every step — a major memory tax when the vocabulary is large (often >100K). `fused_lm_head_token_chunk_size` swaps in a custom fused linear + logprob/entropy kernel that streams through `chunk_size` tokens at a time, avoiding the materialization:
 
 ```toml
 [trainer.model]
-fused_lm_head_token_chunk_size = 1024       # default
+fused_lm_head_token_chunk_size = 8192       # default
 # fused_lm_head_token_chunk_size = "disabled"  # vanilla LM head
 ```
 
-Drop the chunk size further when peak memory is still tight (e.g. with very long sequences); raise it to amortize kernel-launch overhead. SFT training silently disables this (not supported yet). Only available with `model.impl = "custom"`.
+Drop the chunk size when peak memory is still tight (e.g. with very long sequences); raise it to amortize kernel-launch overhead. Only available with `model.impl = "custom"`.
+
+### Skipping Masked Tokens in the LM Head
+
+The loss computations typically only depend on a small subset of tokens: prompts, environment observations, and pack padding aren't scored by any loss. `skip_masked_lm_head_tokens` (on by default) runs the LM head only on the tokens some loss component actually reads. Requires the fused head; ignored with `fused_lm_head_token_chunk_size = "disabled"`, and with `enable_token_export`, which records logprobs for every token of a sequence.
+
+```toml
+[trainer.model]
+skip_masked_lm_head_tokens = true   # default
+```
+
+Head time and its activation memory then scale with the kept fraction, which `perf/lm_head_token_fraction` reports each step. The head is roughly 10% of a dense 8B step and a much larger share when the vocabulary is big relative to the backbone, so the end-to-end win is largest for small models, long prompts, and observation-heavy agentic rollouts.
 
 ## Memory-Tight Recipe
 

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -261,7 +261,10 @@ class ModelConfig(BaseModelConfig):
     """Debugging knobs for the model and distributed training."""
 
     fused_lm_head_token_chunk_size: int | Literal["disabled"] = 8192
-    """Flattened token chunk size for the fused LM head. ``int >= 1`` sets the tokens per LM-head chunk explicitly; ``disabled`` uses the vanilla LM head. SFT training silently disables this (not supported yet)."""
+    """Flattened token chunk size for the fused LM head. ``int >= 1`` sets the tokens per LM-head chunk explicitly; ``disabled`` uses the vanilla LM head."""
+
+    skip_masked_lm_head_tokens: bool = True
+    """Skip the LM-head vocabulary projection (forward and backward) on tokens no loss component reads — prompt tokens, env observations and pack padding. The head is the only layer that can skip them; the backbone still runs on the full sequence because attention needs it. Saves head time and activation memory in proportion to the masked fraction of the batch, and leaves every value the loss reads unchanged. Requires the fused LM head (``fused_lm_head_token_chunk_size``); ignored by the vanilla head, which materializes full-length logits by construction."""
 
     @model_validator(mode="after")
     def trust_remote_code_only_with_hf(self):

src/prime_rl/trainer/model.py

@@ -13,7 +13,7 @@
 import torch._dynamo
 import torch.nn as nn
 from huggingface_hub import snapshot_download
-from jaxtyping import Int
+from jaxtyping import Bool, Int
 from torch import Tensor
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper
 from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader
@@ -1436,6 +1436,10 @@ def forward(
     seq_lens: Int[Tensor, "segments"],
     labels: Int[Tensor, "batch seq"] | None = None,
     temperature: Tensor | None = None,
+    # Tokens the LM head must score. None scores every token; otherwise the head
+    # skips the vocabulary projection (and its backward) everywhere the mask is
+    # False. Only the fused LM head honours it.
+    keep_mask: Bool[Tensor, "batch seq"] | None = None,
     routed_experts: Int[Tensor, "batch seq layers topk"] | None = None,
     # Generic multimodal kwargs (e.g. {"pixel_values": ...,
     # "image_grid_thw": ...} for Qwen3-VL; just {"pixel_values": ...}
@@ -1452,6 +1456,7 @@ def forward(
         "input_ids": input_ids,
         "labels": labels,
         "temperature": temperature,
+        "keep_mask": keep_mask,
     }
 
     if mm_kwargs:

src/prime_rl/trainer/models/layers/lm_head.py

@@ -34,6 +34,28 @@ def _float_and_contiguous(tensor: Tensor | None) -> Tensor | None:
     )
 
 
+def lm_head_keep_index(keep_mask: Tensor) -> Tensor | None:
+
+    index = keep_mask.reshape(-1).nonzero(as_tuple=True)[0]  # Forces CPU<->GPU sync
+    if index.numel() == keep_mask.numel():
+        return None
+    if index.numel() == 0:
+        return torch.zeros(1, dtype=torch.long, device=keep_mask.device)
+    return index
+
+
+def _compact(tensor: Tensor, keep_index: Tensor | None) -> Tensor:
+    """Drop tokens that the head does not use."""
+    return tensor.contiguous() if keep_index is None else tensor.index_select(0, keep_index)
+
+
+def _expand(values: Tensor, keep_index: Tensor | None, num_tokens: int) -> Tensor:
+    """Scatter compacted tokens back into the full token grid - dropped positions are 0 and logprob(0) keeps the importance ratio at exp(0-0)=1"""
+    if keep_index is None:
+        return values
+    return values.new_zeros(num_tokens).index_copy(0, keep_index, values)
+
+
 class FusedOutputLinear(torch.nn.Linear):
     def __init__(self, in_features: int, out_features: int, chunk_size: int):
         super().__init__(in_features, out_features, bias=False)
@@ -44,21 +66,29 @@ def forward(
         hidden_states: torch.Tensor,
         labels: torch.Tensor | None = None,
         temperature: Tensor | None = None,
+        keep_mask: Tensor | None = None,
     ) -> PrimeLmOutput:
         assert labels is not None, "FusedOutputLinear requires labels for chunked logprob computation"
         assert temperature is not None, "FusedOutputLinear requires per-token temperatures"
 
         b, s, h = hidden_states.shape
-        hidden_states = hidden_states.reshape(b * s, h).contiguous()
-        labels = labels.reshape(b * s).contiguous()
-        inv_t = 1.0 / temperature.reshape(b * s).contiguous()  # [N]
+        n = b * s
+        hidden_states = hidden_states.reshape(n, h)
+        labels = labels.reshape(n)
+        inv_t = 1.0 / temperature.reshape(n)  # [N]
+
+        keep_index = lm_head_keep_index(keep_mask) if keep_mask is not None else None
 
         logprobs, entropy = _SequenceChunkedLogProbEntropyFn.apply(
-            hidden_states, self.weight, labels, inv_t, self.chunk_size
+            _compact(hidden_states, keep_index),
+            self.weight,
+            _compact(labels, keep_index),
+            _compact(inv_t, keep_index),
+            self.chunk_size,
         )
 
-        logprobs = logprobs.reshape(b, s)
-        entropy = entropy.reshape(b, s)
+        logprobs = _expand(logprobs, keep_index, n).reshape(b, s)
+        entropy = _expand(entropy, keep_index, n).reshape(b, s)
         return PrimeLmOutput(logprobs=logprobs, entropy=entropy)
 
 
@@ -67,9 +97,15 @@ def __init__(self, in_features: int, out_features: int):
         super().__init__(in_features, out_features, bias=False)
 
     def forward(
-        self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: Tensor | None = None
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        temperature: Tensor | None = None,
+        keep_mask: Tensor | None = None,
     ) -> PrimeLmOutput:
-        # VanillaOutputLinear just returns logits - temperature scaling is done externally in train.py
+        # VanillaOutputLinear just returns logits - temperature scaling is done externally in train.py.
+        # keep_mask is ignored: the caller needs full-length logits, which is the whole [N, V] tensor
+        # the fused head exists to avoid.
         return PrimeLmOutput(logits=super().forward(hidden_states))
 
 
@@ -262,6 +298,7 @@ def new_forward(
         labels: torch.Tensor | None = None,
         logits_to_keep: int = 0,
         temperature: torch.Tensor | None = None,
+        keep_mask: torch.Tensor | None = None,
         **kwargs: object,
     ) -> PrimeLmOutput:
         # For VLM with images, don't create position_ids - let model compute MRoPE internally
@@ -285,6 +322,7 @@ def new_forward(
             hidden_states[:, slice_indices, :],
             labels[:, slice_indices] if labels is not None else None,
             temperature=temperature[:, slice_indices] if temperature is not None else None,
+            keep_mask=keep_mask[:, slice_indices] if keep_mask is not None else None,
         )
 
     # Bind the new forward to the model

src/prime_rl/trainer/models/layers/lm_head_gemma.py

@@ -6,8 +6,11 @@
 
 from prime_rl.trainer.models.layers.lm_head import (
     PrimeLmOutput,
+    _compact,
+    _expand,
     _online_logsumexp_and_weighted_update,
     _patch_model_forward,
+    lm_head_keep_index,
 )
 from prime_rl.utils.logger import get_logger
 
@@ -23,21 +26,30 @@ def forward(
         hidden_states: torch.Tensor,
         labels: torch.Tensor | None = None,
         temperature: Tensor | None = None,
+        keep_mask: Tensor | None = None,
     ) -> PrimeLmOutput:
         assert labels is not None, "GemmaFusedOutputLinear requires labels for chunked logprob computation"
         assert temperature is not None, "GemmaFusedOutputLinear requires per-token temperatures"
 
         b, s, h = hidden_states.shape
-        hidden_states = hidden_states.reshape(b * s, h).contiguous()
-        labels = labels.reshape(b * s).contiguous()
-        inv_t = 1.0 / temperature.reshape(b * s).contiguous()  # [N]
+        n = b * s
+        hidden_states = hidden_states.reshape(n, h)
+        labels = labels.reshape(n)
+        inv_t = 1.0 / temperature.reshape(n)  # [N]
+
+        keep_index = lm_head_keep_index(keep_mask) if keep_mask is not None else None
 
         logprobs, entropy = _GemmaChunkedLogProbEntropyFn.apply(
-            hidden_states, self.weight, labels, inv_t, self.chunk_size, self.softcap
+            _compact(hidden_states, keep_index),
+            self.weight,
+            _compact(labels, keep_index),
+            _compact(inv_t, keep_index),
+            self.chunk_size,
+            self.softcap,
         )
 
-        logprobs = logprobs.reshape(b, s)
-        entropy = entropy.reshape(b, s)
+        logprobs = _expand(logprobs, keep_index, n).reshape(b, s)
+        entropy = _expand(entropy, keep_index, n).reshape(b, s)
         return PrimeLmOutput(logprobs=logprobs, entropy=entropy)
 
 
@@ -47,7 +59,11 @@ def __init__(self, in_features: int, out_features: int, softcap: float):
         self.softcap = softcap
 
     def forward(
-        self, hidden_states: torch.Tensor, labels: torch.Tensor | None = None, temperature: Tensor | None = None
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        temperature: Tensor | None = None,
+        keep_mask: Tensor | None = None,
     ) -> PrimeLmOutput:
         logits = super().forward(hidden_states)
         logits = self.softcap * torch.tanh(logits / self.softcap)

src/prime_rl/trainer/rl/loss.py

@@ -3,7 +3,7 @@
 
 import torch
 from beartype import beartype as typechecker
-from jaxtyping import Bool, Float, Int, jaxtyped
+from jaxtyping import Bool, Float, Int, Shaped, jaxtyped
 from torch import Tensor
 
 from prime_rl.configs.trainer import CustomLossConfig, DefaultLossConfig, IPOLossConfig, LossConfig
@@ -63,12 +63,14 @@ def compute_entropy(shifted_logits: Float[Tensor, "batch seq vocab"]) -> Float[T
     return entropy
 
 
-def shift_tensor_left(t: Float[Tensor, "batch seq"]) -> Float[Tensor, "batch seq"]:
+def shift_tensor_left(t: Shaped[Tensor, "batch seq"]) -> Shaped[Tensor, "batch seq"]:
     """Shifts the tensor one token to the left.
 
-    Used to create labels from input_ids: labels[i] = input_ids[i+1].
-    The last position is padded with 0 (a valid token index) since this value
-    will be shifted off by shift_tensor_right and never used.
+    Used to create labels from input_ids (dtype long) and to move a per-token mask
+    (dtype bool) onto the next-token positions the LM head runs on, so it takes any
+    dtype rather than just float.
+    The last position is padded with 0 (a valid token index, and False for a mask)
+    since this value will be shifted off by shift_tensor_right and never used.
     """
     return torch.cat([t[:, 1:], torch.full((t.shape[0], 1), 0, device=t.device, dtype=t.dtype)], dim=1)
 

src/prime_rl/trainer/rl/train.py

@@ -259,6 +259,20 @@ def train(config: TrainerConfig):
 
     token_exporter = setup_token_exporter(config, parallel_dims, world, logger)
 
+    skip_masked_head = config.model.skip_masked_lm_head_tokens
+    if skip_masked_head and not isinstance(config.model.fused_lm_head_token_chunk_size, int):
+        logger.warning(
+            "Ignoring model.skip_masked_lm_head_tokens: the vanilla LM head returns full-length logits, "
+            "so there is nothing to skip. Set model.fused_lm_head_token_chunk_size to an int to enable it."
+        )
+        skip_masked_head = False
+    if skip_masked_head and config.enable_token_export:
+        logger.warning(
+            "Ignoring model.skip_masked_lm_head_tokens: token export records trainer logprobs and entropy "
+            "for every token of a sequence, including the ones the head would skip."
+        )
+        skip_masked_head = False
+
     gc_handler = GarbageCollection(config.gc.interval) if config.gc else None
 
     logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})")
@@ -325,14 +339,22 @@ def train(config: TrainerConfig):
         local_rl_scale = 0
         local_ce_scale = 0
         local_ref_kl_scale = 0
+        # Tokens the LM head will actually use and score (those in some component's mask)
+        local_head_tokens = 0
+        local_batch_tokens = 0
         for micro_batch in micro_batches:
             mask = micro_batch["loss_mask"]
             rl_w = micro_batch["rl_weights"]
             local_rl_scale += int((mask & (rl_w != 0)).sum()) if rl_w is not None else int(mask.sum())
+            head_mask = mask
             if micro_batch["ce_weights"] is not None:
                 local_ce_scale += int((micro_batch["ce_weights"] != 0).sum())
+                head_mask = head_mask | (micro_batch["ce_weights"] != 0)
             if micro_batch["ref_kl_weights"] is not None:
                 local_ref_kl_scale += int((micro_batch["ref_kl_weights"] != 0).sum())
+                head_mask = head_mask | (micro_batch["ref_kl_weights"] != 0)
+            local_head_tokens += int(head_mask.sum())
+            local_batch_tokens += mask.numel()
         global_scales = torch.tensor(
             [local_rl_scale, local_ce_scale, local_ref_kl_scale], dtype=torch.int64, device="cuda"
         )
@@ -398,6 +420,16 @@ def train(config: TrainerConfig):
 
             labels = shift_tensor_left(input_ids)
 
+            # Tokens the LM head has to use and score: bsasically every token some loss component reads
+            keep_mask = None
+            if skip_masked_head:
+                keep_mask = loss_mask
+                if ce_weights is not None:
+                    keep_mask = keep_mask | (ce_weights != 0)
+                if ref_kl_weights is not None:
+                    keep_mask = keep_mask | (ref_kl_weights != 0)
+                keep_mask = shift_tensor_left(keep_mask)
+
             seq_lens_are_pre_shard = False
 
             if cp_enabled:
@@ -415,6 +447,8 @@ def train(config: TrainerConfig):
                     )
                 seq_lens_are_pre_shard = True
                 labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size)
+                if keep_mask is not None:
+                    keep_mask = shard_for_cp(keep_mask, cp_rank=cp_rank, cp_world_size=cp_size)
                 if routed_experts is not None and not defer_vlm_cp_to_model:
                     routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_size)
 
@@ -444,6 +478,7 @@ def train(config: TrainerConfig):
                     position_ids,
                     labels=labels,
                     temperature=temperatures,
+                    keep_mask=keep_mask,
                     mm_kwargs=mm_kwargs,
                     mm_token_type_ids=mm_token_type_ids,
                     seq_lens=seq_lens,
@@ -662,6 +697,9 @@ def train(config: TrainerConfig):
             "perf/throughput_per_gpu": throughput / world.world_size,
             "perf/mfu": mfu,
             "perf/peak_memory": peak_memory,
+            "perf/lm_head_token_fraction": (
+                local_head_tokens / max(local_batch_tokens, 1) if skip_masked_head else 1.0
+            ),
             "step": progress.step,
         }
         asyncio.run(monitors.log(perf_metrics, step=progress.step))

src/prime_rl/trainer/sft/train.py

@@ -261,6 +261,14 @@ def train(config: SFTConfig):
     dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group()
     cp_size = parallel_dims.cp
 
+    skip_masked_head = config.model.skip_masked_lm_head_tokens
+    if skip_masked_head and not isinstance(config.model.fused_lm_head_token_chunk_size, int):
+        logger.warning(
+            "Ignoring model.skip_masked_lm_head_tokens: the vanilla LM head returns full-length logits, "
+            "so there is nothing to skip. Set model.fused_lm_head_token_chunk_size to an int to enable it."
+        )
+        skip_masked_head = False
+
     def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
         """Forward pass returning (loss_sum, token_count) over unmasked tokens."""
         input_ids = micro_batch["input_ids"].to("cuda", non_blocking=True)
@@ -315,6 +323,7 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
                     seq_lens=seq_lens,
                     labels=target_ids,
                     temperature=temperature,
+                    keep_mask=loss_mask if skip_masked_head else None,
                     mm_kwargs=mm_kwargs,
                     mm_token_type_ids=mm_type_ids,
                     seq_lens_are_pre_shard=seq_lens_are_pre_shard,
@@ -640,6 +649,9 @@ def is_online_eval_step(step: int) -> bool:
             "perf/throughput_per_gpu": throughput / world.world_size,
             "perf/peak_memory": peak_memory,
             "perf/mfu": mfu,
+            "perf/lm_head_token_fraction": (
+                global_token_count_val / num_tokens if skip_masked_head and num_tokens > 0 else 1.0
+            ),
             "step": progress.step,
         }
         asyncio.run(monitors.log(perf_metrics, step=progress.step))