def _gradient_hook(param: Tensor):
    """Called when a parameter's gradient is ready"""
    if param.grad is None:
        return

    bucket_idx = param_to_bucket[param]
    bucket_ready_count[bucket_idx] += 1

    # Check if all parameters in this bucket are ready
    if bucket_ready_count[bucket_idx] == len(param_buckets[bucket_idx]):
        # All-reduce this bucket
        bucket_grads = [p.grad for p in param_buckets[bucket_idx]]

        # For multi-tensor operations, we can reduce them together
        if len(bucket_grads) == 1:
            handle = dist.all_reduce(bucket_grads[0], op=dist.ReduceOp.AVG, async_op=True)
        else:
            # Use multi-tensor all-reduce for efficiency
            handle = dist.all_reduce_coalesced(bucket_grads, op=dist.ReduceOp.AVG, async_op=True)

        bucket_handles[bucket_idx] = handle
Called when a parameter's gradient is ready
train_gpt.py