[BUG] ZeRO++ sharding small parameter raise IndexError

[wuxibin89]

Describe the bug
Our model has a small parameter with shape torch.Size([32]), when enable ZeRO++, it raise following error:

• world_size: 2048
• zero_hpz_partition_size: 16

  File "/usr/local/lib/python3.8/site-packages/accelerate/accelerator.py", line 1344, in prepare
    result = self._prepare_deepspeed(*args)
  File "/usr/local/lib/python3.8/site-packages/accelerate/accelerator.py", line 1851, in _prepare_deepspeed
    engine, optimizer, _, lr_scheduler = ds_initialize(**kwargs)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/__init__.py", line 181, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 313, in __init__
    self.optimizer = self._configure_zero_optimizer(optimizer=None)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1590, in _configure_zero_optimizer
    optimizer = DeepSpeedZeRoOffload(
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 119, in __init__
    self._convert_to_zero_parameters(ds_config, module, mpu)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 194, in _convert_to_zero_parameters
    Init(module=module,
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1017, in __init__
    self._convert_to_zero_parameters(module.parameters(recurse=True))
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1048, in _convert_to_zero_parameters
    self._zero_init_param(param)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1040, in _zero_init_param
    param.partition()
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1375, in partition
    self._partition(param_list, has_been_updated=has_been_updated)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1523, in _partition
    self._partition_param_sec(param)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1696, in _partition_param_sec
    sec_numel).copy_(one_dim_param.narrow(0, secondary_start, sec_numel))
IndexError: start out of range (expected to be in range of [-32, 32], but got 1920)
...
IndexError: start out of range (expected to be in range of [-32, 32], but got 1152)
IndexError: start out of range (expected to be in range of [-32, 32], but got 640)
IndexError: start out of range (expected to be in range of [-32, 32], but got 128)
IndexError: start out of range (expected to be in range of [-32, 32], but got 384)
IndexError: start out of range (expected to be in range of [-32, 32], but got 896)
IndexError: start out of range (expected to be in range of [-32, 32], but got 1792)
IndexError: start out of range (expected to be in range of [-32, 32], but got 1280)
IndexError: start out of range (expected to be in range of [-32, 32], but got 1536)
IndexError: start out of range (expected to be in range of [-32, 32], but got 1408)
IndexError: start out of range (expected to be in range of [-32, 32], but got 256)
IndexError: start out of range (expected to be in range of [-32, 32], but got 1024)
IndexError: start out of range (expected to be in range of [-32, 32], but got 768)

Here is our deepspeed config

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "bf16": {
        "enabled": "auto"
    },
    "communication_data_type": "fp32",
    "gradient_clipping": 1.0,
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "zero_hpz_partition_size": 16,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": 5e8,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e6,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}

I think the problem is that when partitioning second param copy, the param is aligned to dp_world_size instead zero_hpz_partition_size, and then cause torch.narrow raising IndexError.
https://github.com/microsoft/DeepSpeed/blob/v0.14.5/deepspeed/runtime/zero/partition_parameters.py#L1662

In my case, torch.Size([32]) is aligned to 2048, and secondary_partition_size is 2048/16=128, so one_dim_param.narrow(0, 128, 128) is out of range.

@HeyangQin

[wangyuehu]

hi, have you solved this problem?

[amadeusuzx]

Hi did you find a solution to this?
@GuanhuaWang does the zero++ team have any advice for this issue?

[HaoWu1998]

hi, have you solved this problem? I also encountered this problem when training the qwen3.5 model.
here is the deepspeed config:

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "bf16": {
        "enabled": "auto"
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "none",
            "pin_memory": true
        },
        "offload_param": {
            "device": "none",
            "pin_memory": true
        },
        "zero_hpz_partition_size": 16,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1000000000,
        "reduce_bucket_size": 500000000,
        "zero_quantized_weights": false,
        "zero_quantized_gradients": false,
        "stage3_prefetch_bucket_size": 500000000,
        "stage3_param_persistence_threshold": 100000,
        "stage3_max_live_parameters": 1000000000,
        "stage3_max_reuse_distance": 1000000000,
        "stage3_gather_16bit_weights_on_model_save": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

[PKUWZP]

@HaoWu1998 I am happy to look into the problem, can you show your error messages from your experiments?

[HaoWu1998]

[rank29]: Traceback (most recent call last):
[rank29]: File "/root/code/ms-swift/swift/cli/sft.py", line 20, in
[rank29]: sft_main()
[rank29]: File "/root/code/ms-swift/swift/pipelines/train/sft.py", line 353, in sft_main
[rank29]: return SwiftSft(args).main()
[rank29]: ^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/pipelines/train/sft.py", line 29, in init
[rank29]: self._prepare_model_tokenizer()
[rank29]: File "/root/code/ms-swift/swift/ray/base.py", line 168, in wrapper
[rank29]: return func(self, *args, **kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/pipelines/train/sft.py", line 51, in _prepare_model_tokenizer
[rank29]: self.model, self.processor = args.get_model_processor(**kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/arguments/base_args/base_args.py", line 329, in get_model_processor
[rank29]: return get_model_processor(**res)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/register.py", line 625, in get_model_processor
[rank29]: return loader.load()
[rank29]: ^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/register.py", line 474, in load
[rank29]: model, processor = self._get_model_processor(model_dir, config)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/register.py", line 465, in _get_model_processor
[rank29]: model = self.get_model(model_dir, config, processor, self.model_kwargs.copy())
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/models/qwen.py", line 1454, in get_model
[rank29]: return Qwen2VLLoader.get_model(self, model_dir, config, processor, model_kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/models/qwen.py", line 759, in get_model
[rank29]: model = super().get_model(model_dir, config, processor, model_kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/register.py", line 315, in get_model
[rank29]: model = auto_model_cls.from_pretrained(model_dir, config=config, trust_remote_code=True, **model_kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/root/code/ms-swift/swift/model/patcher.py", line 388, in _new_from_pretrained
[rank29]: model = from_pretrained(cls, *args, **kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/transformers/modeling_utils.py", line 4201, in from_pretrained
[rank29]: model = cls(config, *model_args, **model_kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 506, in wrapper
[rank29]: f(module, *args, **kwargs)
[rank29]: File "/usr/local/lib/python3.12/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py", line 1822, in init
[rank29]: self.model = Qwen3_5Model(config)
[rank29]: ^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 506, in wrapper
[rank29]: f(module, *args, **kwargs)
[rank29]: File "/usr/local/lib/python3.12/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py", line 1329, in init
[rank29]: self.language_model = Qwen3_5TextModel._from_config(config.text_config)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/transformers/modeling_utils.py", line 1565, in _from_config
[rank29]: model = cls(config, **kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 506, in wrapper
[rank29]: f(module, *args, **kwargs)
[rank29]: File "/usr/local/lib/python3.12/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py", line 1227, in init
[rank29]: [Qwen3_5DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 506, in wrapper
[rank29]: f(module, *args, **kwargs)
[rank29]: File "/usr/local/lib/python3.12/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py", line 748, in init
[rank29]: self.linear_attn = Qwen3_5GatedDeltaNet(config, layer_idx)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 516, in wrapper
[rank29]: self._post_init_method(module)
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1081, in _post_init_method
[rank29]: self._zero_init_param(param)
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1040, in _zero_init_param
[rank29]: param.partition()
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1375, in partition
[rank29]: self._partition(param_list, has_been_updated=has_been_updated)
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1523, in _partition
[rank29]: self._partition_param_sec(param)
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
[rank29]: ret_val = func(*args, **kwargs)
[rank29]: ^^^^^^^^^^^^^^^^^^^^^
[rank29]: File "/usr/local/lib/python3.12/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1696, in partition_param_sec
[rank29]: sec_numel).copy(one_dim_param.narrow(0, secondary_start, sec_numel))
[rank29]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank29]: IndexError: start out of range (expected to be in range of [-48, 48], but got 52)