Blockwise scaling linear quantization recipe by kwyss-nvidia · Pull Request #1559 · NVIDIA/TransformerEngine

kwyss-nvidia · 2025-03-11T22:20:09Z

Description

Add a recipe configuration to quantizers and GEMM for float8 blockwise scaling.

Quantizer MR: #1513
GEMM MR: #1545

Checked against golden outputs from internal tool.

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

Add recipe and recipe state classes to create quantizers.
Modify linear, layer norm linear, and layer_norm_mlp to use the quantized data paths.
Test for linear and layer_norm.
Revert recent change that breaks test_numerics for all recipes.

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

kwyss-nvidia · 2025-03-12T00:39:59Z

@ptrendx here is a parallel review from my fork with only the recipe changes.
kwyss-nvidia#2

transformer_engine/pytorch/fp8.py

tests/pytorch/distributed/test_numerics.py

transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py

transformer_engine/pytorch/module/linear.py

transformer_engine/common/recipe/recipe_common.cuh

tests/cpp/test_common.cu

tests/pytorch/test_float8_blockwise_gemm_exact.py

transformer_engine/pytorch/tensor/float8_blockwise_tensor.py

Apply MR comment change. Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: kwyss-nvidia <kwyss@nvidia.com>

tests/pytorch/test_float8tensor.py

transformer_engine/common/include/transformer_engine/cast.h

transformer_engine/pytorch/csrc/common.h

transformer_engine/pytorch/csrc/extensions/cast.cpp

transformer_engine/pytorch/module/layernorm_linear.py

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

transformer_engine/common/include/transformer_engine/transformer_engine.h

transformer_engine/common/recipe/__init__.py

transformer_engine/pytorch/distributed.py

transformer_engine/pytorch/module/linear.py

tests/pytorch/test_float8_current_scaling_exact.py

tests/pytorch/test_float8_blockwise_scaling_exact.py

zhongbozhu · 2025-04-09T00:43:51Z

transformer_engine/pytorch/module/linear.py

Float8BlockQuantizer getting added here to set columnwise usage True

But this logic is also not doing anything unless we do fp8 gather, I am okay with keeping it there, just taking a note case in case we forget.

PTAL. Added an assert.

timmoon10 · 2025-04-09T00:53:07Z

/te-ci L1

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

transformer_engine/common/recipe/__init__.py

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

timmoon10

LGTM

transformer_engine/pytorch/csrc/common.h

tests/pytorch/test_float8_current_scaling_exact.py

tests/pytorch/test_float8blockwisetensor.py

Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 · 2025-04-10T00:49:23Z

/te-ci L1

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

Signed-off-by: Xin Yao <yaox12@outlook.com>

yaox12 · 2025-04-10T04:45:42Z

/te-ci pytorch L1

Signed-off-by: Xin Yao <yaox12@outlook.com>

yaox12 · 2025-04-10T05:35:25Z

/te-ci pytorch

yaox12 · 2025-04-10T07:39:48Z

JAX L0/L1 tests passed with build 26726340.
PyTorch L1 tests passed with build 26734338 except for irrelevant failures.
PyTorch L0 tests passed with build 26736222 except for irrelevant failures.

zhongbozhu · 2025-04-10T17:29:08Z

/te-ci pytorch L1

@timmoon10

