Repository navigation
[BUG] Ulysses crashes in Stage 0 #7833
Description
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
Activity
tohtana commentedon Feb 10, 2026
Hi @avicooper1,
Thanks for the detailed report. We re-tested this and can reproduce the exact
get_model_parallel_rankAttributeError onv0.17.5, but we cannot reproduce it on the latest releasev0.18.5(b19987c7) or currentmaster(a44fb58f).Can you try these new versions?
amadhan882 commentedon Feb 13, 2026
Hi @avicooper1 and @tohtana,
I have verified the root cause of the
get_model_parallel_rankcrash by exercising the actual code path indeepspeed/utils/bwc.py.Authentic Reproduction Log (v0.14.2)
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
hasattrcheck indeepspeed/utils/bwc.pyto support non-tensor parallel configurations:tohtana commentedon Feb 17, 2026
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_sizeandbwc_tensor_model_parallel_world_sizeshould 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.
nathon-lee commentedon Mar 20, 2026
@tohtana hi — quick check: since you mentioned this crash doesn’t reproduce on master anymore, do you still want me to proceed with a “guard” PR?
If yes, what behavior do you prefer when
mpudoesn’t provideget_model_parallel_rank/world_size?rank=0andworld_size=1(treat TP as disabled), orAlso, would you rather I implement the guard centrally in
deepspeed/utils/bwc.py, or handle it at the call site (e.g., inbf16_optimizer)?Happy to open a PR following your preferred direction.