Skip to content

Conversation

pstjohn
Copy link
Contributor

@pstjohn pstjohn commented Oct 9, 2025

Description

Currently the LayerNorm layer doesn't inherit from TransformerEngineBaseModule, and so it doesn't have the updated process of using an empty tensor to denote no fp8 metadata in the state_dict. This change adds this logic to the get_extra_state method of BasicOperation.

@pstjohn pstjohn force-pushed the pstjohn/dont-serialize-empty-dict branch from f0ddcbd to 040c3ba Compare October 9, 2025 15:36
@timmoon10
Copy link
Collaborator

/te-ci

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.

We used to support this behavior (see #1320), but it somehow got reverted in the TE 2.0 release (544dd14).

@timmoon10 timmoon10 merged commit dd9433e into NVIDIA:main Oct 9, 2025
38 of 41 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants