From 65c6919426f638872158d592d9ec536360c88082 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Thu, 18 Dec 2025 21:49:22 -0800 Subject: [PATCH 1/2] fix missing layernorm --- megatron/core/models/gpt/fine_grained_callables.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 60094976a9a..6ea3fdc0d8b 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -170,11 +170,14 @@ def forward_impl(self, hidden_states): Returns: The logits or loss depending on whether labels are provided. - - Note: - Final layernorm now has been moved from the post-process stage to the - last decoder layer, so we don't need to run the final layer norm here. """ + + empty_decoder = len(self.gpt_model.decoder.layers) == 0 + layer_norm = self.gpt_model.decoder.final_layernorm + if not self.gpt_model.config.mtp_num_layers and empty_decoder and layer_norm: + hidden_states = layer_norm(hidden_states) + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + # Run GPTModel._postprocess loss = self.gpt_model._postprocess( hidden_states=hidden_states, From 3e78312944019550d22e5993681a4ceb98eba9d0 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 19 Dec 2025 02:20:15 -0800 Subject: [PATCH 2/2] format --- megatron/core/models/gpt/fine_grained_callables.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 6ea3fdc0d8b..741a25326fb 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -176,7 +176,9 @@ def forward_impl(self, hidden_states): layer_norm = self.gpt_model.decoder.final_layernorm if not self.gpt_model.config.mtp_num_layers and empty_decoder and layer_norm: hidden_states = layer_norm(hidden_states) - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=True + ) # Run GPTModel._postprocess loss = self.gpt_model._postprocess(