pytorch / pytorch/executorch

Support for gradients to user inputs in TrainingModule

Open
#13,522 5 comments 0 reactions 1 assignee View on GitHub

@JacobSzwejbka is already working on this.

Since Sep 25, 2025.

module: training
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.