Fix MoE checkpoint save skipping TP ranks when TP > EP * ETP (#4200)

[aarushisingh04]

What does this PR do?

fixes a bug where MoE checkpoints skip saving state dicts for certain TP ranks when TP > EP * ETP.

fixes #4200

Description

• when MoE is used with a topology where TP > EP * ETP, the Expert Data Parallel (EDP) group is larger than the regular Data Parallel (DP) group.
• the previous save condition gated checkpoint generation strictly on edp_rank == 0.
• in this topology some ranks have dp_rank == 0 (they hold a unique dense/attention shard) but edp_rank > 0, so they were incorrectly skipped resulting in missing model_optim_rng.pt files and a crash on reload.

fix: replace edp_rank == 0 with min(dp_rank, edp_rank) == 0 in both save-condition locations. a rank now saves if it is primary for either its dense parameters (DP group) or its expert parameters (EDP group).

scope: this only affects CheckpointType.LEGACY. For modern distributed formats, the ckpt_type != CheckpointType.LEGACY clause short-circuits the check entirely so those paths are unaffected.

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[dimapihtar]

Hi, PR which resolves that issue has been already merged: #4678

Thanks!