Skip to content

[PyTorch] fix input_quantizer usage for save_original_input; fix blockwise FP8 convert_and_update_tensor #1978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 9 commits into from
Aug 6, 2025
Merged

[PyTorch] fix input_quantizer usage for save_original_input; fix blockwise FP8 convert_and_update_tensor #1978

merged 9 commits into from
Aug 6, 2025

Conversation

hxbai
Copy link
Contributor

@hxbai hxbai commented Jul 22, 2025 •
edited
Loading

Description

#1963 added the unnecessary rowwise quantization to blockwise FP8 and MXFP8 in the save_original_input case. This PR removes this extra quantization.

#1952 did not correctly handle Blockwise FP8 tensor update. This PR also fixes this.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@hxbai hxbai mentioned this pull request Jul 22, 2025
Merged
13 tasks
@hxbai hxbai marked this pull request as draft July 22, 2025 15:41
@timmoon10
Copy link
Collaborator

/te-ci pytorch L1

@hxbai hxbai marked this pull request as ready for review July 23, 2025 07:57
@hxbai hxbai force-pushed the save_original_fix branch from da084bb to 368117b Compare July 31, 2025 06:35
@hxbai hxbai changed the title [PyTorch] fix input_quantizer usage in Linear backward for save_original_input [PyTorch] fix input_quantizer usage for save_original_input; fix blockwise FP8 convert_and_update_tensor Aug 1, 2025
@timmoon10
Copy link
Collaborator

/te-ci pytorch L1

timmoon10
timmoon10 approved these changes Aug 1, 2025
View reviewed changes
Copy link
Collaborator

@timmoon10 timmoon10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, pending CI

ksivaman
ksivaman approved these changes Aug 4, 2025
View reviewed changes
Copy link
Member

@ksivaman ksivaman left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ksivaman
Copy link
Member

ksivaman commented Aug 4, 2025

/te-ci pytorch L0 L1

@timmoon10 timmoon10 merged commit de69ca0 into NVIDIA:main Aug 6, 2025
12 of 13 checks passed
nv-akorzh pushed a commit to nv-akorzh/TransformerEngine that referenced this pull request Aug 6, 2025
@hxbai @timmoon10
@pre-commit-ci @ksivaman @nv-akorzh
[PyTorch] fix input_quantizer usage for save_original_input; fix bloc…
79b230d 
…kwise FP8 convert_and_update_tensor (NVIDIA#1978)

* fix input_quantizer in save_original_input bwd

Signed-off-by: Hongxiao Bai <[email protected]>

* fix get shape of blockwise tensor with only compact colwise data

Signed-off-by: Hongxiao Bai <[email protected]>

* fix blockwise FP8 convert_and_update_tensor

Signed-off-by: Hongxiao Bai <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Hongxiao Bai <[email protected]>
Co-authored-by: Tim Moon <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
Signed-off-by: Anton Korzh <[email protected]>
nv-akorzh pushed a commit to nv-akorzh/TransformerEngine that referenced this pull request Aug 7, 2025
@hxbai @timmoon10
@pre-commit-ci @ksivaman @nv-akorzh
[PyTorch] fix input_quantizer usage for save_original_input; fix bloc…
acd5adf 
…kwise FP8 convert_and_update_tensor (NVIDIA#1978)

* fix input_quantizer in save_original_input bwd

Signed-off-by: Hongxiao Bai <[email protected]>

* fix get shape of blockwise tensor with only compact colwise data

Signed-off-by: Hongxiao Bai <[email protected]>

* fix blockwise FP8 convert_and_update_tensor

Signed-off-by: Hongxiao Bai <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Hongxiao Bai <[email protected]>
Co-authored-by: Tim Moon <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
Signed-off-by: Anton Korzh <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@timmoon10 timmoon10 timmoon10 approved these changes

@ksivaman ksivaman ksivaman approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@hxbai @timmoon10 @ksivaman