deepspeedai / deepspeedai/DeepSpeed

[BUG] ZeRO-3 partition does not work in Ulysses SP tutorial

Open
#7,458 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug training
Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

Describe the bug
I'm working on a long sequence training on a 32B model, so I need to use both ZeRO-3 and Ulysses sequence parallel features. I follow the tutorial in Arctic Long Sequence Training (ALST) for HF Transformers integration to build my code. However, the ZeRO-3 partition seems not working in this code: after loading the model, the GPU memory usage on all GPUs is the same as when loading the model on a single GPU, meaning the model parameters have not been sharded across the GPUs.

To Reproduce
I modified the code in the tutorial with minimal changes here:

# tmp/meu.py
from deepspeed.runtime.sequence_parallel.ulysses_sp import UlyssesSPAttentionHF, UlyssesSPDataLoaderAdapter
from deepspeed.runtime.utils import move_to_device
from deepspeed.utils import groups
from torch import tensor
from transformers import AutoModelForCausalLM
import deepspeed
import deepspeed.comm as dist
import torch

model_name_or_path = '/public/share/model/Qwen2.5-32B-Instruct'
max_length = 64
sequence_parallel_size = 2
micro_batch_size = 1

config_dict = {
    "train_micro_batch_size_per_gpu": 1,
    "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": 1.000000e+09, 
        "reduce_bucket_size": "auto", 
        "stage3_prefetch_bucket_size": "auto", 
        "stage3_param_persistence_threshold": "auto", 
        "stage3_max_live_parameters": 1.000000e+09, 
        "stage3_max_reuse_distance": 1.000000e+09, 
        "stage3_gather_16bit_weights_on_model_save": True
    }, 
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 1e-3
        }
    },
    "bf16": {
        "enabled": True,
    },
    "sequence_parallel_size": sequence_parallel_size,
}

dtype = torch.bfloat16

# a simple Dataset
# replace with a real dataset but make sure `position_ids` are returned
input_ids = tensor([[1, 10, 10, 10, 2, 2, 3, 4], [1, 20, 20, 20, 2, 2, 3, 5]], )
position_ids = tensor([[0, 1, 2, 3, 4, 5, 6, 7], [0, 1, 2, 3, 4, 5, 6, 7]])
ds = torch.utils.data.TensorDataset(input_ids, position_ids)
def collate_fn(batch):
    input_ids, position_ids = batch[0]
    return dict(input_ids=input_ids.unsqueeze(0),
                position_ids=position_ids.unsqueeze(0),
                labels=input_ids.unsqueeze(0))

dist.init_distributed(dist_backend='nccl', dist_init_required=True)

# Ulysses injection into HF Transformers
mpu = UlyssesSPAttentionHF.register_with_transformers(
    model_name_or_path=model_name_or_path,
    core_attn_implementation="flash_attention_2",
    sequence_parallel_size=sequence_parallel_size,
    max_length=max_length,
    micro_batch_size=micro_batch_size,
    seq_length_is_variable=True,
)

# Deepspeed setup
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, attn_implementation="flash_attention_2")
model, _, _, _ = deepspeed.initialize(config=config_dict,
                                        model=model,
                                        model_parameters=model.parameters(),
                                        mpu=mpu)

# UlyssesSPDataLoaderAdapter injection
sp_group = groups._get_sequence_parallel_group()
sp_world_size = groups._get_sequence_parallel_world_size()
sp_rank = groups._get_sequence_parallel_rank()
dl = torch.utils.data.DataLoader(ds, batch_size=micro_batch_size, collate_fn=collate_fn)
dl = UlyssesSPDataLoaderAdapter(
    dl,
    sp_rank=sp_rank,
    sp_group=sp_group,
    sp_world_size=sp_world_size,
    device=model.device,
)

