Gradient checkpointing fix #1043
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| from torch_xla.utils.checkpoint import checkpoint | ||
|
|
||
| fn_with_kwargs = partial(fn, **kwargs) | ||
| return checkpoint(fn_with_kwargs, *args) |
There was a problem hiding this comment.
I think this should be the other way around:
| return checkpoint(fn_with_kwargs, *args) | |
| return checkpoint(*args, fn_with_kwargs) |
There was a problem hiding this comment.
No, the first argument must be the function we checkpoint, check here.
tengomucho
left a comment
tengomucho
left a comment
There was a problem hiding this comment.
why do you need this function? Can you provide a test that shows its usage and tests that it works?
It is just a wrapper around the original checkpoint function. The goal is to make it possible to pass keyword arguments to the function we are wrapping with I can add a test, but you mentioned many times that the CI was too long. For such a small change I was not considering it, please confirm you want a test. |
| if self.gradient_checkpointing and self.training: | ||
| hidden_states = checkpoint( | ||
| decoder_layer.__call__, | ||
| hidden_states = checkpoint_with_kwargs( |
There was a problem hiding this comment.
you change this here, only for granite, but not for the other models. I think you should change it everywhere, or nowhere, or explain why you only changed it here.
There was a problem hiding this comment.
I dont undertsand? We change it "everywhere".
What does this PR do?
XLA-friendly gradient checkpointing function that accept keyword arguments.