Repository: pytorch
Issue: Have a different way to check if gradient was computed in the optimizer (not checking for None)
Labels: feature,module: autograd,module: optimizer,triaged

## 🚀 Feature
The feature I want is to be able to check if the gradient was actually computed or not for a specific parameter in the optimizer, instead of checking if the gradient is None (as done currently).

## Motivation
The motivation in my case comes from trying to implement some code that is doing multi-task learning. On each training iteration I'm getting a batch from a single task and compute gradients for that task specific head (and shared backbone for all tasks), while heads for other tasks are not supposed to be updated. This is indeed happening on the first iteration, but after we computed gradients for each task at least one time, each one of them will be updated on *every* iteration, even when no gradients were computed for that task specific head. The reason for that is because most optimizers will only ignore weights if their gradient is None, and if not, they will proceed updating them. Thus, after the gradients were computed ones, they will never be None again (since we set them to zero in optimizer.zero_grad) and on each iteration, the current momentum and weight decay will be applied to the weights for which we didn't really compute the gradients. 

In my opinion this is very dangerous behavior because this is likely not what the user expects to happen, but the model could still work fine and train to some non-trivial accuracy even though each task weights receive additional weight decay and momentum updates. While it might not matter much in some cases, when the tasks are trained for different number of iterations (e.g. one task has 10 times more updates than the other), this might significantly harm performance. And there might be other cases besides multi-task learning, when not every weight is supposed to be updated on each iteration, where this behavior of the optimizer can cause some unexpected results without users even noticing this. 

## Pitch

I am not sure how gradient flow is implemented internally in PyTorch, but conceptually it should be possible (and probably not too complicated) to have a way to check if the gradient was actually computed for each weight after the last call to .backward on the loss. And then this could be used to check if the optimizer should update weights or not, instead of checking for None. 

## Alternatives

As an alternative solution, I can explicitly set gradients for weights that are not supposed to be updated to None on each iteration. However, it seems like that will introduce additional computational overhead on each iteration, because new memory will have to be allocated for the gradient storage. 

Another solution could be to have separate optimizers for each task. However, in this case the optimizer buffers for shared backbone weights will be separate. This means that, for example, momentum will be accumulated separately for each task gradients and not merged together and this might also lead to suboptimal performance. I could also have a separate optimizer for shared weights, but this becomes more complicated when certain tasks share additional weights between each other (for example, part of the head could be shared between two tasks and not shared with other tasks, while backbone is shared across all tasks). 

Another solution could be to set learning rate to zero on each iteration for the parameters that should not be updated. However, in this case, their momentum buffers will still get updated and this will result in different performance.

The solution I think I will end-up doing is to modify the optimizer code to explicitly pass weights that should not be updated on the current iteration. This is also not ideal, because it is cumbersome to check which weights should and should not be updated on the current iteration and I have to look at the names of the weights for this. And I also have to add this new logic for each optimizer I plan to use for my models.

Let me know if there is some other, easy to implement solution for this problem. However, even if that is the case, I think it might still be important to implement it by default in PyTorch or at least mention this in the documentation, or raise some kind of warning for this behavior. Because this problem will not lead to an immediate error and thus it is easy to overlook this and get unexpected results that could potentially harm performance of the models quite significantly.