Must the sp_size be equal to the total_gpus in UlyssesSPAttentionHF?

[NiuMa-1234]

I found the sequence_parallel_size in the provided example of Ulysses (test_ulysses_sp_hf.py) is equal to the world_size( total gpus) , and if the sequence_parallel_size is less than the world_size, the training would encounter an error when going backwards, as shown below:

 self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/opt/conda/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 65, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 522, in backward
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 266, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/function.py", line 289, in apply
    return user_fn(self, *args)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/nn/functional.py", line 343, in backward
    gx = torch.empty_like(grad_outputs[rank])
IndexError: tuple index out of range

And the error is likely caused by that, when executing torch._AllGather, the gradoutput only has sp_world_size items but torch.distributed.get_rank() causes each GPU to choose its own grad from the gradoutput. Therefore raise this index mismatch error.

So is there a must to make sure the sequence_parallel_size be equal to the world_size?

[Mikyx-1]

Yes — sequence_parallel_size < world_size is supported by design; it combines Ulysses sequence parallelism with data parallelism rather than requiring SP to span every rank.

initialize_sequence_parallel() in deepspeed/runtime/sequence_parallel/parallel_state_sp.py only requires world_size % sequence_parallel_size == 0. It builds world_size / sequence_parallel_size independent SP groups, plus one combined "sequence-data-parallel" group spanning every rank (since sequence_parallel_size * data_parallel_size == world_size by construction, that combined group always equals the full world). ZeRO stage 1/2 and stage 3 read sequence_parallel_size and divide gradients by dp_world_size / sequence_parallel_size rather than plain dp_world_size, and deepspeed/utils/groups.py routes ZeRO's reduction/sharding group through this combined group whenever the mpu exposes get_sequence_data_parallel_group() (which the Ulysses mpu does).

I wrote a small repro adapted from tests/unit/ulysses_alst/test_ulysses_sp_hf.py with world_size=4, sequence_parallel_size=2 (i.e. 2 SP groups x dp_size=2) and ran it against current master (deepspeed 0.19.3) on 4 GPUs, for both zero_stage=2 and zero_stage=3:

[rank 0] sp_world_size=2 sp_rank=0 seq_data_parallel_world_size=4 data_parallel_group=None
[rank 0] loss_b=10.425115585327148 - about to call backward()
[rank 0] backward() succeeded
[rank 0] grad_b norm = 0.003700781846418977
[rank 0] optimizer step() succeeded
... (same for ranks 1-3, zero_stage=2)

... (zero_stage=3 also passes, forward/backward/step all succeed)

So on current master, this combination works cleanly for a small Llama model with both ZeRO-2 and ZeRO-3 — no backward-pass error. That means your crash is very likely either:

1. Version-specific — if you're on an older DeepSpeed release, this may have been a bug that's since been fixed on master; worth trying the latest release or master directly, or
2. Config/wiring-specific to your setup — e.g. the mpu not being passed into deepspeed.initialize(), a mismatch in how the dataloader/loss aggregation is combining SP shards, or something particular to your model architecture.

Could you share: your DeepSpeed version, your exact sequence_parallel_size/world_size values, and the full traceback (not just the last few frames)? That'll help pin down which of the two it is. If it turns out to be a regression that's already fixed on master, upgrading should resolve it; if it reproduces on master too, that's a genuine bug worth filing separately with a minimal repro.

[NiuMa-1234]

Thank you for your kind explanation. I found that this bug was fixed in January, as shown in this pr #7809