[Bug] Abnormal `grad_norm` in GRPO LoRA training with DeepSpeed ZeRO-0 (Normal with ZeRO-3)

Describe the bug
I encountered an issue where the grad_norm becomes abnormally large (e.g., ~1656) when performing GRPO training with LoRA using DeepSpeed ZeRO-0. However, when switching to DeepSpeed ZeRO-3, the grad_norm returns to a normal range (e.g., ~0.025).

This behavior suggests a potential issue with gradient synchronization or calculation specific to the ZeRO stage configuration in the GRPO LoRA pipeline.

Your hardware and system info

• Hardware: 4x A800
• ms-swift Version: 3.10.2
• DeepSpeed Version: Verified on both 0.17.6 and 0.18.x
• Model: Qwen3-8B
• Training Type: LoRA (Rank 16)
• Task: RLHF (GRPO)

Reproduction Steps

1. Model: Qwen3-8B
2. Training Method: GRPO + LoRA (Rank 16)
3. Loss Type: BNPO
4. Configuration:
   • Case A (Bug): Run the training command with deepspeed configured to ZeRO-0.
   • Case B (Normal): Run the same command with --deepspeed zero3.

Reproduction Script:

# 4x 80G GPUs
CUDA_VISIBLE_DEVICES=0,1,2,3 \
NPROC_PER_NODE=4 \
swift rlhf \
    --rlhf_type grpo \
    --model_type qwen3 \
    --model <path_to_model> \
    --dataset <path_to_dataset> \
    --external_plugins <plugins> \
    --reward_funcs <funcs> \
    --overlong_filter true \
    --max_completion_length 2048 \
    --learning_rate 1e-6 \
    --torch_dtype bfloat16 \
    --num_train_epochs 1 \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 4 \
    --save_total_limit 32 \
    --save_steps 50 \
    --logging_steps 2 \
    --output_dir output_retro_multi_grpo_3 \
    --gradient_accumulation_steps 32 \
    --warmup_steps 50 \
    --dataloader_num_workers 4 \
    --num_generations 8 \
    --use_vllm true \
    --vllm_gpu_memory_utilization 0.3 \
    --vllm_max_model_len 2560 \
    --sleep_level 0 \
    --offload_model false \
    --offload_optimizer false \
    --gc_collect_after_offload true \
    --tensor_parallel_size 1 \
    --temperature 0.8 \
    --top_p 0.85 \
    --log_completions true \
    --loss_type bnpo \
    --epsilon_high 0.28 \
    --dynamic_sample true \
    --max_resample_times 4 \
    --deepspeed zero0 \
    --train_type lora \
    --target_modules all-linear \
    --lora_rank 16

(Note: Change --deepspeed to zero3 for normal behavior)

Logs Comparison

Abnormal Log (DeepSpeed ZeRO-0):
Note the grad_norm: 1656.27

{
  "loss": 0.01327927,
  "grad_norm": 1656.27481113,
  "learning_rate": 1e-07,
  "completions/mean_length": 830.90234375,
  "completions/min_length": 544.0,
  "completions/max_length": 1752.0,
  "completions/clipped_ratio": 0.0,
  "reward": 1.45605469,
  "reward_std": 0.4302305,
  "frac_reward_zero_std": 0.0,
  "rewards/RetroReward/mean": 1.45605469,
  "rewards/RetroReward/std": 0.51162845,
  "kl": 0.0,
  "clip_ratio/low_mean": 0.0,
  "clip_ratio/low_min": 0.0,
  "clip_ratio/high_mean": 0.0,
  "clip_ratio/high_max": 0.0,
  "clip_ratio/region_mean": 0.0,
  "epoch": 6.895e-05,
  "global_step/max_steps": "1/14503",
  "percentage": "0.01%",
  "elapsed_time": "2m 55s",
  "remaining_time": "29d 10h 53m 37s",
  "memory(GiB)": 52.96,
  "train_speed(iter/s)": 0.005699
}

Normal Log (DeepSpeed ZeRO-3):
Note the grad_norm: 0.025

{
  "loss": 0.01327927,
  "grad_norm": 0.02527279,
  "learning_rate": 1e-07,
  "completions/mean_length": 830.90234375,
  "completions/min_length": 544.0,
  "completions/max_length": 1752.0,
  "completions/clipped_ratio": 0.0,
  "reward": 1.45605469,
  "reward_std": 0.4302305,
  "frac_reward_zero_std": 0.0,
  "rewards/RetroReward/mean": 1.45605469,
  "rewards/RetroReward/std": 0.51162845,
  "kl": 0.0,
  "clip_ratio/low_mean": 0.0,
  "clip_ratio/low_min": 0.0,
  "clip_ratio/high_mean": 0.0,
  "clip_ratio/high_max": 0.0,
  "clip_ratio/region_mean": 0.0,
  "epoch": 6.895e-05,
  "global_step/max_steps": "1/14503",
  "percentage": "0.01%",
  "elapsed_time": "3m 20s",
  "remaining_time": "33d 16h 7m 29s",
  "memory(GiB)": 52.23,
  "train_speed(iter/s)": 0.004985
}

Additional context

• The issue seems specifically tied to the DeepSpeed stage configuration.
• The loss values are identical in the logs provided, but the gradient norm differs significantly, indicating that the forward pass might be consistent, but the backward pass or gradient aggregation is behaving differently under ZeRO-0.

Start by reproducing the reported swift rlhf command with --deepspeed zero0 and then zero3, comparing the gradient-norm calculation and aggregation in the GRPO LoRA training path. Done means identifying and correcting the ZeRO-0 discrepancy so the reported norm is calculated consistently, with a regression test covering both configurations.