Enhance recipe compatibility by negvet · Pull Request #1724 · NVIDIA/TransformerEngine

negvet · 2025-04-28T11:07:23Z

Description

Enables recipe update on the fly. Added a user warning when this happens.

Enables TE1.x checkpoint loading.

Implemented check that is supposed to catch user-defined bug, when there is a mismatch between the recipes in fp8_model_init and fp8_autocast.
Example case to check: recipe is DelayedScaling (DelayedScaling is set in fp8_autocast()), but the weight tensor is MXFP8Tensor (MXFP8BlockScaling is set in fp8_model_init()).

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:

Added check in the base module class
Added test

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

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet · 2025-04-28T11:10:39Z

/te-ci L0

ptrendx · 2025-04-28T17:48:36Z

Maybe add a field in the recipe that has this exepected tensor class? Also, I guess for completeness we should not only deal with QuantizedTensor, but also with the *Base (e.g. Float8TensorBase) classes (I know that they are not used for weights in fp8_model_init).

Maybe add a field in the recipe that has this exepected tensor class?

That is what I also like more. I did not do so due to circular import.
Now, I fixed it with local import, if this is acceptable.
But probably, we need some redesign to better fix this circular import issue (in another PR).

In the future, I would also propose to add a field in the recipe with quantizer class etc.

Also, I guess for completeness we should not only deal with QuantizedTensor, but also with the *Base (e.g. Float8TensorBase) classes (I know that they are not used for weights in fp8_model_init).

Fixed

See 98799f7

ptrendx · 2025-04-28T17:49:31Z

@timmoon10 I vaguealy recall that getattr from nn.Module was slow and that you created a faster function for it at some point, do you remember the details?

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

negvet · 2025-04-29T10:08:54Z

@ksivaman please take a look at e64188b (recipe change warning)

negvet · 2025-04-29T10:14:22Z

/te-ci L0

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet · 2025-04-29T12:23:58Z

@ksivaman, 674a3c9 enables recipe change on the fly.

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet · 2025-04-29T13:18:19Z

TE1.x did not save the recipe in the state.
Based on that, the assumption in 5023a32 is the following: if no recipe in state -> it is TE1.x checkpoint with DelayedScaling.
Does it sound like a safe enough assumption?

With this fix, TE1.x checkpoint can be loaded in ToT TE2.x. (for both TE versions, used the same ToT Megatron-LM)

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet · 2025-04-29T13:37:54Z

/te-ci L0

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

negvet · 2025-04-30T09:26:01Z

/te-ci L0

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman · 2025-05-06T00:25:47Z

/te-ci L0

ptrendx · 2025-05-07T23:16:48Z

 _NUM_MAX_UB_STREAMS = 3
 _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None
 layers_atomic_ring_exchange = []
