Fix torch_dist checkpointing ETP replica_id

[Skylion007]

The replica_id is not correct; expert-tensor-parallelism (ETP) would only save the model weights in 1 tensor replica. After this change, the file size of FSDP ETP models, and pipeline parallelism non-ETP models match.

[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.

[Skylion007]

Ping on this when you have a chance @sbhavani

[hxbai]

Hi @Skylion007 the ETP is for TP splitting of the experts, and it should not be replicated along with it.

Could you please paste out the error log with strict load reported in this issue #1836?

they are half the size they should be with TP2, 1/4 with TP4, 1/8 with TP8 etc

Does the TP here means dense TP or expert TP?

[Skylion007]

Hi @Skylion007 the ETP is for TP splitting of the experts, and it should not be replicated along with it.

Could you please paste out the error log with strict load reported in this issue #1836?

they are half the size they should be with TP2, 1/4 with TP4, 1/8 with TP8 etc

Does the TP here means dense TP or expert TP?

It means ETP since in this context. TP in this PR is ETP because the method is wraped with expert_dist_ckpt_decorator
. I would have renamed the variable here, but I wanted to be consistent and prevent code churn.

Increasing ETP to 2, 4, 8 etc shows the reduction in the size of the checkpoint by roughly the factor of ETP (for an MOE model which consists mostly of MOE layers). I do not have any relevant log files still around as I last ran the broken code in August, and honestly debugged in by roundtripping the checkpoint before I set log_all

Also nit, that decorator should using a TypeVar or typing_extensions.ParamSpec with Callable to prevent type erasure of the method it decorates.

[Skylion007]

@hxbai I'll try to replicate it on a more recent version of Megatron looking at the checkpoints I do have is a result of this bug.
#1778

Where it seems to be deleting layers from the model, but the TFLOPs calculation etc all assumes the layers are still there with decoder-last-layers is set and PP is set to 1. Some of the other checkpoints were on the legacy FSDP path and no longer relevant.