[megatron] fix: support FP8 padding for BSHD actor forward - #6887
[megatron] fix: support FP8 padding for BSHD actor forward#6887gem-mint wants to merge 1 commit into
Conversation
Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: gem-mint <146161229+gem-mint@users.noreply.github.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables FP8 support for the BSHD format. It removes the assertion blocking FP8 in BSHD, propagates the use_fp8_padding parameter to preprocess_bshd, and extracts the sequence length alignment logic into a reusable helper function _align_bshd_max_seqlen_for_fp8. Additionally, a unit test has been added to verify the FP8 padding alignment and roundtrip functionality. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
wuxibin89
changed the title
| position_ids: torch.Tensor, | ||
| sequence_parallel: bool = False, | ||
| pre_process: bool = True, | ||
| use_fp8_padding: bool = False, |
There was a problem hiding this comment.
preprocess_bshd is not used anymore, we now use preprocess_bshd_engine
|
Never mind, I decide not to support fp8 in BSHD format. |
2 participants
What
This is the main-branch companion to #6884, requested by maintainers so the BSHD FP8 actor update fix is available on main as well as release/v0.7.1.
It removes the hard BSHD FP8 assertion in the Megatron actor forward path, propagates
use_fp8_paddinginto regular BSHD preprocessing for the input and logits processor arguments, and aligns BSHD padded sequence length so TransformerEngine FP8 block quantization sees a local token count divisible by 128. The main branch already had equivalent alignment logic inpreprocess_bshd_engine; this patch factors that into a helper and reuses it for regular BSHD too.Why
With
data_format=bshdandfp8=hybrid/e4m3, actor update currently fails before forward with:Removing the assertion alone is not enough: the padded BSHD shape must also satisfy TE FP8 block quantization alignment.
Duplicate-work check
Per
AGENTS.md, I checked for overlapping work before opening this PR:gh issue view 5508 --repo verl-project/verl --commentsgh pr list --repo verl-project/verl --state open --search "5508 in:body" --limit 20gh pr list --repo verl-project/verl --state open --search "bshd fp8" --limit 20gh pr list --repo verl-project/verl --state open --search "actor fp8 padding" --limit 20gh pr list --repo verl-project/verl --state open --search '"fp8 is not supported for bshd format yet"' --limit 20The only same-scope result is #6884, which targets
release/v0.7.1. #6703 appears in some keyword searches but fixes a THD/CP OOB issue and is not a duplicate of this BSHD actor-padding fix.Tests
git diff --checkpassedpython -m py_compile tests/models/test_mcore_bshd_fp8_padding.py verl/models/mcore/util.py verl/models/mcore/model_forward.pypassed using the local bundled Pythonpython -m pytest tests/models/test_mcore_bshd_fp8_padding.py -qcould not run locally because this Windows helper environment does not havepytestinstalledtorchinstalledAI assistance was used to prepare this patch; human review done before marking ready for review.