Skip to content

Fix checkpoint rank selection for Ulysses sequence parallelism #8226

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tohtana merged 2 commits into from
Aug 10, 2026
+132 −24
Merged

Fix checkpoint rank selection for Ulysses sequence parallelism #8226

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6df83ff
Fix checkpoint rank for Ulysses sequence parallelism
sfc-gh-thonguyen Aug 6, 2026
6d5df59
Merge branch 'master' into fix/ulysses-checkpoint-mp-rank
sfc-gh-thonguyen Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 45 additions & 24 deletions deepspeed/runtime/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1790,6 +1790,9 @@ def _configure_distributed_model(self, model):
self.seq_data_parallel_group = groups._get_sequence_data_parallel_group()
self.seq_dp_world_size = groups._get_sequence_data_parallel_world_size()
self.mp_world_size = groups._get_model_parallel_world_size()
self.checkpoint_mp_rank = 0
if self.mp_world_size > 1:
self.checkpoint_mp_rank = self.mpu.get_model_parallel_rank()
self.expert_parallel_group = groups._get_expert_parallel_group_dict()
self.expert_data_parallel_group = groups._get_expert_data_parallel_group_dict()
self.sequence_parallel_size = groups._get_sequence_parallel_world_size()
Expand Down Expand Up @@ -3944,7 +3947,11 @@ def load_moe_state_dict(checkpoint_path,
num_experts=1,
checkpoint_engine=TorchCheckpointEngine(),
autoep_layers=None,
folding_spec=None):
folding_spec=None,
checkpoint_mp_rank=None):
if checkpoint_mp_rank is None:
checkpoint_mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank()

try:
from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer as _AutoEPMoELayer
except ImportError:
Expand All @@ -3970,7 +3977,7 @@ def load_moe_state_dict(checkpoint_path,
-1, # -1 means ignore layer_id
global_expert_id,
tag,
mpu),
checkpoint_mp_rank=checkpoint_mp_rank),
map_location=torch.device('cpu'))

# Updating global -> local expert ids
Expand Down Expand Up @@ -4016,7 +4023,11 @@ def load_moe_state_dict(checkpoint_path,
for local_expert_id in range(num_local_experts):
global_expert_id = expp_rank * num_local_experts + local_expert_id
expert_state_dict = checkpoint_engine.load(DeepSpeedEngine._get_expert_ckpt_name(
checkpoint_path, moe_layer_id, global_expert_id, tag, mpu),
checkpoint_path,
moe_layer_id,
global_expert_id,
tag,
checkpoint_mp_rank=checkpoint_mp_rank),
map_location=torch.device('cpu'))
# print(expert_state_dict.keys())
# Updating global -> local expert ids
Expand All @@ -4040,8 +4051,11 @@ def load_moe_state_dict(checkpoint_path,

for local_expert_id in range(num_local_experts):
global_expert_id = expp_rank * num_local_experts + local_expert_id
expert_ckpt_path = DeepSpeedEngine._get_expert_ckpt_name(checkpoint_path, moe_layer_id,
global_expert_id, tag, mpu)
expert_ckpt_path = DeepSpeedEngine._get_expert_ckpt_name(checkpoint_path,
moe_layer_id,
global_expert_id,
tag,
checkpoint_mp_rank=checkpoint_mp_rank)
if not os.path.exists(expert_ckpt_path):
raise FileNotFoundError(f"Expert checkpoint file not found: {expert_ckpt_path}. "
f"Expected layer_{moe_layer_id} expert_{global_expert_id}.")
Expand Down Expand Up @@ -4138,17 +4152,15 @@ def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16
return zero_ckpt_name

def _get_zero_ckpt_name(self, checkpoints_path, tag):
mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
pp_rank = dist.get_rank(group=self.optimizer.dp_process_group)
bf16_mode = self.bfloat16_enabled()
return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode)
return self._get_rank_zero_ckpt_name(checkpoints_path, tag, self.checkpoint_mp_rank, pp_rank, bf16_mode)

def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None, pp_placeholder=None):
if mp_placeholder is not None:
mp_rank_str = mp_placeholder
else:
mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
mp_rank_str = f"{mp_rank:02d}"
mp_rank_str = f"{self.checkpoint_mp_rank:02d}"

