feat(speculative): support sequence packing in the Domino and JetSpec drafts#3023
Merged
HuiyingLi merged 11 commits intoJul 12, 2026
Merged
Conversation
DFlash builds its dataloader with build_eagle3_dataloader but never passed
packed_sequence_size, so packing was unreachable. Wire it end to end. DFlash's
Q layout is draft blocks only (context is KV), so packing composes cleanly:
- dflash_mask: create_dflash_{sdpa,block}_mask take optional ctx_doc_id /
anchor_doc_id; a block's context prefix is restricted to its anchor's document
(both the dense SDPA mask and the flex BlockMask mask_mod).
- core: forward accepts position_ids / seq_lens / doc_remaining. Anchor sampling
keeps whole blocks inside one document (doc_remaining >= block_size-1); the
context position ids reset per document; the draft block positions are gathered
from the per-document context positions so RoPE stays document-local.
- target: generate_batch runs the frozen target with the block-causal mask (or FA
varlen at batch size 1) so the captured context hidden states do not leak across
documents, and carries the packing metadata to the trainer.
- train_dflash: read packed_sequence_size into both dataloaders, splat the packing
metadata into generate_batch and the trainer, fail fast on incompatible configs
(cp_size>1, FlashAttention target with batch>1), and reject packing on the
JetSpec/Domino subclasses (they override the trainer forward without threading it).
Tests: doc-id construction, SDPA mask document restriction, anchor sampling stays
in-document, per-document draft positions, single-document packed==unpacked
equivalence, a two-document packed training step, target-side document isolation,
and the recipe helpers/gates.
Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
… limitation Address /code-review findings: the flex_attention backend is the DFlash default but every packing test used sdpa, leaving the flex doc-gating branch uncovered. Factor the flex mask_mod into a module-level build_dflash_mask_mod so it can be materialised on CPU via create_mask (FlexAttention's kernel is CUDA-only), and add a test asserting the flex per-token mask equals the SDPA mask and enforces the per-document restriction. Also document that the whole-block-in-document anchor bound (required because _build_block_targets does not encode boundaries) drops packed documents shorter than block_size. Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
… drafts Signed-off-by: khazic <khazzz1c@gmail.com>
…ared block-target builder Signed-off-by: khazic <khazzz1c@gmail.com>
… prologue Signed-off-by: khazic <khazzz1c@gmail.com>
12 tasks
…fake batches The DFlash packing base threads target_batch.position_ids / seq_lens / doc_remaining into the trainer (None off the packing path). The fake target wrappers and stub batches in test_train_dflash.py, test_train_dflash_noanchor_skip.py, test_train_domino.py, and test_train_jetspec.py lacked those attributes, so the unpacked path raised AttributeError. Add the three attributes (None) and accept the packing kwargs, matching the real target batch defaults. Signed-off-by: khazic <khazzz1c@gmail.com>
HuiyingLi
previously approved these changes
Jul 12, 2026
Contributor
|
/claude review |
Comment on lines
+153
to
+160
| """Return the FlexAttention ``mask_mod`` closure encoding the DFlash mask. | ||
|
|
||
| Factored out of :func:`create_dflash_block_mask` so the doc-aware gating can be | ||
| materialised with ``torch.nn.attention.flex_attention.create_mask`` and tested on | ||
| CPU (FlexAttention's kernel is CUDA-only). The closure uses only tensor ops (no | ||
| Python control flow) so ``torch.compile`` can trace it; see the module docstring | ||
| for the mask semantics and the ``ctx_doc_id`` / ``anchor_doc_id`` packing args. | ||
| """ |
Contributor
There was a problem hiding this comment.
build_dflash_mask_mod is a new module-level (public) function that accepts four tensor inputs (anchor_positions, block_keep_mask, ctx_doc_id, anchor_doc_id), but its docstring has no Args section documenting their layouts. The reference to the module docstring covers the mask semantics but not the shape/axis order of these specific tensors, so a caller can't tell the expected layout from the docstring alone. The sibling create_dflash_sdpa_mask / create_dflash_block_mask already document these; please add an Args section here, e.g.:
Args:
anchor_positions: ``[B, N]`` anchor positions per sample (long).
block_keep_mask: ``[B, N]`` per-sample valid-anchor mask (bool).
ctx_len: context length ``S``.
block_size: block size.
num_blocks: number of blocks ``N``.
causal: When True, in-block attention is causal (JetSpec).
ctx_doc_id: ``[B, S]`` long per-context-token document id (packing),
or ``None`` to disable the per-document constraint.
anchor_doc_id: ``[B, N]`` long document id of each anchor; required when
``ctx_doc_id`` is given.
…variants-packing # Conflicts: # nemo_automodel/components/speculative/dflash/core.py # nemo_automodel/components/speculative/dflash/target.py # nemo_automodel/recipes/llm/train_dflash.py
HuiyingLi
approved these changes
Jul 12, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Extends sequence packing to the two DFlash variants, completing the packing coverage item in #2958 for the DFlash family (#3002/#3003/#3004/#3005 covered the other drafts). #3004 gated packing off for these subclasses because their trainer forwards did not thread the packing metadata; this PR threads it and lifts the gate.
DFlashTrainerModule._prepare_block_inputs(with acausalflag for JetSpec), so the packing handling lives in one place. It fails loud when the packing metadata arrives partially (seq_lenswithoutposition_ids/doc_remainingwould otherwise silently drop the in-document anchor constraint)._build_block_targetsgainslabel_startanddoc_remaining: anchor sampling only guarantees offsets up toblock_size - 1stay in the anchor's document, and Domino's shifted label window reaches one past that, so the builder truncates labels at the document boundary.forwardacceptsposition_ids/seq_lens/doc_remaining, delegates the prologue to the base helper, and replaces its hand-rolled label/mask gathering with the shared_build_block_targets(label_start=1 if shift_label else 0, doc_remaining=...). The GRU correction head is position-relative within the block and needs no change.forwardaccepts the packing metadata and calls the shared prologue withcausal=True(the doc-aware mask composes with the block-causal JetSpec mask, added in feat(speculative): support sequence packing in the DFlash draft #3004). The teacher logits come from the packed block-causal target forward, so the gathered teacher distribution (positionsanchor .. anchor+block_size-2) is document-local by the anchor-sampling guarantee.train_domino.py/train_jetspec.pypass the metadata through_run_trainer_step; theNotImplementedErrorpacking gate for subclasses intrain_dflash.pyis removed (the remaining_validate_packing_gateschecks are shared and still apply).Correctness verification
CPU unit tests (96 pass across the dflash/domino/jetspec/packing files), the new ones covering:
shift_labelsettings); a two-document packed step with gradients.shift_label, the one label that crosses into doc1 is dropped (valid_tokensgoes fromblock_sizetoblock_size - 1).GPU recipe smokes (1x A100-80GB, bf16, real Qwen3-4B target, packing enabled, flex_attention backend), both full recipes:
Domino: loss falls 5.0 -> 3.12, block accuracy rises to ~0.18, and the Domino-specific diagnostics (final/base loss split, accept_len ~1.45, lambda_base curriculum decay) are all live under packing.
Domino training log
JetSpec: the forward-KL loss falls 5.3 -> 4.41 and the greedy draft-vs-target accuracy rises to ~0.24 under packing.
JetSpec training log
Scope
Domino and JetSpec online text paths. Packing with context parallelism remains excluded (same shared gate as #3004).