# Normal training loop
for iter, batch in enumerate(dl):
    batch = move_to_device(batch, model.device)

    outputs = model(**batch)
    # as of this writing HF doesn't calculate loss with shift_labels yet and requires us to do it manually (liger does that automatically)
    shift_labels = batch["shift_labels"]
    loss = model.module.loss_function(
        logits=outputs.logits,
        labels=None,
        shift_labels=shift_labels,
        vocab_size=model.module.config.vocab_size,
    )

    # differentiable weighted per-shard-loss aggregation across ranks
    losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
    # special dealing with SFT that has prompt tokens that aren't used in loss computation
    good_tokens = sum((shift_labels != -100).view(-1))
    good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)
    total_loss = sum(losses_per_rank[rank] * good_tokens_per_rank[rank] for rank in range(sp_world_size))
    total_good_tokens = sum(good_tokens_per_rank)
    loss = total_loss / total_good_tokens

    if dist.get_rank() == 0:
        print(f"{iter}: {loss=}")

    model.backward(loss)

This code can run normally, but the VRAM usage is still high, so I cannot train ~16K sequences using the code.

I also tried adding the deepspeed.zero.Init() context in the code when loading the pretrained model:

# Deepspeed setup
with deepspeed.zero.Init(
    remote_device="cpu",
    pin_memory=True,
    mpu=mpu
):
    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, attn_implementation="flash_attention_2")
model, _, _, _ = deepspeed.initialize(config=config_dict,
                                        model=model,
                                        model_parameters=model.parameters(),
                                        mpu=mpu)

but it will lead to this error:

