Skip to content

[Community][Main] fix(moe): Fix theoretical memory calculation of layernorm. - #2434

Merged
Phlip79 merged 2 commits into
NVIDIA:mainfrom
1195343015:patch-2
Jan 30, 2026
Merged

[Community][Main] fix(moe): Fix theoretical memory calculation of layernorm.#2434
Phlip79 merged 2 commits into
NVIDIA:mainfrom
1195343015:patch-2

Conversation

@1195343015

@1195343015 1195343015 commented Nov 30, 2025

Copy link
Copy Markdown
Contributor

The calculation of parameters for q_norm and kv_norm should be multiplied by 2 for LayerNorm (includes both weight and bias) but by 1 for RMSNorm (only weight).

The calculation of parameters for `q_norm` and `kv_norm` should be multiplied by 2.
@copy-pr-bot

copy-pr-bot Bot commented Nov 30, 2025

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Treat different norm size for RMSNorm and LayerNorm
@yuzhongw-nvidia
yuzhongw-nvidia requested a review from a team December 10, 2025 02:32
@yuzhongw-nvidia

Copy link
Copy Markdown
Contributor

/ok to test b3a845a

@github-actions

Copy link
Copy Markdown
Contributor

Thank you for your contribution!

NVIDIA Megatron-LM is currently transitioning to development on Github. We will aim to review your PR after we complete our transition and stabilize our Github development process.

Thank you for your understanding.

@yuzhongw-nvidia yuzhongw-nvidia changed the title [Bugfix] Fix theoretical memory calculation. [Community][Main] fix(moe): Fix theoretical memory calculation of layernorm. Dec 10, 2025

@yuzhongw-nvidia yuzhongw-nvidia left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM. Thanks for your work.

@chtruong814 chtruong814 added the needs-follow-up Issue needs follow-up label Jan 11, 2026
@Phlip79
Phlip79 added this pull request to the merge queue Jan 30, 2026
Merged via the queue into NVIDIA:main with commit 9ad5906 Jan 30, 2026
46 checks passed
@chtruong814 chtruong814 removed the needs-follow-up Issue needs follow-up label Jan 30, 2026
daiyaanarfeen pushed a commit to daiyaanarfeen/Megatron-LM that referenced this pull request Feb 23, 2026
…ernorm. (NVIDIA#2434)

Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
…ernorm. (NVIDIA#2434)

Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
…ernorm. (NVIDIA#2434)

Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[QUESTION] Why doesn't theoretical memory take into account the expert parallelism and the partitioning of experts?

5 participants