[BUG]npu zero3 训练自定义模型时，报错Function SumBackward0 returned an invalid gradient at index 0
Author: Zane-Qbb
Pull request: no
Labels: name: bug

### System Info

**Your hardware and system info**
npu: 晟腾910B4
pytroch/pytroch-npu==2.1.0
transformers==4.46.2 （4.49.*也不行）
accelerate==1.1.1
deepspeed==0.14.*  (0.15.*也不行)
操作系统: pytorch_2.1.0-cann_8.0.rc2-py_3.9-euler_2.10.7:0.0.1


**Additional context**
尝试了限制自定义网络结构部分(group_decode_adapter)不做切片，还是会报错。训练过程打印loss值均正常，require_grads均符合预期（llm, group_decode_adapter等均没有冻结）

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

**Describe the bug**
在910B上，基于Qwen7B做SLAM-omni语音模型的复现，需要使用deepspeed的zero3来启动训练

1.zero3启动训练后报错。
2.zero3启动训练时，如果冻结新增的网络结构部分(group_decode_adapter)，可以正常训练。不冻结则会报错
3.zero2启动训练是正常的，但由于显存限制只能训3B的模型。所以需要使用zero3

**报错信息：**

```
Traceback (most recent call last):
  File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/cli/sft.py", line 5, in <module>
    sft_main()
  File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 268, in sft_main
    return SwiftSft(args).main()
  File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/base.py", line 46, in main
    result = self.run()
  File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 141, in run
    return self.train(trainer)
  File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 193, in train
    trainer.train(trainer.args.resume_from_checkpoint)
  File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/trainers/mixin.py", line 261, in train
    res = super().train(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 2241, in train
    return inner_training_loop(
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 3747, in training_step
    self.accelerator.backward(loss, **kwargs)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/accelerate/accelerator.py", line 2233, in backward
    self.deepspeed_engine_wrapped.backward(loss, **kwargs)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/accelerate/utils/deepspeed.py", line 186, in backward
    self.engine.backward(loss, **kwargs)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2020, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 2259, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/torch/autograd/__init__.py", line 256, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Function SumBackward0 returned an invalid gradient at index 0 - got [0] but expected shape compatible with [12480]
```

**启动会报错的zero3启动脚本：**

```
ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
NPROC_PER_NODE=8 \
swift sft \
--model_type=slam_omni \
--model=/Qwen2.5-3B-for-omni \
--dataset=xxxx.json \
--num_train_epochs=26 \
--train_type=full \
--output_dir=xxxx \
--eval_steps=200000000000 \
--save_steps=2000 \
--device=npu \
--ddp_backend hccl \
--per_device_train_batch_size=1 \
--dataloader_num_workers=5 \
--lazy_tokenize true \
--torch_dtype=float32 \
--check_model=false \
--max_length=3072 \
--learning_rate=1e-4 \
--warmup_steps=1000 \
--lr_scheduler_type=cosine \
--deepspeed=zero3  >log.txt 2>&1 &
```

**基于Qwen扩展新增的网络结构部分：**

```
class Linear_GroupDecodeAdapter(nn.Module):
    def __init__(self, audio_vocab_size, code_layer):
        super(Linear_GroupDecodeAdapter, self).__init__()
        self.audio_vocab_size = audio_vocab_size
        self.code_layer = code_layer
        with Init(data_parallel_group=None, enabled=False):  # 不进行切片
            self.linear = nn.Linear(audio_vocab_size, code_layer * audio_vocab_size)

    def forward(self, logits):
        logits = self.linear(logits)
        return logits
```


**新增部分结构在forward中的使用：**
```
...
logits_a = self.group_decode_adapter(logits[..., self.model_config.vocab_config.padded_text_vocabsize:])
logits_a = torch.split(
    logits_a,
    [
        self.model_config.vocab_config.padded_audio_vocabsize,
        self.model_config.vocab_config.padded_audio_vocabsize,
        self.model_config.vocab_config.padded_audio_vocabsize
    ],
    dim=2
)
logits_a = tuple(tensor.unsqueeze(1) for tensor in logits_a)
logits_a = torch.cat(logits_a, dim=1)
loss_a = slam_loss(logits=logits_a, labels=labels_a) * 3
...
```


### Expected behavior

不报错，正常训练

Comments:
Hi @Zane-Qbb! Can you:

1) Translate the issue
2) Give us a sample file that will let us reproduce the issue?

