[FSDP2 Related]`torch.split_with_sizes_copy` of the GPU does not update the version counter of `out` correctly.

[medivh-xp]

🐛 Describe the bug

I'm reading the source code implementation of FSDP2. I noticed a detail that the unsharded param is obtained through split_with_sizes_copy when needed. Since this operation writes to the memory of the param, I expected it to cause an increase in the version counter, which would further lead to failure to pass the verification of the version counter during the backward pass. However, in fact, it works normally on the GPU.

This confuses me a lot. So I tested split_with_sizes_copy separately and found that its handling of the version counter is inconsistent on the CPU and GPU: on the GPU, the version counter of out is not updated correctly.

import torch

x = torch.ones(4, 2)
y = torch.ones(4, 1)
z = torch.ones(4, 1)
torch.split_with_sizes_copy(x, [1, 1], dim=1, out=[y, z])
print(y._version, z._version)  # 1 1

x = torch.ones(4, 2).cuda()
y = torch.ones(4, 1).cuda()
z = torch.ones(4, 1).cuda()
torch.split_with_sizes_copy(x, [1, 1], dim=1, out=[y, z])
print(y._version, z._version)  # 0 0

Versions

I tested the torch versions including the nightly(2.5.0.dev20240709+cu121) versions since release 2.3.1, and this problem exists in all of them.

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zhaojuanmao @mrshenli @rohan-varma @chauhang

[awgu]

cc: @yifuwang could you take a look when you get the chance? If we fix this, we probably need to use the unsafe version counter context in FSDP2 🤔

[yifuwang]

@albanD hmm I thought the dispatcher/codegen would take care bumping the version. Do you know what might've gone wrong?

[albanD]

The registration in native_functions.yaml seem to hint at the fact that codegen should do it indeed.
@bdhirsh any weird stuff happening for the view _copy function's out= overloads? :p

[bdhirsh]

yup, it looks like we aren't codegen'ining InplaceOrViewType kernels for them - taking a look

[bdhirsh]

@albanD looked around a little more:

(1) it looks like we don't generate ADInplaceOrView kernels for any out= ops that have no returns. I think it's because of this line in the codegen. So I think this issue is less an problem for "auto-generated view_copy.out ops" and more for any out= ops that return no outputs. I know we don't support no-returns ops in autograd today, but is it worth trying to at least register an inplace_or_view kernel for them?

(2) the reason this code happens to bump the VC for CPU and not for CUDA: the cpu impl of split_with_sizes_copy.out doesn't have a dedicated kernel, so it redispatches into a functional op + a call to aten.copy_(), and that copy_ call will end up properly bumping the VC. For cuda though, since we have a dedicated kernel and don't dispatch to aten.copy_(), the VC doesn't get bumped.

example dispatch trace for CPU, where we bump the VC:

 [call] op=[aten::split_with_sizes_copy.out], key=[AutogradCPU]
  [redispatch] op=[aten::split_with_sizes_copy.out], key=[CPU]
   [call] op=[aten::split_with_sizes], key=[ADInplaceOrView]
    [redispatch] op=[aten::split_with_sizes], key=[CPU]
     [call] op=[aten::as_strided], key=[CPU]
     [call] op=[aten::as_strided], key=[CPU]
   [call] op=[aten::copy_], key=[ADInplaceOrView]
    [redispatch] op=[aten::copy_], key=[CPU]
   [call] op=[aten::copy_], key=[ADInplaceOrView]
    [redispatch] op=[aten::copy_], key=[CPU]

and for cuda, where we don't bump the VC:

 [call] op=[aten::split_with_sizes_copy.out], key=[AutogradCUDA]
  [redispatch] op=[aten::split_with_sizes_copy.out], key=[CUDA]
   [call] op=[aten::empty.memory_format], key=[BackendSelect]
    [redispatch] op=[aten::empty.memory_format], key=[CPU]
   [call] op=[aten::to.dtype_layout], key=[BackendSelect]
    [redispatch] op=[aten::to.dtype_layout], key=[CUDA]
     [call] op=[aten::_to_copy], key=[BackendSelect]
      [redispatch] op=[aten::_to_copy], key=[CUDA]
       [call] op=[aten::empty_strided], key=[BackendSelect]
        [redispatch] op=[aten::empty_strided], key=[CUDA]
       [call] op=[aten::copy_], key=[ADInplaceOrView]
        [redispatch] op=[aten::copy_], key=[CUDA]

[albanD]

(1) ho yes we should fix that!
Either fixing codegen to properly bump these or error out when we're given such an op in codegen!