Fix Stage 0 + Ulysses crash: make bwc_tensor_model_parallel_rank() resilient to MP API absence

[nathon-lee]

Title

Fix Stage 0 + Ulysses crash: make bwc_tensor_model_parallel_rank() resilient to MP API absence

Summary

This PR fixes a hard crash when using Ulysses sequence parallelism with ZeRO Stage 0 (BF16_Optimizer).
In this configuration, DeepSpeed calls deepspeed.utils.bwc.bwc_tensor_model_parallel_rank(mpu=...), and the passed mpu object can be deepspeed.runtime.sequence_parallel.parallel_state_sp, which does not implement the deprecated get_model_parallel_rank() API. The current fallback path unconditionally calls mpu.get_model_parallel_rank(), raising AttributeError.

The fix adds a defensive capability check before calling the deprecated API. If the provided mpu does not expose any known tensor/model-parallel rank API, we treat it as “no tensor model parallelism” and return rank 0.

Motivation / Context

• Affected scenario: Ulysses sequence parallel + ZeRO Stage 0
• Failure mode: AttributeError: ... parallel_state_sp has no attribute get_model_parallel_rank
• Root cause: bwc_tensor_model_parallel_rank() falls back to a deprecated API without an hasattr() check.

This change keeps the original priority order intact:

1. get_tensor_model_parallel_rank()
2. get_slice_parallel_rank()
3. get_model_parallel_rank() (deprecated)
4. fallback to 0 if none exist

Changes

• deepspeed/utils/bwc.py
  • Update bwc_tensor_model_parallel_rank() to check hasattr(mpu, "get_model_parallel_rank") before calling it.
  • If mpu provides none of the expected tensor/model-parallel rank APIs, return 0 (no TP).

Why this is safe

• For Megatron / DeepSpeed Topology / any existing MPU that already implements get_tensor_model_parallel_rank() or get_slice_parallel_rank() or get_model_parallel_rank(), behavior is unchanged.
• The new code path only affects the previously-crashing case where the mpu object does not provide any of these methods.

Reproduction

Using the Ulysses ALST tutorial flow, switching ZeRO stage from 3 to 0 triggers the crash during optimizer step when grad norm is computed.

Testing

• Existing unit tests should continue to pass.
• Minimal repro: calling bwc_tensor_model_parallel_rank(mpu=deepspeed.runtime.sequence_parallel.parallel_state_sp) should no longer raise.

References

• DeepSpeed Issue: [BUG] Ulysses crashes in Stage 0 #7833 “Ulysses crashes in Stage 0”

[tohtana]

Hi @nathon-lee,
Thank you for reporting!

I found that we already have a fallback from get_model_parallel_world_size to get_sequence_parallel_world_size.
This was introduced in #7649. Can you make sure that the latest version still raises the error?

[nathon-lee]

Hi @tohtana, thanks for looking into this!

I double-checked master and I think #7649 addresses a different fallback path than the one crashing in Stage 0.

• The crash in [BUG] Ulysses crashes in Stage 0 #7833 happens in deepspeed/utils/bwc.py:bwc_tensor_model_parallel_rank() (called from BF16_Optimizer.get_grads_for_norm() in ZeRO Stage 0).

• In current master, bwc_tensor_model_parallel_rank() falls back to the deprecated API unconditionally:

  else: return mpu.get_model_parallel_rank()

  When mpu is deepspeed.runtime.sequence_parallel.parallel_state_sp (Ulysses), that module does not implement get_model_parallel_rank(), so we still hit:

  AttributeError: ... parallel_state_sp has no attribute 'get_model_parallel_rank'

My understanding is that #7649 added/adjusted APIs and fallbacks around world size (e.g., get_model_parallel_world_size -> get_sequence_parallel_world_size), but it doesn't protect the rank fallback above.

That’s why I opened #7888: it makes bwc_tensor_model_parallel_rank() defensive (checks for get_model_parallel_rank before calling it, and otherwise treats it as “no tensor MP” and returns rank 0).

If you have a preferred behavior (e.g., should we fallback to a sequence-parallel rank instead of 0?), I’m happy to adjust the PR.

[tohtana]

Hi @nathon-lee,
Thank you for the reply. I still don't see why this issue happens with the latest master.
Can you share a small repro?