[QUESTION] `save_checkpoint` with `expert_tensor_parallel_size`

[jeromeku]

Your question
How to save a sharded checkpoint when using MoE parallel folding -- specifically, when ETP != TP? save_checkpoint only supports PP, TP, and EP (as in expert_model_parallel not expert_tensor_parallel)?

[yanring]

Hi @jeromeku, it should be well supported. Did you encounter any issues?

Simply put, we'll handle this in the expert_dist_ckpt_decorator of experts.py.

[jeromeku]

@yanring

Thanks for the response. Is this only the case when using ckpt-format = torch_dist?

I was using ckpt-format=torch which was resulting in an incorrect number of checkpoint files when using different TP / ETP sizes. IIUC, using the plain torch format doesn't go down the sharded_state_dict route so the tp group won't be properly patched per the expert_dist_ckpt_decorator.