Skip to content

[JAX] Scaling Enum Abstracting#1655

Merged
phu0ngng merged 10 commits intoNVIDIA:mainfrom
phu0ngng:scaling_mode_abstract
Apr 9, 2025
Merged

[JAX] Scaling Enum Abstracting#1655
phu0ngng merged 10 commits intoNVIDIA:mainfrom
phu0ngng:scaling_mode_abstract

Conversation

@phu0ngng
Copy link
Collaborator

@phu0ngng phu0ngng commented Apr 8, 2025 •
edited
Loading

Description

This PR discoupled the Scaling Enum value in JAX extension from the Scaling Enum value from TE/Common.

In addition, I also updated part of the NormPrimitive so that it is consistent with others.

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

ksivaman
ksivaman reviewed Apr 8, 2025
View reviewed changes
transformer_engine/jax/csrc/extensions/misc.h Outdated
Copy link
Member

@ksivaman ksivaman Apr 8, 2025

Choose a reason for hiding this comment

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

Is this a typo propagated everywhere or is it meant to be JAXX?

Copy link
Collaborator Author

@phu0ngng phu0ngng Apr 8, 2025

Choose a reason for hiding this comment

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

JAXX = JAX Extension 😎

@phu0ngng
Copy link
Collaborator Author

phu0ngng commented Apr 8, 2025

/te-ci jax L0

jberchtold-nvidia
jberchtold-nvidia reviewed Apr 8, 2025
View reviewed changes
@phu0ngng phu0ngng force-pushed the scaling_mode_abstract branch from af8e8b1 to c3de984 Compare April 8, 2025 17:59
phu0ngng and others added 4 commits April 8, 2025 14:25
@phu0ngng
scaling enum abstract
b9fb508 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@pre-commit-ci @phu0ngng
[pre-commit.ci] auto fixes from pre-commit.com hooks
2e9a07e 
for more information, see https://pre-commit.ci
@phu0ngng
revert unneccessary changes
f6bb911 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
rm NVTE_ from ScalingMode names
7052751 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng phu0ngng force-pushed the scaling_mode_abstract branch from c3de984 to 7052751 Compare April 8, 2025 21:25
@phu0ngng
Copy link
Collaborator Author

phu0ngng commented Apr 8, 2025

/te-ci jax L0

phu0ngng added 2 commits April 8, 2025 18:25
@phu0ngng
rework scaling mode enum in grouped gemm
e6251a8 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
lint fix
2ad7744 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Collaborator Author

phu0ngng commented Apr 9, 2025

/te-ci jax L0

@phu0ngng
rework gemm
1cd0367 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Collaborator Author

phu0ngng commented Apr 9, 2025

/te-ci jax L0

pre-commit-ci bot and others added 3 commits April 9, 2025 13:21
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
5d0a49b 
for more information, see https://pre-commit.ci
@phu0ngng
pass by copy
0a97d52 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
fix norm sharding
d134a3b 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
Copy link
Collaborator Author

phu0ngng commented Apr 9, 2025

/te-ci jax L0

@phu0ngng phu0ngng merged commit 962d9c5 into NVIDIA:main Apr 9, 2025
20 checks passed
@phu0ngng phu0ngng deleted the scaling_mode_abstract branch April 9, 2025 20:25
wdykas pushed a commit to wdykas/TransformerEngine that referenced this pull request Apr 14, 2025
@phu0ngng @wdykas
[JAX] Scaling Enum Abstracting (NVIDIA#1655)
177863f 
* scaling enum abstract

* rm NVTE_ from ScalingMode names

* rework scaling mode enum in grouped gemm

* fix norm sharding

---------

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

Reviewers

@ksivaman ksivaman ksivaman left review comments

@jberchtold-nvidia jberchtold-nvidia jberchtold-nvidia 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

@phu0ngng @ksivaman @jberchtold-nvidia