token_mult_prob_error continuously increases during GRPO training with DeepSeek-R1-0528-Qwen3-8B using Megatron backend, leading to training crash
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
Describe the bug
When training DeepSeek-R1-0528-Qwen3-8B model using GRPO with Megatron backend (TP=4), the token_mult_prob_error metric continuously increases as training progresses, eventually causing the training to crash. This indicates a growing discrepancy between the inference engine (vLLM) and the training engine (Megatron) when processing the same samples.
According to the documentation, token_mult_prob_error should stay close to 1.0, and values exceeding 1.05 indicate potential framework issues. In my case, this metric keeps growing throughout training until the run becomes unstable and crashes.
Note: This is a metric degradation issue observed via WandB monitoring, not a runtime error log.
Steps/Code to reproduce bug
Quick reproduction from checkpoint:
You can reproduce this issue by resuming training from the following checkpoint, which will quickly show the increasing token_mult_prob_error:
Checkpoint path: /lustre/fsw/portfolios/coreai/users/bihu/RL/results/grpo-dapo-deepseek-r1-8B-8n64g-megatron-tp4-lr1e-6/step_450
Full reproduction paths:
Project directory: /lustre/fsw/portfolios/coreai/users/bihu/RL
Launch script: /lustre/fsw/portfolios/coreai/users/bihu/RL/run_grpo_dapo_deepseek_r1_8B.sh
Config file: /lustre/fsw/portfolios/coreai/users/bihu/RL/examples/configs/grpo_dapo_deepseek_r1_8B_8n64g_megatron.yaml.bak_no_fp32logits
Configuration file content (grpo_dapo_deepseek_r1_8B_8n64g_megatron.yaml.bak_no_fp32logits):
defaults: "grpo_math_1B.yaml"grpo: num_prompts_per_step: 32 num_generations_per_prompt: 8loss_fn: reference_policy_kl_penalty: 0.0 ratio_clip_min: 0.2 ratio_clip_max: 0.28 use_importance_sampling_correction: truepolicy: model_name: "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B" train_global_batch_size: 64 train_micro_batch_size: 1 logprob_batch_size: 1 max_total_sequence_length: 16384 dtensor_cfg: enabled: false megatron_cfg: enabled: true converter_type: "Qwen3ForCausalLM" tensor_model_parallel_size: 4 pipeline_model_parallel_size: 1 context_parallel_size: 1 sequence_parallel: false activation_checkpointing: true optimizer: optimizer: "adam" lr: 1.0e-6 min_lr: 1.0e-6 weight_decay: 0.01 bf16: true fp16: false params_dtype: "float32" adam_beta1: 0.9 adam_beta2: 0.999 adam_eps: 1e-8 use_distributed_optimizer: true use_precision_aware_optimizer: false scheduler: lr_decay_style: "constant" sequence_packing: enabled: true generation: max_new_tokens: 15360 vllm_cfg: gpu_memory_utilization: 0.85data: max_input_seq_length: 1024 prompt_file: "examples/prompts/cot.txt" dataset_name: "DAPOMath17K"cluster: gpus_per_node: 8 num_nodes: 8
Launch command:
cd /lustre/fsw/portfolios/coreai/users/bihu/RLbash run_grpo_dapo_deepseek_r1_8B.sh
Or directly with uv:
uv run examples/run_grpo_math.py \ --config=/lustre/fsw/portfolios/coreai/users/bihu/RL/examples/configs/grpo_dapo_deepseek_r1_8B_8n64g_megatron.yaml.bak_no_fp32logits \ cluster.num_nodes=8 \ cluster.gpus_per_node=8
Environment:
8 nodes × 8 GPUs = 64 GPUs
NeMo-RL version: r0.5.0 (container: nemo-rl:r0.5.0-e883ac42)
Model: deepseek-ai/DeepSeek-R1-0528-Qwen3-8B
Dataset: BytedTsinghua-SIA/DAPO-Math-17k (via DAPOMath17K loader)
Expected behavior
The token_mult_prob_error metric should remain stable around 1.0 (or at least below 1.05) throughout training, indicating consistent alignment between the vLLM inference engine and the Megatron training engine. Training should complete without crashing due to numerical instability.
Additional context
WandB metrics screenshot:
Observations:
The token_mult_prob_error starts near 1.0 but gradually increases over hundreds of training steps
Training eventually becomes unstable and crashes
Configuration notes:
Using Megatron backend (not DTensor)
TP=4 with converter_type: "Qwen3ForCausalLM"
params_dtype: "float32" for optimizer parameters
No KL penalty (reference_policy_kl_penalty: 0.0)
Long sequence length (max 16k tokens)
Contributor guide
Assessment
This issue has not been assessed yet.