[bug][algorithm] remove incorrect torch.no_grad() for kl in loss (use_kl_loss=True)#1353
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly addresses a bug where gradients were not flowing through the KL divergence term in the loss function when use_kl_loss=True. The fix involves removing the @torch.no_grad() decorator from the compute_approx_kl function. To maintain correctness in other parts of the code that rely on a gradient-free KL computation, specifically for KL-based reward penalties, a with torch.no_grad() context is added at the call site in apply_reward_kl_penalty. The changes are accurate and well-contained.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
|
GSM8K with and without fix (looks pretty similar, but probably just b/c it's gsm8k): |
|
@erictang000 should be able to see the diff with a large kl penalty. Could also just print the kl loss tensor in the worker before and after the fix (after should have requires grad) |
|
here after setting 
|
…_kl_loss=True) (#1353) Fixes #1340 <!-- devin-review-badge-begin --> --- <a href="https://app.devin.ai/review/novasky-ai/skyrl/pull/1353" target="_blank"> <picture> <source media="(prefers-color-scheme: dark)" srcset="https://static.devin.ai/assets/gh-open-in-devin-review-dark.svg?v=1"> <img src="https://static.devin.ai/assets/gh-open-in-devin-review-light.svg?v=1" alt="Open with Devin"> </picture> </a> <!-- devin-review-badge-end -->
2 participants
Fixes #1340