RLHF SAPO KL=nan

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

I've encountered an issue where during the reinforcement learning phase of training, the KL divergence consistently becomes NaN after a certain number of steps. I followed others' suggestions and set --overlong_filter to false, yet the error still persists. Could you please help me identify where the problem lies? Below are my script and the error details for your reference.
`#!/bin/bash

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export MIN_PIXELS=3136
export MAX_PIXELS=640000
export PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True'
NNODES=5 NODE_RANK=0 MASTER_ADDR=6.5.80.12 MASTER_PORT=29563 NPROC_PER_NODE=8
swift rlhf
--rlhf_type grpo
--model /media/cfs/ea-cvfa-multimodal-reasoning-2/ygq/tmp/V3_stage3.2_SFT/v11-20260315-014930/checkpoint-20-0315
--max_length 16384
--external_plugins /media/cfs/ea-cvfa-multimodal-reasoning-2/ygq/0_work/train_RL/plugin/reward_emb_7b_transfer_sft_base.py
--reward_funcs answer_embedding_accuracy think_embedding_accuracy format
--reward_weights 0.5 0.3 0.2
--use_vllm false
--train_type full
--torch_dtype bfloat16
--dataset /media/cfs/ea-cvfa-multimodal-reasoning-2/ygq/0_work/train_RL/data/thinking_data_RL_0314.jsonl
--max_completion_length 2048
--num_train_epochs 3
--per_device_train_batch_size 16
--per_device_eval_batch_size 16
--learning_rate 1e-6
--gradient_accumulation_steps 2
--steps_per_generation 4
--save_strategy 'steps'
--eval_strategy 'steps'
--eval_steps 10
--save_steps 10
--save_total_limit 100
--logging_steps 1
--output_dir /media/cfs/ea-cvfa-multimodal-reasoning-2/ygq/tmp/V3_stage3.2_RL
--warmup_ratio 0.01
--dataloader_num_workers 64
--num_generations 8
--temperature 1.0
--system "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process here answer here "
--deepspeed zero3
--log_completions true
--num_iterations 1
--async_generate false
--save_only_model false
--freeze_vit true
--freeze_aligner false
--beta 0.001
--attn_impl flash_attn
--overlong_filter false
--truncation_strategy delete
--loss_type sapo
--sequence_parallel_size 4
--soft_max_length 2048
--soft_cache_length 1024
--repetition_penalty 1.1
--max_grad_norm 1.0
# --resume_from_checkpoint /media/cfs/ea-cvfa-multimodal-reasoning-2/ygq/tmp/V3_stage3.2_RL/v13-20260313-011335/checkpoint-25
# --tau_pos 1.0
# --tau_neg 1.05
# --do_train
# --do_eval
# --predict_with_generate
`

error:
`
Train: 14%|█▎ | 47/342 [9:21:59<65:18:54, 797.07s/it]

{'loss': -0.02671666, 'grad_norm': 1.41421356, 'learning_rate': 9.6e-07, 'completions/mean_length': 334.6781311, 'completions/min_length': 122.0, 'completions/max_length': 2048.0, 'completions/clipped_ratio': 0.00078125, 'reward': 0.93877602, 'reward_std': 0.01329423, 'frac_reward_zero_std': 0.040625, 'rewards/AnswerEmbeddingAccuracy/mean': 0.93841827, 'rewards/AnswerEmbeddingAccuracy/std': 0.07926226, 'rewards/ThinkEmbeddingAccuracy/mean': 0.89855623, 'rewards/ThinkEmbeddingAccuracy/std': 0.01878368, 'rewards/Format/mean': 1.0, 'rewards/Format/std': 0.0, 'kl': 0.06128701, 'epoch': 0.41, 'global_step/max_steps': '47/342', 'percentage': '13.74%', 'elapsed_time': '9h 22m 0s', 'remaining_time': '2d 10h 47m 27s', 'memory(GiB)': 121.48, 'train_speed(iter/s)': 0.001394}

Train: 14%|█▎ | 47/342 [9:22:00<65:18:54, 797.07s/it]
Train: 14%|█▎ | 47/342 [9:22:00<65:18:54, 797.07s/it]
Train: 14%|█▍ | 48/342 [9:23:34<47:53:28, 586.42s/it]

{'loss': 0.0, 'grad_norm': 1.41421356, 'learning_rate': 9.6e-07, 'kl': nan, 'epoch': 0.42, 'global_step/max_steps': '48/342', 'percentage': '14.04%', 'elapsed_time': '9h 23m 34s', 'remaining_time': '2d 9h 31m 56s', 'memory(GiB)': 121.48, 'train_speed(iter/s)': 0.001419}

Train: 14%|█▍ | 48/342 [9:23:34<47:53:28, 586.42s/it]
Train: 14%|█▍ | 48/342 [9:23:34<47:53:28, 586.42s/it]/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:112: _assert_async_cuda_kernel: block: [0,0,0], thread: [0,0,0] Assertion probability tensor contains either inf, nan or element < 0 failed.`

How to Reproduce / 如何复现

ms-swift 3.12.dev, B200

Additional Information / 补充信息

No response

Start with the supplied swift rlhf command, especially --loss_type sapo, and reproduce the transition from step 47 to 48 on ms-swift 3.12.dev. Trace the SAPO KL calculation and the CUDA probability assertion using the provided configuration. Done means identifying the NaN trigger and confirming training no longer emits NaN in this reproduction.