[bugfix] fix MTP checkpoint saving with TP > 1

[WhichPaths]

Saving an mcore checkpoint fails for Qwen3.5/3.6 MoE with MTP when TP > 1:

megatron.core.dist_checkpointing.core.CheckpointingException: Invalid access pattern for ShardedTensor(key='language_model.mtp.layers.0.enorm.weight', dtype=torch.bfloat16, local_shape=(2048,), global_shape=(2048,), global_offset=(0,), axis_fragmentations=(1,), replica_id=(0, 0, 0), prepend_axis_num=0, allow_shape_mismatch=False, flattened_range=None): tensor([2], dtype=torch.int32)

On mcore 0.16 the base MultiTokenPredictionLayer doesn't set tp_group, so self.tp_group = getattr(self, 'tp_group', None) leaves it as None. MegatronModule.sharded_state_dict only falls back to the global TP group when the attribute is missing, and make_sharded_tensor_for_checkpoint only falls back when tp_group and dp_cp_group are both None (dp_cp_group always comes in via metadata). So get_pg_rank(None) returns 0 on every TP rank and the replicated enorm/hnorm weights get the same replica_id on both TP ranks, i.e. two main replicas.

Fall back to the global TP group when it's None. Forward doesn't change: the gather/scatter calls that take self.tp_group already resolve None to the same group.

Tested on Qwen3.6-35B-A3B, 2 nodes x 8 H200, TP=2 EP=8, mtp_num_layers=1, megatron-core 0.16.1, ms-swift 4.5.3. With this change the run saves checkpoints at step 6 and step 12 with validate_access_integrity on. Resuming from the step-6 one (--finetune false) gives loss 0.50693 / 0.48621 at steps 7 / 8, vs 0.50688 / 0.48642 in the original run.

The same save path also hits a separate pickling error in ms-swift, fixed in modelscope/ms-swift#10242.