[compat] MTP compat megatron-core main branch / dev branch

[Jintao-Huang]

main: NVIDIA/Megatron-LM#3458
dev: NVIDIA/Megatron-LM#4226

[gemini-code-assist]

Code Review

This pull request adds compatibility support for Megatron Core (mcore) version 0.19.0rc0 and above by conditionally handling different logging signatures of MTPLossLoggingHelper. The review feedback suggests optimizing performance by moving the import of _compute_mtp_acceptance_counts to the module level, avoiding repeated import overhead inside the training loop.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds compatibility support for Megatron-Core (mcore) version 0.19.0rc0 and above in gpt_model.py. It conditionally uses MTPLossLoggingHelper.save_metrics_to_tracker and computes multi-token prediction acceptance counts if the helper method is available, falling back to the original save_loss_to_tracker for older versions. The feedback suggests moving the inline import of _compute_mtp_acceptance_counts to the module level using a try...except block to avoid repeated import overhead inside the training loop and to adhere to PEP 8 guidelines.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

To avoid importing _compute_mtp_acceptance_counts repeatedly inside the training loop, we can conditionally import it at the module level using a try...except ImportError block. This adheres to PEP 8 guidelines which recommend placing imports at the top of the file.

mcore_019 = version.parse(megatron.core.__version__) >= version.parse('0.19.0rc0')

try:
    from megatron.core.transformer.multi_token_prediction import _compute_mtp_acceptance_counts
except ImportError:
    _compute_mtp_acceptance_counts = None

References

1. PEP 8 recommends placing imports at the top of the file, just after any module comments and docstrings, and before module globals and constants. ^(link)

[gemini-code-assist]

Importing _compute_mtp_acceptance_counts inside the training loop (which runs mtp_unroll_steps times per forward pass) introduces unnecessary overhead. By using the module-level imported function, we can avoid this overhead and improve performance.

                    if hasattr(MTPLossLoggingHelper, 'save_metrics_to_tracker') and _compute_mtp_acceptance_counts is not None:
                        # mcore >= 0.19 main branch: save_metrics_to_tracker with correct/total
                        correct, total = _compute_mtp_acceptance_counts(
                            mtp_logits, mtp_labels, loss_mask_, output_layer=None, runtime_gather_output=True)
                        MTPLossLoggingHelper.save_metrics_to_tracker(
                            mtp_loss_for_log,
                            correct,
                            total,
                            mtp_layer_number,
                            self.config.mtp_unroll_steps,
                            avg_group=avg_group)