if self.zero_optimization_partition_weights():
if pp_placeholder is not None:
Expand All @@ -4171,22 +4183,24 @@ def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None, pp_placehol
return ckpt_name

def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank):
mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
ckpt_name = os.path.join(checkpoints_path, str(tag),
f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt')
f'expp_rank_{expp_rank}_mp_rank_{self.checkpoint_mp_rank:02d}_optim_states.pt')
return ckpt_name

@staticmethod
def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None):
mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank()
def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None, checkpoint_mp_rank=None):
if checkpoint_mp_rank is None:
checkpoint_mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank()

if layer_id <= -1:
# Used to support old checkpoint loading
ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag),
f'expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt')
f'expert_{expert_id}_mp_rank_{checkpoint_mp_rank:02d}_model_states.pt')
else:
# Used to support new checkpoint loading
ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag),
f'layer_{layer_id}_expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt')
ckpt_name = os.path.join(
checkpoints_path, '' if tag is None else str(tag), f'layer_{layer_id}_expert_{expert_id}_mp_rank_'
f'{checkpoint_mp_rank:02d}_model_states.pt')
return ckpt_name

def _get_all_ckpt_names(self, checkpoints_path, tag):
Expand Down Expand Up @@ -4373,8 +4387,9 @@ def _load_checkpoint(self,

is_pipe_parallel = isinstance(self.module, PipelineModule)

mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel)
load_path, checkpoint, _ = sd_loader.load(self.mp_world_size,
self.checkpoint_mp_rank,
is_pipe_parallel=is_pipe_parallel)

if checkpoint is None:
return None, None
Expand Down Expand Up @@ -4426,7 +4441,7 @@ def _load_checkpoint(self,
state_dict=checkpoint['module'],
old_moe_load=old_moe_load,
model=self.module,
mpu=self.mpu,
checkpoint_mp_rank=self.checkpoint_mp_rank,
num_experts=self.num_experts,
checkpoint_engine=self.checkpoint_engine,
autoep_layers=autoep_layers,
Expand Down Expand Up @@ -4595,10 +4610,9 @@ def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_si
return zero_ckpt_names

def _get_all_zero_checkpoint_names(self, load_dir, tag, bf16_mode):
mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
zero_ckpt_names = self._get_mp_rank_zero_checkpoint_names(load_dir=load_dir,
tag=tag,
mp_rank=mp_rank,
mp_rank=self.checkpoint_mp_rank,
dp_world_size=self.loaded_checkpoint_dp_world_size,
bf16_mode=bf16_mode)
for i, ckpt_name in enumerate(zero_ckpt_names):
Expand Down Expand Up @@ -4938,7 +4952,11 @@ def autoep_expert_writer() -> bool:
# let save the moe parameters
for global_expert_id, expert_state_dict in experts_state_dict.items():
# save the moe parameters
moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu)
moe_save_path = self._get_expert_ckpt_name(save_dir,
moe_layer_id,
global_expert_id,
tag,
checkpoint_mp_rank=self.checkpoint_mp_rank)
if self.random_ltd_enabled():
expert_state_dict = remove_random_ltd_state_dict(expert_state_dict)
saveable_state_dict = expert_state_dict
Expand Down Expand Up @@ -5035,8 +5053,11 @@ def autoep_expert_writer() -> bool:
zero_partition_count=folding_spec.edp_size,
)

moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag,
self.mpu)
moe_save_path = self._get_expert_ckpt_name(save_dir,
moe_layer_id,
global_expert_id,
tag,
checkpoint_mp_rank=self.checkpoint_mp_rank)
Comment on lines +5056 to +5060

@stas00 stas00 Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there any time where checkpoint_mp_rank arg isn't self.checkpoint_mp_rank?

if so it doesn't need to be passed as an arg and can be pulled from self, no?

@sfc-gh-thonguyen sfc-gh-thonguyen Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In DeepSpeedEngine save paths it is always self.checkpoint_mp_rank. However, _get_expert_ckpt_name is static and is also called by static load_moe_state_dict, including from InferenceEngine, where no DeepSpeedEngine instance exists. The optional argument preserves those legacy/static callers while letting DeepSpeedEngine override the MPU rank for Ulysses. So it cannot generally pull the value from self.

