feat: --trim-backbone-rows — run TTT steps >= 1 only on rows that can still emit loss

examples/configs/README.md

@@ -271,6 +271,8 @@ Common fields:
 | `training.compact_teacher` | `false` | Exact lower-peak-memory teacher projection for offline text EAGLE3. |
 | `training.compact_teacher_chunk_size` | `null` | Positive vocabulary chunk size; requires `compact_teacher: true`. |
 | `training.trim_loss_positions` | `false` | EAGLE3 only. Compute the teacher target_p, draft logits, and loss only at supervised positions (batch size 1, plain KL loss); mathematically equivalent to the full-length path. |
+| `training.trim_backbone_rows` | `false` | Also run the draft backbone at TTT steps >= 1 only on rows that can still emit loss (union across remaining steps); step 0 stays full-length. Requires `trim_loss_positions` and attention backend sdpa/fa/usp. |
+| `training.trim_backbone_rows_max_density` | `0.35` | Per-sample guard for `trim_backbone_rows`: trimmed steps use a costlier per-row attention kernel, so the trim is skipped when the kept rows summed over TTT steps exceed this fraction of the full-length work. Set to `1.0` to always trim. |
 | `training.role` | `all` | Use `all` for offline colocated training; disaggregated entrypoints select `auto`, `producer`, or `consumer`. |
 | `training.seed` | `42` | Run and per-rank RNG seed. |
 | `training.prompt_seed` | `null` | Optional online prompt-shuffle seed. `null` preserves the historical behavior of using `training.seed`. |

specforge/algorithms/contracts.py

@@ -239,6 +239,7 @@ class AlgorithmCapabilities:
     required_batch_size: int | None = None
     supports_compact_teacher: bool = False
     supports_trim_loss_positions: bool = False
+    supports_trim_backbone_rows: bool = False
     supports_vocab_mapping: bool = False
     allows_aux_layer_override: bool = False
 
@@ -256,6 +257,7 @@ def __post_init__(self) -> None:
         for field_name in (
             "supports_compact_teacher",
             "supports_trim_loss_positions",
+            "supports_trim_backbone_rows",
             "supports_vocab_mapping",
             "allows_aux_layer_override",
         ):

specforge/algorithms/eagle3/model.py

@@ -22,6 +22,7 @@
 
 """EAGLE3 training model implementation."""
 
+import logging
 from typing import Callable, List, Optional, Tuple
 
 import torch
@@ -39,6 +40,8 @@
 from specforge.modeling.draft import Eagle3DraftModel
 from specforge.utils import padding
 
+logger = logging.getLogger(__name__)
+
 
 class Eagle3Model(nn.Module):
     pass
