[megatron] fix: MTP loss deadlock when using context parallelism

[xhx1022]

Summary

• Fix deadlock when MTP is enabled with context parallelism (CP > 1)
• get_megatron_mtp_loss internally calls reduce_loss_in_tracker() which does all_reduce over DP+CP group (all CP ranks), but the call was gated by is_mp_src_rank_with_outputs() which requires cp_rank==0, so CP rank>0 never participated in the all_reduce → deadlock
• Changed outer gate from is_mp_src_rank_with_outputs() to mpu.is_pipeline_last_stage(ignore_virtual=True) so all CP/TP ranks participate in the all_reduce, while still only writing metrics on the src rank

[gemini-code-assist]

Code Review

This pull request updates the MTP loss handling in transformer_impl.py to ensure all ranks in the last pipeline stage participate in the reduction process, which is necessary for correct all-reduce operations across DP and CP groups. The feedback identifies a potential risk where a zero value for n_micro_batch could lead to a ZeroDivisionError or IndexError, and suggests adding a guard condition to prevent this.

[ArronHZG]

@arvyanh review

[arvyanh]

@arvyanh review

looks correct, we had a similar fix locally that I forgot to cherry-pick to github

        # MTP loss is gathered across dp and cp
        if (
            self.model_config.mtp.enable
            and mpu.get_tensor_model_parallel_rank() == 0
            and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1
        ):
            # add mtp_losses
            metrics = get_megatron_mtp_loss(n_micro_batch)
            if "metrics" not in losses_reduced[0]:
                losses_reduced[0]["metrics"] = {}
            losses_reduced[0]["metrics"].update(metrics)

The difference here (if self.is_mp_src_rank_with_outputs():) is that only CP=0 reports the loss