replace moe checkpoint dp_world_size with seq_dp_world_size

[wukong1992]

Replace moe checkpoint dp_world_size with seq_dp_world_size to sup moe module with seq parallel.

[stas00]

Why are you proposing to do that?

If you need seq_dp_world_size on resume and for some reason it's not there then probably store it additionally while leaving dp_world_size alone, no?

[wukong1992]

Why are you proposing to do that?

If you need seq_dp_world_size on resume and for some reason it's not there then probably store it additionally while leaving dp_world_size alone, no?

cause no moe save checkpoint use this

DeepSpeed/deepspeed/runtime/engine.py

Line 3908 in 37ad0c0

dp_world_size=self.seq_dp_world_size,

There is no need to redefine seq_dp_world_size here.

[stas00]

Thank you for explaining, then yes, it checks out.

@sfc-gh-truwase, I know you weren't involved in DS-MoE - but shouldn't _save_checkpoint and _save_moe_checkpoint follow the exact same recipe wrt saving state_dict (the

DeepSpeed/deepspeed/runtime/engine.py

Line 3891 in 37ad0c0

state = dict(module=module,

DeepSpeed/deepspeed/runtime/engine.py

Line 3822 in 37ad0c0

state = {

only model_state_dict is what's different - perhaps abstracting this last bit into a helper util and having both functions use it, with the only difference passing the model_state_dict that differs? the moe saving one also saves num_experts but it's not using it on load. I think it's diverging in other aspects since nobody bothered to keep in in sync.

[stas00]

Thank you, @wukong1992