Question regarding register_post_accumulate_grad_hook

I was reading the torch documentation for `register_post_accumulate_grad_hook`, I think they expect params as an input for the hook, but here we use grad for it.

```python
def _allreduce_grads(self, grad):
"""
Performs an all-reduce operation to synchronize gradients across multiple processes.
"""
# No synchronization needed during gradient accumulation, except at the final accumulation step.
if self.require_backward_grad_sync:
dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=pgm.process_group_manager.cp_dp_group)
grad /= pgm.process_group_manager.cp_dp_world_size
return grad
```

I was expecting something like this instead:

```python
def _allreduce_grads(self, param):
if self.require_backward_grad_sync and param.grad is not None:
dist.all_reduce(
param.grad,
op=dist.ReduceOp.SUM,
group=pgm.process_group_manager.cp_dp_group,
)
param.grad /= pgm.process_group_manager.cp_dp_world_size
```

Would this make more sense? Thank you