From e8425958e38545a542dd033a73caf7e99f89eb3f Mon Sep 17 00:00:00 2001 From: julyanghar Date: Tue, 18 Aug 2026 19:31:16 -0500 Subject: [PATCH 1/6] =?UTF-8?q?feat:=20--trim-backbone-rows=20=E2=80=94=20?= =?UTF-8?q?run=20TTT=20steps=20>=3D=201=20only=20on=20rows=20that=20can=20?= =?UTF-8?q?still=20emit=20loss?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to #705 (trim_loss_positions) and the RFC in #706. On prompt-heavy data most positions never contribute gradient, yet the draft backbone re-runs all of them at every TTT step. With training.trim_backbone_rows=true, steps >= 1 forward only R_i = union_{j>=i} { s - j : s in sup, 0 <= s - j < chunk }, the nested union of rows that can still emit loss at this or a later step (hidden states chain between steps, so a row needed later must be forwarded now). Step 0 stays full-length: it provides the K/V context every later row attends to. Mechanics - One shared trimmed-attention path for sdpa/fa/usp (flex is rejected at config validation): grouped-query layout throughout, explicit causal mask from absolute positions (non-contiguous row sets are exactly what flash kernels cannot express), diagonal-private K/V row-aligned via nested-set index maps, per-step activation checkpointing. - Long step-0 K/V (>= 8192) switch to a single fused memory-efficient SDPA call per step: query groups fold into the row axis, diagonal K/V join as extra columns behind a boolean mask, softmax never reaches HBM. - Under USP the trimmed steps bypass ring/Ulysses entirely (per-rank row counts differ): step-0 K/V are all-gathered once with the autograd-aware all_gather (backward = reduce-scatter sum), and enabling is agreed rank-uniformly (all_reduce MIN) so mixed trim/full ranks cannot deadlock. - Guards: requires trim_loss_positions; rejects flex_attention, compact_teacher, mrope drafts, non-contiguous position_ids. Verification - CPU golden tables pin the R_i union math (incl. overlap-tail bound, dead steps, the naive fixed-sup reading is detected by a negative control). - Equivalence vs the untrimmed path: all-ones mask (trivial-equality bound), prompt-heavy/block-boundary/isolated masks, grad-norm check, on sdpa, fa, and 4-rank ring4 USP; the fused SDPA core is additionally forced through every case via a patched threshold. Benchmark (Llama3-8B dims, TTT 7, 16-block 8% supervised mask, 48GB cards) - 8k single-GPU sdpa: baseline OOMs, trimmed runs (38.2 GB, 1.16 s/step) - 32k ring4: 17.9 -> 8.6 GB/rank, 1.68 -> 0.62 s/step - 64k ring4: 32.9 -> 14.9 GB/rank, 4.60 -> 1.81 s/step - 16k single-GPU sdpa: both OOM (step-0's materialized sdpa attention is the binding constraint; fa's step 0 does not have it) Co-Authored-By: Claude Fable 5 --- examples/configs/README.md | 1 + specforge/algorithms/contracts.py | 2 + specforge/algorithms/eagle3/model.py | 136 +++++- specforge/algorithms/eagle3/providers.py | 2 + specforge/algorithms/model_providers.py | 1 + specforge/application/planning.py | 6 + specforge/config/schema.py | 28 ++ specforge/core/eagle3_adapters.py | 3 + specforge/modeling/draft/base.py | 1 + specforge/modeling/draft/llama3_eagle.py | 312 ++++++++++++ specforge/training/strategies/base.py | 3 + .../test_unified_feature_reachability.py | 61 +++ tests/test_runtime/test_trim_backbone_rows.py | 454 ++++++++++++++++++ 13 files changed, 1006 insertions(+), 4 deletions(-) create mode 100644 tests/test_runtime/test_trim_backbone_rows.py diff --git a/examples/configs/README.md b/examples/configs/README.md index 3cbd0cc6a..e7a1afb90 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -271,6 +271,7 @@ 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.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`. | diff --git a/specforge/algorithms/contracts.py b/specforge/algorithms/contracts.py index 1d97da89f..16c6d4734 100644 --- a/specforge/algorithms/contracts.py +++ b/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", ): diff --git a/specforge/algorithms/eagle3/model.py b/specforge/algorithms/eagle3/model.py index d89cf5d0c..5ab8fba66 100644 --- a/specforge/algorithms/eagle3/model.py +++ b/specforge/algorithms/eagle3/model.py @@ -262,6 +262,7 @@ 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, ) -> Tuple[ List[torch.Tensor], List[torch.Tensor], @@ -285,6 +286,10 @@ 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). """ adapter = self._make_adapter() # Step 1: handle vocab size @@ -338,6 +343,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 +419,81 @@ 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 + ) 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 +543,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 +557,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 +743,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 +810,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 +829,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, diff --git a/specforge/algorithms/eagle3/providers.py b/specforge/algorithms/eagle3/providers.py index 835be387c..2ab7cd2cf 100644 --- a/specforge/algorithms/eagle3/providers.py +++ b/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, ), diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index 285433fcd..f729d6a44 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -437,6 +437,7 @@ 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, "compact_teacher": cfg.training.compact_teacher, "compact_teacher_chunk_size": cfg.training.compact_teacher_chunk_size, } diff --git a/specforge/application/planning.py b/specforge/application/planning.py index 44830a927..3f66f20e2 100644 --- a/specforge/application/planning.py +++ b/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, diff --git a/specforge/config/schema.py b/specforge/config/schema.py index 27870d0ca..c8197f40f 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -556,6 +556,15 @@ class TrainingConfig(StrictConfigModel): #: prompt-heavy data. Falls back to the full-length path for batch > 1 or when #: an lk_loss objective is used. trim_loss_positions: bool = False + #: Additionally run the draft backbone at TTT steps >= 1 only on the rows + #: that can still emit loss at that or a later step (their union, since + #: hidden states chain between steps). Step 0 stays full-length -- it + #: provides the cross-position K/V context. Intended for prompt-heavy data + #: (small supervised fraction): the trimmed steps use explicit masked-matmul + #: attention, so with dense supervision they degenerate to eager full + #: attention. Requires trim_loss_positions and an attention backend in + #: {sdpa, fa, usp}. + trim_backbone_rows: bool = False #: DFlash-family objective/model knobs. num_anchors: int = Field(default=512, gt=0) loss_decay_gamma: Optional[float] = None @@ -625,6 +634,25 @@ def _validate_training_shape(self): "training.sp_ulysses_size/sp_ring_size require " "training.attention_backend=usp" ) + if self.trim_backbone_rows: + if not self.trim_loss_positions: + raise ValueError( + "training.trim_backbone_rows requires " + "training.trim_loss_positions=true (it builds on the same " + "supervised-row bookkeeping)" + ) + if self.attention_backend not in ("sdpa", "fa", "usp"): + raise ValueError( + "training.trim_backbone_rows supports attention_backend " + "sdpa/fa/usp only (flex_attention's cache semantics differ)" + ) + if self.compact_teacher: + raise ValueError( + "training.trim_backbone_rows is incompatible with " + "training.compact_teacher (the compact-teacher branch " + "bypasses the trim bookkeeping, so the flag would " + "silently do nothing)" + ) return self diff --git a/specforge/core/eagle3_adapters.py b/specforge/core/eagle3_adapters.py index dfcac43ef..9550566af 100644 --- a/specforge/core/eagle3_adapters.py +++ b/specforge/core/eagle3_adapters.py @@ -28,6 +28,9 @@ class StepState(BackboneStepState): class BackendAdapter: + #: sequence-parallel world size; UspAdapter overrides in __init__. + sp_world_size = 1 + def __init__(self, model: "OnlineEagle3Model"): self.m = model diff --git a/specforge/modeling/draft/base.py b/specforge/modeling/draft/base.py index be073fa0e..2c4adfb9a 100644 --- a/specforge/modeling/draft/base.py +++ b/specforge/modeling/draft/base.py @@ -103,6 +103,7 @@ def backbone( position_ids: torch.Tensor, past_key_values: Optional[Cache] = None, use_cache: bool = True, + trim_rows_ctx: Optional[dict] = None, ) -> torch.Tensor: """ The backbone of the draft model. diff --git a/specforge/modeling/draft/llama3_eagle.py b/specforge/modeling/draft/llama3_eagle.py index 52da46ce8..555ccaf01 100644 --- a/specforge/modeling/draft/llama3_eagle.py +++ b/specforge/modeling/draft/llama3_eagle.py @@ -658,6 +658,22 @@ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): .contiguous() ) + def _trim_step0_kv(self, cache_hidden): + """Step-0 K/V for the trimmed path, as [B, KVH, {G|1}, L0, D]. + + The sdpa TTT branch caches step-0 K/V already expanded to the full + head count ([B, H, L, D], repeat_kv layout: head h maps to kv head + h // G), so viewing it grouped is free. + """ + k0, v0 = cache_hidden[0][0], cache_hidden[1][0] + bsz, _, l0, hd = k0.shape + kvh = self.num_key_value_heads + groups = self.num_key_value_groups + return ( + k0.view(bsz, kvh, groups, l0, hd), + v0.view(bsz, kvh, groups, l0, hd), + ) + def forward( self, hidden_states: torch.Tensor, @@ -667,7 +683,12 @@ def forward( past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, + trim_rows_ctx: Optional[dict] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if trim_rows_ctx is not None: + return _trimmed_rows_attention( + self, hidden_states, cache_hidden, trim_rows_ctx + ) bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) @@ -804,7 +825,12 @@ def forward( past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, + trim_rows_ctx: Optional[dict] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if trim_rows_ctx is not None: + raise NotImplementedError( + "trim_backbone_rows is not supported with flex_attention" + ) bsz, q_len, _ = hidden_states.size() past_seen_tokens = ( @@ -1291,6 +1317,13 @@ def __init__(self, config): ): _raise_standard_flash_attn_unavailable() + def _trim_step0_kv(self, cache_hidden): + """fa caches step-0 K/V as [B, L, KVH, D]; keep KV heads un-repeated + and let the grouped matmul broadcast over query groups.""" + k0 = cache_hidden[0][0].transpose(1, 2).unsqueeze(2) # [B, KVH, 1, L, D] + v0 = cache_hidden[1][0].transpose(1, 2).unsqueeze(2) + return k0, v0 + def forward( self, hidden_states: torch.Tensor, @@ -1300,7 +1333,12 @@ def forward( past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, + trim_rows_ctx: Optional[dict] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if trim_rows_ctx is not None: + return _trimmed_rows_attention( + self, hidden_states, cache_hidden, trim_rows_ctx + ) bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) @@ -1384,6 +1422,35 @@ def __init__(self, config): self.gather_idx = 1 self.use_sync = False + def _trim_step0_kv(self, cache_hidden): + """Gather step-0 K/V to GLOBAL length once for the trimmed steps. + + Trimmed steps bypass Ulysses/ring entirely: per-rank row counts differ + (all-to-all needs equal lengths) and flash cannot express causal masks + over non-contiguous rows. ``torch.distributed.nn.functional.all_gather`` + is used deliberately instead of ``specforge.distributed.Gather``: its + backward is the reduce-scatter SUM of every rank's gradient, which is + the correct semantics here because each rank computes a different loss + from the gathered K/V (Gather's world-size-scaled slice backward + assumes identical per-rank computation). KV heads stay un-repeated; + the grouped matmul broadcasts over query groups. + """ + import torch.distributed.nn.functional as dist_nn + + # step-0 entries are post-a2a: [B, C*ulysses, H/ulysses, D] + k0 = cache_hidden[0][0] + v0 = cache_hidden[1][0] + if self.sp_ulysses_degree > 1: + k0 = torch.cat(dist_nn.all_gather(k0, group=self.ulysses_pg), dim=2) + v0 = torch.cat(dist_nn.all_gather(v0, group=self.ulysses_pg), dim=2) + if self.sp_ring_degree > 1: + k0 = torch.cat(dist_nn.all_gather(k0, group=self.ring_pg), dim=1) + v0 = torch.cat(dist_nn.all_gather(v0, group=self.ring_pg), dim=1) + return ( + k0.transpose(1, 2).unsqueeze(2), # [B, KVH, 1, L_global, D] + v0.transpose(1, 2).unsqueeze(2), + ) + def forward( self, hidden_states: torch.Tensor, @@ -1393,9 +1460,15 @@ def forward( past_key_values: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, + trim_rows_ctx: Optional[dict] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: from yunchang.comm import SeqAllToAll4D + if trim_rows_ctx is not None: + return _trimmed_rows_attention( + self, hidden_states, cache_hidden, trim_rows_ctx + ) + bsz, q_len, _ = hidden_states.size() local_q_len = q_len @@ -1567,6 +1640,241 @@ def forward(self, hidden_states): return self.weight * hidden_states.to(input_dtype) +def _trimmed_attention_core(qg, k0, v0, positions, scale, neg_inf, *diag_kv): + """Pure-local attention math for one trimmed step (checkpoint-friendly). + + Recomputed in backward via activation checkpointing: the [n_i, L0]-sized + fp32 softmax tensors would otherwise stay resident once per trimmed step, + which at long global lengths dwarfs every saving the trimming made. + Collectives are deliberately kept OUTSIDE (k0/v0 arrive pre-gathered): + a collective inside a recomputed region would re-issue during backward in + rank-dependent order and deadlock. + + Rows are processed in chunks of ``_TRIM_ROW_CHUNK`` so the transient + [rows, L0] fp32 attention block stays bounded at long global lengths + (at 64k global and full rows it would otherwise be a >10 GiB single + allocation). Per-row math is independent, so chunking is exact. + """ + n = qg.shape[3] + if n <= _TRIM_ROW_CHUNK: + return _trimmed_attention_rows(qg, k0, v0, positions, scale, neg_inf, *diag_kv) + outs = [] + for s in range(0, n, _TRIM_ROW_CHUNK): + e = min(s + _TRIM_ROW_CHUNK, n) + chunk_diag = [d[:, :, s:e] for d in diag_kv] + outs.append( + _trimmed_attention_rows( + qg[:, :, :, s:e], + k0, + v0, + positions[:, s:e], + scale, + neg_inf, + *chunk_diag, + ) + ) + return torch.cat(outs, dim=3) + + +_TRIM_ROW_CHUNK = 512 + +#: Above this step-0 K/V length the trimmed step switches from the exact +#: eager core to one fused memory-efficient SDPA call (no [rows, L] tensor is +#: ever materialized). Short lengths keep the eager core, whose math is +#: bit-aligned with the native sdpa TTT branch. +_TRIM_SDPA_MIN_K0 = 8192 + + +def _trimmed_attention_sdpa(q, k0, v0, positions, *diag_kv): + """Long-k0 core: one fused memory-efficient SDPA call per step. + + ``q`` is [B, KVH, G, m, D]; ``k0``/``v0`` are [B, KVH, L0, D] + (un-expanded); diagonal entries are [B, KVH, m, D]. Query groups are + folded into the row axis so q and k share ``num_heads`` (the fused + kernels require equal head counts), and the diagonal-private K/V join as + extra key/value columns visible only to their own row via the boolean + mask. The kernel keeps softmax internal, so nothing [rows, L0]-sized in + fp32 ever reaches HBM. + """ + bsz, kvh, groups, m, d = q.shape + half = len(diag_kv) // 2 + diag_ks, diag_vs = diag_kv[:half], diag_kv[half:] + k_ext = torch.cat((k0, *diag_ks), dim=2) if diag_ks else k0 + v_ext = torch.cat((v0, *diag_vs), dim=2) if diag_vs else v0 + col = torch.arange(k0.shape[2], device=q.device) + allow = col[None, :] <= positions.view(-1, 1) # [m, L0] + if diag_ks: + eye = torch.eye(m, dtype=torch.bool, device=q.device) + allow = torch.cat([allow] + [eye] * len(diag_ks), dim=1) + allow = allow.unsqueeze(0).expand(groups, m, -1).reshape(groups * m, -1) + qf = q.reshape(bsz, kvh, groups * m, d) + out = nn.functional.scaled_dot_product_attention( + qf, k_ext, v_ext, attn_mask=allow[None, None] + ) + return out.view(bsz, kvh, groups, m, d) + + +def _trimmed_attention_rows(qg, k0, v0, positions, scale, neg_inf, *diag_kv): + attn_weights = torch.matmul(qg, k0.transpose(-1, -2)) * scale + col = torch.arange(k0.shape[-2], device=qg.device) + causal = col[None, :] > positions.view(-1, 1) # [rows, L0] + attn_weights = attn_weights.masked_fill(causal[None, None, None], neg_inf) + half = len(diag_kv) // 2 + diag_ks, diag_vs = diag_kv[:half], diag_kv[half:] + for ke in diag_ks: + w = (qg * ke.unsqueeze(2)).sum(-1) * scale # [B, KVH, G, rows] + attn_weights = torch.cat((attn_weights, w[..., None]), dim=-1) + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + qg.dtype + ) + L0 = k0.shape[-2] + attn_output = torch.matmul(attn_weights[..., :L0], v0) # [B, KVH, G, rows, D] + for i, ve in enumerate(diag_vs): + attn_output = attn_output + attn_weights[..., L0 + i, None] * ve.unsqueeze(2) + return attn_output + + +def _trimmed_rows_attention(attn, hidden_states, cache_hidden, ctx): + """TTT steps >= 1 with backbone-row trimming (``--trim-backbone-rows``). + + ``hidden_states`` holds only the surviving rows R_i (the union of every row + that can still emit loss at this or a later step). The math mirrors the + sdpa TTT branch: full cross attention against step-0 K/V plus one + diagonal-private term per later step. It is shared by every backend + because it is plain matmul math -- the causal mask is built from the rows' + absolute positions, so non-contiguous row sets (which flash kernels cannot + express) are fine. Assumes batch=1 with no padding mask (the position-built + causal mask is the only masking applied). Everything runs in + grouped-query layout so the step-0 + K/V stay at ``num_key_value_heads`` and are never materialized per query + head. + + ``ctx`` carries per-step fields set by the TTT loop: + positions LongTensor [1, n_i] absolute RoPE positions (global under USP) + diag_sels dict: cache entry e -> positions of R_i within R_e + per-forward constants: ``k0_len`` (number of step-0 K/V rows; global length + under USP), and persistent fields ``k0``/``v0`` ([B, KVH, {G|1}, L0, D]), + prepared once via the backend's ``_trim_step0_kv``. + + Later steps' K/V are appended un-repeated as [B, KVH, n_i, D]; they are + only ever consumed by this function (via ``diag_sels``), never by the + step-0 native paths. + """ + if isinstance(attn.rotary_emb, LlamaMutiRotaryEmbedding): + # Multimodal RoPE carries 3D position ids and a different rotary call + # signature; the trimmed path's absolute-position bookkeeping does not + # support it yet. Fail fast rather than mis-slice positions. + raise NotImplementedError( + "trim_backbone_rows does not support multimodal RoPE (mrope) drafts" + ) + bsz, q_len, _ = hidden_states.size() + lck = len(cache_hidden[0]) + kvh = attn.num_key_value_heads + groups = attn.num_key_value_groups + + query_states = attn.q_proj(hidden_states).view( + bsz, q_len, attn.num_heads, attn.head_dim + ) + key_states = attn.k_proj(hidden_states).view(bsz, q_len, kvh, attn.head_dim) + value_states = attn.v_proj(hidden_states).view(bsz, q_len, kvh, attn.head_dim) + + # RoPE at the rows' absolute positions. The cos/sin cache must cover the + # largest absolute position, not q_len (the rows are a sparse subset). + cos, sin = attn.rotary_emb(query_states, seq_len=ctx["k0_len"] + lck) + cos, sin = cos.to(query_states.device), sin.to(query_states.device) + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin, ctx["positions"] + lck, unsqueeze_dim=2 + ) + + key_states = key_states.transpose(1, 2) # [B, KVH, n_i, D] + value_states = value_states.transpose(1, 2) + cache_hidden[0] = cache_hidden[0] + [key_states] + cache_hidden[1] = cache_hidden[1] + [value_states] + + if ctx.get("k0") is None: + ctx["k0"], ctx["v0"] = attn._trim_step0_kv(cache_hidden) + k0, v0 = ctx["k0"], ctx["v0"] # [B, KVH, {G|1}, L0, D] + scale = 1.0 / math.sqrt(attn.head_dim) + + # Grouped-query layout: [B, n, H, D] -> [B, KVH, G, n, D] + qg = query_states.view(bsz, q_len, kvh, groups, attn.head_dim).permute( + 0, 2, 3, 1, 4 + ) + + # Row-align the diagonal-private K/V (cheap index_selects, kept outside + # the checkpointed region so the core is pure math on plain tensors). + diag_ks, diag_vs = [], [] + n_entries = len(cache_hidden[0]) + for e in range(1, n_entries): + ke = cache_hidden[0][e] + ve = cache_hidden[1][e] + if e < n_entries - 1: + sel = ctx["diag_sels"][e] + ke = ke.index_select(2, sel) + ve = ve.index_select(2, sel) + diag_ks.append(ke) + diag_vs.append(ve) + + # The heavy part (cross attention against step-0 K/V + diagonal terms) is + # activation-checkpointed: its [n_i, L0] fp32 softmax tensors are the + # dominant residency and are cheap to recompute. + neg_inf = torch.finfo(qg.dtype).min + if ctx["k0_len"] >= _TRIM_SDPA_MIN_K0: + # k0 arrives as [B, KVH, {G|1}, L0, D]; index 0 of dim 2 recovers the + # un-expanded [B, KVH, L0, D] view in both layouts for free. + k0f, v0f = k0[:, :, 0], v0[:, :, 0] + if torch.is_grad_enabled() and qg.requires_grad: + from torch.utils.checkpoint import checkpoint + + attn_output = checkpoint( + _trimmed_attention_sdpa, + qg, + k0f, + v0f, + ctx["positions"], + *diag_ks, + *diag_vs, + use_reentrant=False, + ) + else: + attn_output = _trimmed_attention_sdpa( + qg, k0f, v0f, ctx["positions"], *diag_ks, *diag_vs + ) + elif torch.is_grad_enabled() and qg.requires_grad: + # Checkpoint PER ROW-CHUNK: backward recomputes one chunk at a time, + # so the [rows, L0] fp32 transients stay bounded in both directions. + from torch.utils.checkpoint import checkpoint + + n = qg.shape[3] + outs = [] + for s in range(0, n, _TRIM_ROW_CHUNK): + e = min(s + _TRIM_ROW_CHUNK, n) + chunk_diag = [d[:, :, s:e] for d in (*diag_ks, *diag_vs)] + outs.append( + checkpoint( + _trimmed_attention_rows, + qg[:, :, :, s:e], + k0, + v0, + ctx["positions"][:, s:e], + scale, + neg_inf, + *chunk_diag, + use_reentrant=False, + ) + ) + attn_output = outs[0] if len(outs) == 1 else torch.cat(outs, dim=3) + else: + attn_output = _trimmed_attention_core( + qg, k0, v0, ctx["positions"], scale, neg_inf, *diag_ks, *diag_vs + ) + + attn_output = attn_output.permute(0, 3, 1, 2, 4).reshape( + bsz, q_len, attn.head_dim * attn.num_heads + ) + return attn.o_proj(attn_output) + + class LlamaDecoderLayer(nn.Module): def __init__(self, config, attention_backend: str = "sdpa"): super().__init__() @@ -1605,6 +1913,7 @@ def forward( past_key_values: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, + trim_rows_ctx: Optional[dict] = None, ) -> Tuple[ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] ]: @@ -1637,6 +1946,7 @@ def forward( past_key_values=past_key_values, output_attentions=output_attentions, use_cache=use_cache, + trim_rows_ctx=trim_rows_ctx, ) hidden_states = residual + hidden_states @@ -1785,6 +2095,7 @@ def backbone( position_ids: torch.Tensor, past_key_values: Optional[Cache] = None, use_cache: bool = True, + trim_rows_ctx: Optional[dict] = None, ) -> torch.Tensor: return self.midlayer( input_emb=input_embeds, @@ -1795,4 +2106,5 @@ def backbone( past_key_values=past_key_values, output_attentions=False, use_cache=False, + trim_rows_ctx=trim_rows_ctx, ) diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index 4d262cec9..dae1f3b50 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -173,6 +173,7 @@ def __init__( target_head: Optional[nn.Module] = None, ploss_decay: float = 0.8, trim_loss_positions: bool = False, + trim_backbone_rows: bool = False, compact_teacher: bool = False, compact_teacher_chunk_size: Optional[int] = None, ) -> None: @@ -180,6 +181,7 @@ def __init__( self.target_head = target_head self.ploss_decay = ploss_decay self.trim_loss_positions = trim_loss_positions + self.trim_backbone_rows = trim_backbone_rows self.compact_teacher = compact_teacher self.compact_teacher_chunk_size = compact_teacher_chunk_size if compact_teacher: @@ -320,6 +322,7 @@ def forward_loss( else None ), trim_loss_positions=self.trim_loss_positions, + trim_backbone_rows=self.trim_backbone_rows, **compact_kwargs, ) weights = [self.ploss_decay**i for i in range(len(plosses))] diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index 680069324..f8a4b7846 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -59,6 +59,7 @@ def test_compact_teacher_reaches_the_eagle3_step_provider(self): resolved.algorithm.providers.step.options(cfg), { "trim_loss_positions": True, + "trim_backbone_rows": False, "compact_teacher": True, "compact_teacher_chunk_size": 2048, }, @@ -81,6 +82,66 @@ def test_trim_loss_positions_rejects_non_eagle3_strategy(self): ): resolve_run(cfg) + def test_trim_backbone_rows_reaches_the_eagle3_step_provider(self): + cfg = Config.model_validate( + { + **OFFLINE_EAGLE3, + "training": { + "attention_backend": "sdpa", + "trim_loss_positions": True, + "trim_backbone_rows": True, + }, + } + ) + resolved = resolve_run(cfg) + + self.assertEqual( + resolved.algorithm.providers.step.options(cfg)["trim_backbone_rows"], + True, + ) + + def test_trim_backbone_rows_requires_trim_loss_positions(self): + with self.assertRaisesRegex( + ValueError, "requires\\s+training.trim_loss_positions" + ): + Config.model_validate( + { + **OFFLINE_EAGLE3, + "training": { + "attention_backend": "sdpa", + "trim_backbone_rows": True, + }, + } + ) + + def test_trim_backbone_rows_rejects_flex_attention(self): + with self.assertRaisesRegex(ValueError, "sdpa/fa/usp only"): + Config.model_validate( + { + **OFFLINE_EAGLE3, + "training": { + "attention_backend": "flex_attention", + "trim_loss_positions": True, + "trim_backbone_rows": True, + }, + } + ) + + def test_trim_backbone_rows_rejects_compact_teacher(self): + with self.assertRaisesRegex(ValueError, "incompatible with"): + Config.model_validate( + { + **OFFLINE_EAGLE3, + "training": { + "attention_backend": "sdpa", + "trim_loss_positions": True, + "trim_backbone_rows": True, + "compact_teacher": True, + "compact_teacher_chunk_size": 2048, + }, + } + ) + def test_loader_and_profiler_options_reach_the_canonical_trainer(self): eagle = resolve_run(Config.model_validate(OFFLINE_EAGLE3)) dflash = resolve_run( diff --git a/tests/test_runtime/test_trim_backbone_rows.py b/tests/test_runtime/test_trim_backbone_rows.py new file mode 100644 index 000000000..52aeda09f --- /dev/null +++ b/tests/test_runtime/test_trim_backbone_rows.py @@ -0,0 +1,454 @@ +# coding=utf-8 +"""trim_backbone_rows: TTT steps >= 1 forward only rows that can still emit loss. + +Layers (CI keeps guarding the union-row math even without GPUs): + +1. ``TestBackboneRowsGolden`` (CPU) -- hand-derived literal tables for + ``R_i = union_{j >= i} { s - j : s in sup, 0 <= s - j < chunk_len }``: + the nested-union semantics, the overlap-tail bound, dead-tail padding and + non-USP back-compat. (The naive "reuse ``sup`` at every step" reading is + wrong -- loss rows shift left once per TTT step -- and these tables pin the + corrected math.) +2. ``TestEquivTrimBackboneSdpa`` (one GPU) -- per-step loss parity, backbone + trim ON vs OFF (both with trim_loss_positions on): an all-supervised mask + where R_i covers every row (trivial-equality boundary: any row/position/ + mask error is exposed without tolerance cover), plus prompt-heavy, + block-boundary, and isolated-point masks, plus a total-grad-norm check. +3. ``TestEquivTrimBackboneFa`` -- same masks with attention_backend=fa + (step 0 native flash, trimmed steps on the shared manual path). +4. ``TestEquivTrimBackboneUspFourRank`` -- ring-4 offline pipeline: the + trimmed steps bypass ring/Ulysses and attend an all-gathered step-0 K/V, + so per-rank row counts may differ freely; masks include supervision + straddling every rank boundary and a tail-only rank (dead steps + ranks + falling back to the untrimmed path). +""" + +import json +import os +import shutil +import tempfile +import unittest +from unittest import mock + +import torch + +CUDA = torch.cuda.is_available() +NGPU = torch.cuda.device_count() if CUDA else 0 +WORLD_SIZE = 4 +SEQ = 48 +TTT = 3 + + +def _has_standard_flash_attention() -> bool: + try: + from flash_attn import flash_attn_varlen_func # noqa: F401 + from flash_attn.bert_padding import pad_input, unpad_input # noqa: F401 + from flash_attn.flash_attn_interface import ( # noqa: F401 + _flash_attn_varlen_backward, + ) + except Exception: + return False + return True + + +class TestBackboneRowsGolden(unittest.TestCase): + """Hand-derived expected outputs for the B-level row sets (CPU only).""" + + def _mk(self, mask_list, seed=1, vocab=32, draft=8): + g = torch.Generator().manual_seed(seed) + ids = torch.randperm(vocab, generator=g)[:draft].sort().values + t2d = torch.zeros(vocab, dtype=torch.bool) + t2d[ids] = True + L = len(mask_list) + lm = torch.tensor(mask_list, dtype=torch.long).view(1, L, 1) + tgt = torch.randn(1, L, vocab, generator=g) + return tgt, t2d, lm + + def test_union_with_tail_bound(self): + # C=6, k=3, mask=[0,0,0,0,1,1,1,0] -> sup={4,5,6} + # loss rows: j0={4,5} j1={3,4,5} j2={2,3,4} + # R_1 = j1 U j2 = {2,3,4,5}; R_2 = j2 = {2,3,4} + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 1, 1, 1, 0]) + p = _build_trim_pack(tgt, t2d, lm, length=3, chunk_len=6, backbone_rows=True) + self.assertEqual(p["b_rows_steps"][1].tolist(), [2, 3, 4, 5]) + self.assertEqual(p["b_rows_steps"][2].tolist(), [2, 3, 4]) + # step-1 hidden comes from the full-length step-0 output + # (R_i nested in R_{i-1}: step-i prev-selection = b_diag_sels[i][i-1]) + # loss rows inside the compact row space + self.assertEqual(p["b_loss_sel"][1].tolist(), [1, 2, 3]) # {3,4,5} in R_1 + self.assertEqual(p["b_loss_sel"][2].tolist(), [0, 1, 2]) # {2,3,4} in R_2 + # diagonal map: R_2 inside R_1 + self.assertEqual(p["b_diag_sels"][2][1].tolist(), [0, 1, 2]) + + def test_allones_covers_every_row(self): + # mask = all ones over 8 slots, C=6, k=3 -> R_1 = R_2 = {0..5} + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([1] * 8) + p = _build_trim_pack(tgt, t2d, lm, length=3, chunk_len=6, backbone_rows=True) + self.assertEqual(p["b_rows_steps"][1].tolist(), [0, 1, 2, 3, 4, 5]) + self.assertEqual(p["b_rows_steps"][2].tolist(), [0, 1, 2, 3, 4, 5]) + + def test_isolated_points_shrinking_chain(self): + # non-USP: C=L=8, k=3, mask=[1,0,0,1,0,0,0,0] -> sup={0,3} + # loss rows: j0={0,3} j1={2} j2={1} + # R_1 = {1,2}; R_2 = {1} + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([1, 0, 0, 1, 0, 0, 0, 0]) + p = _build_trim_pack(tgt, t2d, lm, length=3, backbone_rows=True) + self.assertEqual(p["b_rows_steps"][1].tolist(), [1, 2]) + self.assertEqual(p["b_rows_steps"][2].tolist(), [1]) + self.assertEqual(p["b_diag_sels"][2][1].tolist(), [0]) + + def test_dead_tail_keeps_nested_padding(self): + # sup so late that the last step has no reachable row: the padded row + # must come from the previous set (nested chain survives). + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 1, 0, 0, 0, 0, 0, 0]) + # sup={1}: loss rows j0={1} j1={0} j2={} (1-2 < 0) + p = _build_trim_pack(tgt, t2d, lm, length=3, backbone_rows=True) + self.assertEqual(p["b_rows_steps"][1].tolist(), [0]) + self.assertEqual(p["b_rows_steps"][2].tolist(), [0]) # padded from R_1 + self.assertEqual(p["nrows_steps"][2], 0) + + def test_off_flag_adds_no_keys(self): + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 1, 1, 1, 0]) + p = _build_trim_pack(tgt, t2d, lm, length=3, chunk_len=6) + self.assertNotIn("b_rows_steps", p) + + +def _masks_single(): + m = {} + allones = torch.ones(SEQ, dtype=torch.long) + m["allones"] = allones + prompt_heavy = torch.zeros(SEQ, dtype=torch.long) + prompt_heavy[SEQ // 2 :] = 1 + m["prompt_heavy"] = prompt_heavy + blocks = torch.zeros(SEQ, dtype=torch.long) + blocks[[10, 11, 12, 13, 22, 23, 24, 25, 34, 35, 36, 37]] = 1 + m["blocks"] = blocks + isolated = torch.zeros(SEQ, dtype=torch.long) + isolated[[0, 9, 21, 33, 45]] = 1 + m["isolated"] = isolated + return m + + +class _SingleRankEquivBase(unittest.TestCase): + backend = "sdpa" + allones_tol = 1e-6 + + @classmethod + def _build(cls, workdir): + from tests.test_runtime import _fixtures as fx + from specforge.algorithms.eagle3.model import OnlineEagle3Model + from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig + from specforge.modeling.target.target_head import TargetHead + + fx.build_single_rank_distributed(port="29881") + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + fx.write_draft_config(os.path.join(workdir, "draft.json")) + fx.write_target_head_dir(os.path.join(workdir, "target")) + fx.write_vocab_mapping(os.path.join(workdir, "vm.pt")) + cfg = AutoDraftModelConfig.from_file(os.path.join(workdir, "draft.json")) + dm = AutoDraftModel.from_config( + cfg, attention_backend=cls.backend, torch_dtype=torch.bfloat16 + ).cuda() + dm.load_vocab_mapping(os.path.join(workdir, "vm.pt")) + dm.freeze_embedding() + model = OnlineEagle3Model( + draft_model=dm, length=TTT, attention_backend=cls.backend + ).cuda() + head = TargetHead.from_pretrained( + os.path.join(workdir, "target"), lm_head_key="lm_head.weight" + ) + return model, head + + def _batch(self, head, loss_mask): + from tests.test_runtime import _fixtures as fx + + g = torch.Generator().manual_seed(11) + input_ids = torch.randint(0, fx.V, (1, SEQ), generator=g) + hidden = torch.randn(1, SEQ, fx.H, generator=g).to(torch.bfloat16) + aux = torch.randn(1, SEQ, 3 * fx.H, generator=g).to(torch.bfloat16) + input_ids, target_hidden, lm = head.preprocess( + input_ids, hidden, loss_mask.view(1, SEQ) + ) + target = head(target_hidden.cuda()) + return dict( + input_ids=input_ids.cuda(), + attention_mask=torch.ones(1, SEQ, device="cuda"), + loss_mask=lm.cuda(), + target=target, + hidden_states=aux.cuda(), + ) + + def _step_losses(self, model, batch, trim_backbone): + with torch.no_grad(): + plosses, *_ = model( + trim_loss_positions=True, + trim_backbone_rows=trim_backbone, + **batch, + ) + return [float(p.item()) for p in plosses] + + def _grad_norm(self, model, batch, trim_backbone): + model.zero_grad(set_to_none=True) + plosses, *_ = model( + trim_loss_positions=True, trim_backbone_rows=trim_backbone, **batch + ) + loss = sum(0.8**i * plosses[i] for i in range(len(plosses))) + loss.backward() + total = torch.zeros((), device="cuda", dtype=torch.float64) + for p in model.parameters(): + if p.grad is not None: + total += (p.grad.double() ** 2).sum() + model.zero_grad(set_to_none=True) + return float(total.sqrt().item()) + + def _run_case(self, name, mask, tol_fn): + workdir = tempfile.mkdtemp(prefix="trim_b_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + torch.use_deterministic_algorithms(True, warn_only=True) + model, head = self._build(workdir) + model.train() + batch = self._batch(head, mask) + off = self._step_losses(model, batch, False) + on = self._step_losses(model, batch, True) + self.assertEqual(len(off), TTT) + for j, (a, b) in enumerate(zip(off, on)): + self.assertLessEqual( + abs(a - b), tol_fn(a), msg=f"{name} step{j}: off={a} on={b}" + ) + g_off = self._grad_norm(model, batch, False) + g_on = self._grad_norm(model, batch, True) + self.assertLessEqual( + abs(g_on - g_off), 5e-3 * g_off, msg=f"{name} grads: {g_off} vs {g_on}" + ) + + def test_allones_exact(self): + self._run_case( + "allones", _masks_single()["allones"], lambda a: self.allones_tol + ) + + def test_sdpa_core_forced(self): + # Force the long-k0 SDPA core (normally k0_len >= 8192) through the + # same all-ones + adversarial comparisons; measured residual is ~1e-5 + # (fused-kernel vs native op ordering), far below any discrete error. + from specforge.modeling.draft import llama3_eagle as le + + with mock.patch.object(le, "_TRIM_SDPA_MIN_K0", 1): + masks = _masks_single() + for name in ("allones", "isolated"): + with self.subTest(mask=name): + self._run_case(name, masks[name], lambda a: 5e-4) + + def test_adversarial_masks(self): + masks = _masks_single() + for name in ("prompt_heavy", "blocks", "isolated"): + with self.subTest(mask=name): + self._run_case(name, masks[name], lambda a: max(1e-3 * abs(a), 1e-4)) + + +@unittest.skipUnless(CUDA, "requires one CUDA device") +class TestEquivTrimBackboneSdpa(_SingleRankEquivBase): + backend = "sdpa" + allones_tol = 1e-6 + + +@unittest.skipUnless( + CUDA and _has_standard_flash_attention(), + "requires CUDA and flash-attn", +) +class TestEquivTrimBackboneFa(_SingleRankEquivBase): + backend = "fa" + # step 0 stays on flash while the reference path keeps flash at every + # step; the trimmed steps change implementation, so bit-level equality is + # not expected even for the all-ones mask. + allones_tol = 1e-5 + + +def _write_usp_workdir(workdir): + from tests.test_runtime import _fixtures as fx + + fx.write_draft_config(os.path.join(workdir, "draft.json")) + fx.write_target_head_dir(os.path.join(workdir, "target")) + fx.write_vocab_mapping(os.path.join(workdir, "vocab_mapping.pt")) + masks = {} + m1 = torch.ones(SEQ, dtype=torch.long) + m1[-1] = 0 + masks["allones"] = m1 + m3 = torch.zeros(SEQ, dtype=torch.long) + m3[[10, 11, 12, 13, 22, 23, 24, 25, 34, 35, 36, 37]] = 1 + masks["boundary"] = m3 + m4 = torch.zeros(SEQ, dtype=torch.long) + m4[[12, 13]] = 1 + masks["tailonly"] = m4 + g = torch.Generator().manual_seed(11) + base_input = torch.randint(0, fx.V, (SEQ,), generator=g) + base_hid = torch.randn(1, SEQ, fx.H, generator=g).to(torch.bfloat16) + base_aux = torch.randn(1, SEQ, 3 * fx.H, generator=g).to(torch.bfloat16) + for name, lm in masks.items(): + d = os.path.join(workdir, f"features_{name}") + os.makedirs(d, exist_ok=True) + torch.save( + { + "input_ids": base_input.clone(), + "loss_mask": lm.clone(), + "hidden_state": base_hid.clone(), + "aux_hidden_state": base_aux.clone(), + }, + os.path.join(d, "0000.ckpt"), + ) + + +def _usp_worker(rank, world_size, port, workdir): + from tests.test_runtime import _fixtures as fx + + fx.init_rank_distributed( + rank, world_size, tp_size=1, sp_ulysses_size=1, sp_ring_size=4, port=str(port) + ) + try: + import torch.distributed as dist + + from specforge.algorithms.builtin import builtin_algorithm_registry + from specforge.algorithms.eagle3.model import OnlineEagle3Model + from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig + from specforge.modeling.target.target_head import TargetHead + from specforge.runtime.data_plane import FeatureDataLoader, LocalFeatureStore + + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + torch.use_deterministic_algorithms(True, warn_only=True) + cfg = AutoDraftModelConfig.from_file(os.path.join(workdir, "draft.json")) + dm = AutoDraftModel.from_config( + cfg, attention_backend="usp", torch_dtype=torch.bfloat16 + ).cuda() + dm.load_vocab_mapping(os.path.join(workdir, "vocab_mapping.pt")) + dm.freeze_embedding() + model = OnlineEagle3Model( + draft_model=dm, length=TTT, attention_backend="usp" + ).cuda() + model.train() + target_head = TargetHead.from_pretrained( + os.path.join(workdir, "target"), lm_head_key="lm_head.weight" + ) + algorithm = builtin_algorithm_registry().resolve("eagle3") + provider = algorithm.providers.offline_for("text") + + results = {} + for case in ("allones", "boundary", "tailonly"): + refs = provider.build_reader( + os.path.join(workdir, f"features_{case}"), + run_id=f"trimb-{case}", + ttt_length=TTT, + max_len=SEQ, + ).read() + loader = FeatureDataLoader( + LocalFeatureStore(f"trimb-{case}-{rank}"), + refs=refs, + batch_size=1, + collate_fn=provider.build_collator(), + per_sample_transform=provider.build_normalizer( + SEQ, ttt_length=TTT, use_usp_preprocess=True + ), + strategy=algorithm.name, + ) + batch = next(iter(loader)) + + def step_losses(trim_backbone): + strat = algorithm.providers.step.build( + model, + target_head=target_head, + trim_loss_positions=True, + trim_backbone_rows=trim_backbone, + ) + with torch.no_grad(): + out = strat.forward_loss(batch) + return [float(p.item()) for p in out.metrics["plosses"]] + + results[case] = {"off": step_losses(False), "on": step_losses(True)} + + gathered = [None] * world_size + dist.all_gather_object(gathered, results) + if rank == 0: + with open(os.path.join(workdir, "results.json"), "w") as fh: + json.dump(gathered, fh) + dist.barrier() + finally: + from specforge.distributed import destroy_distributed + + destroy_distributed() + + +def _usp_worker_forced_sdpa(rank, world_size, port, workdir): + from specforge.modeling.draft import llama3_eagle as le + + le._TRIM_SDPA_MIN_K0 = 1 # force the long-k0 SDPA core in every worker + _usp_worker(rank, world_size, port, workdir) + + +@unittest.skipUnless( + CUDA and NGPU >= WORLD_SIZE and _has_standard_flash_attention(), + "requires four CUDA devices and the standard flash-attn USP interfaces", +) +class TestEquivTrimBackboneUspFourRank(unittest.TestCase): + def test_sdpa_core_matches_on_ring4(self): + import torch.multiprocessing as mp + + workdir = tempfile.mkdtemp(prefix="trim_b_usp_sdpa_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + _write_usp_workdir(workdir) + mp.spawn( + _usp_worker_forced_sdpa, + args=(WORLD_SIZE, 29886, workdir), + nprocs=WORLD_SIZE, + join=True, + ) + with open(os.path.join(workdir, "results.json")) as fh: + gathered = json.load(fh) + for case in ("allones", "boundary", "tailonly"): + for rank, res in enumerate(gathered): + for j, (a, b) in enumerate(zip(res[case]["off"], res[case]["on"])): + tol = 5e-4 if case == "allones" else max(1e-3 * abs(a), 2e-4) + self.assertLessEqual( + abs(a - b), + tol, + msg=f"sdpa-core {case} rank{rank} step{j}: off={a} on={b}", + ) + + def test_backbone_trim_matches_on_ring4(self): + import torch.multiprocessing as mp + + workdir = tempfile.mkdtemp(prefix="trim_b_usp_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + _write_usp_workdir(workdir) + mp.spawn( + _usp_worker, + args=(WORLD_SIZE, 29885, workdir), + nprocs=WORLD_SIZE, + join=True, + ) + with open(os.path.join(workdir, "results.json")) as fh: + gathered = json.load(fh) + for case in ("allones", "boundary", "tailonly"): + for rank, res in enumerate(gathered): + off, on = res[case]["off"], res[case]["on"] + self.assertEqual(len(off), TTT) + for j, (a, b) in enumerate(zip(off, on)): + tol = 2e-4 if case == "allones" else max(1e-3 * abs(a), 1e-4) + self.assertLessEqual( + abs(a - b), + tol, + msg=f"{case} rank{rank} step{j}: off={a} on={b}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) From 3aa6f7c03302459c31b667cf1c8bc46d63556fe7 Mon Sep 17 00:00:00 2001 From: julyanghar Date: Tue, 18 Aug 2026 19:38:33 -0500 Subject: [PATCH 2/6] style: isort import order in trim_backbone_rows tests Co-Authored-By: Claude Fable 5 --- tests/test_runtime/test_trim_backbone_rows.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_runtime/test_trim_backbone_rows.py b/tests/test_runtime/test_trim_backbone_rows.py index 52aeda09f..d665a76f5 100644 --- a/tests/test_runtime/test_trim_backbone_rows.py +++ b/tests/test_runtime/test_trim_backbone_rows.py @@ -145,10 +145,10 @@ class _SingleRankEquivBase(unittest.TestCase): @classmethod def _build(cls, workdir): - from tests.test_runtime import _fixtures as fx from specforge.algorithms.eagle3.model import OnlineEagle3Model from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig from specforge.modeling.target.target_head import TargetHead + from tests.test_runtime import _fixtures as fx fx.build_single_rank_distributed(port="29881") torch.manual_seed(0) From 175b3c8c8c4bd2476698994c1f5d36bbdfc836bc Mon Sep 17 00:00:00 2001 From: julyanghar Date: Tue, 18 Aug 2026 20:07:29 -0500 Subject: [PATCH 3/6] test: cover trim_backbone_rows in builtin provider contract fixtures Co-Authored-By: Claude Fable 5 --- tests/test_algorithms/test_builtin_providers.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index dc92b1d91..256506a27 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -166,6 +166,7 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): training = SimpleNamespace( attention_backend="flex_attention", trim_loss_positions=True, + trim_backbone_rows=False, compact_teacher=True, compact_teacher_chunk_size=1024, lambda_base_start=0.75, @@ -217,6 +218,7 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): "eagle3_kl_scale", "eagle3_kl_decay", "eagle3_trim_loss_positions", + "eagle3_trim_backbone_rows", "eagle3_compact_teacher", }, "peagle": { @@ -264,6 +266,7 @@ def test_step_runtime_config_binds_options_contract_and_missing_key_policy(self) training=SimpleNamespace( attention_backend="flex_attention", trim_loss_positions=False, + trim_backbone_rows=False, compact_teacher=False, compact_teacher_chunk_size=None, ) @@ -299,6 +302,7 @@ def test_step_runtime_config_binds_options_contract_and_missing_key_policy(self) ( ("compact_teacher", False), ("compact_teacher_chunk_size", None), + ("trim_backbone_rows", False), ("trim_loss_positions", False), ), ) From b027e2eb681df9e797f84905b7c74d896d12aefd Mon Sep 17 00:00:00 2001 From: julyanghar Date: Tue, 18 Aug 2026 20:55:55 -0500 Subject: [PATCH 4/6] test: widen fa all-ones tolerance for cross-hardware kernel noise 1e-5 was calibrated on one GPU model; CI's hardware measures 1.2e-5. 1e-4 still sits ~250x below the smallest discrete semantic error (~loss/C). Co-Authored-By: Claude Fable 5 --- tests/test_runtime/test_trim_backbone_rows.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_runtime/test_trim_backbone_rows.py b/tests/test_runtime/test_trim_backbone_rows.py index d665a76f5..a2284a051 100644 --- a/tests/test_runtime/test_trim_backbone_rows.py +++ b/tests/test_runtime/test_trim_backbone_rows.py @@ -271,7 +271,10 @@ class TestEquivTrimBackboneFa(_SingleRankEquivBase): # step 0 stays on flash while the reference path keeps flash at every # step; the trimmed steps change implementation, so bit-level equality is # not expected even for the all-ones mask. - allones_tol = 1e-5 + # 1e-4 absorbs cross-hardware kernel noise (1.2e-5 measured on CI's GPUs + # vs <1e-5 locally) while staying ~250x below the smallest possible + # discrete row/denominator error (~loss/C). + allones_tol = 1e-4 def _write_usp_workdir(workdir): From 0a619a205924446b34d0408e87f73195224aa943 Mon Sep 17 00:00:00 2001 From: julyanghar Date: Thu, 20 Aug 2026 17:14:27 -0500 Subject: [PATCH 5/6] feat(eagle3): skip backbone-row trimming when supervision is too dense Row trimming replaces flash attention with an explicit masked-matmul kernel that costs more per row, so it only pays while the kept-row sets are small. Measured here (8192-token chunk, TTT=7, RTX 6000 Ada): flash attention runs at ~16.7 us/row against 28-62 us/row for the trimmed kernel, so step time breaks even near 40% density and degrades to 2.4x slower at full supervision, where no row can be dropped at all. On ShareGPT that regime is the common case, not the exception: samples under 4096 tokens are 99.7% of the corpus and their median supervised fraction is 0.92. Compare the two sides directly -- trimmed rows sum_i |R_i| against full rows (k-1) * chunk -- and keep the B path only while the former stays under trim_backbone_rows_max_density (default 0.35, set 1.0 to always trim). The decision is per sample rather than per step: |R_i| barely varies across steps for contiguous supervision, and a mixed full/trimmed unroll would need per-backend K/V layout conversion plus an extra head-dimension gather under USP. Under USP the ranks agree by MIN, for the same reason the enable flag does: a mixed decision would mismatch collectives and hang. Existing equivalence tests pin the guard off so they keep exercising the trimmed kernels; new tests cover both sides of the threshold, that the guard actually fires, and that the decision is rank-uniform on 4 ranks. Co-Authored-By: Claude Fable 5 --- examples/configs/README.md | 1 + specforge/algorithms/eagle3/model.py | 56 +++++++ specforge/algorithms/model_providers.py | 1 + specforge/config/schema.py | 9 ++ specforge/training/strategies/base.py | 3 + .../test_algorithms/test_builtin_providers.py | 3 + .../test_unified_feature_reachability.py | 1 + tests/test_runtime/test_trim_backbone_rows.py | 152 +++++++++++++++++- 8 files changed, 218 insertions(+), 8 deletions(-) diff --git a/examples/configs/README.md b/examples/configs/README.md index e7a1afb90..22a818baa 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -272,6 +272,7 @@ Common fields: | `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`. | diff --git a/specforge/algorithms/eagle3/model.py b/specforge/algorithms/eagle3/model.py index 5ab8fba66..09710ac9c 100644 --- a/specforge/algorithms/eagle3/model.py +++ b/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 @@ -263,6 +266,7 @@ def forward( 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], @@ -290,6 +294,9 @@ def forward( >= 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 @@ -471,6 +478,55 @@ def forward( 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): b_active = b_enabled and idx >= 1 if b_active: diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index f729d6a44..56b75ae68 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -438,6 +438,7 @@ 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, } diff --git a/specforge/config/schema.py b/specforge/config/schema.py index c8197f40f..1f75e05e1 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -565,6 +565,15 @@ class TrainingConfig(StrictConfigModel): #: attention. Requires trim_loss_positions and an attention backend in #: {sdpa, fa, usp}. trim_backbone_rows: bool = False + #: Per-step guard for trim_backbone_rows. A trimmed step replaces flash + #: attention with an explicit masked-matmul kernel that costs more per row, + #: so trimming only pays while the kept-row set is small: measured on a + #: 8192-token sequence the step time breaks even around 40% density and + #: degrades to 2.4x slower at full supervision. Steps whose kept-row count + #: exceeds this fraction of the sequence therefore run the full-length path + #: instead. |R_i| is non-increasing in i, so the full-length steps are + #: always a prefix. Set to 1.0 to disable the guard (always trim). + trim_backbone_rows_max_density: float = Field(default=0.35, gt=0.0, le=1.0) #: DFlash-family objective/model knobs. num_anchors: int = Field(default=512, gt=0) loss_decay_gamma: Optional[float] = None diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index dae1f3b50..a72fddd39 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -174,6 +174,7 @@ def __init__( ploss_decay: float = 0.8, trim_loss_positions: bool = False, trim_backbone_rows: bool = False, + trim_backbone_rows_max_density: float = 0.35, compact_teacher: bool = False, compact_teacher_chunk_size: Optional[int] = None, ) -> None: @@ -182,6 +183,7 @@ def __init__( self.ploss_decay = ploss_decay self.trim_loss_positions = trim_loss_positions self.trim_backbone_rows = trim_backbone_rows + self.trim_backbone_rows_max_density = trim_backbone_rows_max_density self.compact_teacher = compact_teacher self.compact_teacher_chunk_size = compact_teacher_chunk_size if compact_teacher: @@ -323,6 +325,7 @@ def forward_loss( ), trim_loss_positions=self.trim_loss_positions, trim_backbone_rows=self.trim_backbone_rows, + trim_backbone_rows_max_density=(self.trim_backbone_rows_max_density), **compact_kwargs, ) weights = [self.ploss_decay**i for i in range(len(plosses))] diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index 256506a27..0602639af 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -167,6 +167,7 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): attention_backend="flex_attention", trim_loss_positions=True, trim_backbone_rows=False, + trim_backbone_rows_max_density=0.35, compact_teacher=True, compact_teacher_chunk_size=1024, lambda_base_start=0.75, @@ -267,6 +268,7 @@ def test_step_runtime_config_binds_options_contract_and_missing_key_policy(self) attention_backend="flex_attention", trim_loss_positions=False, trim_backbone_rows=False, + trim_backbone_rows_max_density=0.35, compact_teacher=False, compact_teacher_chunk_size=None, ) @@ -303,6 +305,7 @@ def test_step_runtime_config_binds_options_contract_and_missing_key_policy(self) ("compact_teacher", False), ("compact_teacher_chunk_size", None), ("trim_backbone_rows", False), + ("trim_backbone_rows_max_density", 0.35), ("trim_loss_positions", False), ), ) diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index f8a4b7846..aab2a1c30 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -60,6 +60,7 @@ def test_compact_teacher_reaches_the_eagle3_step_provider(self): { "trim_loss_positions": True, "trim_backbone_rows": False, + "trim_backbone_rows_max_density": 0.35, "compact_teacher": True, "compact_teacher_chunk_size": 2048, }, diff --git a/tests/test_runtime/test_trim_backbone_rows.py b/tests/test_runtime/test_trim_backbone_rows.py index a2284a051..550b29d3a 100644 --- a/tests/test_runtime/test_trim_backbone_rows.py +++ b/tests/test_runtime/test_trim_backbone_rows.py @@ -23,7 +23,9 @@ falling back to the untrimmed path). """ +import io import json +import logging import os import shutil import tempfile @@ -142,6 +144,7 @@ def _masks_single(): class _SingleRankEquivBase(unittest.TestCase): backend = "sdpa" allones_tol = 1e-6 + max_density = 1.0 @classmethod def _build(cls, workdir): @@ -194,6 +197,11 @@ def _step_losses(self, model, batch, trim_backbone): plosses, *_ = model( trim_loss_positions=True, trim_backbone_rows=trim_backbone, + # Pin the density guard off: these cases (allones in particular) + # sit above the default threshold, so with the guard active they + # would fall back to the full path and stop exercising the + # trimmed kernels while still passing. + trim_backbone_rows_max_density=self.max_density, **batch, ) return [float(p.item()) for p in plosses] @@ -201,7 +209,10 @@ def _step_losses(self, model, batch, trim_backbone): def _grad_norm(self, model, batch, trim_backbone): model.zero_grad(set_to_none=True) plosses, *_ = model( - trim_loss_positions=True, trim_backbone_rows=trim_backbone, **batch + trim_loss_positions=True, + trim_backbone_rows=trim_backbone, + trim_backbone_rows_max_density=self.max_density, + **batch, ) loss = sum(0.8**i * plosses[i] for i in range(len(plosses))) loss.backward() @@ -277,6 +288,57 @@ class TestEquivTrimBackboneFa(_SingleRankEquivBase): allones_tol = 1e-4 +@unittest.skipUnless(CUDA, "requires one CUDA device") +class TestDensityGateTrimSide(_SingleRankEquivBase): + """Budget wide open: every mask must still take (and match on) the B path.""" + + backend = "sdpa" + allones_tol = 1e-6 + max_density = 1.0 + + def test_gate_open_matches_full(self): + masks = _masks_single() + for name in ("blocks", "prompt_heavy", "isolated"): + with self.subTest(mask=name): + self._run_case(name, masks[name], lambda a: max(1e-3 * abs(a), 1e-4)) + + +@unittest.skipUnless(CUDA, "requires one CUDA device") +class TestDensityGateFullSide(_SingleRankEquivBase): + """Budget closed: the guard drops to the A-level path, results unchanged. + + The guard exists because trimmed steps use a costlier per-row kernel, so it + must be able to turn the B path off without changing what training sees. + Both sides of the threshold are therefore checked against the same + untrimmed reference. + """ + + backend = "sdpa" + allones_tol = 1e-6 + max_density = 0.01 + + def test_gate_closed_matches_full(self): + masks = _masks_single() + for name in ("blocks", "prompt_heavy", "isolated", "allones"): + with self.subTest(mask=name): + self._run_case(name, masks[name], lambda a: max(1e-3 * abs(a), 1e-4)) + + def test_gate_closed_actually_skips_trim(self): + # Guard against the silent-pass failure mode: a gate that never fires + # would make the case above vacuous, so assert the model said so. + workdir = tempfile.mkdtemp(prefix="trim_b_gate_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + model, head = self._build(workdir) + model.train() + batch = self._batch(head, _masks_single()["blocks"]) + with self.assertLogs("specforge.algorithms.eagle3.model", level="INFO") as cm: + self._step_losses(model, batch, True) + self.assertTrue( + any("trim_backbone_rows skipped" in line for line in cm.output), + f"gate should have closed, got: {cm.output}", + ) + + def _write_usp_workdir(workdir): from tests.test_runtime import _fixtures as fx @@ -293,6 +355,19 @@ def _write_usp_workdir(workdir): m4 = torch.zeros(SEQ, dtype=torch.long) m4[[12, 13]] = 1 masks["tailonly"] = m4 + # Ranks land on opposite sides of the density budget: the first quarter is + # fully supervised (its local row sets are near-full, so that rank alone + # would refuse to trim) while the rest carries one supervised position per + # quarter. Without a rank-uniform decision the ranks would run different + # code paths and their collectives would not line up. + m5 = torch.zeros(SEQ, dtype=torch.long) + m5[: SEQ // 4] = 1 + # Every rank's shard needs an adjacent supervised pair of its own: the data + # filter drops shards without one, which empties the loader on those ranks. + for q in (1, 2, 3): + base = q * (SEQ // 4) + m5[[base + 2, base + 3]] = 1 + masks["rank_split"] = m5 g = torch.Generator().manual_seed(11) base_input = torch.randint(0, fx.V, (SEQ,), generator=g) base_hid = torch.randn(1, SEQ, fx.H, generator=g).to(torch.bfloat16) @@ -346,7 +421,7 @@ def _usp_worker(rank, world_size, port, workdir): provider = algorithm.providers.offline_for("text") results = {} - for case in ("allones", "boundary", "tailonly"): + for case in ("allones", "boundary", "tailonly", "rank_split"): refs = provider.build_reader( os.path.join(workdir, f"features_{case}"), run_id=f"trimb-{case}", @@ -365,18 +440,40 @@ def _usp_worker(rank, world_size, port, workdir): ) batch = next(iter(loader)) - def step_losses(trim_backbone): + def step_losses(trim_backbone, density=1.0, capture=False): strat = algorithm.providers.step.build( model, target_head=target_head, trim_loss_positions=True, trim_backbone_rows=trim_backbone, + trim_backbone_rows_max_density=density, ) - with torch.no_grad(): - out = strat.forward_loss(batch) - return [float(p.item()) for p in out.metrics["plosses"]] - - results[case] = {"off": step_losses(False), "on": step_losses(True)} + # capture=True records whether this rank's gate closed, by + # listening for the model's own skip log. + buf = io.StringIO() + handler = logging.StreamHandler(buf) + target = logging.getLogger("specforge.algorithms.eagle3.model") + if capture: + target.addHandler(handler) + target.setLevel(logging.INFO) + try: + with torch.no_grad(): + out = strat.forward_loss(batch) + finally: + if capture: + target.removeHandler(handler) + losses = [float(p.item()) for p in out.metrics["plosses"]] + return (losses, buf.getvalue()) if capture else losses + + entry = {"off": step_losses(False), "on": step_losses(True)} + if case == "rank_split": + # Ranks sit on opposite sides of the budget here; the decision + # must still come out identical everywhere or the collectives + # would diverge. Capture the flag each rank actually used. + gated, out = step_losses(True, density=0.35, capture=True) + entry["gated"] = gated + entry["trim_flag"] = 0 if "trim_backbone_rows skipped" in out else 1 + results[case] = entry gathered = [None] * world_size dist.all_gather_object(gathered, results) @@ -401,6 +498,45 @@ def _usp_worker_forced_sdpa(rank, world_size, port, workdir): CUDA and NGPU >= WORLD_SIZE and _has_standard_flash_attention(), "requires four CUDA devices and the standard flash-attn USP interfaces", ) +class TestDensityGateUspAgreement(unittest.TestCase): + """The density decision must be identical on every rank. + + Per-rank chunks carry different supervision, so the row-count budget can + land on opposite sides of the threshold across ranks. A rank that trims + while another runs the full path issues a different sequence of + collectives, which hangs the job rather than producing a wrong number -- + so this asserts the agreed flag, not just the losses. + """ + + def test_gate_flag_is_rank_uniform(self): + import torch.multiprocessing as mp + + workdir = tempfile.mkdtemp(prefix="trim_b_gate_usp_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + _write_usp_workdir(workdir) + mp.spawn( + _usp_worker, + args=(WORLD_SIZE, 29887, workdir), + nprocs=WORLD_SIZE, + join=True, + ) + with open(os.path.join(workdir, "results.json")) as fh: + gathered = json.load(fh) + flags = [res["rank_split"]["trim_flag"] for res in gathered] + self.assertNotIn(-1, flags, "gate instrumentation did not fire on some rank") + self.assertEqual( + len(set(flags)), 1, f"ranks disagreed on the density gate: {flags}" + ) + for rank, res in enumerate(gathered): + off, gated = res["rank_split"]["off"], res["rank_split"]["gated"] + for j, (a, b) in enumerate(zip(off, gated)): + self.assertLessEqual( + abs(a - b), + max(1e-3 * abs(a), 2e-4), + msg=f"rank_split rank{rank} step{j}: off={a} gated={b}", + ) + + class TestEquivTrimBackboneUspFourRank(unittest.TestCase): def test_sdpa_core_matches_on_ring4(self): import torch.multiprocessing as mp From b604de822155ab5e53c76cc406f12ab20da44db0 Mon Sep 17 00:00:00 2001 From: julyanghar Date: Thu, 20 Aug 2026 22:38:52 -0500 Subject: [PATCH 6/6] docs: fix stale schema comment (guard is per sample, not per step) Co-Authored-By: Claude Fable 5 --- specforge/config/schema.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/specforge/config/schema.py b/specforge/config/schema.py index 1f75e05e1..78f9c95cb 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -565,14 +565,16 @@ class TrainingConfig(StrictConfigModel): #: attention. Requires trim_loss_positions and an attention backend in #: {sdpa, fa, usp}. trim_backbone_rows: bool = False - #: Per-step guard for trim_backbone_rows. A trimmed step replaces flash + #: Cost guard for trim_backbone_rows. A trimmed step replaces flash #: attention with an explicit masked-matmul kernel that costs more per row, - #: so trimming only pays while the kept-row set is small: measured on a - #: 8192-token sequence the step time breaks even around 40% density and - #: degrades to 2.4x slower at full supervision. Steps whose kept-row count - #: exceeds this fraction of the sequence therefore run the full-length path - #: instead. |R_i| is non-increasing in i, so the full-length steps are - #: always a prefix. Set to 1.0 to disable the guard (always trim). + #: so trimming only pays while the kept-row sets are small: measured on a + #: 8192-token sequence the step time breaks even around 40% supervision + #: density and degrades to 2.4x slower at full supervision. The guard + #: compares the trimmed work sum_i |R_i| against this fraction of the full + #: work (k-1) * chunk and, when it is exceeded, runs the whole sample on + #: the full-length path -- the decision is per sample, not per step, so a + #: single unroll never mixes trimmed and full steps. Set to 1.0 to disable + #: the guard (always trim). trim_backbone_rows_max_density: float = Field(default=0.35, gt=0.0, le=1.0) #: DFlash-family objective/model knobs. num_anchors: int = Field(default=512, gt=0)