[BUG] DeepSpeed Ulysses zero3 compatibility

Describe the bug
Training a hf model (llama 3.1 with peft) on long context with sequence_parallel_size > 1 works only up until zero stage 2.
If I set "stage" to 3 I get the following error:

[rank1]:   File "/root/miniconda3/envs/finetuning/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py", line 1464, in partition_grads
[rank1]:     grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel())
[rank1]: RuntimeError: start (0) + length (8388608) exceeds dimension size (4194304).

I also had to disable this assertion when switching over from zero 1 to 3:

assert train_batch == micro_batch * grad_acc * self.world_size

So maybe there is an issue with the world_size definition when running zero3 (though even when fixing this to the correct world size and device_mesh the same error occurs)?

To Reproduce
Running the example from:
DeepSpeedExamples/post_training/sequence_parallelism/test_ulysses.py
with:

 "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        },
        "overlap_comm": True,
        "contiguous_gradients": True,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": True
    },

on the hf pr: https://github.com/huggingface/transformers/pull/32305

Expected behavior
ZeRo-3 should work as stated in the official blog post.

ds_report output

DeepSpeed general environment info:
torch install path ............... ['/root/miniconda3/envs/finetuning/lib/python3.10/site-packages/torch']
torch version .................... 2.4.1+cu121
deepspeed install path ........... ['/root/miniconda3/envs/finetuning/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.15.1, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 2.4, cuda 12.1
shared memory (/dev/shm) size .... 321.31 GB

System info:

• OS: [e.g. Ubuntu 18.04]
• GPU count and types [e.g. two machines with x8 A100s each]
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version
• Any other relevant info about your setup

Launcher context
I am using the deepspeed launcher.

Thanks for the help!
Even if this not officially supported I would be thankful for some pointers, so I can implement something on my own.
For context:
We want to train a 70B model on seq length of 60k. 8B already works with Ulysses, but without zero-3 I think 70B is impossible on a single node.

Reproduce the failure with DeepSpeedExamples/post_training/sequence_parallelism/test_ulysses.py using the configuration and Hugging Face PR linked in the issue. Start from deepspeed/runtime/zero/stage3.py at partition_grads and compare the world_size and device_mesh assumptions; done means ZeRO-3 runs the Ulysses example without the grad partition-size error.