[training, test] feat: add unit tests for padding_utils helpers - #3786
Closed
lonexreb wants to merge 1 commit into
Closed
[training, test] feat: add unit tests for padding_utils helpers#3786lonexreb wants to merge 1 commit into
lonexreb wants to merge 1 commit into
Conversation
padding_utils.py exports three small, pure-tensor helpers used to align variable-length batches to a target sequence length with an upper cap: pad_or_truncate_2d_to_len pad_or_truncate_pos_to_len pad_or_truncate_attn_to_len They operate on CPU tensors only and have no GPU or distributed dependencies, but were uncovered by unit tests. This commit adds 20 tests across three classes covering: - 2D pad: None passthrough, padding to target with default and custom pad_value, truncation to max_cap, identity pass-through between target and cap, and float-dtype pad value. - Position-id pad: None passthrough, monotonic extension shape and values, batch broadcasting, dtype preservation, truncation to max_cap, and identity pass-through. - Attention mask pad: None passthrough, 2D bool/int mask padding with correct pad value, 2D truncation and pass-through, 4D square-mask padding in both seq dims, 4D truncation and pass-through, and a 3D-mask ValueError. Follows the pattern of recently-merged unit-test PRs NVIDIA-NeMo#3650 and NVIDIA-NeMo#3695. No production changes. Signed-off-by: lonexreb <reach2shubhankar@gmail.com>
yaoyu-33
added
area:training
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
yaoyu-33
removed
the
needs-review
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.
Summary
`src/megatron/bridge/training/utils/padding_utils.py` exports three pure-tensor helpers that align variable-length batches to a target sequence length with an upper cap:
They are CPU-only, no torch.distributed dependencies, and previously had no direct unit test coverage.
What this PR adds
20 tests in `tests/unit_tests/training/utils/test_padding_utils.py`:
`TestPadOrTruncate2dToLen` (6)
`TestPadOrTruncatePosToLen` (6)
`TestPadOrTruncateAttnToLen` (8)
Why
Follows the pattern of recently-merged unit-test PRs #3650 (pg_utils) and #3695 (FLOPs calculator).
Test plan
🤖 Generated with Claude Code