From eba9a4c07c9665c78ef3fb75af0579844607bf1d Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 5 Jun 2026 16:13:11 +0800 Subject: [PATCH] fix(layer_wise): tag MTP-stage word_embeddings as is_embedding_or_output_parameter The is_embedding_or_output_parameter attribute (used by decoupled_lr, the Muon LayerWise optimizer, and FSDP) was only set on the pre_process stage's word embedding. On an MTP stage (mtp_process=True, pre_process=False) the duplicated word_embeddings copy was left untagged, so the LayerWise optimizer treated it as a Muon-managed parameter instead of routing it through the DistributedOptimizer, duplicating its optimizer state and inflating peak memory. Tag the MTP-stage embedding too, mirroring the pre_process path. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../core/models/common/language_module/language_module.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 9d8eddbdda6..d238e2130d9 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import os from typing import Optional, Tuple @@ -202,7 +202,8 @@ def setup_embeddings_and_output_layer(self) -> None: # Mark embedding and output layer for decoupled_lr and other features. # This is the original Megatron attribute used by decoupled_lr, Muon, FSDP, etc. - if self.pre_process and hasattr(self, 'embedding'): + # Include MTP-stage embedding too: it is a duplicated copy of the pre_process embedding + if (self.pre_process or getattr(self, 'mtp_process', False)) and hasattr(self, 'embedding'): self.embedding.word_embeddings.weight.is_embedding_or_output_parameter = True if ( self.post_process