From b27639bff6456bed118a45fe0d13881cf5aa0045 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:09:53 -0700 Subject: [PATCH 1/6] perf(speculative): compute draft attention with a FlexAttention block-sparse mask The DFlash draft handed SDPA a materialized [B, 1, Q, KV] additive float mask. That mask disqualifies every fused backend -- FlashAttention rejects arbitrary masks and caps head_dim at 256 -- so PyTorch fell through to the cutlass memory-efficient backend, whose only kernel is **sm80**. Every DFlash-family drafter was therefore running an Ampere attention kernel on Blackwell, at roughly 2.5% of peak. Kernel attribution put that one backward kernel at 59% of the whole training step. The predicate is block-sparse by construction -- each query block sees a context PREFIX up to its anchor plus its own block on the diagonal, about 33% dense -- so it expresses exactly as a FlexAttention BlockMask, and the generated Triton kernel skips the ~67% of tiles that are entirely masked. Measured on PDX B300, 3 nodes (1 vLLM serve + 2 trainer = 16 DP ranks), Gemma-4-E4B DSpark, bs 4 x seq 4096, median over 16 ranks x 20 steady-state steps: step 1.690 -> 0.711 s (2.38x) attention bwd 0.997 -> 0.094 s attention fwd 0.124 -> 0.011 s Two portability constraints are baked into dflash_flex_attention.py because both were found the hard way on torch 2.11 / sm103: * head_dim > 256 must pin 32x32 kernel tiles. 64x32 and 64x64 give "misaligned address"; 128x32 gives "unspecified launch failure". * The BlockMask block size must divide the kernel's BLOCK_M/BLOCK_N. head_dim > 256 pins 32x32 so 64 is safe, but head_dim <= 256 lets Inductor autotune up to BLOCK_M=128, where a 64-wide mask block raises "Q and KV block size must be divisible by BLOCK_M and BLOCK_N". Gemma-4 dodged this only because head_dim 512 pins the tiles; every head_dim <= 256 drafter (Qwen3, Kimi) would have hit it. Also note enable_gqa=True is deliberately NOT used: forward matches, but backward took 568 ms -- 10x the pre-repeated path and 2.5x worse than the SDPA baseline it replaces. K/V are always pre-repeated instead. The branch sits inside _build_draft_attention_mask rather than at its three call sites, so DFlash, DSpark and Domino all benefit and the existing monkeypatch-spy tests keep observing a single mask build per forward. Off by default (dflash_use_flex_attention). OPERATIONAL PREREQUISITE: TRITON_CACHE_DIR must be node-local. FlexAttention compiles Triton templates on the first backward, and pointing that cache at a shared filesystem makes concurrent ranks race -- 16 ranks on lustre died with "OSError: [Errno 14] Bad address" reading a .json mid-write. Signed-off-by: Hao Guo --- modelopt/torch/speculative/config.py | 15 ++ .../torch/speculative/dflash/dflash_model.py | 1 + .../plugins/dflash_flex_attention.py | 192 ++++++++++++++++++ .../torch/speculative/plugins/hf_dflash.py | 18 ++ .../speculative/plugins/modeling_dflash.py | 59 +++--- .../speculative/plugins/test_hf_dflash.py | 121 +++++++++++ 6 files changed, 379 insertions(+), 27 deletions(-) create mode 100644 modelopt/torch/speculative/plugins/dflash_flex_attention.py diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 10af2366f9e..8df353b2daa 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -170,6 +170,21 @@ class DFlashConfig(ModeloptBaseConfig): description="Whether to use torch.compile on DFlash forward/loss methods.", ) + dflash_use_flex_attention: bool = ModeloptField( + default=False, + description=( + "Compute the draft's attention with torch FlexAttention over a block-sparse " + "BlockMask instead of handing SDPA a materialized [B, 1, Q, KV] float mask. " + "The mask is only ~33% dense (each query block sees a context PREFIX up to its " + "anchor, plus its own block on the diagonal), and the dense form additionally " + "forces SDPA onto the cutlass memory-efficient backend -- whose only kernel is " + "sm80 -- because the fused backends reject arbitrary masks and cap head_dim at " + "256. Measured 4.0x on the draft attention (225 ms -> 56 ms fwd+bwd) at the " + "Gemma-4-E4B DSpark shape on B300, within bf16 rounding of the dense path. " + "Requires torch >= 2.5. Off by default." + ), + ) + dflash_swa_window_size: int | None = ModeloptField( default=None, description=( diff --git a/modelopt/torch/speculative/dflash/dflash_model.py b/modelopt/torch/speculative/dflash/dflash_model.py index 24f2143bf84..7d90329d4bb 100644 --- a/modelopt/torch/speculative/dflash/dflash_model.py +++ b/modelopt/torch/speculative/dflash/dflash_model.py @@ -49,4 +49,5 @@ def modify(self, config): self.dflash_report_acc = config.dflash_report_acc self.dflash_use_torch_compile = config.dflash_use_torch_compile self.dflash_swa_window_size = config.dflash_swa_window_size + self.dflash_use_flex_attention = config.dflash_use_flex_attention self.dflash_export_rope_scaling = config.dflash_export_rope_scaling diff --git a/modelopt/torch/speculative/plugins/dflash_flex_attention.py b/modelopt/torch/speculative/plugins/dflash_flex_attention.py new file mode 100644 index 00000000000..8ea8b12a5a6 --- /dev/null +++ b/modelopt/torch/speculative/plugins/dflash_flex_attention.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Block-sparse FlexAttention path for the DFlash/DSpark draft. + +The draft's attention mask is dense in shape but sparse in content. Query block ``b`` +with anchor ``a_b`` attends to + + context: kv < a_b -- a prefix, and anchors are sorted, so + across blocks this is a staircase + draft: kv in [S + Bb, S + Bb + B) -- the block diagonal, ``block_size`` wide + +Handing SDPA a materialized ``[B, 1, Q, KV]`` float mask makes it compute all of it, +and disqualifies every fused backend on the way: PyTorch's FlashAttention kernels +reject arbitrary masks *and* cap ``head_dim`` at 256, while a Gemma-4 draft's +``global_head_dim`` is 512. What is left is the cutlass memory-efficient backend, +whose only kernel is **sm80** -- an Ampere kernel on Blackwell. Profiling the +Gemma-4-E4B DSpark run on B300 put that single backward kernel at 59% of the whole +training step, running at roughly 2.5% of the GPU's bf16 peak. + +FlexAttention instead takes the mask as a predicate, compiles it into the kernel, and +skips fully-masked tiles. Measured on B300 at the production shape +(q[4,16,4096,512], kv[4,1->16,8192,512], 512 anchors x block 8, 33% mask density): + + dense-mask SDPA fwd 24.80 ms fwd+bwd 225.37 ms + flex (this file) fwd 6.84 ms fwd+bwd 56.40 ms 4.0x + +Both are within bf16 rounding of each other on out/dq/dk/dv, including the +fully-masked rows that invalid blocks produce. + +Two non-obvious requirements, both established by measurement rather than docs: + +* ``head_dim`` 512 overflows shared memory at FlexAttention's default tiles (263 KB + required against a 232 KB limit), so the tiles are pinned above ``head_dim`` 256. + Of the shapes that fit, only 32x32 runs at all on torch 2.11 / sm103: 64x32 and + 64x64 fault with "misaligned address" and 128x32 with "unspecified launch + failure". At ``head_dim`` <= 256 the library defaults are far better than anything + pinned (9.8x over SDPA at 256), so they are left alone. +* ``enable_gqa=True`` must NOT be used. Its forward matches the pre-repeated path + exactly, but its backward takes 568 ms -- 10x slower, and 2.5x worse than the SDPA + baseline it is meant to replace. K/V are repeated to the query head count first, + which is what HF's sdpa path does anyway. +""" + +import torch + +__all__ = ["build_draft_block_mask", "flex_attention_forward", "is_block_mask"] + +# head_dim > 256 cannot use FlexAttention's default tiles (SMEM overflow); 32x32 with +# two pipeline stages is the only combination measured to both fit and run. +_LARGE_HEAD_DIM_KERNEL_OPTIONS = { + "BLOCK_M": 32, + "BLOCK_N": 32, + "BLOCK_M1": 32, + "BLOCK_N1": 32, + "BLOCK_M2": 32, + "BLOCK_N2": 32, + "num_stages": 2, + "num_warps": 4, +} +_MAX_DEFAULT_TILE_HEAD_DIM = 256 + +# Block granularity of the BlockMask itself. Finer granularity tracks the true 33% +# density more closely (64 -> 65.7% sparsity vs 128 -> 64.2%) and measured 58.0 ms +# against 61.4 ms; 32 is marginally better again but doubles the metadata for ~1 ms. +_PINNED_TILE_MASK_BLOCK_SIZE = 64 +_DEFAULT_TILE_MASK_BLOCK_SIZE = 128 + + +def _mask_block_size(head_dim): + """Mask granularity, which is coupled to the kernel tiles and cannot be chosen freely. + + FlexAttention requires the BlockMask's block size to be divisible by the kernel's + BLOCK_M/BLOCK_N, and raises "Q and KV block size must be divisible by BLOCK_M and + BLOCK_N" otherwise. So: + + * head_dim > 256 pins 32x32 tiles (SMEM), so 64 is safe -- and finer granularity + tracks the true mask density better (65.7% sparsity vs 64.2% at 128, measured + 58.0 ms vs 61.4 ms). + * head_dim <= 256 leaves the tiles to Inductor's autotuner, which may pick up to + 128. Only FlexAttention's own default of 128 is guaranteed compatible with every + choice it can make. + """ + if head_dim > _MAX_DEFAULT_TILE_HEAD_DIM: + return _PINNED_TILE_MASK_BLOCK_SIZE + return _DEFAULT_TILE_MASK_BLOCK_SIZE + +_flex_attention_compiled = None +_create_block_mask_compiled = None + + +def _flex_ops(): + """Resolve and compile the FlexAttention entry points once per process.""" + global _flex_attention_compiled, _create_block_mask_compiled + if _flex_attention_compiled is None: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + # dynamic=False: shapes are fixed by (batch, num_anchors, block_size, seq_len) + # within a run, and dynamic shapes measurably deoptimize the generated kernel. + _flex_attention_compiled = torch.compile(flex_attention, dynamic=False) + _create_block_mask_compiled = torch.compile(create_block_mask, dynamic=False) + return _flex_attention_compiled, _create_block_mask_compiled + + +def is_block_mask(mask) -> bool: + """True if ``mask`` is a FlexAttention ``BlockMask`` rather than a dense tensor.""" + if mask is None or torch.is_tensor(mask): + return False + try: + from torch.nn.attention.flex_attention import BlockMask + except ImportError: + return False + return isinstance(mask, BlockMask) + + +def build_draft_block_mask( + seq_len, anchor_positions, block_keep_mask, n_blocks, block_size, window, device, + head_dim, +): + """BlockMask equivalent of ``HFDFlashModel._build_draft_attention_mask``. + + Same predicate, expressed for FlexAttention instead of materialized. Rebuilt every + step because anchors are resampled on every forward; the compiled builder costs + ~0.1 ms, against ~1.5 ms to materialize the dense mask it replaces. + """ + _, create_block_mask = _flex_ops() + bsz = anchor_positions.shape[0] + q_len = n_blocks * block_size + kv_len = seq_len + q_len + # Indexed inside mask_mod, which runs under vmap -- keep them on-device and integral. + anchors = anchor_positions.to(device=device, dtype=torch.int32) + keep = block_keep_mask.to(device=device, dtype=torch.bool) + + def mask_mod(b, h, q_idx, kv_idx): + q_block = q_idx // block_size + anchor = anchors[b, q_block] + is_ctx = kv_idx < seq_len + ctx_ok = is_ctx & (kv_idx < anchor) + if window is not None: + # Same sliding window as the dense path: measured against the query's REAL + # position (anchor + position-in-block), not its index in the draft block. + ctx_ok = ctx_ok & (kv_idx > anchor + (q_idx % block_size) - window) + draft_ok = (~is_ctx) & (q_block == (kv_idx - seq_len) // block_size) + return (ctx_ok | draft_ok) & keep[b, q_block] + + return create_block_mask( + mask_mod, bsz, None, q_len, kv_len, device=device, + BLOCK_SIZE=_mask_block_size(head_dim), + ) + + +def _repeat_kv(x, n_rep): + """HF's ``repeat_kv``: [B, n_kv, S, D] -> [B, n_kv * n_rep, S, D].""" + if n_rep == 1: + return x + b, h, s, d = x.shape + return x[:, :, None].expand(b, h, n_rep, s, d).reshape(b, h * n_rep, s, d) + + +def flex_attention_forward(query, key, value, block_mask, scaling): + """FlexAttention with the draft's BlockMask. Returns ``[B, q_len, n_heads, head_dim]``. + + The layout matches what HF's ``sdpa_attention_forward`` returns so the caller's + ``reshape(bsz, q_len, -1)`` is unchanged. + """ + flex_attention, _ = _flex_ops() + head_dim = query.shape[-1] + n_rep = query.shape[1] // key.shape[1] + kernel_options = ( + _LARGE_HEAD_DIM_KERNEL_OPTIONS if head_dim > _MAX_DEFAULT_TILE_HEAD_DIM else None + ) + attn_output = flex_attention( + query, + _repeat_kv(key, n_rep), + _repeat_kv(value, n_rep), + block_mask=block_mask, + scale=scaling, + kernel_options=kernel_options, + ) + return attn_output.transpose(1, 2).contiguous() diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index b025b3088a2..7ffd735553a 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -667,6 +667,24 @@ def _build_draft_attention_mask( q_len = n_blocks * block_size kv_len = seq_len + q_len + # FlexAttention consumes the same predicate as a BlockMask and skips the ~67% of + # tiles that are entirely masked; DFlashAttention routes a BlockMask to + # flex_attention_forward. See dflash_flex_attention.py for why this matters so + # much here (dense mask -> sm80 memory-efficient kernel -> 59% of the step). + if getattr(self, "dflash_use_flex_attention", False): + from .dflash_flex_attention import build_draft_block_mask + + return build_draft_block_mask( + seq_len, + anchor_positions, + block_keep_mask, + n_blocks, + block_size, + window, + device, + head_dim=self.dflash_module.layers[0].self_attn.head_dim, + ) + q_indices = torch.arange(q_len, device=device).view(1, 1, -1, 1) kv_indices = torch.arange(kv_len, device=device).view(1, 1, 1, -1) q_block_ids = q_indices // block_size diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 9b6c459c1a2..a162eb236a1 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -185,6 +185,36 @@ def _get_attn_fn(self): self._attn_fn = ALL_ATTENTION_FUNCTIONS.get(impl, ALL_ATTENTION_FUNCTIONS["sdpa"]) return self._attn_fn + def _attend(self, q, k, v, attention_mask, bsz, q_len): + """Run attention and project, routing a FlexAttention BlockMask to the flex kernel. + + ``attention_mask`` is either the dense additive [B, 1, Q, KV] tensor (HF attention + dispatch) or a BlockMask carrying the same predicate block-sparsely. + """ + from .dflash_flex_attention import flex_attention_forward, is_block_mask + + if is_block_mask(attention_mask): + dropout = 0.0 if not self.training else self.attention_dropout + if dropout: + raise ValueError( + "FlexAttention path does not support attention_dropout > 0 " + f"(got {dropout}); unset dflash_use_flex_attention." + ) + attn_output = flex_attention_forward(q, k, v, attention_mask, self.scaling) + else: + attn_fn = self._get_attn_fn() + attn_output, _ = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + ) + return self.o_proj(attn_output.reshape(bsz, q_len, -1)) + def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None): """Forward with KV injection. @@ -218,20 +248,7 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) - # Use HF's attention dispatch (handles GQA internally) - attn_fn = self._get_attn_fn() - attn_output, _ = attn_fn( - self, - q, - k, - v, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - sliding_window=self.sliding_window, - ) - attn_output = attn_output.reshape(bsz, q_len, -1) - return self.o_proj(attn_output) + return self._attend(q, k, v, attention_mask, bsz, q_len) class DFlashGemma4Attention(DFlashAttention): @@ -353,19 +370,7 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) - attn_fn = self._get_attn_fn() - attn_output, _ = attn_fn( - self, - q, - k, - v, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - sliding_window=self.sliding_window, - ) - attn_output = attn_output.reshape(bsz, q_len, -1) - return self.o_proj(attn_output) + return self._attend(q, k, v, attention_mask, bsz, q_len) class DFlashDecoderLayer(nn.Module): diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index ca0c6f779a8..ef6831a1d96 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -250,3 +250,124 @@ def test_offline_forward_self_logit_distillation_recomputes_logits(self, offline assert hasattr(output, "logits") assert output.logits is not None assert torch.isfinite(output.loss).item() + + +class TestDFlashFlexAttentionGPU: + """FlexAttention path must match the dense-mask SDPA path it replaces. + + The block-sparse BlockMask encodes exactly the predicate + ``_build_draft_attention_mask`` materializes, so switching implementations may only + move results by float rounding -- not by a masked position becoming visible, and not + at the fully-masked query rows that invalid blocks produce. + + These build a wider model than the rest of this file: ``get_tiny_llama`` defaults to + hidden_size 32 over 16 heads, i.e. head_dim 2, and FlexAttention's Triton templates + need head_dim >= 16 (below that Inductor finds no valid config and raises). 128/2 + heads gives head_dim 64, the smallest standard size. + """ + + SEQ = 64 + BASE_KWARGS = { + "num_hidden_layers": 4, + "hidden_size": 128, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "intermediate_size": 64, + "max_position_embeddings": 256, + "vocab_size": 64, + } + + @classmethod + def _model(cls, **cfg_overrides): + model = get_tiny_llama(**cls.BASE_KWARGS) + config = get_dflash_config() + config.update(cfg_overrides) + mtsp.convert(model, [("dflash", config)]) + return model.cuda().train() + + @classmethod + def _pair(cls, **cfg_overrides): + """A dense model and a flex model with identical draft weights.""" + dense = cls._model(**cfg_overrides) + flex = cls._model(dflash_use_flex_attention=True, **cfg_overrides) + flex.dflash_module.load_state_dict(dense.dflash_module.state_dict()) + return dense, flex + + @classmethod + def _inputs(cls, bsz=2): + input_ids = torch.randint(0, cls.BASE_KWARGS["vocab_size"], (bsz, cls.SEQ), device="cuda") + attention_mask = torch.ones(bsz, cls.SEQ, dtype=torch.long, device="cuda") + return input_ids, attention_mask + + def test_mask_builder_returns_block_mask(self): + """With the flag on, the mask builder hands back a BlockMask, not a dense tensor.""" + pytest.importorskip("torch.nn.attention.flex_attention") + from modelopt.torch.speculative.plugins.dflash_flex_attention import is_block_mask + + model = self._model(dflash_use_flex_attention=True) + mask = model._build_draft_attention_mask( + self.SEQ, + torch.tensor([[4, 8]], device="cuda"), + torch.tensor([[True, True]], device="cuda"), + 2, + torch.float32, + torch.device("cuda"), + window=None, + ) + assert is_block_mask(mask) + assert not torch.is_tensor(mask) + + @pytest.mark.parametrize("window", [None, 8]) + def test_matches_dense_mask_path(self, window): + """Loss agrees with the dense path to bf16 tolerance, with and without SWA.""" + pytest.importorskip("torch.nn.attention.flex_attention") + overrides = {} if window is None else {"dflash_swa_window_size": window} + dense, flex = self._pair(**overrides) + input_ids, attention_mask = self._inputs() + + # Anchors are resampled from the RNG on every forward, so both models must draw + # from the same seed or they would see different anchors, not different kernels. + torch.manual_seed(1234) + out_dense = dense(input_ids=input_ids, attention_mask=attention_mask) + torch.manual_seed(1234) + out_flex = flex(input_ids=input_ids, attention_mask=attention_mask) + + torch.testing.assert_close(out_flex.loss, out_dense.loss, rtol=2e-2, atol=2e-2) + + def test_matches_dense_mask_path_with_invalid_blocks(self): + """Fully-masked query rows (invalid blocks) must not diverge either. + + A row of all -inf is the one place the two kernels could legitimately disagree + (softmax of nothing), and answer_only_loss produces such rows whenever a sample + has fewer valid anchors than the batch maximum. + """ + pytest.importorskip("torch.nn.attention.flex_attention") + dense, flex = self._pair() + input_ids, attention_mask = self._inputs() + # Row 1 keeps far fewer supervised positions than row 0, so its trailing blocks + # come back with block_keep_mask False. + labels = input_ids.clone() + labels[1, : self.SEQ - BLOCK_SIZE] = -100 + + torch.manual_seed(7) + out_dense = dense(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + torch.manual_seed(7) + out_flex = flex(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + + assert torch.isfinite(out_flex.loss), "flex produced a non-finite loss" + torch.testing.assert_close(out_flex.loss, out_dense.loss, rtol=2e-2, atol=2e-2) + + def test_backward_produces_finite_grads(self): + """The flex path is differentiable and its grads are finite.""" + pytest.importorskip("torch.nn.attention.flex_attention") + flex = self._model(dflash_use_flex_attention=True) + input_ids, attention_mask = self._inputs() + + flex(input_ids=input_ids, attention_mask=attention_mask).loss.backward() + grads = [ + p.grad + for p in flex.dflash_module.parameters() + if p.requires_grad and p.grad is not None + ] + assert grads, "no draft gradients were produced" + assert all(torch.isfinite(g).all() for g in grads) From dd5b7fcd8ca996499f6e6e0707e18441b7777cce Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:11:39 -0700 Subject: [PATCH 2/6] perf(speculative): project teacher rows on demand and collapse the DSpark metric syncs Three independent inefficiencies in the DSpark training step, all of which grew out of materializing or synchronizing more than the loss actually consumes. * The offline/streaming path built full-sequence [B, seq, vocab] base logits and then gathered a few thousand rows out of them. DFlashBaseModelOutput now carries the post-final-norm base hidden and takes defer_lm_head, so the loss projects only the rows it reads. At the Gemma-4-E4B shape that removes roughly 17 GB/step of memory traffic. * The Markov bias was added out-of-place, allocating a third [B, N, bs, vocab] tensor (8.6 GB at this shape). It now folds in place when the uncorrected logits are dead, which is the case whenever dflash_report_acc is off -- their only consumer is the base_accuracy diagnostic. Numerically identical: a.add_(b) and a + b compute the same elementwise rounded sum, and neither lm_head nor markov_w2 saves its OUTPUT for backward. * Five per-step .item() calls were collapsed into a single torch.stack([...]).tolist(). Each .item() is a full device synchronize, and five of them chop up the window in which DDP's all-reduce can hide behind backward. This one change was worth 39 ms of NCCL exposure on its own. compute_markov_latent() is split out of compute_markov_bias() so the rank-r argument to markov_w2 is reachable without the vocab-wide projection; _rnn_step returns the latent rather than a bias, and both call sites are updated. Measured on PDX B300, 3 nodes (1 vLLM serve + 2 trainer = 16 DP ranks), Gemma-4-E4B DSpark, bs 4 x seq 4096, median over 16 ranks x 20 steady-state steps: 0.711 -> 0.629 s/step. Signed-off-by: Hao Guo --- .../torch/speculative/plugins/hf_dspark.py | 161 ++++++++++++++---- .../speculative/plugins/modeling_dflash.py | 27 ++- .../speculative/plugins/modeling_dspark.py | 41 +++-- 3 files changed, 177 insertions(+), 52 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 4183a63f1d1..70aa1306f9f 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -78,7 +78,44 @@ __all__ = ["HFDSparkModel"] -def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024): +def _tvd_chunk(a, b): + """Per-token TVD for one row chunk: ``(softmax(a) - softmax(b)).abs().sum(-1)``.""" + return ( + (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1) + ) + + +# torch.compile of _tvd_chunk, built once per process and reused. Eager, this chain is +# six separate passes over [chunk, vocab] tensors (two bf16->fp32 casts, two softmaxes, a +# subtract, an abs) that are 1 GB each at the Gemma-4 shape; Inductor fuses it into a +# couple of kernels. Profiling put softmax + the surrounding elementwise ops at ~50% of +# the training step once FlexAttention removed the attention bottleneck. +_compiled_tvd_chunk = None +_tvd_compile_failed = False + + +def _get_tvd_chunk(use_compile: bool): + """Return the TVD chunk fn, compiled when asked for and when compilation succeeds. + + Deliberately NOT wrapped in ``torch._dynamo.config.suppress_errors = True`` (which the + Eagle plugin sets globally): that turns a compile failure into a silent fallback to + eager, i.e. a performance feature that reports success while doing nothing. Here a + failure is warned about once and then remembered. + """ + global _compiled_tvd_chunk, _tvd_compile_failed + if not use_compile or _tvd_compile_failed: + return _tvd_chunk + if _compiled_tvd_chunk is None: + try: + _compiled_tvd_chunk = torch.compile(_tvd_chunk, dynamic=False, fullgraph=True) + except Exception as exc: + _tvd_compile_failed = True + logger.warning("torch.compile of the DSpark TVD chunk failed (%s); using eager.", exc) + return _tvd_chunk + return _compiled_tvd_chunk + + +def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None): """Total-variation distance ||softmax(a)-softmax(b)||_1 / ... per token, memory-lean. Materializing both [N, vocab] float32 softmax tensors at once OOMs at large @@ -86,13 +123,9 @@ def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024): gradient-checkpoint each chunk so the wide softmaxes are recomputed in backward rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``. - """ - - def _chunk(a, b): - return ( - (torch.softmax(a.float(), dim=-1) - torch.softmax(b.float(), dim=-1)).abs().sum(dim=-1) - ) + """ + _chunk = chunk_fn or _tvd_chunk outs = [] for i in range(0, final_logits.size(0), chunk_size): a, b = final_logits[i : i + chunk_size], teacher_logits[i : i + chunk_size] @@ -141,9 +174,17 @@ def get_exporter(self): return DSparkExporter(self) - def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_positions, n_blocks): + def _apply_markov_head( + self, hidden, backbone_logits, input_ids, anchor_positions, n_blocks, inplace=False + ): """Add the Markov transition bias to the backbone base logits. + ``inplace`` folds the bias into ``backbone_logits`` instead of allocating a third + [B, N, bs, vocab] tensor (8.6 GB at the Gemma-4 shape). Safe for autograd -- neither + ``lm_head`` nor ``markov_w2`` saves its OUTPUT for backward -- but it leaves + ``backbone_logits`` holding the corrected logits, so the caller must not still need + the uncorrected ones (they are only used for the ``base_accuracy`` metric). + Returns ``(final_logits [B, N, bs, V], confidence_logits [B, N, bs] | None)``. """ bsz, seq_len = input_ids.shape @@ -160,7 +201,7 @@ def _apply_markov_head(self, hidden, backbone_logits, input_ids, anchor_position prev_ids = torch.gather(input_ids.unsqueeze(1).expand(-1, n_blocks, -1), 2, prev_idx) bias = self.dflash_module.compute_markov_bias(prev_ids, hidden4d) - final4d = base4d + bias + final4d = base4d.add_(bias) if inplace else base4d + bias confidence_logits = None if self.dflash_module.use_confidence_head: @@ -176,13 +217,24 @@ def _compute_dspark_loss( anchor_positions, block_keep_mask, loss_mask, - target_model_logits, + target_model_logits=None, + teacher_hidden=None, ): """Compute the three-term DSpark loss (CE + TVD + confidence BCE) and metrics. Uses next-token (shift_label) alignment: block position k predicts the token at anchor+k+1; the aligned target distribution is the base model's own next-token distribution at position anchor+k (= label index - 1). + + The teacher distribution can arrive two ways. ``teacher_hidden`` ([B, seq, H], the + base model's post-final-norm hidden) is preferred: only the N*block_size rows the + loss actually reads are gathered and projected, so the full-sequence + [B, seq, vocab] tensor is never built and never gathered out of -- at the Gemma-4 + shape that is ~17 GB/step of memory traffic. ``target_model_logits`` is the + fallback for producers that hand over logits directly. + + ``backbone_logits`` may be None, which skips the ``base_accuracy`` diagnostic (and + with it a second full-vocab argmax); see ``dflash_report_acc``. """ bsz, seq_len = input_ids.shape bs = self.dflash_block_size @@ -217,7 +269,7 @@ def _compute_dspark_loss( weight_mask = weight_mask * decay flat_final = final_logits.reshape(-1, vocab) - flat_base = backbone_logits.reshape(-1, vocab) + flat_base = None if backbone_logits is None else backbone_logits.reshape(-1, vocab) flat_targets = target_ids.reshape(-1) flat_weights = weight_mask.reshape(-1) valid_count = flat_weights.sum() + 1e-6 @@ -225,12 +277,27 @@ def _compute_dspark_loss( # Aligned target distribution: base-model logits that predict token anchor+k+1 # sit at position anchor+k (= label index - 1). teacher_indices = (safe_label_indices - 1).clamp(min=0) - teacher_logits = torch.gather( - target_model_logits.unsqueeze(1).expand(-1, n_blocks, -1, -1), - 2, - teacher_indices.unsqueeze(-1).expand(-1, -1, -1, vocab), - ) - flat_teacher = teacher_logits.reshape(-1, vocab).detach() + with torch.no_grad(): + if teacher_hidden is not None: + hdim = teacher_hidden.size(-1) + gathered_hidden = torch.gather( + teacher_hidden.unsqueeze(1).expand(-1, n_blocks, -1, -1), + 2, + teacher_indices.unsqueeze(-1).expand(-1, -1, -1, hdim), + ) + flat_teacher = self._base_model_lm_head(gathered_hidden.reshape(-1, hdim)) + else: + if target_model_logits is None: + raise ValueError( + "DSpark loss needs the base distribution: pass teacher_hidden " + "(preferred) or target_model_logits." + ) + flat_teacher = torch.gather( + target_model_logits.unsqueeze(1).expand(-1, n_blocks, -1, -1), + 2, + teacher_indices.unsqueeze(-1).expand(-1, -1, -1, vocab), + ).reshape(-1, vocab) + flat_teacher = flat_teacher.detach() if valid_count <= 1.0: loss = flat_final.sum() * 0.0 @@ -243,7 +310,11 @@ def _compute_dspark_loss( # Term 2: total-variation distance between the corrected draft and target. # Chunked + checkpointed to avoid materializing two [N, vocab] softmaxes at once. - l1_per_token = _tvd_per_token(flat_final, flat_teacher) + l1_per_token = _tvd_per_token( + flat_final, + flat_teacher, + chunk_fn=_get_tvd_chunk(getattr(self, "dflash_use_torch_compile", False)), + ) l1_loss = (l1_per_token * flat_weights).sum() / valid_count # Term 3: confidence head BCE against the analytical accept rate c* = 1 - 0.5*TVD. @@ -264,19 +335,25 @@ def _compute_dspark_loss( with torch.no_grad(): eval_count = binary_eval_mask.sum() + 1e-6 keep = binary_eval_mask > 0.5 - accuracy = ( - ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count - ).item() - base_accuracy = ( - ((flat_base.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count - ).item() + acc = ((flat_final.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count + base_acc = ( + acc.new_zeros(()) + if flat_base is None + else ((flat_base.argmax(dim=-1) == flat_targets) & keep).sum().float() / eval_count + ) + # ONE device sync for all five scalars. Each .item() is a full synchronize, and + # five of them per step chop up the window in which DDP's all-reduce can hide + # behind backward -- measured comm exposure is 18% of the step post-FlexAttention. + acc_v, base_acc_v, ce_v, l1_v, conf_v = torch.stack( + [acc, base_acc, ce_loss.detach(), l1_loss.detach(), confidence_loss.detach()] + ).tolist() metrics = { - "ce_loss": ce_loss.detach().item(), - "l1_loss": l1_loss.detach().item(), - "confidence_loss": float(confidence_loss.detach().item()), - "base_accuracy": base_accuracy, + "ce_loss": ce_v, + "l1_loss": l1_v, + "confidence_loss": conf_v, + "base_accuracy": base_acc_v, } - return loss, accuracy, metrics + return loss, acc_v, metrics def forward( self, @@ -336,9 +413,14 @@ def forward( self._base_model_norm, self._base_model_lm_head, need_logits=True, + # Hand back the normed hidden instead of full-sequence logits; the loss + # projects only the rows it reads. Producers that supply base_model_logits + # directly still come back with logits and take the fallback path. + defer_lm_head=True, ) target_hidden = base_outputs.target_hidden target_model_logits = base_outputs.logits + teacher_hidden = base_outputs.base_hidden else: # Call the inner base model directly (NOT super().forward(), which during # training runs the full DFlash pipeline). Compute target-model logits via @@ -349,7 +431,9 @@ def forward( attention_mask=attention_mask, output_hidden_states=True, ) - target_model_logits = self._base_model_lm_head(base_out.last_hidden_state) + # lm_head is applied per-row inside the loss, not over the whole sequence. + teacher_hidden = base_out.last_hidden_state + target_model_logits = None offset = 1 selected = [base_out.hidden_states[lid + offset] for lid in self.target_layer_ids] target_hidden = torch.cat(selected, dim=-1) # [B, seq, num_layers * H] @@ -399,19 +483,30 @@ def forward( ) # 6. Backbone logits → Markov correction → three-term loss. + # dflash_report_acc gates the base_accuracy diagnostic (the draft's accuracy BEFORE + # the Markov correction). With it off, the uncorrected logits are dead after the + # bias is added, so the bias folds in place and a second full-vocab argmax is + # skipped -- two [N, vocab] tensors' worth of traffic per step. + report_base_acc = getattr(self, "dflash_report_acc", True) backbone_logits = self._base_model_lm_head(hidden).reshape(bsz, n_blocks, block_size, -1) final_logits, confidence_logits = self._apply_markov_head( - hidden, backbone_logits, input_ids, anchor_positions, n_blocks + hidden, + backbone_logits, + input_ids, + anchor_positions, + n_blocks, + inplace=not report_base_acc, ) loss, accuracy, metrics = self._compute_dspark_loss( - backbone_logits, + backbone_logits if report_base_acc else None, final_logits, confidence_logits, input_ids, anchor_positions, block_keep_mask, loss_mask, - target_model_logits, + target_model_logits=target_model_logits, + teacher_hidden=teacher_hidden, ) return ModelOutput(loss=loss, logits=None, train_acc=[[accuracy]], dspark_metrics=metrics) diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index a162eb236a1..3302844e32c 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -65,10 +65,19 @@ class DFlashBaseModelOutput: target_hidden: torch.Tensor # concatenated hidden states from target layers [B, seq, N*H] logits: torch.Tensor | None = None # base model logits [B, seq, vocab] + # Post-final-norm base hidden [B, seq, H], i.e. lm_head's input. Consumers that only + # need the base distribution at a handful of positions project THIS at those rows + # instead of materialising (and then gathering out of) full-sequence logits. + base_hidden: torch.Tensor | None = None @classmethod def from_offline_dict( - cls, d: dict, base_model_norm=None, base_model_lm_head=None, need_logits=False + cls, + d: dict, + base_model_norm=None, + base_model_lm_head=None, + need_logits=False, + defer_lm_head=False, ): """Construct from a dict of pre-computed base model outputs (offline training). @@ -85,19 +94,25 @@ def from_offline_dict( to lm_head would be a corrupt distillation target). """ logits = d.get("base_model_logits") + base_hidden = None if need_logits and logits is None: + out_hiddens = d.get("base_model_hidden_states") + if out_hiddens is None: + raise KeyError("base_model_hidden_states") + base_hidden = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm) + if defer_lm_head: + # Caller will project only the rows it needs; skip the full-sequence + # [B, seq, vocab] materialisation entirely. + return cls(target_hidden=d["aux_hidden_states"], base_hidden=base_hidden) if base_model_lm_head is None: raise ValueError( "need_logits=True but base_model_lm_head is None; cannot reconstruct logits." ) - out_hiddens = d.get("base_model_hidden_states") - if out_hiddens is None: - raise KeyError("base_model_hidden_states") - out_hiddens = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm) - logits = base_model_lm_head(out_hiddens) + logits = base_model_lm_head(base_hidden) return cls( target_hidden=d["aux_hidden_states"], logits=logits, + base_hidden=base_hidden, ) diff --git a/modelopt/torch/speculative/plugins/modeling_dspark.py b/modelopt/torch/speculative/plugins/modeling_dspark.py index aa1a4181bbd..1dbce2768c9 100644 --- a/modelopt/torch/speculative/plugins/modeling_dspark.py +++ b/modelopt/torch/speculative/plugins/modeling_dspark.py @@ -142,44 +142,58 @@ def prev_token_embeddings(self, prev_ids: torch.Tensor) -> torch.Tensor: """Look up the Markov embedding ``W1[x_{k-1}]`` of the teacher-forced prev tokens.""" return self.markov_w1(prev_ids.long()) - def compute_markov_bias(self, prev_ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor: - """Compute the transition bias ``B_k`` added to the backbone base logits. + def compute_markov_latent(self, prev_ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor: + """Rank-``r`` state the transition bias is projected from, BEFORE ``markov_w2``. + + Every head type ends in ``markov_w2()``; this returns that argument so a + caller can defer the vocab-wide projection (which materialises a + ``[B, N, block_size, vocab]`` tensor) to the point where it is actually needed. Args: prev_ids: Teacher-forced previous-token ids per block position [B, N, block_size]. hidden: Backbone hidden states [B, N, block_size, H] (used by gated/rnn heads). Returns: - Logit bias [B, N, block_size, vocab]. + Latent [B, N, block_size, r]. """ prev_emb = self.prev_token_embeddings(prev_ids) # [B, N, bs, r] if self.markov_head_type == "vanilla": - return self.markov_w2(prev_emb) + return prev_emb if self.markov_head_type == "gated": gate = torch.sigmoid(self.gate_proj(torch.cat([hidden, prev_emb], dim=-1))) - return self.markov_w2(gate.to(prev_emb.dtype) * prev_emb) + return gate.to(prev_emb.dtype) * prev_emb # rnn: unroll the gated recurrence over the block dimension. block_size = prev_ids.shape[-1] leading = prev_emb.shape[:-2] # [B, N] state = torch.zeros(*leading, self.markov_rank, device=prev_emb.device, dtype=hidden.dtype) - biases = [] + latents = [] for k in range(block_size): - state, bias = self._rnn_step(state, prev_emb[..., k, :], hidden[..., k, :]) - biases.append(bias) - return torch.stack(biases, dim=-2) + state, latent = self._rnn_step(state, prev_emb[..., k, :], hidden[..., k, :]) + latents.append(latent) + return torch.stack(latents, dim=-2) + + def compute_markov_bias(self, prev_ids: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor: + """Compute the transition bias ``B_k`` added to the backbone base logits. + + Returns: + Logit bias [B, N, block_size, vocab]. + """ + return self.markov_w2(self.compute_markov_latent(prev_ids, hidden)) def _rnn_step(self, state, prev_emb, hidden): - """One GRU-like recurrent step. Returns (new_state [.., r], bias [.., vocab]).""" + """One GRU-like recurrent step. Returns (new_state [.., r], latent [.., r]). + + The latent is ``markov_w2``'s input, not the bias -- see compute_markov_latent. + """ z = torch.cat([state, prev_emb, hidden], dim=-1) gate_raw, candidate_raw, output_raw = self.joint_proj(z).chunk(3, dim=-1) gate = torch.sigmoid(gate_raw) candidate = torch.tanh(candidate_raw) new_state = gate * state + (1.0 - gate) * candidate - bias = self.markov_w2(torch.tanh(output_raw)) - return new_state, bias + return new_state, torch.tanh(output_raw) def markov_step(self, prev_token: torch.Tensor, hidden: torch.Tensor, state=None): """One autoregressive Markov step (inference): bias for a single position. @@ -204,7 +218,8 @@ def markov_step(self, prev_token: torch.Tensor, hidden: torch.Tensor, state=None state = torch.zeros( prev_emb.shape[0], self.markov_rank, device=prev_emb.device, dtype=hidden.dtype ) - state, bias = self._rnn_step(state, prev_emb, hidden) + state, latent = self._rnn_step(state, prev_emb, hidden) + bias = self.markov_w2(latent) return bias, state def compute_confidence_logits( From 1e288e548965c7828c8e9eb1c8ad8ca0f0df0740 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:12:04 -0700 Subject: [PATCH 3/6] perf(speculative): chunk the DSpark TVD with Tensor.split instead of slicing _tvd_per_token chunked its row axis with final_logits[i : i + chunk_size]. Slicing and splitting agree on every value, so no numerical test can tell them apart -- only the backward graph can, and there the difference is asymptotic. A slice per chunk creates one SliceBackward0 each, and every one of those allocates a zero tensor of the FULL [N, vocab] shape and scatters its own chunk's gradient into it: O(n_chunks * N * vocab). Tensor.split creates a single SplitBackward0 whose backward is one cat: O(N * vocab). At the Gemma-4-E4B shape (N = 4 * 512 * 8 = 16384 rows, vocab = 262144, so 8 GiB per [N, vocab] bf16 tensor, 16 chunks) that 16x amplification was 93.2 ms/step -- 15% of the step and the second-largest item in the profile after the DDP all-reduce. It had been flat at 97 +/- 1 ms across three earlier optimisation rounds, invisible because it was never broken out of the elementwise bucket. Measured on PDX B300, 3 nodes (1 vLLM serve + 2 trainer = 16 DP ranks), bs 4 x seq 4096, median over 16 ranks x 20 steady-state steps: 0.629 -> 0.541 s/step, with -85.7 ms of the -87.9 ms landing in backward and forward unchanged. In isolation at the production shape, backward went 198.0 -> 107.9 ms while peak memory was unchanged (41.9 -> 42.1 GiB). Bitwise identical, verified elementwise at the production shape: the returned per-token TVD, grad(hidden) and grad(markov_w2.weight) all compare exactly equal. End to end over 40 logged steps, max |this - previous| on the loss is 0.018, against a 0.025 spread among three earlier runs -- i.e. inside the pipeline's own run-to-run noise. The new tests do two things the existing suite could not. TestTvdPerTokenChunking parametrises chunk_size over values that divide N and values that leave a ragged tail: every existing DSpark test runs 96 rows against the default chunk_size of 1024, so all of them fit in a single chunk and no chunk-boundary bug could ever fail CI. test_chunks_via_split_not_slice walks the autograd graph and asserts SplitBackward is present and SliceBackward is not -- pinning the optimisation itself, since it has no observable numerical signature. Signed-off-by: Hao Guo --- .../torch/speculative/plugins/hf_dspark.py | 28 ++++++- .../speculative/plugins/test_hf_dspark.py | 80 ++++++++++++++++++- 2 files changed, 105 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 70aa1306f9f..e38afe8e9bb 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -124,11 +124,35 @@ def _tvd_per_token(final_logits, teacher_logits, chunk_size=1024, chunk_fn=None) rather than held — peak memory ~ chunk_size*vocab instead of N*vocab. The math is identical to ``(softmax(final)-softmax(teacher)).abs().sum(-1)``. + Chunking goes through ``Tensor.split``, NOT ``final_logits[i : i + chunk_size]``. + Both produce the same views over the same rows, so the forward values are + identical — but the backward graphs are not, and the difference is large. A slice + per chunk creates one ``SliceBackward0`` each, and every one of those allocates a + zero tensor of the FULL [N, vocab] shape and scatters its own chunk's gradient + into it, so the cost is O(n_chunks * N * vocab). ``split`` creates a single + ``SplitBackward0`` whose backward is one ``cat``, i.e. O(N * vocab). + + At the Gemma-4-E4B shape (N = bsz 4 * n_blocks 512 * block_size 8 = 16384, + vocab = 262144 -> 8 GiB per [N, vocab] bf16 tensor, 16 chunks) that is not a + micro-optimization: measured on a B300, backward went 198.0 ms -> 107.9 ms, and + kernel attribution on the training profile put 93.2 ms/step -- 15% of a 614 ms + step, the second-largest item after the DDP all-reduce -- on ``SliceBackward0``. + Peak memory is unchanged (41.9 -> 42.1 GiB). Outputs and gradients are bitwise + identical: verified elementwise at the production shape for the returned + per-token TVD, grad(hidden) and grad(markov_w2.weight). + + Do NOT substitute ``torch.chunk`` or ``torch.tensor_split``. Those split into a + fixed NUMBER of pieces and so pick different boundaries (ceil(N/k)), which changes + the shapes handed to ``chunk_fn`` -- and ``chunk_fn`` may be a + ``torch.compile(..., dynamic=False)`` build that recompiles per shape. """ _chunk = chunk_fn or _tvd_chunk outs = [] - for i in range(0, final_logits.size(0), chunk_size): - a, b = final_logits[i : i + chunk_size], teacher_logits[i : i + chunk_size] + for a, b in zip( + final_logits.split(chunk_size, dim=0), + teacher_logits.split(chunk_size, dim=0), + strict=True, + ): if torch.is_grad_enabled() and a.requires_grad: outs.append(torch.utils.checkpoint.checkpoint(_chunk, a, b, use_reentrant=False)) else: diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 6a0f884d33f..9d421451cc7 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -34,7 +34,7 @@ import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel -from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel +from modelopt.torch.speculative.plugins.hf_dspark import HFDSparkModel, _tvd_per_token from modelopt.torch.speculative.plugins.modeling_dflash import DFlashModule from modelopt.torch.speculative.plugins.modeling_dspark import DSparkModule @@ -311,3 +311,81 @@ def test_export_config_has_dspark_fields(self, tmp_path): assert dc["shift_label"] is True assert "mask_token_id" in dc assert "target_layer_ids" in dc + + +class TestTvdPerTokenChunking: + """``_tvd_per_token`` must be chunk-size-invariant, and must chunk via ``split``. + + No other DSpark test reaches a second chunk: the tiny fixture yields N = 2 * 12 * 4 + = 96 rows against the default ``chunk_size=1024``, so every existing test runs in a + single chunk and any chunk-boundary bug -- a mis-ordered ``cat``, a ragged tail + handled wrong, a row/label misalignment -- passes CI silently. These tests call the + helper directly so they can force many chunks on a tiny tensor. + """ + + @staticmethod + def _run(chunk_size, n=12, vocab=32): + """Forward + backward through ``_tvd_per_token`` from a fixed seed.""" + torch.manual_seed(0) + final = torch.randn(n, vocab, requires_grad=True) + teacher = torch.randn(n, vocab) + out = _tvd_per_token(final, teacher, chunk_size=chunk_size) + # Weight the rows unequally, so a cat that reassembles the chunks in the wrong + # order cannot cancel out in the reduction. + (out * torch.arange(1, n + 1, dtype=out.dtype)).sum().backward() + return out.detach(), final.grad.detach() + + # 12 rows: 5 and 7 leave a ragged last chunk (5+5+2, 7+5); 13 and 1024 exceed n. + @pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 7, 11, 12, 13, 1024]) + def test_chunk_size_invariant(self, chunk_size): + ref_out, ref_grad = self._run(12) # one chunk == the unchunked reference + out, grad = self._run(chunk_size) + assert torch.equal(out, ref_out), f"TVD value changed at chunk_size={chunk_size}" + assert torch.equal(grad, ref_grad), f"TVD grad changed at chunk_size={chunk_size}" + + def test_chunks_via_split_not_slice(self): + """Pin the optimization itself, not just its result. + + Slicing and splitting agree on every value, so no numerical test can tell them + apart -- only the backward graph can. Slicing costs O(n_chunks * N * vocab) + because each ``SliceBackward0`` zero-fills a full [N, vocab] tensor; at the + Gemma-4 shape that was 93.2 ms/step. This test is what stops the loop being + "simplified" back to ``final_logits[i : i + chunk_size]``. + """ + final = torch.randn(8, 4, requires_grad=True) + teacher = torch.randn(8, 4) + out = _tvd_per_token(final, teacher, chunk_size=2) + + # `alive` is load-bearing, not debris: accessing `.next_functions` hands back a + # FRESH python wrapper for each node every time, so a node we do not hold a + # reference to is freed the moment we pop it -- and CPython happily reuses that + # address for a later, different node. Without `alive` the id() check reports a + # false "already visited" and the walk truncates after three nodes, missing the + # Split/Slice node entirely (verified: both variants returned the same + # ['AbsBackward0', 'CatBackward0', 'SumBackward1']). + seen, visited, alive, stack = set(), set(), [], [out.grad_fn] + while stack: + fn = stack.pop() + if fn is None or id(fn) in visited: + continue + visited.add(id(fn)) + alive.append(fn) + seen.add(type(fn).__name__) + stack.extend(nxt for nxt, _ in fn.next_functions) + + assert any(name.startswith("SplitBackward") for name in seen), ( + f"expected a SplitBackward node in the graph, saw: {sorted(seen)}" + ) + assert not any(name.startswith("SliceBackward") for name in seen), ( + f"chunking regressed to per-chunk slicing, saw: {sorted(seen)}" + ) + + def test_no_grad_path_matches_grad_path(self): + """The ``requires_grad=False`` branch skips checkpointing; values must not move.""" + torch.manual_seed(0) + final = torch.randn(12, 32) + teacher = torch.randn(12, 32) + with torch.no_grad(): + plain = _tvd_per_token(final, teacher, chunk_size=5) + grad_out = _tvd_per_token(final.requires_grad_(True), teacher, chunk_size=5) + assert torch.equal(plain, grad_out.detach()) From 31b25143fd8346ee8f36fb49eb4727f9950c9536 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 28 Aug 2026 13:41:05 -0700 Subject: [PATCH 4/6] perf(speculative): pin the draft's block count so the compiled attention has one shape _sample_anchor_positions sized its output by the batch: min(num_anchors, valid_counts.max() - 1), i.e. the longest answer in the batch. n_blocks sets the draft's q_len and kv_len, and both the block-mask builder and flex_attention are torch.compile(..., dynamic=False), so every distinct value was a distinct shape and a fresh compile. Each cost 170-300 s, and new ones kept arriving, because any batch whose answers are all shorter than the cap mints a width never seen before. That is not a warm-up cost, it is the run. A 3-node Gemma-4-E4B job spent 1360 of its first 1510 training seconds frozen in recompilation. Reconstructing the consumption rate from the serve's request log shows the trainer alternating between 0.6 s/step bursts and multi-minute stalls -- seven of them in ~107 steps, the last still arriving at step 102. max_n is now min(num_anchors, max_anchor + 1), a function of the config and the sequence length alone. The bound the old code computed survives as `cap`, a 0-dim tensor applied before the sort rather than as a slice width, so the sampled anchors are bitwise what they were. Widening the slice and sorting after would instead pull anchors from past the bound into the front of the row and silently move which positions are trained on, while every shape assertion still passed. The surplus columns carry keep=False, which every consumer already handles because rows shorter than the batch maximum have always produced them. Dropping the .item() also takes a host sync out of every step. Measured on 3 nodes (1 serve, 2 trainer, bs 4 x 4096 over 16 ranks), 40 steps: before after compile events 4 1 compile seconds 772.1 173.5 wall (s) 799.1 204.7 3.90x The 20-step mean loss is unchanged against a production run of the same recipe -- 3.211 / 2.660 / 2.621 / 2.614 before, 3.211 / 2.661 / 2.623 / 2.601 after, inside this pipeline's 0.025 run-to-run spread. Steady-state step time moves 0.534 -> 0.556 s: short batches now pay for 512 blocks they used to skip, which is the trade. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/speculative/plugins/hf_dflash.py | 47 +++-- .../speculative/plugins/test_hf_dflash.py | 172 ++++++++++++++++++ 2 files changed, 208 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 7ffd735553a..c15ca2a1722 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -584,6 +584,24 @@ def _sample_anchor_positions(self, seq_len, loss_mask, device): Returns (anchor_positions [B, N], block_keep_mask [B, N]). + ``N`` is fixed by the config and the sequence length, never by the batch + contents. It used to be ``min(num_anchors, valid_counts.max() - 1)``, i.e. a + function of the longest answer in the batch. ``n_blocks`` sets the draft's + ``q_len`` and ``kv_len``, and both the block-mask builder and the attention are + ``torch.compile(..., dynamic=False)``, so every distinct value cost a fresh + 170-300 s recompile -- and new values kept appearing indefinitely, since any + batch whose answers are all shorter than the cap mints one. A 3-node production + run spent 1360 of its first 1510 training seconds frozen in recompilation + against a 0.6 s/step steady state. + + Which anchors get sampled is unchanged: ``cap`` below still applies the old + data-dependent bound, just on-device and before the sort rather than as a slice + width. The surplus columns carry ``keep=False``, which every consumer already + handles because rows shorter than the batch maximum have always produced them: + the losses normalize by the weight sum derived from ``block_keep_mask``, and + ``mask_mod`` ands the same flag in, so FlexAttention skips those tiles instead + of computing them. + TODO: Fix the random seed per epoch (change between epochs) so that anchor positions are deterministic within an epoch. This would allow caching the derived masks and position IDs across steps while preserving the same data augmentation @@ -596,27 +614,34 @@ def _sample_anchor_positions(self, seq_len, loss_mask, device): valid = loss_mask[:, : max_anchor + 1] > 0.5 valid_counts = valid.sum(dim=1) - max_n = min(num_anchors, int(valid_counts.max().item()) - 1) - if max_n <= 0: - # No valid anchors — return empty - anchors = torch.zeros(bsz, 1, dtype=torch.long, device=device) - keep = torch.zeros(bsz, 1, dtype=torch.bool, device=device) - return anchors, keep + # Static. Bounded by max_anchor + 1 as well as num_anchors so that short sequences + # (unit tests, short-context recipes) cannot ask for more columns than exist. + max_n = min(num_anchors, max_anchor + 1) + + # The bound the old code computed, kept on-device so the shapes above stay static + # and the .item() sync is gone. valid_counts <= max_anchor + 1 by construction, so + # clamping to max_n leaves this equal to min(num_anchors, valid_counts.max() - 1). + cap = (valid_counts.max() - 1).clamp(min=0, max=max_n) indices = torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1) - masked_indices = torch.where(valid, indices, torch.tensor(seq_len + 1, device=device)) + fill = torch.tensor(seq_len + 1, device=device) + masked_indices = torch.where(valid, indices, fill) random_vals = torch.rand(bsz, max_anchor + 1, device=device) random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device)) _, sorted_idx = random_vals.sort(dim=1) gathered = torch.gather(masked_indices, 1, sorted_idx) - anchors = gathered[:, :max_n].sort(dim=1).values - keep = torch.arange(max_n, device=device).unsqueeze(0) < valid_counts.unsqueeze(1).clamp( - max=max_n - ) + # Blank past `cap` BEFORE sorting, not after. Sorting a wider slice would pull + # anchors from beyond the old bound into the front of the row and silently change + # which positions are trained on; blanking first reproduces the old + # ``gathered[:, :cap].sort()`` exactly and leaves the surplus as fill. + cols = torch.arange(max_n, device=device).unsqueeze(0) + anchors = torch.where(cols < cap, gathered[:, :max_n], fill).sort(dim=1).values + + keep = cols < torch.minimum(valid_counts, cap).unsqueeze(1) anchors = torch.where(keep, anchors, torch.tensor(0, dtype=torch.long, device=device)) return anchors, keep diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ab5c5a57d21..336bec17e96 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -912,3 +912,175 @@ def test_multi_turn_masks_only_assistant(self, tiny_tokenizer): # User/system content should NOT appear in unmasked tokens assert "You are helpful" not in decoded assert "How are you?" not in decoded + + +def _legacy_sample_anchor_positions(model, seq_len, loss_mask, device): + """Verbatim copy of the pre-static implementation, kept as the semantic reference.""" + bs = model.dflash_block_size + bsz = loss_mask.shape[0] + max_anchor = max(seq_len - bs, 0) + num_anchors = getattr(model, "_num_anchors", 512) + + valid = loss_mask[:, : max_anchor + 1] > 0.5 + valid_counts = valid.sum(dim=1) + max_n = min(num_anchors, int(valid_counts.max().item()) - 1) + + if max_n <= 0: + return ( + torch.zeros(bsz, 1, dtype=torch.long, device=device), + torch.zeros(bsz, 1, dtype=torch.bool, device=device), + ) + + indices = torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1) + masked_indices = torch.where(valid, indices, torch.tensor(seq_len + 1, device=device)) + random_vals = torch.rand(bsz, max_anchor + 1, device=device) + random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device)) + _, sorted_idx = random_vals.sort(dim=1) + gathered = torch.gather(masked_indices, 1, sorted_idx) + anchors = gathered[:, :max_n].sort(dim=1).values + keep = torch.arange(max_n, device=device).unsqueeze(0) < valid_counts.unsqueeze(1).clamp( + max=max_n + ) + anchors = torch.where(keep, anchors, torch.tensor(0, dtype=torch.long, device=device)) + return anchors, keep + + +class TestAnchorSamplingStaticShape: + """n_blocks is fixed by the config, and the anchors it picks are still the legacy ones. + + A data-dependent n_blocks made the draft's q_len/kv_len data-dependent, and both the + block-mask builder and the attention are compiled with ``dynamic=False`` -- so every + new batch shape cost a multi-minute recompile, indefinitely. These tests pin both + halves of the fix: the shape no longer moves with the data, and what the model trains + on did not change. + """ + + DEVICE = torch.device("cpu") + + @staticmethod + def _model(num_anchors): + model = get_tiny_llama(num_hidden_layers=4) + config = get_dflash_config(block_size=BLOCK_SIZE) + config["dflash_num_anchors"] = num_anchors + mtsp.convert(model, [("dflash", config)]) + return model + + @staticmethod + def _loss_mask(lengths, seq_len=SEQ_LEN): + mask = torch.zeros(len(lengths), seq_len) + for row, n in enumerate(lengths): + mask[row, :n] = 1.0 + return mask + + # Each case puts the legacy bound somewhere different: under the cap, ragged across + # rows, at 1, and in the two degenerate cases the old code special-cased. + @pytest.mark.parametrize( + "lengths", + [(13, 13), (13, 5), (9, 9), (5, 3), (2, 1), (1, 1), (0, 0), (13, 0)], + ) + def test_matches_legacy_sampling(self, lengths): + """Same seed in, same anchors and same keep mask out -- bitwise.""" + model = self._model(num_anchors=8) + loss_mask = self._loss_mask(lengths) + + torch.manual_seed(1234) + legacy_anchors, legacy_keep = _legacy_sample_anchor_positions( + model, SEQ_LEN, loss_mask, self.DEVICE + ) + torch.manual_seed(1234) + anchors, keep = model._sample_anchor_positions(SEQ_LEN, loss_mask, self.DEVICE) + + n_old = legacy_keep.shape[1] + assert keep.shape[1] >= n_old + assert torch.equal(keep[:, :n_old], legacy_keep) + assert torch.equal(anchors[:, :n_old], legacy_anchors) + # Everything past the legacy bound is inert padding. + assert not keep[:, n_old:].any() + assert not anchors[:, n_old:].any() + + def test_sort_is_truncated_before_it_widens(self): + """The surplus columns must not pull unsampled anchors to the front of the row. + + Slicing to the static width *after* sorting would do exactly that: anchors past + the legacy bound would sort in among the kept ones and silently move which + positions are trained on, while every shape assertion still passed. + """ + model = self._model(num_anchors=8) + # One long row, so the legacy bound (valid_counts.max() - 1) sits below the static + # width and there really are surplus columns to get wrong. + loss_mask = self._loss_mask((6, 6)) + torch.manual_seed(7) + legacy_anchors, legacy_keep = _legacy_sample_anchor_positions( + model, SEQ_LEN, loss_mask, self.DEVICE + ) + torch.manual_seed(7) + anchors, keep = model._sample_anchor_positions(SEQ_LEN, loss_mask, self.DEVICE) + assert legacy_keep.shape[1] < keep.shape[1], "fixture no longer exercises truncation" + kept = anchors[keep] + assert torch.equal(kept, legacy_anchors[legacy_keep]) + + def test_shape_is_independent_of_batch_contents(self): + """The whole point: one shape, therefore one compile.""" + model = self._model(num_anchors=8) + shapes = { + model._sample_anchor_positions(SEQ_LEN, self._loss_mask(lengths), self.DEVICE)[0].shape + for lengths in [(13, 13), (13, 5), (9, 9), (5, 3), (2, 1), (1, 1), (0, 0)] + } + assert len(shapes) == 1, f"n_blocks still varies with the batch: {shapes}" + + def test_shape_is_min_of_num_anchors_and_sequence(self): + anchors, _ = self._model(num_anchors=4)._sample_anchor_positions( + SEQ_LEN, self._loss_mask((13, 13)), self.DEVICE + ) + assert anchors.shape[1] == 4, "num_anchors should bind here" + anchors, _ = self._model(num_anchors=512)._sample_anchor_positions( + SEQ_LEN, self._loss_mask((13, 13)), self.DEVICE + ) + assert anchors.shape[1] == SEQ_LEN - BLOCK_SIZE + 1, "the sequence should bind here" + + def test_sampling_does_not_sync_on_the_batch(self): + """The old bound read valid_counts back to the host, stalling the pipeline. + + Parsed rather than grepped: the prose explaining the removal names the call it + removed, so a substring search over the source matches its own docstring. + """ + import ast + import inspect + import textwrap + + tree = ast.parse( + textwrap.dedent(inspect.getsource(hf_dflash.HFDFlashModel._sample_anchor_positions)) + ) + called = { + node.func.attr + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + } + assert not called & {"item", "tolist"}, f"host sync in anchor sampling: {called}" + + def test_trailing_padding_blocks_do_not_change_the_loss(self): + """The padding the static shape introduces is weightless in the loss and accuracy. + + Mathematically exact -- block_keep_mask zeroes those rows in both the numerator + and the normalizer -- but the reduction is over more elements, so compare at + float tolerance rather than bitwise. + """ + model = self._model(num_anchors=8) + vocab, bsz, n_blocks, pad = 32, 1, 2, 3 + torch.manual_seed(0) + input_ids = torch.randint(0, vocab, (bsz, SEQ_LEN)) + loss_mask = torch.ones(bsz, SEQ_LEN) + logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab) + anchors = torch.tensor([[0, BLOCK_SIZE]])[:, :n_blocks] + keep = torch.ones(bsz, n_blocks) + + base_loss, base_acc = model._compute_loss(logits, input_ids, anchors, keep, loss_mask) + padded_loss, padded_acc = model._compute_loss( + torch.cat([logits, torch.randn(bsz, pad * BLOCK_SIZE, vocab)], dim=1), + input_ids, + torch.cat([anchors, torch.zeros(bsz, pad, dtype=anchors.dtype)], dim=1), + torch.cat([keep, torch.zeros(bsz, pad)], dim=1), + loss_mask, + ) + torch.testing.assert_close(padded_loss, base_loss) + assert padded_acc == pytest.approx(base_acc) From 2ec89ec35925e72c89a627ccefc011f06f3334a3 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sat, 29 Aug 2026 04:38:54 -0700 Subject: [PATCH 5/6] fix(speculative): fail closed when the streaming /done check cannot be reached The RDMA hidden-states pool is a ring: the producer reuses slots, and the consumer detects a mid-read overwrite by asking /done whether the slot's generation still matches the one its bytes were written under. That check is the only guard against a lapped slot -- the token_ids comparison right after it comes from the server's per-request record, not from the bytes just read, so it passes on a mis-slotted read -- and it defaulted to `valid = True` whenever the /done call itself raised. The two mistakes are not symmetric. Discarding a good sample costs one resample. Keeping a lapped one trains the draft on another prompt's activations, silently: nothing downstream can tell, and the loss curve absorbs it. A /done that errors proves nothing about the slot either, so it now takes the same path an explicit lap does. The transport failure gets its own message rather than reusing the lap warning: a wave of unreachable /done is a sidecar problem, a wave of laps is ring pressure, and the two call for different fixes (a sick server vs. more pool_slots). Not theoretical -- a 3-node Gemma-4-E4B run logged 62 lapped-slot resamples over 80 steps at the launch script's default pool_slots=16, and 16 at pool_slots=128, so the guard fires in normal operation. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../plugins/hf_streaming_dataset.py | 16 +++- .../plugins/test_hf_streaming_dataset.py | 82 +++++++++++++++++++ 2 files changed, 94 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index 0f9713948fd..1a81dbdf519 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py @@ -536,14 +536,22 @@ def _fetch(self, sample: dict) -> EagleFetchPayload | None: time.sleep(0.0002) agent.release_xfer_handle(h) hidden_states = view.clone() # copy out before /done so the gen check brackets the read - # /done frees the slot + reports valid; valid=False -> ring lapped us mid-read, bytes - # stale -> resample. A failed /done can't prove staleness, so default valid=True. + # /done frees the slot and reports whether the ring lapped us mid-read (stale bytes). + # It fails CLOSED: a /done that errors cannot prove the slot is still ours either, and + # the two mistakes are not symmetric. Discarding a good sample costs one resample; + # keeping a lapped one trains the draft on another prompt's activations, and nothing + # downstream catches that -- the token_ids check below compares the server's + # per-request record, not the bytes just read, so it passes on a mis-slotted read. + # A rising resample rate is at least visible in the logs; silent corruption is not. try: valid = self._http_rdma.get( f"http://{host}:{port}/done", params={"req_id": rid} ).json()["valid"] - except Exception: - valid = True + except Exception as exc: + # Deliberately distinct from the lap message below: a wave of these is a sidecar + # problem, not ring pressure, and the two call for different fixes. + warn_rank_0(f"[streaming] /done failed for {sample['cid']} ({exc!r}); resampling") + return None if not valid: warn_rank_0(f"[streaming] slot lapped mid-read for {sample['cid']}; resampling") return None diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index efd77267b3c..43d641bbf39 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py @@ -430,6 +430,88 @@ def test_lapped_slot_is_treated_as_miss(monkeypatch): ds[0] +def _done_raising_handler(seq, n_layers, hidden, *, fail_first_n, calls): + """Sidecar handler whose /done raises a transport error for its first ``fail_first_n`` + calls, then behaves normally. ``calls`` accumulates the paths hit.""" + inner = _rdma_sidecar_handler(seq, n_layers, hidden) + state = {"done": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls.append(request.url.path) + if request.url.path == "/done": + state["done"] += 1 + if state["done"] <= fail_first_n: + raise httpx.ConnectError("simulated sidecar failure") + return inner(request) + + return handler + + +def test_failed_done_is_treated_as_a_miss(monkeypatch): + """A /done that errors must not be read as "the slot is still ours". + + The gen check /done performs is the only thing between a lapped ring slot and training + on another prompt's activations: the token_ids comparison that follows compares the + server's per-request record, not the bytes just read, so it passes on a mis-slotted + read. Trusting the read when /done is unreachable trades a visible resample for silent + corruption, so the failure must fail closed like an explicit lap does. + """ + seq, n_layers, hidden = 8, 3, 16 + calls: list[str] = [] + _mock_rdma( + monkeypatch, + _done_raising_handler(seq, n_layers, hidden, fail_first_n=10**6, calls=calls), + ) + + ds = EagleVllmStreamingDataset( + entries=[{"conversation_id": "c-0", "messages": [{"role": "user", "content": "x"}]}], + tokenizer=_tokenizer_returning(seq), + config=EagleVllmStreamingConfig( + server_urls="http://mock:8000", + model="mock-model", + max_seq_len=seq, + fail_after_consecutive_skips=100, + ), + ) + with pytest.raises(RuntimeError, match="no fetchable sample"): + ds[0] + assert "/done" in calls, "fixture never reached /done" + + +def test_failed_done_resamples_instead_of_returning_the_read(monkeypatch): + """The discard is a resample, not a hard failure: the next entry is fetched and returned. + + Pins that failing closed costs one extra round trip rather than breaking the epoch -- + the cheap half of the asymmetry that justifies discarding. + """ + seq, n_layers, hidden = 8, 3, 16 + calls: list[str] = [] + _mock_rdma( + monkeypatch, + _done_raising_handler(seq, n_layers, hidden, fail_first_n=1, calls=calls), + ) + + ds = EagleVllmStreamingDataset( + entries=[ + {"conversation_id": f"c-{i}", "messages": [{"role": "user", "content": "x"}]} + for i in range(2) + ], + tokenizer=_tokenizer_returning(seq), + config=EagleVllmStreamingConfig( + server_urls="http://mock:8000", + model="mock-model", + max_seq_len=seq, + fail_after_consecutive_skips=100, + ), + ) + + batch = ds[0] + assert batch["base_model_hidden_states"].shape == (seq, hidden) + # Two prompts posted: the first read was thrown away, the second is what came back. + assert calls.count("/v1/completions") == 2 + assert calls.count("/done") == 2 + + def test_oversize_server_response_raises(monkeypatch): """If the server captured more tokens than max_seq_len (its connector max_tokens > our recv buffer), reading would silently truncate the slice; fail loud instead so the From a1e5ae6ad10086e81dd8a5a6df815d0b7221b4ab Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sat, 29 Aug 2026 05:27:15 -0700 Subject: [PATCH 6/6] perf(speculative): compile the draft stack and stop DDP hunting unused parameters The draft's decoder loop was where the step's small-kernel tail lived: a profiled Gemma-4-E4B step ran ~1500 pointwise launches, 403 aten::copy_ and 205 device-to-device memcpys, individually microseconds and collectively about a third of the step. Nothing but a compiler fixes that shape of cost, and dflash_use_torch_compile -- despite promising "torch.compile on DFlash forward/loss methods" -- only ever reached the DSpark TVD chunk. It now also wraps DFlashModule's body. dynamic=False is affordable only because the block count is pinned; while n_blocks tracked the batch this would have recompiled per width at minutes each. The lazy rotary construction stays outside the compiled region (it mutates the module), and the compiled body is used in training only: generation runs the same module at a varying length, which under dynamic=False would mint one compile per length. The recipe's ddp_find_unused_parameters can now be false, which needed one fix first. Torch reports it "did not find any unused parameters", but that is a statement about the batches seen: on a batch whose loss terms are all skipped, _compute_dspark_loss returned `flat_final.sum() * 0`, and the confidence head hangs off its own projection rather than final_logits, so it received nothing and DDP would abort the run. That branch now walks dflash_module.parameters() the way forward()'s no-valid-anchor early return already did. Both paths are pinned by tests, since the failure mode is a crash hours into a run. Measured on 3 nodes, 80 steps, against the same recipe without the change: before after step time (median) 0.5497 s 0.4927 s -10.4% forward 0.1305 s 0.1005 s -23% backward 0.3478 s 0.3048 s -12% one-time compile 173.5 s 201.3 s The extra 28 s of compile pays for itself after ~490 steps, so an 80-step profile actually gets slower (228.8 s -> 243.3 s) while the 88k-step run this targets goes 13.5 h -> 12.1 h. The 20-step mean loss is unchanged: 3.209 / 2.661 / 2.621 / 2.625 against 3.209 / 2.661 / 2.621 / 2.588, the first three windows equal to three decimals. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/torch/speculative/config.py | 6 +- .../torch/speculative/plugins/hf_dflash.py | 5 + .../torch/speculative/plugins/hf_dspark.py | 9 +- .../speculative/plugins/modeling_dflash.py | 31 +++- .../dspark_gemma4_e4b.yaml | 8 +- .../speculative/plugins/test_hf_dspark.py | 135 ++++++++++++++++++ 6 files changed, 190 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 8df353b2daa..29625191ef3 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -167,7 +167,11 @@ class DFlashConfig(ModeloptBaseConfig): dflash_use_torch_compile: bool = ModeloptField( default=True, - description="Whether to use torch.compile on DFlash forward/loss methods.", + description=( + "Whether to torch.compile the compute-heavy parts of DFlash training: the draft " + "decoder stack and the DSpark TVD chunk. Costs a one-time compile on the first " + "training step; the draft's shapes are static, so it does not recur." + ), ) dflash_use_flex_attention: bool = ModeloptField( diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index c15ca2a1722..4d9c99f72c1 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -568,6 +568,11 @@ def modify(self, config): self.is_quantized = False self._num_anchors = self.dflash_num_anchors + # Opt-in Inductor fusion of the draft stack; see DFlashModule._body for why this is + # only affordable now that the block count is static. + self.dflash_module._dflash_compile_stack = bool( + getattr(self, "dflash_use_torch_compile", False) + ) def _build_draft_module(self, dflash_config): """Build the draft module. Subclasses override to use an augmented module.""" diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index e38afe8e9bb..9aea6037fe0 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -324,7 +324,14 @@ def _compute_dspark_loss( flat_teacher = flat_teacher.detach() if valid_count <= 1.0: - loss = flat_final.sum() * 0.0 + # Every draft parameter must receive a gradient, not just the ones behind + # final_logits: the confidence head hangs off compute_confidence_logits, which + # this branch skips, so `flat_final.sum() * 0` alone leaves it unused and DDP + # with find_unused_parameters=False aborts the run on the first such batch. + # Mirrors the same guard in forward()'s no-valid-anchor early return. + loss = flat_final.sum() * 0.0 + sum( + p.sum() for p in self.dflash_module.parameters() + ) * 0.0 metrics = {"ce_loss": 0.0, "l1_loss": 0.0, "confidence_loss": 0.0, "base_accuracy": 0.0} return loss, 0.0, metrics diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 3302844e32c..cc1cc7b37f5 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -573,9 +573,38 @@ def _init_weights(self, config): def forward(self, noise_embedding, target_hidden, position_ids, attention_mask=None): """Forward with feature fusion, KV injection, and position embeddings.""" + # Lazy rotary construction mutates the module, so it stays outside the compiled + # region below: Dynamo would either graph-break on it or bake in the first call's + # state. + self._maybe_init_rotary_emb(device=noise_embedding.device) + return self._body()(noise_embedding, target_hidden, position_ids, attention_mask) + + def _body(self): + """Return the draft stack, Inductor-compiled on first use when asked for. + + The layer loop is where the step's small-kernel tail lives: at the Gemma-4-E4B + shape a profiled step ran ~1500 pointwise launches, 403 ``aten::copy_`` and 205 + device-to-device memcpys, individually microseconds and collectively about a third + of the step. Fusing them is a compiler's job. + + ``dynamic=False`` is only affordable because ``_sample_anchor_positions`` pins the + block count; while n_blocks tracked the batch, every new width cost a fresh + multi-minute compile. + + Training only. Generation (AR validation, drafting) runs the same module at a + different and varying shape, which under ``dynamic=False`` would mint a compile per + length -- exactly the trade the pinned block count was introduced to avoid. + """ + if not self.training or not getattr(self, "_dflash_compile_stack", False): + return self._forward_body + if getattr(self, "_compiled_body", None) is None: + self._compiled_body = torch.compile(self._forward_body, dynamic=False) + return self._compiled_body + + def _forward_body(self, noise_embedding, target_hidden, position_ids, attention_mask): + """Feature fusion, rotary selection, the decoder stack, and the final norm.""" hidden_states = noise_embedding target_hidden = self.hidden_norm(self.fc(target_hidden)) - self._maybe_init_rotary_emb(device=hidden_states.device) per_kind = { kind: emb(hidden_states, position_ids) for kind, emb in getattr(self, "rotary_emb_by_kind", {}).items() diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml index dc849cb1ba7..899bf148c48 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml @@ -103,7 +103,13 @@ training: bf16: true tf32: true remove_unused_columns: false - ddp_find_unused_parameters: true + # false: every draft parameter gets a gradient on every batch, including the + # no-valid-anchor batches that skip the three loss terms -- forward() and the + # valid_count <= 1 branch both walk dflash_module.parameters() to keep DDP fed. True + # costs an extra autograd-graph traversal per iteration for nothing (torch itself warns + # it "did not find any unused parameters"). Precondition for this recipe: + # dflash_confidence_head_alpha > 0, otherwise the confidence head is never in the loss. + ddp_find_unused_parameters: false ddp_timeout: 1800 report_to: none diff --git a/tests/unit/torch/speculative/plugins/test_hf_dspark.py b/tests/unit/torch/speculative/plugins/test_hf_dspark.py index 9d421451cc7..d33071ad128 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dspark.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dspark.py @@ -389,3 +389,138 @@ def test_no_grad_path_matches_grad_path(self): plain = _tvd_per_token(final, teacher, chunk_size=5) grad_out = _tvd_per_token(final.requires_grad_(True), teacher, chunk_size=5) assert torch.equal(plain, grad_out.detach()) + + +def _draft_args(model, n_blocks=2, bsz=1): + """Build the draft module's inputs the way the training forward does. + + Worth going through the model's own helpers: the draft attends over the context as + keys but only its own blocks as queries, so target_hidden is seq_len long while + noise_embedding is n_blocks * block_size, and position_ids spans both. Hand-rolled + shapes silently disagree inside apply_rotary_pos_emb. + """ + m = model.dflash_module + dt = m.fc.weight.dtype # the draft carries the base model's dtype, not fp32 + torch.manual_seed(0) + input_ids = torch.randint(1, model.dflash_config.vocab_size, (bsz, SEQ_LEN)) + anchors = (torch.arange(n_blocks).unsqueeze(0) * BLOCK_SIZE).expand(bsz, -1).contiguous() + keep = torch.ones(bsz, n_blocks, dtype=torch.bool) + noise = model._build_noise_embedding(input_ids, anchors, keep, n_blocks).to(dt) + target = torch.randn(bsz, SEQ_LEN, m.fc.in_features, dtype=dt) + pos = model._build_position_ids(SEQ_LEN, anchors, input_ids.device) + mask = model._build_draft_attention_mask( + SEQ_LEN, anchors, keep, n_blocks, dt, input_ids.device, window=None + ) + return input_ids, anchors, keep, (noise, target, pos, mask) + + +def _dspark_model(use_compile=False, **cfg_kwargs): + model = get_tiny_llama(num_hidden_layers=4) + cfg = _get_dspark_config(**cfg_kwargs) + cfg["dflash_use_torch_compile"] = use_compile + mtsp.convert(model, [("dflash", cfg)]) + model.train() + return model + + +class TestDraftStackCompile: + """The draft stack is Inductor-compiled only when asked for, and only while training.""" + + def test_flag_off_keeps_the_eager_body(self): + m = _dspark_model(use_compile=False).dflash_module + assert m._body() == m._forward_body + + def test_eval_keeps_the_eager_body(self): + """Generation runs this module at a varying length; under dynamic=False that would + mint one compile per length -- the cost the pinned block count exists to avoid.""" + m = _dspark_model(use_compile=True).dflash_module + m.eval() + assert m._body() == m._forward_body + + def test_compiled_body_is_built_once(self): + m = _dspark_model(use_compile=True).dflash_module + first = m._body() + assert first != m._forward_body + assert m._body() is first, "recompiled on a later step" + + def test_compiled_matches_eager(self): + """Same weights, same inputs, both bodies -- same hidden states.""" + model = _dspark_model(use_compile=True) + m = model.dflash_module + _, _, _, args = _draft_args(model) + m._maybe_init_rotary_emb(device=args[0].device) # normally done by forward() + + with torch.no_grad(): + ref = m._forward_body(*args) + got = m._body()(*args) + assert m._body() != m._forward_body, "fixture did not actually compile" + # Inductor may fuse and reassociate, so this is agreement to the dtype's precision, + # not bitwise equality. + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) + + +class TestDdpGradientCoverage: + """Every draft parameter must get a gradient on every batch, degenerate ones included. + + This is the precondition for ddp_find_unused_parameters=false: DDP aborts the run the + first time a parameter that joined the reduction produces no gradient. The confidence + head is the one at risk -- it hangs off its own projection and never appears in + final_logits, so the branches that skip the loss terms have to reach it deliberately. + """ + + @staticmethod + def _model(): + return _dspark_model(use_confidence_head=True, confidence_alpha=1.0) + + @staticmethod + def _ungraded(model): + return [ + n + for n, p in model.dflash_module.named_parameters() + if p.requires_grad and p.grad is None + ] + + def test_batch_with_no_valid_anchor(self): + """Nothing is a training target, so forward() returns before building the draft.""" + model = self._model() + torch.manual_seed(0) + input_ids = torch.randint(1, model.dflash_config.vocab_size, (2, SEQ_LEN)) + out = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + labels=torch.full_like(input_ids, -100), + ) + out.loss.backward() + assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}" + + def test_loss_branch_with_zero_total_weight(self): + """Anchors exist but every label position is masked, so the three terms are skipped. + + Driven through the real module rather than synthetic logits: a gradient assertion + is only meaningful if the graph actually reaches the parameters. + """ + model = self._model() + m = model.dflash_module + n_blocks, bsz = 2, 1 + input_ids, anchors, _, args = _draft_args(model, n_blocks=n_blocks, bsz=bsz) + hidden = m(*args) + vocab = model.dflash_config.vocab_size + backbone_logits = torch.randn(bsz, n_blocks * BLOCK_SIZE, vocab, dtype=hidden.dtype) + # The teacher distribution is gathered before the degenerate branch is reached, so + # it has to be present even though this batch contributes nothing to the loss. + target_model_logits = torch.randn(bsz, SEQ_LEN, vocab, dtype=hidden.dtype) + final_logits, confidence_logits = model._apply_markov_head( + hidden, backbone_logits, input_ids, anchors, n_blocks + ) + loss, _, _ = model._compute_dspark_loss( + backbone_logits, + final_logits, + confidence_logits, + input_ids, + anchors, + torch.ones(bsz, n_blocks), # blocks are kept ... + torch.zeros(bsz, SEQ_LEN), # ... but no label position carries weight + target_model_logits=target_model_logits, + ) + loss.backward() + assert not self._ungraded(model), f"no gradient for {self._ungraded(model)}"