diff --git a/examples/configs/README.md b/examples/configs/README.md index 451e9ac6c..fcdf8b1e6 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -233,6 +233,7 @@ Common fields: | `training.resume_from` | `null` | Full-run checkpoint/run root: draft, optimizer/scheduler, counters, data position, and RNG. Mutually exclusive with `model.draft_checkpoint_path`. | | `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` | 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.role` | `all` | Use `all` for local offline training; disaggregated entrypoints select `auto`, `producer`, or `consumer`. | | `training.seed` | `42` | Run and per-rank RNG seed. | diff --git a/specforge/algorithms/eagle3/model.py b/specforge/algorithms/eagle3/model.py index c16022401..77493aa4b 100644 --- a/specforge/algorithms/eagle3/model.py +++ b/specforge/algorithms/eagle3/model.py @@ -117,6 +117,7 @@ def __init__( lk_loss_type: Optional[str] = None, kl_scale: float = 1.0, kl_decay: float = 1.0, + trim_loss_positions: bool = False, ): """ Args: @@ -125,6 +126,11 @@ def __init__( lk_loss_type: LK loss objective type. One of {"lambda", "alpha"}. kl_scale: Initial KL weight scale for lambda LK loss. kl_decay: Decay factor for adaptive KL weight in lambda LK loss. + trim_loss_positions: when set, the teacher target_p, draft logits and + loss are computed only at supervised (loss-masked) positions rather + than over the full sequence. Mathematically equivalent after + rescaling the mean denominator; off by default. Automatically falls + back to the full-length path for batch > 1 or when an lk_loss is used. """ super().__init__() self.draft_model = draft_model @@ -133,6 +139,7 @@ def __init__( self.lk_loss_type = lk_loss_type self.kl_scale = kl_scale self.kl_decay = kl_decay + self.trim_loss_positions = trim_loss_positions def _make_adapter(self) -> BackendAdapter: if self.attention_backend == "usp": @@ -149,6 +156,8 @@ def _acc_and_loss( position_mask: torch.Tensor, loss_mask: torch.Tensor, adapter: BackendAdapter, + loss_scale: float = 1.0, + full_positions: Optional[int] = None, ) -> Tuple[ torch.Tensor, torch.Tensor, @@ -183,8 +192,14 @@ def _acc_and_loss( reduce_metrics_fn=adapter.reduce_metrics, reduce_loss_fn=adapter.reduce_loss, ) + if loss_scale != 1.0: + # The trimmed loss kernel averages over n_sup supervised positions, but + # the full-length semantics average over L; rescale to recover it. Only + # valid when lk_loss_type is None (a plain KL loss is linearly scalable). + loss = loss * loss_scale loss_denom = torch.tensor( - logits.shape[0] * logits.shape[1], + logits.shape[0] + * (full_positions if full_positions is not None else logits.shape[1]), device=logits.device, dtype=torch.float32, ) @@ -291,18 +306,61 @@ def forward( chunk_size=compact_teacher_chunk_size, ) del target_hidden_for_compact + trim_pack = None else: - ( - target_p_padded, - target_p_on_draft_padded, - target_token_ids_padded, - position_mask, - ) = _compute_target_p_padded( - target=target, - t2d=self.draft_model.t2d, - loss_mask=loss_mask, - length=self.length, + # A-level trim: with batch==1 and no lk_loss, compute the teacher only at + # supervised positions; fall back to the full path otherwise. Under USP + # the backbone keeps running on this rank's own chunk (usp_chunk_size = + # local_len - ttt_length); the local buffer's ttt_length overlap tail may + # only act as teacher positions for own-chunk rows, never emit loss rows + # itself (those rows belong to the next rank), so the per-step row sets + # are bounded by chunk_len. + # chunk_len must come from the SAME source the full path uses for its + # slicing/normalization: the hidden-state sequence length (the loss + # kernel means over backbone rows). loss_mask can carry an extra + # zero-padded slot in the offline pipeline, so deriving from it would + # be off by one (wrong rows, wrong denominator, and under USP a + # backbone length that disagrees with full-path ranks). + if self.attention_backend == "usp": + trim_chunk_len = hidden_states.shape[1] - self.length + else: + trim_chunk_len = hidden_states.shape[1] + _trim_ok = ( + self.trim_loss_positions + and self.lk_loss_type is None + and loss_mask.shape[0] == 1 + and trim_chunk_len > 0 + and int(loss_mask.sum().item()) > 0 ) + trim_pack = None + if _trim_ok: + # Returns None when no supervised position can reach any row + # (e.g. supervision only beyond the reachable window); then we + # fall through to the full path below. + trim_pack = _build_trim_pack( + target, + self.draft_model.t2d, + loss_mask, + self.length, + chunk_len=trim_chunk_len, + ) + if trim_pack is not None: + target_p_padded = None + target_p_on_draft_padded = None + target_token_ids_padded = None + position_mask = trim_pack["position_mask_sup"] + else: + ( + target_p_padded, + target_p_on_draft_padded, + target_token_ids_padded, + position_mask, + ) = _compute_target_p_padded( + target=target, + t2d=self.draft_model.t2d, + loss_mask=loss_mask, + length=self.length, + ) del target torch.cuda.empty_cache() @@ -362,33 +420,55 @@ def forward( raise ValueError(f"Unknown attention backend: {self.attention_backend}") for idx in range(self.length): - state = adapter.step_view( - idx=idx, - ttt_length=self.length, - global_input_ids=global_input_ids, - attention_mask=attention_mask, - loss_mask=loss_mask, - position_ids=position_ids, - hidden_states=hidden_states, - target_p_padded=target_p_padded, - target_p_on_draft_padded=target_p_on_draft_padded, - target_token_ids_padded=target_token_ids_padded, - position_mask=position_mask, - seq_length=seq_length, - ) + if 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 + # supervised rows go through logits/loss below. + state = None + if self.attention_backend == "usp": + _c = trim_pack["full_len"] + step_input_ids = global_input_ids[:, :_c] + step_hidden = hidden_states[:, :_c, :] + step_attn = attention_mask[:, :_c] + step_pos = position_ids[:, : _c * adapter.sp_ulysses_degree] + else: + step_input_ids = global_input_ids + step_hidden = hidden_states + step_attn = attention_mask + step_pos = position_ids + else: + state = adapter.step_view( + idx=idx, + ttt_length=self.length, + global_input_ids=global_input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + position_ids=position_ids, + hidden_states=hidden_states, + target_p_padded=target_p_padded, + target_p_on_draft_padded=target_p_on_draft_padded, + target_token_ids_padded=target_token_ids_padded, + position_mask=position_mask, + seq_length=seq_length, + ) + step_input_ids = state.input_ids + step_hidden = state.hidden_states + step_attn = state.attention_mask + step_pos = state.position_ids is_last = idx == self.length - 1 # Step 5.1: embed the input ids - inputs_embeds = self.draft_model.embed_input_ids(state.input_ids) + inputs_embeds = self.draft_model.embed_input_ids(step_input_ids) inputs_embeds = inputs_embeds.to(hidden_states.dtype) # Step 5.2: run the draft model backbone hidden_states_out = self.draft_model.backbone( input_embeds=inputs_embeds, - hidden_states=state.hidden_states, + hidden_states=step_hidden, cache_hidden=cache_hidden, - attention_mask=state.attention_mask, - position_ids=state.position_ids, + attention_mask=step_attn, + position_ids=step_pos, past_key_values=past_key_values, use_cache=True, ) @@ -396,27 +476,65 @@ def forward( # update hidden states for next step hidden_states = hidden_states_out - # Step 5.4: get logits - logits = self.draft_model.compute_logits(hidden_states) - - # Step 5.5 + 5.6: metric and loss - ( - acc, - acceptance_rate, - loss, - correct, - denom, - metric_loss, - loss_denom, - ) = self._acc_and_loss( - logits=logits, - target_p=state.target_p, - target_p_on_draft=state.target_p_on_draft, - target_token_ids=state.target_token_ids, - position_mask=state.position_mask, - loss_mask=state.loss_mask, - adapter=adapter, - ) + # Step 5.4 + 5.5 + 5.6: logits, metric and loss + if trim_pack is not None: + # A-level: only the rows that can carry loss at this step go through + # norm + lm_head. Rows shift down by one per step (rows = sup - idx) + # while the teacher/mask stay pinned at the supervised positions. + rows_j = trim_pack["rows_steps"][idx] + keep_j = trim_pack["keep_steps"][idx] + nrows_j = trim_pack["nrows_steps"][idx] + logits = self.draft_model.compute_logits( + hidden_states.index_select(1, rows_j) + ) + pm_j = trim_pack["position_mask_sup"].index_select(1, keep_j) + lm_j = trim_pack["loss_mask_sup"].index_select(1, keep_j) + if nrows_j == 0: + # Dead step: no own-chunk row carries loss here (e.g. all local + # supervised positions sit in the USP overlap tail at this + # depth). The pack padded one dummy entry; zeroing its masks + # makes the contribution exactly zero while every rank still + # runs the same kernels and collective calls. + pm_j = torch.zeros_like(pm_j) + lm_j = torch.zeros_like(lm_j) + ( + acc, + acceptance_rate, + loss, + correct, + denom, + metric_loss, + loss_denom, + ) = self._acc_and_loss( + logits=logits, + target_p=trim_pack["target_p_c"].index_select(1, keep_j), + target_p_on_draft=trim_pack["on_draft_c"].index_select(1, keep_j), + target_token_ids=trim_pack["token_ids_c"].index_select(1, keep_j), + position_mask=pm_j, + loss_mask=lm_j, + adapter=adapter, + loss_scale=nrows_j / trim_pack["full_len"], + full_positions=trim_pack["full_len"], + ) + else: + logits = self.draft_model.compute_logits(hidden_states) + ( + acc, + acceptance_rate, + loss, + correct, + denom, + metric_loss, + loss_denom, + ) = self._acc_and_loss( + logits=logits, + target_p=state.target_p, + target_p_on_draft=state.target_p_on_draft, + target_token_ids=state.target_token_ids, + position_mask=state.position_mask, + loss_mask=state.loss_mask, + adapter=adapter, + ) acces.append(acc) acceptance_rates.append(acceptance_rate) plosses.append(loss) @@ -516,3 +634,137 @@ def _compute_metric_counts(logits, target_token_ids, loss_mask, d2t): ).sum() denom = loss_mask.sum().clamp_min(1e-6) return correct, denom + + +def _compute_target_p_eager(target, t2d, loss_mask, row_chunk=256): + """Uncompiled variant of the teacher target_p computation. + + Kept uncompiled because the supervised-row count varies per batch, which would + trigger repeated torch.compile recompilation. Mathematically identical to the + compiled path; chunks over rows to bound the transient full-vocab fp32 + activation (which can reach several GB when the row count is large). + """ + tps, tpds, toks, pms = [], [], [], [] + n = target.shape[1] + for s in range(0, n, row_chunk): + t = target[:, s : s + row_chunk].float() + ids = t.argmax(-1) + tm = t2d[ids][..., None].int() + pms.append(tm * loss_mask[:, s : s + row_chunk]) + dth = t[..., t2d] + tps.append(F.softmax(dth, dim=2).detach()) + lse = torch.logsumexp(t, dim=-1, keepdim=True) + tpds.append(torch.exp(dth - lse).detach()) + toks.append(ids.detach()) + return ( + torch.cat(tps, 1), + torch.cat(tpds, 1), + torch.cat(toks, 1), + torch.cat(pms, 1), + ) + + +def _build_trim_pack(target, t2d, loss_mask, length, chunk_len=None): + """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 + ``padding(..., left=False)`` to ``position_mask`` / ``loss_mask`` once per TTT + step, so at step j the mask seen at row p is ``mask[p + j]``; meanwhile + ``step_view`` slices the padded teacher so row p is supervised by the teacher at + absolute position ``p + j``. A row therefore contributes at step j iff + ``p + j`` is supervised, i.e. ``p = s - j`` for some supervised position s. + + So the rows shift *down* by one per step while the teacher/mask positions stay + pinned at the supervised set: + + step j: rows = {s - j : s in sup, s >= j} teacher/mask at those s + + That makes the teacher cheap: it only ever has to be evaluated at ``sup`` + (no sliding window), and each step just drops the entries whose row would fall + off the front of the sequence. + + ``chunk_len`` is the number of rows the backbone actually produces per step. + On single-rank backends it equals the full local length L (default). Under USP + the backbone runs on this rank's own chunk (``usp_chunk_size = L - + ttt_length``) while the local buffer keeps a ``ttt_length`` overlap tail: + tail positions may act as teachers for own-chunk rows at deeper steps, but + tail rows belong to the next rank and must never emit loss here — hence the + additional ``s - j < chunk_len`` bound on the row set, and ``full_len`` (the + loss denominator) becomes ``chunk_len`` to match the full path's + mean-over-chunk semantics. + + A step whose row set comes out empty (all supervised positions out of reach + at that depth) is padded with one dummy entry and reported with + ``nrows_steps[j] == 0``; the caller zeroes its masks so the step contributes + exactly zero loss while kernel launches and collective calls stay aligned + across ranks. + + batch == 1 only (online training uses batch == 1 per rank); the caller falls + back to the full-length path otherwise. + + Returns a dict with, per step j: ``rows_steps[j]`` (row indices into the draft + hidden states), ``keep_steps[j]`` (which supervised entries survive), + ``nrows_steps[j]`` (real row count, 0 for dead steps), plus the teacher + tables evaluated once at ``sup`` and the mask values at ``sup``. + """ + with torch.no_grad(): + B, L = loss_mask.shape[0], loss_mask.shape[1] + assert B == 1, "trim path requires batch==1" + if chunk_len is None: + chunk_len = L + sup = loss_mask.view(-1).nonzero(as_tuple=False).squeeze(-1) # [n_sup] + # Positions beyond chunk_len + length - 2 can never supervise any row at + # any step (would need j >= length); dropping them keeps every later + # index within the hidden/target sequence range even when loss_mask is + # longer than the hidden states (offline pipelines pad it by one). + sup = sup[sup < chunk_len + length - 1] + if sup.numel() == 0: + # Nothing reachable at any step; tell the caller to use the full path. + return None + # Teacher is only ever needed at the supervised positions themselves. + target_sel = target[:, sup] # [1, n_sup, V_target] + lm_sel = loss_mask[:, sup] + target_p_c, on_draft_c, token_ids_c, _ = _compute_target_p_eager( + target_sel, t2d, lm_sel + ) + pm_sup = _compute_position_mask_at(target, t2d, loss_mask, sup) + lm_sup = loss_mask.view(-1)[sup].view(1, -1, 1) + + rows_steps, keep_steps, nrows_steps = [], [], [] + 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 n == 0: + # Dead step: pad with one dummy entry; the caller zeroes its + # masks so it contributes nothing. + keep = pad_idx + rows = pad_idx + else: + rows = sup[keep] - j + rows_steps.append(rows) + keep_steps.append(keep) + nrows_steps.append(n) + return dict( + sup=sup, + rows_steps=rows_steps, + keep_steps=keep_steps, + nrows_steps=nrows_steps, + target_p_c=target_p_c, + on_draft_c=on_draft_c, + token_ids_c=token_ids_c, + position_mask_sup=pm_sup, + loss_mask_sup=lm_sup, + full_len=chunk_len, + ) + + +def _compute_position_mask_at(target, t2d, loss_mask, sup): + """position_mask = t2d[argmax(target)] * loss_mask, computed only at sup positions.""" + with torch.no_grad(): + tsel = target[:, sup] + ids = tsel.float().argmax(-1) + tm = t2d[ids][..., None].int() + return tm * loss_mask[:, sup] diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index 61eb1acea..157fdf5d8 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -252,6 +252,7 @@ def build_eagle3_model( lk_loss_type=cfg.training.lk_loss_type, kl_scale=cfg.training.kl_scale, kl_decay=cfg.training.kl_decay, + trim_loss_positions=cfg.training.trim_loss_positions, ).to(device=_device(), dtype=_torch_dtype(cfg)) needs_target_head = cfg.mode == "offline" or ( cfg.deployment.mode == "disaggregated" and cfg.training.role == "consumer" diff --git a/specforge/config/schema.py b/specforge/config/schema.py index cabaa975d..864bb6bca 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -507,6 +507,12 @@ class TrainingConfig(StrictConfigModel): lk_loss_type: Optional[Literal["lambda", "alpha"]] = None kl_scale: float = 1.0 kl_decay: float = 1.0 + #: Compute the teacher target_p, draft logits and loss only at supervised + #: (loss-masked) positions instead of over the full sequence. Mathematically + #: equivalent (the mean denominator is rescaled) and saves memory/compute on + #: 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 #: DFlash-family objective/model knobs. num_anchors: int = Field(default=512, gt=0) loss_decay_gamma: Optional[float] = None diff --git a/tests/test_runtime/test_equiv_trim_loss_positions.py b/tests/test_runtime/test_equiv_trim_loss_positions.py new file mode 100644 index 000000000..6b13dd983 --- /dev/null +++ b/tests/test_runtime/test_equiv_trim_loss_positions.py @@ -0,0 +1,89 @@ +# coding=utf-8 +"""Equivalence: trim_loss_positions must not change the training loss. + +A-level position trimming computes the teacher target_p, the draft logits and the +loss only at supervised (loss-masked) positions instead of over the full sequence. +It is mathematically equivalent to the full-length path (the mean denominator is +rescaled from n_sup back to the full length). This test runs the identical forward +with trimming off and on and asserts the per-step losses match within bf16 +tolerance. + +GPU-only, matching the other EAGLE3 equivalence tests in this directory. +""" + +import os +import shutil +import tempfile +import unittest + +import torch + +CUDA = torch.cuda.is_available() + + +@unittest.skipUnless(CUDA, "trim_loss_positions equivalence requires CUDA") +class TestEquivTrimLossPositions(unittest.TestCase): + def test_trim_loss_positions_matches_full(self): + torch.manual_seed(0) + from tests.test_runtime import _fixtures as fx + + fx.build_single_rank_distributed(port="29567") + + workdir = tempfile.mkdtemp(prefix="equiv_trim_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + + model, target_head = fx.build_eagle3(workdir, ttt=3) + model.eval() + + # One offline sample gives us (input_ids, target, loss_mask, hidden_state). + feature_dir = os.path.join(workdir, "features") + os.makedirs(feature_dir, exist_ok=True) + fx.write_offline_files(feature_dir, n=1, seq=16) + batch = torch.load( + os.path.join(feature_dir, sorted(os.listdir(feature_dir))[0]), + map_location="cpu", + ) + + # `hidden_state` is the target-side capture that the target head turns into + # the teacher distribution; `aux_hidden_state` is the draft backbone input. + input_ids, target, loss_mask = target_head.preprocess( + batch["input_ids"].unsqueeze(0), + batch["hidden_state"], + batch["loss_mask"].unsqueeze(0), + ) + target = target_head(target.cuda()) + hidden_states = batch["aux_hidden_state"].cuda() + input_ids = input_ids.cuda() + + # Prompt-heavy mask so trimming is non-trivial: first half unsupervised. + loss_mask = loss_mask.cuda().clone() + loss_mask[:, : loss_mask.shape[1] // 2] = 0 + attention_mask = torch.ones_like(input_ids) + + @torch.no_grad() + def step_losses(trim: bool): + model.trim_loss_positions = trim + plosses, *_ = model( + input_ids=input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + target=target, + hidden_states=hidden_states, + ) + return [float(p.item()) for p in plosses] + + full = step_losses(False) + trimmed = step_losses(True) + + self.assertEqual(len(full), len(trimmed)) + for i, (a, b) in enumerate(zip(full, trimmed)): + tol = 5e-3 * max(abs(a), abs(b)) + 1e-4 + self.assertLessEqual( + abs(a - b), + tol, + msg=f"step {i}: full={a} trimmed={b} (tol={tol})", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_runtime/test_equiv_trim_usp.py b/tests/test_runtime/test_equiv_trim_usp.py new file mode 100644 index 000000000..d67959d90 --- /dev/null +++ b/tests/test_runtime/test_equiv_trim_usp.py @@ -0,0 +1,296 @@ +# coding=utf-8 +"""trim_loss_positions under USP sequence parallelism. + +Three layers, so CI keeps guarding the USP row-selection math even on hosts +without four GPUs: + +1. ``TestTrimPackGolden`` (CPU) -- hand-derived literal tables for the per-step + row sets: the overlap-tail bound (``s - j < chunk_len``), dead-step padding, + the unreachable-supervision fallback, and non-USP back-compat. +2. ``TestTrimLossAnalytic`` (one GPU) -- with zero logits and normalized + teachers every masked row's loss is exactly ``ln(draft_vocab)``, so the + trim-scaled and full-shaped losses must both hit a closed-form constant. +3. ``TestEquivTrimUspFourRank`` (four GPUs + flash-attn, like + ``test_equiv_4rank``) -- per-step loss parity, trim ON vs OFF, on a real + ring-4 offline pipeline across three adversarial masks: all-supervised + (trivial-equality boundary: any row/denominator error is exposed without + tolerance cover), supervision straddling every rank boundary + (overlap-tail-as-teacher), and supervision only inside one rank's overlap + tail (dead steps plus ranks with no supervision at all, which fall back to + the full path -- the mixed-path collective-alignment hazard). +""" + +import json +import math +import os +import shutil +import tempfile +import unittest + +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 TestTrimPackGolden(unittest.TestCase): + """Hand-derived expected outputs for _build_trim_pack (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_usp_tail_bound(self): + # C=6, ttt=2, local_len=8, supervised {4,5,6}; 6 is an overlap-tail + # position: a legal teacher, never a loss row. + # step0: {s : s-0 < 6} = {4,5} -> rows [4,5], n=2 + # step1: all of {4,5,6} -> rows [3,4,5], n=3 + 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=2, chunk_len=6) + self.assertEqual(p["full_len"], 6) + self.assertEqual(p["sup"].tolist(), [4, 5, 6]) + self.assertEqual(p["rows_steps"][0].tolist(), [4, 5]) + self.assertEqual(p["keep_steps"][0].tolist(), [0, 1]) + self.assertEqual(p["nrows_steps"][0], 2) + self.assertEqual(p["rows_steps"][1].tolist(), [3, 4, 5]) + self.assertEqual(p["nrows_steps"][1], 3) + + def test_usp_dead_step(self): + # Supervision only at {6,7} (pure tail). Position 7 can never reach a + # row (needs j >= 2) and is filtered; 6 is unreachable at step 0. + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 0, 0, 1, 1]) + p = _build_trim_pack(tgt, t2d, lm, length=2, chunk_len=6) + self.assertEqual(p["sup"].tolist(), [6]) + self.assertEqual(p["nrows_steps"][0], 0) # dead step + self.assertEqual(p["rows_steps"][0].tolist(), [0]) # padded dummy + self.assertEqual(p["rows_steps"][1].tolist(), [5]) + self.assertEqual(p["nrows_steps"][1], 1) + + def test_unreachable_supervision_falls_back(self): + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 0, 0, 0, 1]) + self.assertIsNone(_build_trim_pack(tgt, t2d, lm, length=2, chunk_len=6)) + + def test_non_usp_backcompat(self): + # chunk_len=None -> C=L: the pre-USP semantics, plus full_len now comes + # from the row count rather than the (possibly padded) mask length. + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([1, 0, 0, 0, 1, 0, 0, 1]) + p = _build_trim_pack(tgt, t2d, lm, length=2) + self.assertEqual(p["full_len"], 8) + self.assertEqual(p["rows_steps"][0].tolist(), [0, 4, 7]) + self.assertEqual(p["rows_steps"][1].tolist(), [3, 6]) + self.assertEqual(p["keep_steps"][1].tolist(), [1, 2]) + self.assertEqual(p["nrows_steps"][1], 2) + + +@unittest.skipUnless(CUDA, "loss kernel is a Triton kernel") +class TestTrimLossAnalytic(unittest.TestCase): + """Zero logits + normalized teachers => masked row loss == ln(D) exactly.""" + + def test_trim_and_full_hit_closed_form(self): + from specforge.core.loss import LogSoftmaxLoss + + D = 64 + g = torch.Generator().manual_seed(2) + + def one_hot_rows(n): + t = torch.zeros(1, n, D, device="cuda") + t[0, torch.arange(n), torch.randint(0, D, (n,), generator=g)] = 1.0 + return t + + # Trim-shaped: 3 selected rows, all masked-in; kernel mean == ln(D); + # rescaled by nrows/C = 3/6 -> 3*ln(64)/6. + kernel = LogSoftmaxLoss.apply( + torch.zeros(1, 3, D, device="cuda"), + one_hot_rows(3), + torch.ones(1, 3, 1, device="cuda"), + ) + self.assertAlmostEqual(kernel.item(), math.log(D), places=5) + self.assertAlmostEqual((kernel * (3 / 6)).item(), 2.0794415417, places=5) + + # Full-shaped: 6 rows, 3 masked-in -> same constant with no rescale. + pm = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda").view(1, 6, 1) + full = LogSoftmaxLoss.apply( + torch.zeros(1, 6, D, device="cuda"), one_hot_rows(6), pm + ) + self.assertAlmostEqual(full.item(), 2.0794415417, places=5) + + +def _write_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"), + ) + return list(masks) + + +def _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"trimusp-{case}", + ttt_length=TTT, + max_len=SEQ, + ).read() + loader = FeatureDataLoader( + LocalFeatureStore(f"trimusp-{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)) + strat = algorithm.providers.step.build(model, target_head=target_head) + + def step_losses(trim): + model.trim_loss_positions = trim + with torch.no_grad(): + out = strat.forward_loss(batch) + return [float(p.item()) for p in out.metrics["plosses"]] + + results[case] = {"full": step_losses(False), "trim": 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() + + +@unittest.skipUnless( + CUDA and NGPU >= WORLD_SIZE and _has_standard_flash_attention(), + "requires four CUDA devices and the standard flash-attn USP interfaces", +) +class TestEquivTrimUspFourRank(unittest.TestCase): + def test_trim_matches_full_per_step_on_ring4(self): + import torch.multiprocessing as mp + + workdir = tempfile.mkdtemp(prefix="trim_usp_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + _write_workdir(workdir) + mp.spawn( + _worker, + args=(WORLD_SIZE, 29871, 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): + full, trim = res[case]["full"], res[case]["trim"] + self.assertEqual(len(full), TTT) + for j, (a, b) in enumerate(zip(full, trim)): + if case == "allones": + # expected near-bit-equal; 1e-6 is ~4 orders below the + # smallest possible discrete error (one row's worth, + # ~loss/C) while allowing 1-2 ulp of fp32 noise + tol = 1e-6 + else: + tol = max(1e-3 * abs(a), 1e-4) + self.assertLessEqual( + abs(a - b), + tol, + msg=f"{case} rank{rank} step{j}: full={a} trim={b}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2)