feat(datasets): add cp=1 shared-prefix prefix-tree attention for rollouts - #2564
Merged
HuiyingLi merged 5 commits intoJun 17, 2026
Merged
Conversation
…outs Fold shared-prefix RL rollouts (one prompt -> N completions) into a deduplicated flat layout plus a block-sparse prefix-tree AttnMaskSpec, and hand the per-step spec to the magi attention backend (cp=1) via set_active_attn_spec. Each completion attends FULL to the shared prompt and CAUSAL to itself, so the prompt is encoded once instead of N times. - prefix_tree.py: fold_shared_prefix_rollouts, prefix_tree_collate_fn, build_mock_rollout_dataset - train_ft.py: activate the batch's AttnMaskSpec each step (self-clearing) - example config + unit tests + a GPU magi parity script Refs: verl RFC #6401, Automodel NVIDIA-NeMo#2385 Signed-off-by: khazic <khazzz1c@gmail.com>
Backend-agnostic correctness check that realizes the folded AttnMaskSpec as a flex_attention mask (pure torch, runs on Ampere) and compares per-token output against an independent dense-mask SDPA reference. Validates the fold+spec logic where no Hopper GPU / magi build is available. Signed-off-by: khazic <khazzz1c@gmail.com>
Address /code-review + /simplify findings: - Labels are now pre-shifted (token t predicts t+1, completion's last token masked). This repo's MaskedCrossEntropy does no internal shift and the standard dataset path hands pre-shifted labels, so the prior predict-self labels were a degenerate identity objective. Drop the speculative completion_labels param to keep the shift unambiguous. - Move set_active_attn_spec activation into MagiState.prepare_llm_batch so the recipe stays magi-agnostic. - Move build_mock_rollout_dataset to mock_prefix_tree.py (mock_* convention). - Slim FoldedRollouts to the consumed fields; fuse the fold's per-completion passes; factor the shared parity-oracle mask into _prefix_tree_reference. - Example uses Qwen2.5-0.5B (a registered custom model) so backend.attn=magi wires the custom attn_func that honors the AttnMaskSpec; a plain HF model would route through the HF magi backend and ignore the prefix-tree spec. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| @@ -0,0 +1,164 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
Contributor
There was a problem hiding this comment.
LGTM — the fold logic, label shifting, position IDs, and spec construction all look correct. Good test coverage across basic, single-completion, empty-prompt, and validation-error cases, plus the independent dense-mask parity checks in the functional tests.
The import-linter independence contract forbids components.datasets importing components.distributed. The fold no longer imports AttnMaskSpec: it emits the tree structure (node_lengths, sample_paths) on the batch under 'prefix_tree', and MagiState.prepare_llm_batch builds and activates the AttnMaskSpec from it. Parity scripts (outside the contract) build the spec themselves. Signed-off-by: khazic <khazzz1c@gmail.com>
The prefix-tree AttnMaskSpec is handed to the attn_func out-of-band because HF's attention interface has a fixed signature and cannot receive a custom mask argument. Only the custom-model magi attn_func reads the active spec; the HF magi backend uses the plain causal varlen key and would silently drop the prefix-tree mask. Guard the HF dispatch path so it fails loudly with config guidance instead of training on the wrong mask. Signed-off-by: khazic <khazzz1c@gmail.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
HuiyingLi
approved these changes
Jun 17, 2026
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
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 does this PR do ?
Adds cp=1 multi-turn / shared-prefix prefix-tree attention for RL rollouts on the MagiAttention backend. A group of "one prompt + N sampled completions" is folded into a single deduplicated flat sequence
[prompt | c_0 | c_1 | ...]plus a block-sparse prefix-tree mask: every completion attends FULL to the shared prompt and CAUSAL to itself, so the prompt is encoded once instead of N times (verl RFC #6401 / Automodel #2385). Scope is cp=1 only (no context-parallel dispatch).Changelog
components/datasets/llm/prefix_tree.py:fold_shared_prefix_rollouts()(deduplicated flatinput_ids/ pre-shifted next-tokenlabels/ per-completionposition_ids, building the mask via the existingAttnMaskSpec.prefix_tree) andprefix_tree_collate_fn()(local_batch_size=1, attachesmagi_attn_spec).components/distributed/magi_attn_utils.py:MagiState.prepare_llm_batchactivates the per-stepAttnMaskSpecout-of-band viaset_active_attn_spec(self-clearing each step). Only the custom-model magi attn_func reads the active spec, so the recipe stays magi-agnostic.components/datasets/llm/mock_prefix_tree.py: deterministic mock rollout data for smoke runs.examples/llm_finetune/qwen/qwen25_magi_prefix_tree_rollouts.yaml: cp=1 example on Qwen2.5-0.5B (a registered custom model, sobackend.attn=magiwires the custom attn_func that honors the spec).tests/functional_tests/attention/(flex_attention for Ampere, magi FFA for Hopper) that compare per-token output against an independently built dense-mask SDPA oracle.Correctness
The fold + spec + mask realization is validated on an A100 via the backend-agnostic flex_attention parity (output vs an independent dense-mask SDPA reference, built from the rollout structure, not from the spec):
Prefix-tree unit tests pass locally. The Hopper-only magi FFA kernel parity (
prefix_tree_magi_parity.py) consumes the same spec and is left for an H-card run; it exercises the third-party SandAI kernel, not the logic this PR adds, which the A100 check already covers.Before your PR is "Ready for review"
Pre checks:
Additional Information