fix: fix temperature-related issues

[zhandaz]

What does this PR do ?

Fix the temperature-related accuracy issues in both the dtensor path and the mcore path.

As reported in #887 and some other researchers, training GRPO (some variants of Qwen 2.5 7B, AceMath) does not work.

uv run python -u NeMo-Skills/nemo_skills/training/nemo_rl/start_grpo.py \
    ++policy.model_name=models/final_hf_checkpoint_1220 \
    ++cluster.gpus_per_node=8 \
    ++cluster.num_nodes=${NUM_ACTOR_NODES} \
    +data.train_data_path=data/deepseek_r1/rl_data/acemath_rl/shuffled_acemath_rl_920k.jsonl \
    +data.val_data_path=data/deepseek_r1/rl_data/acemath_rl/shuffled_acemath_rl_920k.jsonl \
    ++checkpointing.checkpoint_dir=results/${TASK_NAME}_${EXP_SUFFIX}/checkpoints \
    ++logger.log_dir=results/${TASK_NAME}_${EXP_SUFFIX}/training-logs \
    ++logger.wandb_enabled=True \
    ++logger.wandb.project=${WANDB_PROJECT} \
    ++logger.wandb.name=${WANDB_NAME} \
    ++logger.wandb.id=${WANDB_NAME} \
    ++policy.train_global_batch_size=1024 \
    ++policy.train_micro_batch_size=1 \
    ++policy.max_total_sequence_length=8192 \
    ++policy.logprob_batch_size=4 \
    ++policy.sequence_packing.enabled=True \
    ++policy.sequence_packing.train_mb_tokens=8192 \
    ++policy.sequence_packing.logprob_mb_tokens=8192 \
    ++policy.sequence_packing.algorithm="modified_first_fit_decreasing" \
    ++policy.sequence_packing.sequence_length_round=64 \
    ++policy.generation.vllm_cfg.tensor_parallel_size=4 \
    ++policy.generation.vllm_cfg.gpu_memory_utilization=0.8 \
    ++policy.generation.vllm_cfg.pipeline_parallel_size=1 \
    ++policy.generation.temperature=0.6 \
    ++policy.optimizer.kwargs.lr=1e-06 \
    ++loss_fn.reference_policy_kl_penalty=0.001 \
    ++policy.dtensor_cfg.enabled=True \
    ++policy.dtensor_cfg.tensor_parallel_size=2 \
    ++policy.dtensor_cfg.sequence_parallel=False \
    ++policy.dtensor_cfg.activation_checkpointing=True \
    ++data.prompt.prompt_config=qwen/math-cot \
    ++data.prompt.prompt_template=qwen-instruct \
    ++grpo.num_prompts_per_step=128 \
    ++grpo.num_generations_per_prompt=8

1. Reproduce the errors

As we can see in the images below and the wandb, the rewards decrease a lot after around 50 steps, and the token_mult_prob_error (meaning) is large.

2. Fix the temperature error in the DTensor Policy Worker

We found the reason could be: in this PR #660, it scales wrt the temperature according to the vllm engine version (as vllm V1 engine does not return the final logprobs wrt sampling parameters). This fix is actually wrong. We should follow the pattern in #316 to scale logits according to the temperatures anyway. Otherwise, it may lead to off-policy. The inference and the training behavior do not match.

By removing the if in

RL/nemo_rl/models/policy/dtensor_policy_worker.py

Lines 471 to 472 in df31c1b

	if not is_vllm_v1_engine_enabled():
	logits.div_(self.cfg["generation"]["temperature"])

we can see that the reward is back to normal, and token_mult_prob_error is much reasonable (though it still has one large spike that is larger than 2).

3. Fix the returned logprobs in vLLM V1

We know that the returned logprobs from vLLM V1 is not correct for post-training because it is the raw logits (before applying any sampling paramters). Therefore, the token_mult_prob_error is actually an inaccurate indicator. The convergence should not be affected since use_importance_sampling_correction and use_on_policy_kl_approximation are False, and the vllm returned logprobs are not involved in the real training (just used in the metrics calculation).

To confirm that, we do a file patch over the vLLM to return the logprobs after temperature. The reward curve remains almost the same in case 2, and the token_mult_prob_error is almost around 1 (which is expected).

Related

Top-p and Top-k fixing are not supported yet. It is not quite elegant to patch vllm to support top-p and top-k as it involves patching multiple files and many code blocks. We have more discussion in #773 and #69.

Issues

Fixes #902, #887

[github-actions]

❌ Submodule Fast-Forward Check Failed

Check based on commit: e9523c2 (PR #935 from zhanda/debug-accuracy)

❌ Submodules that need attention:

NeMo: ❌ PR branch is BEHIND main branch
TARGET (main branch): https://github.com/NVIDIA/NeMo/commits/5c42641e344a487c7ca5b253a7483f0af8ef40e6/
CURRENT (PR #935 from zhanda/debug-accuracy): https://github.com/NVIDIA/NeMo/commits/aaefedd1d13f4ccd5cd06a19e06f1df33589a235/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

[zhandaz]

Have rerun the experiments after merging the main. There is a small spike. Should be fine.

If there is no further comment, should be ready to merge.

[wangshangsam]

LGTM!

I don't have further comments but I just have a quick question - does it mean that, once vLLM 0.11 is released, we don't need to patch
raw_logprobs = self.compute_logprobs(logits) into
raw_logprobs = self.compute_logprobs(self.apply_temperature(logits.to(torch.float32), sampling_metadata.temperature)) entirely?

[zhandaz]

I don't have further comments but I just have a quick question - does it mean that, once vLLM 0.11 is released, we don't need to patch raw_logprobs = self.compute_logprobs(logits) into raw_logprobs = self.compute_logprobs(self.apply_temperature(logits.to(torch.float32), sampling_metadata.temperature)) entirely?

Yes! All these logics can be greatly simplified once the final_logprobs can be returned by vllm. Hopefully it is available in vllm=0.11.0.

[github-actions]

ℹ️ File Synchronization Check

Check based on commit: 8769667 (PR #935 from zhanda/debug-accuracy)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

• nemo_rl/models/policy/dtensor_policy_worker.py
• nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

ℹ️ File Synchronization Check

Check based on commit: 1dd2d46 (PR #935 from zhanda/debug-accuracy)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

• nemo_rl/models/policy/dtensor_policy_worker.py
• nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

ℹ️ File Synchronization Check

Check based on commit: 1dd2d46 (PR #935 from zhanda/debug-accuracy)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

• nemo_rl/models/policy/dtensor_policy_worker.py
• nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[zhandaz]

The reward seems similar for mcore path after this patch. However, the token_mult_prob_error issue is not mitigated much in this case. I am testing with deepseek-ai/DeepSeek-R1-Distill-Qwen-7B.

[github-actions]

ℹ️ File Consistency Check

Check based on commit: 589d0d8 (PR #935 from zhanda/debug-accuracy)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

• nemo_rl/models/policy/dtensor_policy_worker.py
• nemo_rl/models/policy/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}