(modularity) [sjtu_hansenyu@gpu08 modularity2]$ deepspeed --num_gpus 2 tmp/meu.py 
[2025-07-30 22:46:51,824] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-07-30 22:46:55,710] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-07-30 22:46:58,111] [WARNING] [runner.py:220:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
[2025-07-30 22:46:58,111] [INFO] [runner.py:610:main] cmd = /public/home/sjtu_hansenyu/miniconda3/envs/modularity/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMCwgMV19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None tmp/meu.py
[2025-07-30 22:47:00,730] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-07-30 22:47:04,121] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-07-30 22:47:05,954] [INFO] [launch.py:146:main] WORLD INFO DICT: {'localhost': [0, 1]}
[2025-07-30 22:47:05,954] [INFO] [launch.py:152:main] nnodes=1, num_local_procs=2, node_rank=0
[2025-07-30 22:47:05,954] [INFO] [launch.py:163:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0, 1]})
[2025-07-30 22:47:05,954] [INFO] [launch.py:164:main] dist_world_size=2
[2025-07-30 22:47:05,954] [INFO] [launch.py:168:main] Setting CUDA_VISIBLE_DEVICES=0,1
[2025-07-30 22:47:05,955] [INFO] [launch.py:256:main] process 23069 spawned with command: ['/public/home/sjtu_hansenyu/miniconda3/envs/modularity/bin/python', '-u', 'tmp/meu.py', '--local_rank=0']
[2025-07-30 22:47:05,956] [INFO] [launch.py:256:main] process 23070 spawned with command: ['/public/home/sjtu_hansenyu/miniconda3/envs/modularity/bin/python', '-u', 'tmp/meu.py', '--local_rank=1']
[2025-07-30 22:47:08,397] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-07-30 22:47:08,455] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-07-30 22:47:11,485] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-07-30 22:47:11,706] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-07-30 22:47:13,224] [INFO] [comm.py:676:init_distributed] cdb=None
[2025-07-30 22:47:13,224] [INFO] [comm.py:707:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2025-07-30 22:47:13,649] [INFO] [comm.py:676:init_distributed] cdb=None
You are attempting to use Flash Attention 2.0 without specifying a torch dtype. This might lead to unexpected behaviour
[2025-07-30 22:47:14,374] [INFO] [partition_parameters.py:366:__exit__] finished initializing model - num_params = 1, num_elems = 0.78B
[rank0]: Traceback (most recent call last):
[rank0]:   File "/public/home/sjtu_hansenyu/workspace/modularity2/tmp/meu.py", line 81, in <module>
[rank0]:     model = AutoModelForCausalLM.from_pretrained(model_name_or_path, attn_implementation="flash_attention_2")
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 571, in from_pretrained
[rank0]:     return model_class.from_pretrained(
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/modeling_utils.py", line 279, in _wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4342, in from_pretrained
[rank0]:     model = cls(config, *model_args, **model_kwargs)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 529, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 742, in __init__
[rank0]:     self.model = Qwen2Model(config)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 529, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 453, in __init__
[rank0]:     self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 539, in wrapper
[rank0]:     self._post_init_method(module)
[rank0]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1152, in _post_init_method
[rank0]:     param.data = param.data.to(self.local_device)
[rank0]: NotImplementedError: Cannot copy out of meta tensor; no data!
You are attempting to use Flash Attention 2.0 without specifying a torch dtype. This might lead to unexpected behaviour
[rank1]: Traceback (most recent call last):
[rank1]:   File "/public/home/sjtu_hansenyu/workspace/modularity2/tmp/meu.py", line 81, in <module>
[rank1]:     model = AutoModelForCausalLM.from_pretrained(model_name_or_path, attn_implementation="flash_attention_2")
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 571, in from_pretrained
[rank1]:     return model_class.from_pretrained(
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/modeling_utils.py", line 279, in _wrapper
[rank1]:     return func(*args, **kwargs)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4342, in from_pretrained
[rank1]:     model = cls(config, *model_args, **model_kwargs)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 529, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 742, in __init__
[rank1]:     self.model = Qwen2Model(config)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 529, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 453, in __init__
[rank1]:     self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 539, in wrapper
[rank1]:     self._post_init_method(module)
[rank1]:   File "/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1152, in _post_init_method
[rank1]:     param.data = param.data.to(self.local_device)
[rank1]: NotImplementedError: Cannot copy out of meta tensor; no data!
[rank0]:[W730 22:47:14.018809811 ProcessGroupNCCL.cpp:1496] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
[2025-07-30 22:47:16,967] [INFO] [launch.py:319:sigkill_handler] Killing subprocess 23069
[2025-07-30 22:47:16,968] [INFO] [launch.py:319:sigkill_handler] Killing subprocess 23070
[2025-07-30 22:47:17,007] [ERROR] [launch.py:325:sigkill_handler] ['/public/home/sjtu_hansenyu/miniconda3/envs/modularity/bin/python', '-u', 'tmp/meu.py', '--local_rank=1'] exits with return code = 1

Expected behavior
Model parameters correctly partitioned across GPUs.

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
--------------------------------------------------
async_io ............... [NO] ....... [OKAY]
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]
fp_quantizer ........... [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `dlvsym'
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `dlopen'
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `dlclose'
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `dlerror'
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `dlsym'
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `shm_open'
/public/software/compiler/gnu/11.3.0/bin/../lib/gcc/x86_64-pc-linux-gnu/11/../../../../x86_64-pc-linux-gnu/bin/ld: /public/home/sjtu_hansenyu/.local/cuda/lib64/libcufile.so: undefined reference to `shm_unlink'
collect2: error: ld returned 1 exit status
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.6
 [WARNING]  using untested triton version (3.2.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 ............... ['/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/torch']
torch version .................... 2.6.0+cu124
deepspeed install path ........... ['/public/home/sjtu_hansenyu/miniconda3/envs/modularity/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.17.2, unknown, unknown
torch cuda version ............... 12.4
torch hip version ................ None
nvcc version ..................... 12.2
deepspeed wheel compiled w. ...... torch 2.4, cuda 12.1
shared memory (/dev/shm) size .... 1007.64 GB

System info (please complete the following information):

  • OS: CentOS Linux 7
  • GPU count and types: single node, NVIDIA A800 x8
  • Python version: 3.10.15

Launcher context
deepspeed launching

Docker context
None

Additional context
None

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with tmp/meu.py and the Ulysses ALST tutorial, comparing the normal model-loading path with the deepspeed.zero.Init() path. Read the ZeRO partitioning traceback around partition_parameters.py and the Transformers Qwen2 model construction. Done means the reproduction initializes successfully and ZeRO-3 parameters are sharded so GPU memory is reduced across ranks.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.