@@ -262,6 +265,8 @@ def forward(
         target_head_weight: Optional[torch.Tensor] = None,
         compact_teacher_chunk_size: int = DEFAULT_VOCAB_CHUNK_SIZE,
         trim_loss_positions: bool = False,
+        trim_backbone_rows: bool = False,
+        trim_backbone_rows_max_density: float = 0.35,
     ) -> Tuple[
         List[torch.Tensor],
         List[torch.Tensor],
@@ -285,6 +290,13 @@ def forward(
                 states in draft-vocab space and ``target`` is ignored.
             trim_loss_positions: compute the teacher, draft logits and loss only at
                 supervised positions when the batch/objective supports it.
+            trim_backbone_rows: additionally run the draft backbone at TTT steps
+                >= 1 only on the rows that can still emit loss at that or a later
+                step (requires trim_loss_positions; step 0 stays full-length as it
+                provides the cross-position K/V context).
+            trim_backbone_rows_max_density: steps whose kept-row set exceeds this
+                fraction of the chunk run the full-length path instead -- trimming
+                costs more per row, so it stops paying once few rows are dropped.
         """
         adapter = self._make_adapter()
         # Step 1: handle vocab size
@@ -338,6 +350,7 @@ def forward(
                     loss_mask,
                     self.length,
                     chunk_len=trim_chunk_len,
+                    backbone_rows=trim_backbone_rows,
                 )
             if trim_pack is not None:
                 target_p_padded = None
@@ -413,8 +426,130 @@ def forward(
         else:
             raise ValueError(f"Unknown attention backend: {self.attention_backend}")
 
+        b_enabled = trim_pack is not None and "b_rows_steps" in trim_pack
+        if trim_backbone_rows and self.attention_backend == "usp":
+            # Backbone-row trimming changes the backbone's collective pattern
+            # (one K/V all-gather instead of per-step ring attention), so the
+            # decision must be rank-uniform: a rank whose pack fell back to
+            # None (no reachable supervision) would otherwise keep issuing
+            # ring collectives that trimmed ranks never join -> deadlock.
+            # Agree by MIN across the sequence-parallel group; on disagreement
+            # every rank drops to the A-level path, whose backbone collectives
+            # are identical to the full path's.
+            flag = torch.tensor(1 if b_enabled else 0, device=hidden_states.device)
+            torch.distributed.all_reduce(
+                flag, op=torch.distributed.ReduceOp.MIN, group=adapter.sp_group
+            )
+            if int(flag.item()) == 0:
+                # b_enabled alone gates every B-level read below; the pack's
+                # b_* entries are simply never consulted again.
+                b_enabled = False
+        if b_enabled and position_ids is not None and position_ids.dim() == 3:
+            # Multimodal RoPE carries [axes, batch, seq] position ids; the
+            # row-selection below would slice the batch axis. Fail fast here
+            # (the attention-level guard cannot be reached for this layout).
+            raise NotImplementedError(
+                "trim_backbone_rows does not support multimodal RoPE (mrope) "
+                "drafts (3D position_ids)"
+            )
+        if b_enabled:
+            # The trimmed path builds its causal mask from position VALUES
+            # (row at position p attends step-0 keys 0..p), which matches the
+            # full path's index-based mask only when positions are the row
+            # indices themselves (plain arange locally, or the collator's
+            # contiguous global arange under USP). Packed sequences with
+            # per-segment position resets would silently diverge -- reject.
+            _pos = position_ids[:, : trim_pack["full_len"]]
+            if not bool((_pos.diff(dim=-1) == 1).all()):
+                raise NotImplementedError(
+                    "trim_backbone_rows requires contiguous ascending "
+                    "position_ids (packed/segment-reset positions are not "
+                    "supported)"
+                )
+        # Persistent across steps: the trimmed-attention path stashes the
+        # (possibly all-gathered) step-0 K/V here at the first trimmed step.
+        # k0_len is the step-0 K/V row count (global length under USP).
+        trim_ctx = (
+            {
+                "k0": None,
+                "v0": None,
+                "k0_len": trim_pack["full_len"] * adapter.sp_world_size,
+            }
+            if b_enabled
+            else None
+        )
+        # Cost guard. A trimmed step swaps flash attention for an explicit
+        # masked-matmul kernel that costs more per row, so trimming only pays
+        # while the kept-row sets are small. Measured on this repo (8192-token
+        # chunk, TTT=7, RTX 6000 Ada): flash attention costs ~16.7 us/row and
+        # the trimmed kernel 28-62 us/row, so the step time breaks even near
+        # 40% density and degrades to 2.4x slower at full supervision, where
+        # no row can be dropped at all. Compare the two sides directly:
+        #   trimmed rows  sum_i |R_i|   vs   full rows  (k-1) * chunk
+        # and keep the B path only while the former stays under the budget.
+        # The decision is per sample, not per step: |R_i| barely varies across
+        # steps for contiguous supervision (it shrinks by one row per
+        # supervised run per step), and a mixed full/trimmed unroll would need
+        # per-backend K/V layout conversion plus an extra head-dimension
+        # gather under USP -- new collectives for a safety guard is a bad
+        # trade.
+        if b_enabled:
+            chunk = trim_pack["full_len"]
+            rows_sum = sum(
+                int(trim_pack["b_rows_steps"][i].numel()) for i in range(1, self.length)
+            )
+            budget = trim_backbone_rows_max_density * (self.length - 1) * chunk
+            trim_pays = rows_sum <= budget
+            if self.attention_backend == "usp":
+                # Same reason as the b_enabled agreement above: per-rank row
+                # counts differ, so ranks can land on opposite sides of the
+                # budget, and a mixed trimmed/full unroll across ranks would
+                # mismatch collectives. Agree by MIN -- running the full path
+                # is always correct, it only forgoes savings.
+                flag = torch.tensor(1 if trim_pays else 0, device=hidden_states.device)
+                torch.distributed.all_reduce(
+                    flag, op=torch.distributed.ReduceOp.MIN, group=adapter.sp_group
+                )
+                trim_pays = bool(flag.item())
+            if not trim_pays:
+                # Surface it: the user asked for row trimming and did not get
+                # it, and the reason (dense supervision) is a property of their
+                # data, not a bug.
+                logger.info(
+                    "trim_backbone_rows skipped for this sample: kept rows "
+                    "%d exceed the budget %d (%.2f x %d steps x %d chunk); "
+                    "the trimmed attention kernel costs more per row than "
+                    "flash attention, so trimming this sample would be slower",
+                    rows_sum,
+                    int(budget),
+                    trim_backbone_rows_max_density,
+                    self.length - 1,
+                    chunk,
+                )
+                b_enabled = False
         for idx in range(self.length):
-            if trim_pack is not None:
+            b_active = b_enabled and idx >= 1
+            if b_active:
+                # B-level: this step's backbone runs only R_i = the union of all
+                # rows that can still emit loss at this or a later step. hidden
+                # chains through per-step row maps (R_i is nested in R_{i-1});
+                # input_ids/positions follow the absolute rows so RoPE and the
+                # causal mask stay position-correct inside the trimmed path.
+                rows_b = trim_pack["b_rows_steps"][idx]
+                # R_i nested in R_{i-1}: positions of R_i inside the previous
+                # step's compact output are exactly the last diagonal map.
+                prev_sel = (
+                    rows_b if idx == 1 else trim_pack["b_diag_sels"][idx][idx - 1]
+                )
+                step_hidden = hidden_states.index_select(1, prev_sel)
+                step_input_ids = global_input_ids.index_select(1, rows_b)
+                step_pos = position_ids[:, : trim_pack["full_len"]].index_select(
+                    1, rows_b
+                )
+                step_attn = None  # the trimmed path builds its own causal mask
+                trim_ctx["positions"] = step_pos
+                trim_ctx["diag_sels"] = trim_pack["b_diag_sels"][idx]
+            elif trim_pack is not None:
                 # A-level: the teacher tables are already compacted to supervised
                 # positions; the backbone runs exactly the same inputs as the full
                 # path (per-rank chunk under USP, full length otherwise) and only
@@ -464,6 +599,7 @@ def forward(
                 position_ids=step_pos,
                 past_key_values=past_key_values,
                 use_cache=True,
+                **({"trim_rows_ctx": trim_ctx} if b_active else {}),
             )
 
             # update hidden states for next step
@@ -477,8 +613,14 @@ def forward(
                 rows_j = trim_pack["rows_steps"][idx]
                 keep_j = trim_pack["keep_steps"][idx]
                 nrows_j = trim_pack["nrows_steps"][idx]
+                if b_active:
+                    # backbone output is compacted to R_i rows; select loss rows
+                    # by their position inside R_i rather than absolute index.
+                    loss_sel = trim_pack["b_loss_sel"][idx]
+                else:
+                    loss_sel = rows_j
                 logits = self.draft_model.compute_logits(
-                    hidden_states.index_select(1, rows_j)
+                    hidden_states.index_select(1, loss_sel)
                 )
                 pm_j = trim_pack["position_mask_sup"].index_select(1, keep_j)
                 lm_j = trim_pack["loss_mask_sup"].index_select(1, keep_j)
@@ -657,7 +799,9 @@ def _compute_target_p_eager(target, t2d, loss_mask, row_chunk=256):
     )
 
 
-def _build_trim_pack(target, t2d, loss_mask, length, chunk_len=None):
+def _build_trim_pack(
+    target, t2d, loss_mask, length, chunk_len=None, backbone_rows=False
+):
     """A-level trim (--trim-loss-positions): keep only the rows that can carry loss.
 
     Derivation of the per-step row set. On the full-length path the loop applies
@@ -722,13 +866,15 @@ def _build_trim_pack(target, t2d, loss_mask, length, chunk_len=None):
         )
         lm_sup = loss_mask.view(-1)[sup].view(1, -1, 1)
 
-        rows_steps, keep_steps, nrows_steps = [], [], []
+        rows_steps, keep_steps, nrows_steps, raw_rows = [], [], [], []
         pad_idx = torch.zeros(1, dtype=sup.dtype, device=sup.device)
         for j in range(length):
             keep = (
                 ((sup >= j) & (sup - j < chunk_len)).nonzero(as_tuple=False).squeeze(-1)
             )
             n = int(keep.numel())
+            if backbone_rows:
+                raw_rows.append(sup[keep] - j)
             if n == 0:
                 # Dead step: pad with one dummy entry; the caller zeroes its
                 # masks so it contributes nothing.
@@ -739,7 +885,45 @@ def _build_trim_pack(target, t2d, loss_mask, length, chunk_len=None):
             rows_steps.append(rows)
             keep_steps.append(keep)
             nrows_steps.append(n)
+
+        extra = {}
+        if backbone_rows and length > 1:
+            # B-level (--trim-backbone-rows): the rows step i must forward are
+            #   R_i = union_{j >= i} { s - j : s in sup, 0 <= s - j < chunk_len }
+            # -- every row that can still emit loss at this or a later step
+            # (hidden chains between steps, so a row needed at step j must be
+            # forwarded at every step i <= j). The sets are nested
+            # (R_i \subseteq R_{i-1}), which makes the per-step row maps plain
+            # searchsorted lookups. Step 0 always runs full-length: it produces
+            # the K/V context every later row attends to.
+            b_rows_steps, b_diag_sels, b_loss_sel = {}, {}, {}
+            prev = None
+            for i in range(1, length):
+                nonempty = [r for r in raw_rows[i:] if r.numel()]
+                if nonempty:
+                    uni = torch.unique(torch.cat(nonempty))
+                else:
+                    # Dead tail: keep one row from the previous set so the
+                    # nested chain (and kernel/collective alignment across
+                    # ranks) survives; its loss is zero-masked by nrows == 0.
+                    uni = (prev[:1] if prev is not None else pad_idx).clone()
+                b_rows_steps[i] = uni
+                b_diag_sels[i] = {
+                    e: torch.searchsorted(b_rows_steps[e], uni) for e in range(1, i)
+                }
+                b_loss_sel[i] = (
+                    torch.searchsorted(uni, rows_steps[i])
+                    if nrows_steps[i] > 0
+                    else pad_idx
+                )
+                prev = uni
+            extra = dict(
+                b_rows_steps=b_rows_steps,
+                b_diag_sels=b_diag_sels,
+                b_loss_sel=b_loss_sel,
+            )
         return dict(
+            **extra,
             sup=sup,
             rows_steps=rows_steps,
             keep_steps=keep_steps,

specforge/algorithms/eagle3/providers.py

@@ -69,6 +69,7 @@ def resume_contract(config, draft_model, training_model):
         "eagle3_kl_scale": float(training_model.kl_scale),
         "eagle3_kl_decay": float(training_model.kl_decay),
         "eagle3_trim_loss_positions": bool(config.training.trim_loss_positions),
+        "eagle3_trim_backbone_rows": bool(config.training.trim_backbone_rows),
         "eagle3_compact_teacher": bool(config.training.compact_teacher),
         "eagle3_compact_teacher_chunk_size": (
             config.training.compact_teacher_chunk_size
@@ -162,6 +163,7 @@ def algorithm_spec() -> AlgorithmSpec:
             attention_backends={"sdpa", "flex_attention", "fa", "usp"},
             supports_compact_teacher=True,
             supports_trim_loss_positions=True,
+            supports_trim_backbone_rows=True,
             supports_vocab_mapping=True,
             allows_aux_layer_override=True,
         ),

specforge/algorithms/model_providers.py

@@ -437,6 +437,8 @@ def build_dspark_model(
 def eagle3_strategy_kwargs(cfg: Config) -> Dict[str, Any]:
     return {
         "trim_loss_positions": cfg.training.trim_loss_positions,
+        "trim_backbone_rows": cfg.training.trim_backbone_rows,
+        "trim_backbone_rows_max_density": (cfg.training.trim_backbone_rows_max_density),
         "compact_teacher": cfg.training.compact_teacher,
         "compact_teacher_chunk_size": cfg.training.compact_teacher_chunk_size,
     }

specforge/application/planning.py

@@ -126,6 +126,12 @@ def _validate_algorithm_capabilities(
             "training.trim_loss_positions"
         )
 
+    if training.trim_backbone_rows and not capabilities.supports_trim_backbone_rows:
+        raise ValueError(
+            f"algorithm {algorithm.name!r} does not support "
+            "training.trim_backbone_rows"
+        )
+
 
 def _validate_training_topology(
     cfg: Config,