+_QUANTIZED_WEIGHT_TENSOR_TYPES = (


We recently merged the change to introduce QuantizedTensorBase class so you should be able to remove this.

Fixed in 7c2f5eb

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet · 2025-05-09T13:39:51Z

/te-ci L0

ptrendx · 2025-05-12T18:21:48Z

+from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor
+from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor
+from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensor


Actually you can't do it that way - this recipe class is in common, so is used in the JAX version as well. You will need to do the import and setting in the __post_init__ method as well (with try/except block).

Right, fixed in d6ea981

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

negvet · 2025-05-13T10:59:27Z

/te-ci L0

ksivaman · 2025-05-15T06:39:56Z

    def __post_init__(self) -> None:
        assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported."
+        try:
+            from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor


Not a fan of adding framework specific logic in the core lib at all, even if it is guarded with try-except. So far we haven't had to do this and probably shouldn't introduce it here unless a strong reason exists to do so.

How about we reverse the logic and instead of including tensor info in the recipe, we include some recipe info in the tensor, and that should achieve the same goal of being able to match the two together, right?

ksivaman · 2025-05-15T06:44:13Z

+            if isinstance(tensor, QuantizedTensorBase) and not isinstance(
+                tensor, recipe.expected_tensor_class
+            ):


Suggested change

if isinstance(tensor, QuantizedTensorBase) and not isinstance(

tensor, recipe.expected_tensor_class

):

if isinstance(tensor, QuantizedTensorBase) and not tensor.recipe_name() == recipe.__class__.__name__:

Something like this to address my above comment maybe? This would require some hardcoding, i.e. each of the tensor types would be associated with a fixed "string" as recipe class name but I think that should be ok? In my example MXFP8Tensor.recipe_name() would return MXFP8BlockScaling. We could think of better ideas.

@negvet @ptrendx

This is a good approach, although string-based comparison is not that robust (what if name is changed, inheritance is not naturally supported etc.).

What about the following update of your proposal:

# In each tensor class (e.g., Float8Tensor) @classmethod def get_compatible_recipe_types(cls): return [DelayedScaling, Float8CurrentScaling]

and then we just check isinstance(recipe, tuple(tensor.__class__.get_compatible_recipe_types()))

@ptrendx

But this is still a workaround, fitted for this particular case.
Another one could be a static recipe_to_tensor_map defined in the pytorch module.
Although if there is a stronger demand for a framework-specific logic in the recipe, we might consider to define abstract base recipe classes in the common module and move the framework-specific implementations to their respective modules (class DelayedScaling(BaseDelayedScaling)). But not sure if we need that.

The framework specific logic lives in the quantizers, the recipe itself makes sense in common.

I like the get_compatible_recipe_types idea. Some early nitpicks to help review later:

@classmethod def get_compatible_recipes(cls): return (DelayedScaling, Float8CurrentScaling)

and we can check isinstance(recipe, tensor.get_compatible_recipes())

ksivaman

https://github.com/NVIDIA/TransformerEngine/pull/1724/files#r2090393437

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet · 2025-05-19T17:53:31Z

As discussed with @ptrendx, allowed quantizer to know about the recipe, rather than a tensor.
See a2910d7

negvet · 2025-05-19T18:07:01Z

/te-ci L0

ksivaman

LGTM!

* Check tensor-recipe compatibility Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Tensor class in recipe, checking for *Base Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Extend recipe __repr__ with recipe_type Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Warn about recipe change Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable dynamic recipe change: clear fp8 workspace Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * TE 1.x checkpoint compatibility Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Disable warning for recipe wrappers Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Test recipe change Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use QuantizedTensorBase Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Fix circular import Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Revert previous circular import fix Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * Fix pytorch imports in common Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Let quantizer know about the recipe Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix imports Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> --------- Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Check tensor-recipe compatibility

86e1502

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

ptrendx reviewed Apr 28, 2025

View reviewed changes

negvet added 3 commits April 29, 2025 08:11

Tensor class in recipe, checking for *Base

98799f7

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

Extend recipe __repr__ with recipe_type

67bb47f

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

Warn about recipe change

e64188b

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet force-pushed the et/recipe_compat branch from ecb8b96 to e64188b Compare April 29, 2025 10:02

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

eac9fe1

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

negvet marked this pull request as ready for review April 29, 2025 10:14

Enable dynamic recipe change: clear fp8 workspace

674a3c9

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

TE 1.x checkpoint compatibility

5023a32

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

Disable warning for recipe wrappers

243b6f4

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet changed the title ~~Check tensor-recipe compatibility~~ Enhance recipe compatibility Apr 29, 2025

negvet requested a review from ptrendx April 29, 2025 16:35

negvet and others added 2 commits April 30, 2025 09:23

Test recipe change

400e206

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

01a48da

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

ptrendx and others added 2 commits May 1, 2025 13:26

Merge branch 'main' into et/recipe_compat

60cb9f1

Merge branch 'main' into et/recipe_compat

1f02262

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ptrendx reviewed May 7, 2025

View reviewed changes

negvet added 3 commits May 9, 2025 12:23

merge main

5047291

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

Use QuantizedTensorBase

7c2f5eb

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

Fix circular import

5cdda54

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

Revert previous circular import fix

463ad23

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

negvet requested a review from ptrendx May 9, 2025 15:54

ptrendx reviewed May 12, 2025

View reviewed changes

negvet and others added 2 commits May 13, 2025 08:09

Fix pytorch imports in common

d6ea981

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

c5c27a6

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

ksivaman reviewed May 15, 2025

View reviewed changes

ksivaman requested changes May 15, 2025

View reviewed changes

negvet and others added 4 commits May 19, 2025 17:33

Let quantizer know about the recipe

a2910d7

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

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

7d7841a

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

Merge branch 'main' into et/recipe_compat

81b1acc

Fix imports

e357606

Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com>

ksivaman approved these changes May 19, 2025

View reviewed changes

ksivaman merged commit 730fd11 into NVIDIA:main May 19, 2025
27 of 31 checks passed

timmoon10 mentioned this pull request Jun 12, 2025

[PyTorch] Add support for FP8 current scaling in operation-based API #1858

Merged

14 tasks

Conversation

negvet commented Apr 28, 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!

negvet commented Apr 28, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet Apr 29, 2025 • 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!

negvet Apr 29, 2025 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet commented Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

negvet commented Apr 29, 2025

Uh oh!

negvet commented Apr 29, 2025

Uh oh!

negvet commented Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

negvet commented Apr 29, 2025

Uh oh!

negvet commented Apr 30, 2025

Uh oh!

ksivaman commented May 6, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet commented May 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

negvet commented May 13, 2025

Uh oh!

ksivaman May 15, 2025 • 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!

ksivaman May 15, 2025 • 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!

ksivaman May 15, 2025 • 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!

negvet May 16, 2025 • 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!

negvet May 16, 2025 • 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!

Choose a reason for hiding this comment

Uh oh!

ksivaman left a comment

Choose a reason for hiding this comment

Uh oh!

negvet commented May 19, 2025

Uh oh!

negvet commented May 19, 2025

negvet commented Apr 28, 2025 •

edited

Loading

negvet Apr 29, 2025 •

edited

Loading

negvet Apr 29, 2025 •

edited

Loading

negvet commented Apr 29, 2025 •

edited

Loading

negvet commented Apr 29, 2025 •

edited

Loading

ksivaman May 15, 2025 •

edited

Loading

ksivaman May 15, 2025 •

edited

Loading

ksivaman May 15, 2025 •

edited

Loading

negvet May 16, 2025 •

edited

Loading

negvet May 16, 2025 •

edited

Loading