SIGN IN SIGN UP

[FSDP2] Guard _unsharded_param access in to_accumulated_grad_if_needed (#194058)

Fixes an `AttributeError` in `FSDPParam.to_accumulated_grad_if_needed`.

### The bug

```python
def to_accumulated_grad_if_needed(self) -> None:
    # Access `_unsharded_param` to bypass the sharded state check since we
    # prefer to reshard before upcasting the gradient to save memory
    if (
        self.reduce_dtype is None
        or self._unsharded_param.grad is None      # <-- unguarded
        or self._unsharded_param.grad.dtype == self.reduce_dtype
    ):
        return
```

`_unsharded_param` is created by `init_unsharded_param` and dropped by
`free_unsharded_param`, so it is not always present. `init_unsharded_param`
guards its own access with `hasattr("_unsharded_param")` for that reason; this
call site does not, and a parameter that has not been all-gathered raises

```
AttributeError: 'FSDPParam' object has no attribute '_unsharded_param'.
Did you mean: 'unsharded_param'?
```

### How it shows up

Fine-tuning a 9.65B vision-language model with `fully_shard` and a mixed-precision
policy whose `reduce_dtype` differs from the parameter dtype. Gradient
accumulation then walks parameters that took no part in the forward pass, in this
case a vision tower under text-only batches, and those were never gathered.

It reproduces:

* with optimizer offload enabled and disabled;
* with `reshard_after_forward=False`.

The last one is the informative case: the attribute is not being freed after a
gather, it is never created at all for those parameters.

### The fix

Read the attribute through `getattr` and fold "not gathered" into the condition
the method already returns early for. A parameter that was never gathered has no
unsharded gradient to upcast, so behaviour is unchanged for every parameter that
does have one, and the local variable also avoids three repeated attribute
lookups on the hot path.

### Validation

Two full fine-tunes on the same code, data and configuration, one of which
previously failed at this line every time:

| model | before | after |
|---|---|---|
| 4.66B | trains (never hit the path) | unchanged, step 35, loss 0.209, grad-norm 0.68 |
| 9.65B | AttributeError at step 1, three times | trains, step 2, loss 0.234, grad-norm 7.24 |

The 9.65B loss and grad-norm land in the same range as the 4.66B run that was
never affected, so the guard is not masking a gradient that should have been
upcast.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/194058
Approved by: https://github.com/weifengpy
K
khazic committed
2fd2f338d8437898d4667cb28f4a63fb03044749
Parent: 89f89a3
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/23/2026, 7:50:21 AM