* Add GEMM logic for blockwise quantized tensors. GEMM test cases included in pytorch integration. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update NVTE_BLOCK_SCALING for GEMM. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Gate feature on CUDA 12.9 Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Gemm typo. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Remove unecessary type converter change. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Reflect epilogue availability and test supported epilogues. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * GEMM simplifications from recipe branch. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Format py code. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update GEMM DGelu tests to match support depending on output dtype. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Force pow2Scales in GEMM Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Add GEMM test to pytorch test suite. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Add copyright to GEMM test. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update import for GEMM test. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Add license. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update test gemm supported predicate. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Use sgemm like interfaces and naming. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Rewrite GEMM comment. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * MR Feedback. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Recipe setup for Linear modules. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Use 12.9 feature test. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Run against tensor dumps from internal library. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update FIXME to TODO with linked issue. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update full recompute feature to save recipe. The recompute context uses the same recipe and fp8 settings as the original fwd pass. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * MR Feedback. Avoid reusing quantizer objects. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update logic in module. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Format py. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update for PP bug. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update test numerics. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update force_power_of_2 scales in the recipe. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update usage method to satisfy upstream changes. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * fix subchannel recipe in distributed test with bf16 gather Signed-off-by: zhongboz <zhongboz@nvidia.com> * Edit and cleanup BF16 gather code. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update test import. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * support columnwise only mode to 1D quantize kernel Signed-off-by: zhongboz <zhongboz@nvidia.com> * Format and move enum Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Skip alloc. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * try async bf16 gather Signed-off-by: zhongboz <zhongboz@nvidia.com> * Format python code. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Document and type code. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update pytorch lint errors. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Dont set high precision dtype. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Add test for sanity and CG; fix CG for sequential? Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Keep make_quantizers API stable Update num_quantizers instead to pass cuda_graph tests. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Fix import name. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Rename recipe method. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Skip grouped linear sanity test. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Set usage before BF16 gather. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * refactor for nvte_quantize_v2 Signed-off-by: zhongboz <zhongboz@nvidia.com> * Format code. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Cleanup nvte_quantize_v2 Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Test fp32 scales. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Disable CUDA graph. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Simplify layernorm linear Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Cleanup layernorm linear. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * LayerNorm linear bwd gather logic. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Communication updates. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update transformer_engine/pytorch/ops/op.py Apply MR comment change. Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: kwyss-nvidia <kwyss@nvidia.com> * Lint fix. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * MR feedback. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Enable cuda graph tests. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Reduce chance of spurious failure and reword. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Review suggestions from @timmoon10 Signed-off-by: Tim Moon <tmoon@nvidia.com> * Update CPP tests. Signed-off-by: Keith Wyss <kwyss@nvidia.com> * Update common.h Signed-off-by: Xin Yao <yaox12@outlook.com> * Update test_float8blockwisetensor.py Signed-off-by: Xin Yao <yaox12@outlook.com> --------- Signed-off-by: Keith Wyss <kwyss@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: kwyss-nvidia <kwyss@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Xin Yao <yaox12@outlook.com> Co-authored-by: zhongboz <zhongboz@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Xin Yao <yaox12@outlook.com>

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch 2 times, most recently from c42f005 to 16cc925 Compare March 11, 2025 22:23

kwyss-nvidia mentioned this pull request Mar 12, 2025

Mirror of linear changes kwyss-nvidia/TransformerEngine#2

Open

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch from 16cc925 to db411ce Compare March 12, 2025 00:51

kwyss-nvidia changed the title ~~Draft: Blockwise scaling linear quantization recipe~~ Blockwise scaling linear quantization recipe Mar 12, 2025

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch 2 times, most recently from 8bb7d63 to 365a4d9 Compare March 13, 2025 00:06

yaox12 reviewed Mar 13, 2025

View reviewed changes

transformer_engine/pytorch/fp8.py Outdated Show resolved Hide resolved

transformer_engine/pytorch/fp8.py Outdated Show resolved Hide resolved

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch 5 times, most recently from 51fbe41 to 78c194d Compare March 17, 2025 17:33

yaox12 mentioned this pull request Mar 19, 2025

How can we integrate the DeepGEMM Fp8 GEMM implementation in TE's block-wise scaling？ #1509

Open

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch from 78c194d to 8f4f0f0 Compare March 19, 2025 22:43

BestJuly reviewed Mar 28, 2025

View reviewed changes

tests/pytorch/distributed/test_numerics.py Outdated Show resolved Hide resolved

transformer_engine/pytorch/tensor/_internal/float8_blockwise_tensor_base.py Outdated Show resolved Hide resolved

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch 5 times, most recently from 22828fe to 413331d Compare April 1, 2025 23:23

zhongbozhu reviewed Apr 2, 2025

View reviewed changes

transformer_engine/pytorch/module/linear.py Outdated Show resolved Hide resolved

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch 2 times, most recently from dbb6f63 to 7bfea52 Compare April 2, 2025 02:22

