Skip to content

[JAX] TensorUsage + FP8 GEMM with all layouts handling on BW#1844

Merged
phu0ngng merged 6 commits intoNVIDIA:mainfrom
phu0ngng:tensor_usage
Jun 18, 2025
Merged

[JAX] TensorUsage + FP8 GEMM with all layouts handling on BW#1844
phu0ngng merged 6 commits intoNVIDIA:mainfrom
phu0ngng:tensor_usage

Conversation

@phu0ngng
Copy link
Copy Markdown
Collaborator

@phu0ngng phu0ngng commented Jun 3, 2025 •
edited
Loading

Description

TensorUsage + FP8 GEMM with all layouts handling on BW.

Verified that no NT enforcements by JAX, i.e. no additional transpose.

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

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

denera
denera requested changes Jun 3, 2025
View reviewed changes
Comment thread transformer_engine/jax/quantize/device_utils.py Outdated
jberchtold-nvidia
jberchtold-nvidia reviewed Jun 3, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@jberchtold-nvidia jberchtold-nvidia left a comment

Choose a reason for hiding this comment

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

Overall looks good, some small comments. I like the TensorUsage concept

Comment thread transformer_engine/jax/layernorm_dense.py Outdated
Copy link
Copy Markdown
Collaborator

@jberchtold-nvidia jberchtold-nvidia Jun 3, 2025

Choose a reason for hiding this comment

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

Should we keep the check for whether we're training or not? So it'd be something like this

casted_ln_out.get_tensor(TensorUsage.LHS_TRANS) if quantizer_set.x.is_2x2x() else None

because on Hopper this layout wouldn't exist when doing 1x for inference, right?

Copy link
Copy Markdown
Collaborator Author

@phu0ngng phu0ngng Jun 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

No, we don't need to do that here.
Whether it's training or not should be handled in the get_tensor() method, and currently, we don't support it yet.

Copy link
Copy Markdown
Collaborator

@jberchtold-nvidia jberchtold-nvidia Jun 4, 2025

Choose a reason for hiding this comment

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

Ah I see, that will make it easier if it's automatic. So the idea is based on whether it's x or kernel, we know which usage (LHS/RHS) will be used in the forward and the backward, so we can automatically know to remove it when doing inference (not currently, but we have all the required info to support it in future)?

Comment thread transformer_engine/jax/quantize/scaling_modes.py Outdated
@phu0ngng phu0ngng added the 2.5.0 label Jun 6, 2025
@phu0ngng phu0ngng requested a review from huanghua1994 June 13, 2025 17:43
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 13, 2025

Adding @huanghua1994 to review the changes in grouped_gemm.

@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 13, 2025

/te-ci JAX L0

jberchtold-nvidia
jberchtold-nvidia previously approved these changes Jun 13, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@jberchtold-nvidia jberchtold-nvidia left a comment

Choose a reason for hiding this comment

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

LGTM

@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 16, 2025

/te-ci JAX L0

@phu0ngng phu0ngng force-pushed the tensor_usage branch 3 times, most recently from b411402 to 0644636 Compare June 16, 2025 18:21
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 16, 2025

/te-ci JAX L0

@phu0ngng
TensorUsage + FP8 GEMM with all layouts handling on BW
1f5d486 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 16, 2025

/te-ci JAX L0

1 similar comment
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 17, 2025

/te-ci JAX L0

@phu0ngng
fix for quantizer 2x2x
b56031f 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 17, 2025

/te-ci JAX L0

@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 17, 2025

/te-ci JAX L0

@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 17, 2025

/te-ci JAX L0

@phu0ngng
fix github build
450eee1 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Copy Markdown
Collaborator Author

phu0ngng commented Jun 17, 2025

/te-ci JAX L0

@phu0ngng phu0ngng merged commit 3a298e6 into NVIDIA:main Jun 18, 2025
23 checks passed
@phu0ngng phu0ngng deleted the tensor_usage branch June 18, 2025 11:47
KshitijLakhani pushed a commit that referenced this pull request Jun 27, 2025
@phu0ngng @KshitijLakhani
[JAX] TensorUsage + FP8 GEMM with all layouts handling on BW (#1844)
82bff47 
* TensorUsage + FP8 GEMM with all layouts handling on BW

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@denera denera denera approved these changes

@jberchtold-nvidia jberchtold-nvidia jberchtold-nvidia approved these changes

@huanghua1994 huanghua1994 Awaiting requested review from huanghua1994

Assignees

No one assigned

Labels

2.5.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@phu0ngng @denera @jberchtold-nvidia