[BUG] Ulysses crashes in Stage 0

[avicooper1]

Describe the bug
Training with Ulysses sequence parallel in Stages 1/2/3 works out of the box. Using the same training pipeline with Stage 0 crashes with the following error stack trace:

File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2470, in step
  self._take_model_step(lr_kwargs)
File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2365, in _take_model_step
  self.optimizer.step()
File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
  return func(*args, **kwargs)
File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py", line 278, in step
  non_expert_grads_for_norm, expert_grads_for_norm = self.get_grads_for_norm()
File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
  return func(*args, **kwargs)
File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/deepspeed/runtime/bf16_optimizer.py", line 394, in get_grads_for_norm
  tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu)
File "/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/deepspeed/utils/bwc.py", line 42, in bwc_tensor_model_parallel_rank
  return mpu.get_model_parallel_rank()
AttributeError: module 'deepspeed.runtime.sequence_parallel.parallel_state_sp' has no attribute 'get_model_parallel_rank'

As noted, the Ulysses code does not contain a method for get_model_parallel_rank.

Looking through the stack trace, the optimizer class for Stage 0 differs from that of Stage 1 (which makes sense, Stage 1 parallelizes the optimizer,) ie. deepspeed.runtime.bf16_optimizer.BF16_Optimizer and deepspeed.runtime.zero.stage_1_and_2.DeepSpeedZeroOptimizer respectively. This code divergence means that only the Stage 0 optimizer calls this problematic line.

To Reproduce
I followed the tutorial: Arctic Long Sequence Training (ALST) for HF Transformers integration

Changing the Stage from 3 to 0 caused the error.

Expected behavior
It is my assumption that enabling sequence parallel is independent of enabling any model related parallelism (ie. stages 1/2/3.) For example, in a case where I have an extremely small model that fits easily on to 1 GPU (1B model) but the context is extremely long and cannot fit on to a single GPU.

Is there a technical limitation (ie. the optimizer needs to be shared in Stage 1) in order for sequence parallel to work? Or can the right methods be added to make it work for Stage 0 as well?

Next steps
If it is the case that this is not a technical limitation but instead a bug, I am happy to implement the necessary changes, with some direction on how to make these changes.

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
dc ..................... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
 [WARNING]  FP Quantizer is using an untested triton version (3.6.0), only 2.3.(0, 1) and 3.0.0 are known to be compatible with these kernels
fp_quantizer ........... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.10
 [WARNING]  using untested triton version (3.6.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/torch']
torch version .................... 2.10.0+cu128
deepspeed install path ........... ['/home/acooper/KosupaLLM/.venv/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.17.5, unknown, unknown
torch cuda version ............... 12.8
torch hip version ................ None
nvcc version ..................... 12.3
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
shared memory (/dev/shm) size .... 12.00 GB

[tohtana]

Hi @avicooper1,

Thanks for the detailed report. We re-tested this and can reproduce the exact get_model_parallel_rank AttributeError on v0.17.5, but we cannot reproduce it on the latest release v0.18.5 (b19987c7) or current master (a44fb58f).

Can you try these new versions?

[amadhan882]

Hi @avicooper1 and @tohtana,

I have verified the root cause of the get_model_parallel_rank crash by exercising the actual code path in deepspeed/utils/bwc.py.

Authentic Reproduction Log (v0.14.2)

Attempting authentic call to deepspeed.utils.bwc...

--- EXACT TRACEBACK LOG ---
Traceback (most recent call last):
  File "repro.py", line 23, in <module>
    bwc.bwc_tensor_model_parallel_rank(mpu=sp_module)
  File "/usr/local/lib/python3.12/dist-packages/deepspeed/utils/bwc.py", line 42, in bwc_tensor_model_parallel_rank
    return mpu.get_model_parallel_rank()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'deepspeed.runtime.sequence_parallel.parallel_state_sp' has no attribute 'get_model_parallel_rank'

Root Cause Analysis

In Stage 0 (BF16_Optimizer), the engine calls the bwc (Backward Compatibility) utility to fetch the model parallel rank for gradient scaling. However, bwc.py (Line 42) blindly calls .get_model_parallel_rank() without a safety check.

Because Ulysses partitions along the sequence dimension and not the tensor dimension, this attribute is correctly absent. This creates a conflict whenever Ulysses is combined with the Stage 0 optimizer.

Proposed Fix

Add a defensive hasattr check in deepspeed/utils/bwc.py to support non-tensor parallel configurations:

# deepspeed/utils/bwc.py
def bwc_tensor_model_parallel_rank(mpu=None):
    if mpu and hasattr(mpu, 'get_model_parallel_rank'):
        return mpu.get_model_parallel_rank()
    return 0 

[tohtana]

Thank you for the investigation, @amadhan882. We really appreciate your help.

I confirm that the master branch still doesn't have such a check. I think get_model_parallel_world_size and bwc_tensor_model_parallel_world_size should have the same issue.
Could you open a PR to improve it? Though this issue doesn't happen with master, it would still be good to have such a guard.