Skip attention mask materialization in RL training - #5885
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
848e52f to
7a4206d
Compare
| pad_mask_loss, | ||
| create_attention_mask=True): | ||
| """Build masks and position id for left to right model.""" | ||
| assert create_attention_mask or not reset_attention_mask, \ |
There was a problem hiding this comment.
I think we need a docstring now explaining what reset/create attention mask do and why we need both.
There was a problem hiding this comment.
Done, what do you think?
yobibyte
left a comment
yobibyte
left a comment
There was a problem hiding this comment.
I would run some 20 steps of an environment before/after for this to compare, but I do not have an argument why. Just intuition. Also left some comments.
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
I did that here, although the "after" in this run was multiple changes packaged together: https://wandb.ai/adlr/megatron-rl?nw=9vg5y1477c9 Do you want an actual before/after without any confounding variables, or were you asking for just a sanity check? |
sanity check, looks ok |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31657412435 |
5 participants
What does this PR do?
Adds a
create_attention_maskargument toget_ltor_masks_and_position_idsthat can be toggled toFalse, skipping the creation of the attention mask when it's unnecessary.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.