Update base for Update on "[DTensor][conv] add DTensor convolution_backward op support for case where the input Tensor has requires_grad=False"

Browse source

## Summary
DTensor `convolution_backward` op throws exception when the input Tensor has `requires_grad=False` which happens if the conv layer is the first layer in the model. 

ATEN convolution_backward op Usually returns 3 Tensors (grad_input, grad_weight, grad_bias) and the `grad_input` is actually an Optional[Tensor] which can be `None` in the case mentioned above.

However, the DTensor sharding propagation rule and corresponding TP conv backward implementation both assume that the `grad_input` would be existent. 

## Fix
allow the `grad_input` to be `None` for `convolution_backward` op.

## Test
`pytest test/distributed/tensor/test_convolution_ops.py`

## Follow-up
The current implementation of DTensor conv op also ignores `output_mask` and this may need further care.


cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o tianyu-l

[ghstack-poisoned]

...

This commit is contained in:

Xilun Wu 2025-02-07 11:39:42 -08:00

parent 3de0c70c94

commit 6ef57f056c

Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL

Show stats Download patch file Download diff file Expand all files Collapse all files

Diff content is not available