Skip to content

Don't pickle an empty dict in LayerNorm and pt base modules - #2253

Merged
timmoon10 merged 2 commits into
NVIDIA:mainfrom
pstjohn:pstjohn/dont-serialize-empty-dict
Oct 9, 2025
Merged

Don't pickle an empty dict in LayerNorm and pt base modules#2253
timmoon10 merged 2 commits into
NVIDIA:mainfrom
pstjohn:pstjohn/dont-serialize-empty-dict

Conversation

@pstjohn

Copy link
Copy Markdown
Contributor

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.

Signed-off-by: Peter St. John <pstjohn@nvidia.com>
@pstjohn
pstjohnforce-pushed the pstjohn/dont-serialize-empty-dict branch from f0ddcbd to 040c3baCompareOctober 9, 2025 15:36
@timmoon10

Copy link
Copy Markdown
Member

/te-ci

@timmoon10timmoon10 left a comment

Copy link
Copy Markdown
Member

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:mainOct 9, 2025
38 of 41 checks passed
Sign up for freeto 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

@pstjohn@timmoon10@ksivaman