chore: CI proxy for #2564 - #2565
Closed
HuiyingLi wants to merge 5 commits into
Closed
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 #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
Author
|
/ok to test 548787c |
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>
Contributor
Author
|
/ok to test 9ce5f1b |
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>
Contributor
Author
|
/ok to test 6f329f7 |
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 #2564 (548787c) so internal CI runs under the internal-contributor queue; results post back to #2564 via the shared SHA. Source of truth: #2564. Closed once CI completes.