kunlunl reviewed Apr 2, 2025

View reviewed changes

transformer_engine/common/recipe/recipe_common.cuh Outdated Show resolved Hide resolved

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch from 7bfea52 to 44843bb Compare April 2, 2025 17:34

zhongbozhu reviewed Apr 2, 2025

View reviewed changes

tests/cpp/test_common.cu Outdated Show resolved Hide resolved

zhongbozhu reviewed Apr 2, 2025

View reviewed changes

tests/pytorch/test_float8_blockwise_gemm_exact.py Show resolved Hide resolved

zhongbozhu reviewed Apr 2, 2025

View reviewed changes

transformer_engine/pytorch/tensor/float8_blockwise_tensor.py Show resolved Hide resolved

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch from d7775fc to b62d555 Compare April 8, 2025 23:35

Update transformer_engine/pytorch/ops/op.py

196cd6d

Apply MR comment change. Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: kwyss-nvidia <kwyss@nvidia.com>

timmoon10 reviewed Apr 8, 2025

View reviewed changes

zhongbozhu reviewed Apr 9, 2025

View reviewed changes

transformer_engine/pytorch/module/layernorm_linear.py Outdated Show resolved Hide resolved

Lint fix.

67e790b

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch from 8fc753d to 67e790b Compare April 9, 2025 00:05

ptrendx reviewed Apr 9, 2025

View reviewed changes

transformer_engine/common/include/transformer_engine/transformer_engine.h Outdated Show resolved Hide resolved

transformer_engine/common/recipe/__init__.py Outdated Show resolved Hide resolved

timmoon10 reviewed Apr 9, 2025

View reviewed changes

zhongbozhu reviewed Apr 9, 2025

View reviewed changes

MR feedback.

ea9e46b

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

kwyss-nvidia force-pushed the kwyss/subchannel_recipe_linear branch from 6948759 to ea9e46b Compare April 9, 2025 01:32

Enable cuda graph tests.

324792b

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

yaox12 reviewed Apr 9, 2025

View reviewed changes

transformer_engine/common/recipe/__init__.py Outdated Show resolved Hide resolved

Reduce chance of spurious failure and reword.

54e7279

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

BestJuly mentioned this pull request Apr 9, 2025

Plans for block-wise FP8 quantization during training? #1411

Closed

timmoon10 approved these changes Apr 10, 2025

View reviewed changes

transformer_engine/pytorch/csrc/common.h Show resolved Hide resolved

tests/pytorch/test_float8_current_scaling_exact.py Outdated Show resolved Hide resolved

tests/pytorch/test_float8blockwisetensor.py Outdated Show resolved Hide resolved

timmoon10 added 2 commits April 10, 2025 00:48

Review suggestions from @timmoon10

0bf7844

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Merge branch 'main' into kwyss/subchannel_recipe_linear

62662ae

kwyss-nvidia and others added 2 commits April 9, 2025 17:56

Update CPP tests.

7efac72

Signed-off-by: Keith Wyss <kwyss@nvidia.com>

Update common.h

c3ee3d8

Signed-off-by: Xin Yao <yaox12@outlook.com>

Update test_float8blockwisetensor.py

59cb49c

Signed-off-by: Xin Yao <yaox12@outlook.com>

timmoon10 merged commit a8f0fe0 into NVIDIA:main Apr 10, 2025
19 of 23 checks passed

Comments

Conversation

kwyss-nvidia commented Mar 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

kwyss-nvidia commented Mar 12, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhongbozhu Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

kwyss-nvidia Apr 9, 2025

Choose a reason for hiding this comment

Uh oh!

timmoon10 commented Apr 9, 2025

Uh oh!

Uh oh!

timmoon10 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Apr 10, 2025

Uh oh!

yaox12 commented Apr 10, 2025

Uh oh!

yaox12 commented Apr 10, 2025

Uh oh!

yaox12 commented Apr 10, 2025

Uh oh!

zhongbozhu commented Apr 10, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

kwyss-nvidia commented Mar 11, 2025 •

edited

Loading

timmoon10 left a comment •

edited

Loading