chore: CI proxy for #3004#3016
Open
HuiyingLi wants to merge 6 commits into
Open
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>
Contributor
Author
|
/ok to test ea320ca |
Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
Author
|
/ok to test c6aa976 |
…rget batches The DFlash packing change threads target_batch.position_ids / seq_lens / doc_remaining into the trainer (None off the packing path). The fake target wrappers in test_train_dflash.py and test_train_dflash_noanchor_skip.py returned a SimpleNamespace lacking those attributes, so the unpacked path raised AttributeError. Add the three attributes (None) and accept the packing kwargs, matching DFlashTargetBatch's real defaults. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
Author
|
/ok to test 2d0dbcd |
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
2 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.
CI-only proxy PR. Do not merge / do not review. Points at the exact head commit of #3004 (ea320ca) so internal CI runs under the internal-contributor queue; results post back to #3004 via the shared SHA. Source of truth: #3004. Close once CI completes.