feat(gpt): add output postprocess hook

megatron/core/models/common/model_chunk_schedule_plan.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
 from contextlib import nullcontext
-from typing import Optional
+from typing import Any, Callable, Optional
 
 import torch
 from torch import Tensor
@@ -282,6 +282,9 @@ def __init__(
         runtime_gather_output: Optional[bool] = None,
         loss_mask: Optional[Tensor] = None,
         padding_mask=None,
+        *,
+        output_processor: Optional[Callable[..., Tensor]] = None,
+        output_processor_context: Optional[Any] = None,
     ):
         """Initialize the schedule plan of all Transformer layers' sub-modules.
 
@@ -299,6 +302,10 @@ def __init__(
             extra_block_kwargs: Additional keyword arguments for blocks.
             runtime_gather_output: Whether to gather output at runtime.
             loss_mask (torch.Tensor): Used to mask out some portions of the loss
+            output_processor (Callable): Custom postprocess hook to run instead of the
+                default logits/loss path.
+            output_processor_context (Any): User-defined context object forwarded to
+                `output_processor`.
 
         Returns:
             The model chunk schedule plan.
@@ -324,6 +331,8 @@ def __init__(
         self._model_chunk_state.padding_mask = padding_mask
         self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
         self._model_chunk_state.runtime_gather_output = runtime_gather_output
+        self._model_chunk_state.output_processor = output_processor
+        self._model_chunk_state.output_processor_context = output_processor_context
         self._model_chunk_state.model = model
         self._model_chunk_state.context = None
         self._model_chunk_state.context_mask = None

megatron/core/models/gpt/fine_grained_callables.py

@@ -228,6 +228,8 @@ def forward_impl(self, hidden_states):
             sequence_len_offset=self.chunk_state.sequence_len_offset,
             runtime_gather_output=self.chunk_state.runtime_gather_output,
             extra_block_kwargs=self.chunk_state.extra_block_kwargs,
+            output_processor=self.chunk_state.output_processor,
+            output_processor_context=self.chunk_state.output_processor_context,
         )
 
         # For now, 1f1b only supports fp16 module

megatron/core/models/gpt/gpt_model.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from collections import OrderedDict
-from typing import Dict, Literal, Optional
+from typing import Any, Callable, Dict, Literal, Optional
 
 import torch
 from torch import Tensor
@@ -491,6 +491,8 @@ def forward(
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
         padding_mask: Optional[Tensor] = None,
+        output_processor: Optional[Callable[..., Tensor]] = None,
+        output_processor_context: Optional[Any] = None,
     ) -> Tensor:
         """Forward function of the GPT Model This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -504,6 +506,10 @@ def forward(
             padding_mask (Tensor, optional): Padding mask for MoE routing.
                 Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
                 Only used for MoE layers to exclude padding tokens from routing computations.
+            output_processor (Callable, optional): Custom postprocess hook that receives
+                decoder hidden states and output-layer helpers, then returns the model output.
+            output_processor_context (Any, optional): User-defined context object forwarded to
+                `output_processor`.
         """
         if self.config.fine_grained_activation_offloading:
             self.preprocess_for_fine_grained_offloading()
@@ -563,6 +569,8 @@ def forward(
             runtime_gather_output=runtime_gather_output,
             extra_block_kwargs=extra_block_kwargs,
             inference_context=inference_context,
+            output_processor=output_processor,
+            output_processor_context=output_processor_context,
         )
 
     def _postprocess(
@@ -584,6 +592,8 @@ def _postprocess(
         runtime_gather_output=None,
         extra_block_kwargs=None,
         inference_context=None,
+        output_processor=None,
+        output_processor_context=None,
     ):
         """Postprocesses decoder hidden states to generate logits or compute loss.
 
@@ -650,6 +660,26 @@ def _postprocess(
                 )
         sequence_parallel_override = False
 
+        if output_processor is not None:
+            return output_processor(
+                hidden_states=hidden_states,
+                output_layer=self.output_layer,
+                output_weight=output_weight,
+                labels=labels,
+                loss_mask=loss_mask,
+                input_ids=input_ids,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                decoder_input=decoder_input,
+                inference_context=inference_context,
+                packed_seq_params=packed_seq_params,
+                runtime_gather_output=runtime_gather_output,
+                context=output_processor_context,
+                compute_language_model_loss=self.compute_language_model_loss,
+                scale_logits=self._scale_logits,
+                config=self.config,
+            )
+
         if in_inference_mode and inference_context.config.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
@@ -719,6 +749,9 @@ def build_schedule_plan(
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
         padding_mask: Optional[Tensor] = None,
+        *,
+        output_processor: Optional[Callable[..., Tensor]] = None,
+        output_processor_context: Optional[Any] = None,
     ):
         """Builds a computation schedule plan for the model.
 
@@ -745,6 +778,10 @@ def build_schedule_plan(
                 Parameters for inference. Defaults to None.
             loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
             padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
+            output_processor (Callable, optional): Custom postprocess hook to run in the
+                schedule-plan postprocess node instead of the default logits/loss path.
+            output_processor_context (Any, optional): User-defined context object forwarded to
+                `output_processor`.
 
         Returns:
             TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -767,6 +804,8 @@ def build_schedule_plan(
             runtime_gather_output,
             loss_mask,
             padding_mask,
+            output_processor=output_processor,
+            output_processor_context=output_processor_context,
         )
 
     def sharded_state_dict(

tests/unit_tests/models/test_gpt_model.py

@@ -7,6 +7,7 @@
 
 import pytest
 import torch
+import torch.nn.functional as F
 from packaging import version
 from pytest import approx
 from transformer_engine.pytorch.fp8 import check_fp8_support
@@ -112,6 +113,121 @@ def test_post_process_forward(self):
         assert logits.shape[1] == sequence_length
         assert logits.shape[2] == self.gpt_model.vocab_size
 
+    @pytest.mark.internal
+    def test_output_processor_forward(self):
+        config: TransformerConfig = self.gpt_model.config
+        sequence_length = self.gpt_model.max_sequence_length
+        micro_batch_size = 2
+
+        self.gpt_model.cuda()
+        self.gpt_model.eval()
+
+        data = list(range(sequence_length))
+        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        labels = (input_ids + 1) % self.gpt_model.vocab_size
+        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool
+        ).cuda()
+
+        context = {"selected_token_positions": torch.tensor([0, 2], device="cuda")}
+        seen = {}
+
+        def output_processor(**kwargs):
+            seen.update(kwargs)
+            logits, _ = kwargs["output_layer"](
+                kwargs["hidden_states"],
+                weight=kwargs["output_weight"],
+                runtime_gather_output=kwargs["runtime_gather_output"],
+            )
+            logits = kwargs["scale_logits"](logits).transpose(0, 1).contiguous().float()
+            token_logprobs = F.log_softmax(logits, dim=-1).gather(
+                dim=-1, index=kwargs["labels"].unsqueeze(-1)
+            )
+            token_logprobs = token_logprobs.squeeze(-1)
+            return token_logprobs.index_select(1, kwargs["context"]["selected_token_positions"])
+
+        with torch.no_grad():
+            logits = self.gpt_model.forward(
+                input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask
+            ).float()
+            expected = F.log_softmax(logits, dim=-1).gather(dim=-1, index=labels.unsqueeze(-1))
+            expected = expected.squeeze(-1).index_select(1, context["selected_token_positions"])
+
+            output = self.gpt_model.forward(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                output_processor=output_processor,
+                output_processor_context=context,
+            )
+
+        assert torch.allclose(output, expected)
+        assert seen["context"] is context
+        assert seen["output_layer"] is self.gpt_model.output_layer
+        assert seen["output_weight"] is None
+        assert seen["labels"] is labels
+        assert seen["runtime_gather_output"] is None
+        assert seen["config"] is config
+        assert seen["compute_language_model_loss"].__self__ is self.gpt_model
+        assert seen["scale_logits"].__self__ is self.gpt_model
+
+    @pytest.mark.internal
+    def test_build_schedule_plan_threads_output_processor(self):
+        sequence_length = self.gpt_model.max_sequence_length
+        micro_batch_size = 2
+
+        self.gpt_model.cuda()
+
+        data = list(range(sequence_length))
+        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        labels = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool
+        ).cuda()
+
+        context = {"loss_mask": torch.ones_like(labels)}
+
+        seen = {}
+
+        def output_processor(**kwargs):
+            seen.update(kwargs)
+            return kwargs["hidden_states"].sum()
+
+        schedule_plan = self.gpt_model.build_schedule_plan(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            output_processor=output_processor,
+            output_processor_context=context,
+        )
+
+        assert schedule_plan.state.output_processor is output_processor
+        assert schedule_plan.state.output_processor_context is context
+
+        schedule_plan.state.decoder_input = None
+        schedule_plan.state.rotary_pos_emb = None
+        schedule_plan.state.rotary_pos_cos = None
+        schedule_plan.state.rotary_pos_sin = None
+        schedule_plan.state.sequence_len_offset = None
+
+        hidden_states = torch.randn(
+            sequence_length,
+            micro_batch_size,
+            self.gpt_model.config.hidden_size,
+            device="cuda",
+            requires_grad=True,
+        )
+        output = schedule_plan.post_process.forward(hidden_states)
+
+        assert output.item() == pytest.approx(hidden_states.sum().item())
+        assert seen["context"] is context
+        assert seen["labels"] is labels
+        assert seen["output_layer"] is self.gpt_model.output_layer
+
 
 def test_get_mlp_module_spec_interface():
     # Get the function signature