Fix PP inference correctness in megatron_generate and megatron_importer

Two bugs caused pipeline-parallel inference to produce garbage output:

1. megatron_generate/megatron_prefill used get_forward_backward_func() (the
   training pipeline scheduler), which is not designed for inference. Rewrote
   both functions to use explicit P2P communication via
   recv_from_prev_pipeline_rank_ / send_to_next_pipeline_rank, matching the
   pattern from run_mcore_inference.

2. import_mcore_gpt_from_hf loads HF weights into stage 0's embedding but
   never updates the output_layer on the last PP stage when
   share_embeddings_and_output_weights=True. After import, call
   model.setup_embeddings_and_output_layer() to re-run the all-reduce that
   syncs the output layer from stage 0 to the last stage.

Also parametrize the megatron_generate test to cover both TP and PP.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

CHANGELOG.rst

@@ -22,6 +22,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix Megatron utility functions for generation (with pipeline parallelism) and MMLU score evaluation (10x speedup).
 - Fix Minitron pruning (``mcore_minitron``) for MoE models. Importance estimation hooks were incorrectly registered for MoE modules and NAS step was hanging before this.
 - Fix TRT support for remote autotuning in ONNX Autotune from 10.16+ to 10.15+ and fix TRT versioning check to the ``trtexec`` version instead of the TRT Python API when using ``trtexec`` backend.
 

modelopt/torch/export/plugins/megatron_importer.py

@@ -747,6 +747,17 @@ def _import_state_dict(self):
         if hasattr(model, "output_layer") and not model.share_embeddings_and_output_weights:
             self.rules["output_layer"](model.output_layer)
 
+        # For PP with shared embedding/output weights, re-sync the output layer on the last
+        # pipeline stage from stage 0's (now HF-loaded) embedding. At model init,
+        # setup_embeddings_and_output_layer() zeros out the last stage's weight and all-reduces
+        # from stage 0. After importing HF weights into stage 0's embedding, that sync is stale,
+        # so we re-run it here.
+        if (
+            model.share_embeddings_and_output_weights
+            and model.config.pipeline_model_parallel_size > 1
+        ):
+            model.setup_embeddings_and_output_layer()
+
         # MTP
         if hasattr(model, "mtp"):
             layer_pbar.set_description("Importing MTP")

modelopt/torch/utils/logging.py

@@ -105,8 +105,9 @@ def no_stdout():
 
 def print_rank_0(*args, **kwargs):
     """Prints only on the master process."""
+    kwargs.setdefault("flush", True)
     if dist.is_master():
-        print(*args, **kwargs, flush=True)
+        print(*args, **kwargs)
 
 
 def warn_rank_0(message, *args, **kwargs):