> Hi [@Zane-Qbb](https://github.com/Zane-Qbb)! Can you:
> 
> 1. Translate the issue
> 2. Give us a sample file that will let us reproduce the issue?

Thank you for your reply. The issue has been identified and resolved. However, to provide a reference for others who may encounter similar problems, I will translate the issue into English:



System Info
Your hardware and system info
NPU: Ascend 910B4
PyTorch/PyTorch-NPU==2.1.0
Transformers==4.46.2 (4.49.* also doesn't work)
Accelerate==1.1.1
DeepSpeed==0.14.* (0.15.* also doesn't work)
Operating System: pytorch_2.1.0-cann_8.0.rc2-py_3.9-euler_2.10.7:0.0.1

Additional context
Tried to restrict the custom network structure part (group_decode_adapter) from slicing, but it still reports an error. The loss values printed during training are normal, and require_grads are as expected (llm, group_decode_adapter, etc. are not frozen).

Who can help?
No response

Information

The official example scripts

My own modified scripts
Tasks

An officially supported task in the examples folder (such as GLUE/SQuAD, ...)

My own task or dataset (give details below)
Reproduction
Describe the bug
On the 910B, based on Qwen7B to reproduce the SLAM-omni speech model, it is necessary to use DeepSpeed's zero3 to start training.

Zero3 reports an error after starting training.
When starting training with zero3, if the newly added network structure part (group_decode_adapter) is frozen, training can proceed normally. If not frozen, it reports an error.
Zero2 training starts normally, but due to memory limitations, only the 3B model can be trained. Therefore, zero3 is needed.
Error message:

Traceback (most recent call last):
File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/cli/sft.py", line 5, in <module>
sft_main()
File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 268, in sft_main
return SwiftSft(args).main()
File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/base.py", line 46, in main
result = self.run()
File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 141, in run
return self.train(trainer)
File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 193, in train
trainer.train(trainer.args.resume_from_checkpoint)
File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/trainers/mixin.py", line 261, in train
res = super().train(*args, **kwargs)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 2241, in train
return inner_training_loop(
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 3747, in training_step
self.accelerator.backward(loss, **kwargs)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/accelerate/accelerator.py", line 2233, in backward
self.deepspeed_engine_wrapped.backward(loss, **kwargs)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/accelerate/utils/deepspeed.py", line 186, in backward
self.engine.backward(loss, **kwargs)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2020, in backward
self.optimizer.backward(loss, retain_graph=retain_graph)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 2259, in backward
self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
scaled_loss.backward(retain_graph=retain_graph)
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/torch/_tensor.py", line 492, in backward
torch.autograd.backward(
File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/torch/autograd/init.py", line 256, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: Function SumBackward0 returned an invalid gradient at index 0 - got [0] but expected shape compatible with [12480]
The zero3 startup script that reports an error:

ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
NPROC_PER_NODE=8
swift sft
--model_type=slam_omni
--model=/Qwen2.5-3B-for-omni
--dataset=xxxx.json
--num_train_epochs=26
--train_type=full
--output_dir=xxxx
--eval_steps=200000000000
--save_steps=2000
--device=npu
--ddp_backend hccl
--per_device_train_batch_size=1
--dataloader_num_workers=5
--lazy_tokenize true
--torch_dtype=float32
--check_model=false
--max_length=3072
--learning_rate=1e-4
--warmup_steps=1000
--lr_scheduler_type=cosine
--deepspeed=zero3 >log.txt 2>&1 &
The newly added network structure part based on Qwen:

class Linear_GroupDecodeAdapter(nn.Module):
def init(self, audio_vocab_size, code_layer):
super(Linear_GroupDecodeAdapter, self).init()
self.audio_vocab_size = audio_vocab_size
self.code_layer = code_layer
with Init(data_parallel_group=None, enabled=False): # No slicing
self.linear = nn.Linear(audio_vocab_size, code_layer * audio_vocab_size)

def forward(self, logits):
    logits = self.linear(logits)
    return logits
Usage of the newly added structure in forward:

...
loglogits_a = self.group_decode_adapter(logits[..., self.model_config.vocab_config.padded_text_vocabsize:])
logits_a = torch.split(
logits_a,
[
self.model_config.vocab_config.padded_audio_vocabsize,
self.model_config.vocab_config.padded_audio_vocabsize,
self.model_config.vocab_config.padded_audio_vocabsize
],
dim=2
)
logits_a = tuple(tensor.unsqueeze(1) for tensor in logits_a)
logits_a = torch.cat(logits_a, dim=1)
loss_a = slam_loss(logits=logits_a, labels=labels_a) * 3
...
Expected behavior
No errors, normal training

> > Hi [@Zane-Qbb](https://github.com/Zane-Qbb)! Can you:
> > 
> > 1. Translate the issue
> > 2. Give us a sample file that will let us reproduce the issue?
> 
> Thank you for your reply. The issue has been identified and resolved. However, to provide a reference for others who may encounter similar problems, I will translate the issue into English:
> 
> System Info Your hardware and system info NPU: Ascend 910B4 PyTorch/PyTorch-NPU==2.1.0 Transformers==4.46.2 (4.49.* also doesn't work) Accelerate==1.1.1 DeepSpeed==0.14.* (0.15.* also doesn't work) Operating System: pytorch_2.1.0-cann_8.0.rc2-py_3.9-euler_2.10.7:0.0.1
> 
> Additional context Tried to restrict the custom network structure part (group_decode_adapter) from slicing, but it still reports an error. The loss values printed during training are normal, and require_grads are as expected (llm, group_decode_adapter, etc. are not frozen).
> 
> Who can help? No response
> 
> Information
> 
> The official example scripts
> 
> My own modified scripts Tasks
> 
> An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
> 
> My own task or dataset (give details below) Reproduction Describe the bug On the 910B, based on Qwen7B to reproduce the SLAM-omni speech model, it is necessary to use DeepSpeed's zero3 to start training.
> 
> Zero3 reports an error after starting training. When starting training with zero3, if the newly added network structure part (group_decode_adapter) is frozen, training can proceed normally. If not frozen, it reports an error. Zero2 training starts normally, but due to memory limitations, only the 3B model can be trained. Therefore, zero3 is needed. Error message:
> 
> Traceback (most recent call last): File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/cli/sft.py", line 5, in sft_main() File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 268, in sft_main return SwiftSft(args).main() File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/base.py", line 46, in main result = self.run() File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 141, in run return self.train(trainer) File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/llm/train/sft.py", line 193, in train trainer.train(trainer.args.resume_from_checkpoint) File "/home/ma-user/modelarts/user-job-dir/ms-swift-master/swift/trainers/mixin.py", line 261, in train res = super().train(*args, **kwargs) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 2241, in train return inner_training_loop( File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 2548, in _inner_training_loop tr_loss_step = self.training_step(model, inputs, num_items_in_batch) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/transformers/trainer.py", line 3747, in training_step self.accelerator.backward(loss, **kwargs) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/accelerate/accelerator.py", line 2233, in backward self.deepspeed_engine_wrapped.backward(loss, **kwargs) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/accelerate/utils/deepspeed.py", line 186, in backward self.engine.backward(loss, **kwargs) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn ret_val = func(*args, **kwargs) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2020, in backward self.optimizer.backward(loss, retain_graph=retain_graph) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn ret_val = func(*args, **kwargs) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 2259, in backward self.loss_scaler.backward(loss.float(), retain_graph=retain_graph) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward scaled_loss.backward(retain_graph=retain_graph) File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/torch/_tensor.py", line 492, in backward torch.autograd.backward( File "/home/ma-user/anaconda3/envs/PyTorch-2.1.0/lib/python3.9/site-packages/torch/autograd/init.py", line 256, in backward Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass RuntimeError: Function SumBackward0 returned an invalid gradient at index 0 - got [0] but expected shape compatible with [12480] The zero3 startup script that reports an error:
> 
> ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 NPROC_PER_NODE=8 swift sft --model_type=slam_omni --model=/Qwen2.5-3B-for-omni --dataset=xxxx.json --num_train_epochs=26 --train_type=full --output_dir=xxxx --eval_steps=200000000000 --save_steps=2000 --device=npu --ddp_backend hccl --per_device_train_batch_size=1 --dataloader_num_workers=5 --lazy_tokenize true --torch_dtype=float32 --check_model=false --max_length=3072 --learning_rate=1e-4 --warmup_steps=1000 --lr_scheduler_type=cosine --deepspeed=zero3 >log.txt 2>&1 & The newly added network structure part based on Qwen:
> 
> class Linear_GroupDecodeAdapter(nn.Module): def init(self, audio_vocab_size, code_layer): super(Linear_GroupDecodeAdapter, self).init() self.audio_vocab_size = audio_vocab_size self.code_layer = code_layer with Init(data_parallel_group=None, enabled=False): # No slicing self.linear = nn.Linear(audio_vocab_size, code_layer * audio_vocab_size)
> 
> def forward(self, logits): logits = self.linear(logits) return logits Usage of the newly added structure in forward:
> 
> ... loglogits_a = self.group_decode_adapter(logits[..., self.model_config.vocab_config.padded_text_vocabsize:]) logits_a = torch.split( logits_a, [ self.model_config.vocab_config.padded_audio_vocabsize, self.model_config.vocab_config.padded_audio_vocabsize, self.model_config.vocab_config.padded_audio_vocabsize ], dim=2 ) logits_a = tuple(tensor.unsqueeze(1) for tensor in logits_a) logits_a = torch.cat(logits_a, dim=1) loss_a = slam_loss(logits=logits_a, labels=labels_a) * 3 ... Expected behavior No errors, normal training



**Solved because the customized forward is written:

dummy_loss = sum(p.sum() for p in self.group_decode_adapter.parameters()) * 0.0
loss = loss + dummy_loss

This calculation is intended to enable all structures of group_decode_adapter to participate in loss calculation when some structures of group_decode_adapter are not used. In this way, backpropagation is normal during distributed training. This operation does not seem to be required under zero3.**