Support for gradients to user inputs in TrainingModule
@JacobSzwejbka is already working on this.
Since Sep 25, 2025.
- Dominant language
- Python
- Stars
- 5k
- Forks
- 1.2k
- Avg merge
- 2d 10h
- Merged PRs (30d)
- 581
Description
Currently, TrainingModule doesn't seem to support retrieving gradients wrt user input that are part of the joint fwd+bwd method execution outputs, if one or more of the inputs have requires_grad=True during _export_forward_backward.
I need to access to input grads for my use-case and have added rudimentary support on my end by simply accessing the list of outputs based on param_index and fqn_list.size(), and populating a std::vector with the input grads in-order. Perhaps, supporting this similar to named_parameters or named_gradients would also require changes on torch.export side to include the names and meta/fqn methods to retrieve the names and populate a std::unordered_map..(?)
Adding official support for this in the TrainingModule would be of great help!
cc @JacobSzwejbka
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.