saveable = expert_state_dict
if self.checkpoint_engine.preserves_storage_sharing():
saveable = clone_tensors_for_torch_save(expert_state_dict)
Expand Down
66 changes: 66 additions & 0 deletions tests/unit/sequence_parallelism/test_ulysses.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import torch.nn.functional as F
import deepspeed.comm as dist
from deepspeed import initialize
import deepspeed.runtime.sequence_parallel.parallel_state_sp as sp_mpu
from transformers import AutoModel
from unit.common import DistributedTest
from deepspeed.sequence.layer import _SeqAllToAll
Expand All @@ -19,6 +20,71 @@
#Use mesh device to create data and sequence parallel group


class TestUlyssesCheckpointLoad(DistributedTest):
world_size = 2

def test_load_non_sequence_parallel_checkpoint(self, tmpdir):
config = {
"train_batch_size": self.world_size,
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-3
}
},
"zero_optimization": {
"stage": 1
}
}
hidden_dim = 4
source_model = SimpleModel(hidden_dim)
with torch.no_grad():
for parameter in source_model.parameters():
parameter.fill_(1.5)
expected = {name: tensor.detach().cpu().clone() for name, tensor in source_model.state_dict().items()}
source_engine, _, _, _ = initialize(model=source_model,
model_parameters=source_model.parameters(),
config=config)

checkpoint_dir = str(tmpdir)
tag = "no_sp"
source_engine.save_checkpoint(checkpoint_dir, tag=tag)
dist.barrier()

sp_mpu.initialize_sequence_parallel(self.world_size)
target_model = SimpleModel(hidden_dim)
target_engine, _, _, _ = initialize(model=target_model,
model_parameters=target_model.parameters(),
config=config,
mpu=sp_mpu)
assert target_engine.mp_world_size == 1
assert target_engine.checkpoint_mp_rank == 0
assert sp_mpu.get_model_parallel_rank() == dist.get_rank()

target_engine.load_checkpoint(checkpoint_dir,
tag=tag,
load_module_only=True,
load_optimizer_states=False,
load_lr_scheduler_states=False)

for name, tensor in target_engine.module.state_dict().items():
assert torch.equal(tensor.detach().cpu(), expected[name])

sp_tag = "with_sp"
target_engine.save_checkpoint(checkpoint_dir, tag=sp_tag)
dist.barrier()

resumed_model = SimpleModel(hidden_dim)
resumed_engine, _, _, _ = initialize(model=resumed_model,
model_parameters=resumed_model.parameters(),
config=config,
mpu=sp_mpu)
resumed_engine.load_checkpoint(checkpoint_dir, tag=sp_tag)

for name, tensor in resumed_engine.module.state_dict().items():
assert torch.equal(tensor.detach().cpu(), expected[name])


class TestUlyssesUtils(DistributedTest):
world_size = 4

Expand Down
21 changes: 21 additions & 0 deletions tests/unit/v1/moe/test_autoep_checkpoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -412,6 +412,27 @@ def test_autoep_metadata_schema_validation(self):
"moe_layer_id": 0
}])

def test_legacy_mpu_checkpoint_rank_argument(self):
from deepspeed.runtime.engine import DeepSpeedEngine

class LegacyMPU:

def get_model_parallel_rank(self):
return 3

mpu = LegacyMPU()
expert_path = DeepSpeedEngine._get_expert_ckpt_name("/fake", 1, 2, "tag", mpu)
assert expert_path.endswith("layer_1_expert_2_mp_rank_03_model_states.pt")
override_path = DeepSpeedEngine._get_expert_ckpt_name("/fake", 1, 2, "tag", mpu, checkpoint_mp_rank=0)
assert override_path.endswith("layer_1_expert_2_mp_rank_00_model_states.pt")

DeepSpeedEngine.load_moe_state_dict(checkpoint_path="/fake",
tag="tag",
state_dict={},
old_moe_load=False,
model=nn.Linear(1, 1),
mpu=mpu)


class TestAutoEPZero3UniversalCheckpoint(DistributedTest):
world_size = 2
Expand Down
Loading