[RLlib] - Fix numerical overflow in gradient clipping for (many) large gradients #45055
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…ring values or summing many large values results in +/- infinity. As we clip by multiplying with a clipping coefficient instead of overriding values inside of the gradients tensors this modification allows to clip very large gradients. Signed-off-by: Simon Zehnder <[email protected]>
Signed-off-by: Simon Zehnder <[email protected]>
Signed-off-by: Simon Zehnder <[email protected]>
Signed-off-by: Simon Zehnder <[email protected]>
sven1977
reviewed
May 2, 2024
| @@ -94,6 +95,29 @@ def test_copy_torch_tensors(self): | |||
| all(copied_tensor.detach().numpy() == tensor.detach().cpu().numpy()) | |||
| ) | |||
|
|
|||
| def test_large_gradients_clipping(self): | |||
There was a problem hiding this comment.
Nice! Thanks for creating this important test case!
simonsays1980
requested review from
avnishn,
ArturNiederfahrenhorst,
maxpumperla and
kouroshHakha
as code owners
sven1977
reviewed
May 2, 2024
rllib/utils/torch_utils.py
Outdated
| if torch.logical_or(total_norm.isnan(), total_norm.isinf()): | ||
| raise RuntimeError( | ||
| f"The total norm of order {norm_type} for gradients from " | ||
| "`parameters` is non-finite, so it cannot be clipped. " | ||
| ) | ||
| clip_coef = grad_clip / (total_norm + 1e-6) | ||
| clip_coef = grad_clip / torch.maximum( | ||
| torch.tensor(grad_clip), total_norm + 1e-6 |
There was a problem hiding this comment.
Two questions:
- Would this
torch.tensor()pose a danger when on another device? GPU? - Can we add a comment here (or enhance the one below) explaining why we compute the coeff like this? What are the expected final values of the coeff (between 0.0 and 1.0)?
There was a problem hiding this comment.
Good catch! Yes, let's put the tensor on the device we have extracted before to be save. I will add then a note why we want the coefficient not larger than 1.0.
sven1977
approved these changes
May 2, 2024
There was a problem hiding this comment.
Looks great! Thanks for fixing this very important piece of code @simonsays1980 . Just 2 nits and questions before we can merge.
Signed-off-by: Simon Zehnder <[email protected]>
sven1977
reviewed
May 2, 2024
| clip_coef = grad_clip / (total_norm + 1e-6) | ||
| # We do want the coefficient to be in between 0.0 and 1.0, therefore | ||
| # if the global_norm is smaller than the clip value, we use the clip value | ||
| # as normalization constant. |
There was a problem hiding this comment.
Perfect, thanks for making this comment much more clear!
|
Looks great, let's merge once tests pass ... |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Why are these changes needed?
Large gradients and many of these could lead to numerical overflow when computing their l2-norm in
torch_utils.clip_gradients(using the "global_norm"). This is counterproductive as a user wants to clip such gradients and instead runs into numerical overflow because of clipping gradients.This PR proposes small changes to turn
infandneginfvalues returned from norms to10e8and-10e8, respectively. This does not harm gradients themselves (if these for example were alreadyinf/neginfb/c we clip gradients by multiplication and not overriding values).Related issue number
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.