Transfer data directly to the device

[will-cromar]

Take two. See the notes from the original PR #5752

New changes:

• Use the actual byte strides of the input at::Tensor instead of the strides in the xla::Shape.
• Downcast the input at::Tensor if the target type differs from the actual type.
• Release the GIL in XlaDataToTensors, since at::Tensor destruction after TransferToServer can deadlock with TransferFromServer.

[will-cromar]

The issue here is that I was calculating the byte_strides for an at::Tensor based on the given xla::Shape, which may not match. There's still an unresolved issue here that PjRtClient::BufferFromHostBuffer relies on the target xla::PrimitiveType to determine the size of each element in the source.

In the numerous special cases in GetDevicePrimitiveType where we silently modify types, this means that BufferFromHostBuffer will read the wrong length. For example, it will read the first 32 bits of a float64 on TPU, since we silently downcast float64 to float32. Some of these cases are now irrelevant, but I'm not sure that all of them are. cc @JackCaoG to weigh in.

In cases where the source type does not match the output type, I believe we'll still have to "stage" the data in an xla::Literal or cast the at::Tensor. The conversion has to happen somewhere, but I don't think we can do it just by modifying BufferFromHostBuffer's data pointer and byte_strides.

[will-cromar]

I have everything working locally now. Separating the more tedious changes here into #5777.

After this PR, we'll still have to make an intermediate copy if the input tensor type does not match the target type. I added a counter to capture this overhead, since it may have a performance impact. For what it's worth, casting the at::Tensor directly still seems to be faster than copying to an xla::Literal.

Getting around the copy is simple: just create the tensors such that the type matches what it will be on the device. So if you want a bf16 on device, make a bf16 tensor.

[will-cromar]

The TPU CI is currently hanging on SPMDLoadPlannerTest.test_resolve_and_commit_sharded_tensor. I can reproduce this locally, but not if I step through the lines with pdb. That suggests it may be hanging due to a race condition. This test does not look particularly related to this PR.

[will-cromar]

@jonb377 and I were able to figure out where the deadlock is. The hang is caused by a GIL deadlock when we try to retrieve data from the device before a transfer finishes, and the transfer is the only thing keeping an at::Tensor alive. If anything, I'm surprised only one test case caught this bug. Marking this PR as "do not merge" until we find a fix.

Here's what's happening:

1. The main thread moves some data to the TPU. This immediately kicks off an async thread to handle the transfer (which is kind of the point of the whole PR).
2. The main thread continues to call local_shards, which calls down to XlaDataToTensors, which calls TransferToServer, which has to block until the transfer from (1) is complete. Note that the main thread is holding the GIL.
3. The async transfer thread from (1) kicks off, holding a callback that holds the only shared_ptr to an AtenSource, which holds an at::Tensor. This is required to ensure that the at::Tensor is kept alive until the transfer is complete.
4. The actual transfer of data from host to device is complete, firing the callback that holds the only shared_ptr to the AtenSource.
5. The callback functionally does nothing except delete the shared_ptr<AtenSource>. Because it is the last owner of the underlying at::Tensor, this destroys the at::Tensor, which destroys the c10::TensorImpl.
6. c10::TensorImpl's destructor waits to acquire the GIL.
7. Since TransferToServer fires the callback synchronously, it cannot return until the GIL is acquired, even though the actual transfer to the TPU is complete.
8. See 2.

Here's the relevant stack trace through TransferToServer from gdb:

...
#4  take_gil (ceval=0x7f210c228048 <_PyRuntime+584>, tstate=0x7f1a40002bb0) at Python/ceval_gil.h:206
#5  0x00007f210c0c3655 in PyEval_AcquireThread (tstate=0x7f1a40002bb0) at Python/ceval.c:316
#6  0x00007f2107849f05 in pybind11::gil_scoped_acquire::gil_scoped_acquire() () from /usr/local/lib/python3.8/site-packages/torch/lib/libtorch_python.so
...
#8  0x00007f21073f932b in c10::impl::PyObjectSlot::maybe_destroy_pyobj() () from /usr/local/lib/python3.8/site-packages/torch/lib/libc10.so
...
#11 0x00007f21073e8669 in c10::TensorImpl::~TensorImpl() () from /usr/local/lib/python3.8/site-packages/torch/lib/libc10.so
#12 0x00007f1e55ad83b3 in std::_Sp_counted_ptr_inplace<torch_xla::runtime::AtenSource, std::allocator<torch_xla::runtime::AtenSource>, (__gnu_cxx::_Lock_policy)2>::_M_dispose() () from /usr/local/lib/python3.8/site-packages/torch_xla-2.2.0+git30abda8-py3.8-linux-x86_64.egg/_XLAC.cpython-38-x86_64-linux-gnu.so
#13 0x00007f1e55935b3a in std::_Sp_counted_base<(__gnu_cxx::_Lock_policy)2>::_M_release() () from /usr/local/lib/python3.8/site-packages/torch_xla-2.2.0+git30abda8-py3.8-linux-x86_64.egg/_XLAC.cpython-38-x86_64-linux-gnu.so
#14 0x00007f1e55d9d3c4 in std::_Function_handler<void (), torch_xla::runtime::PjRtComputationClient::TransferToServer(absl::lts_20230125::Span<std::shared_ptr<torch_xla::runtime::TensorSource const> const>)::{lambda()#1}>::_M_manager(std::_Any_data&, std::_Any_data const&, std::_Manager_operation) () from /usr/local/lib/python3.8/site-packages/torch_xla-2.2.0+git30abda8-py3.8-linux-x86_64.egg/_XLAC.cpython-38-x86_64-linux-gnu.so
#15 0x00007f1e5a81ef92 in std::_Function_handler<void (PJRT_Error*), xla::PjRtCApiClient::BufferFromHostBufferInternalImpl(void const*, xla::PrimitiveType, absl::lts_20230125::Span<long const>, std::optional<absl::lts_20230125::Span<long const> >, xla::PjRtClient::HostBufferSemantics, std::function<void ()>, std::variant<xla::PjRtDevice*, xla::PjRtMemorySpace*>, xla::Layout const*)::{lambda(PJRT_Error*)#1}>::_M_manager(std::_Any_data&, std::_Any_data const&, std::_Manager_operation) () from /usr/local/lib/python3.8/site-packages/torch_xla-2.2.0+git30abda8-py3.8-linux-x86_64.egg/_XLAC.cpython-38-x86_64-linux-gnu.so

[will-cromar]

I fixed a similar GIL deadlock bug about a year ago in #4504. In that case, the solution was to release the GIL during TransferFromServer, but it looks like I inserted that fix into a different code path than SPMD uses.

[will-cromar]

I wrapped the GIL release and data transfer into a new utility, ReleaseGilAndTransferData. GetTensors and XlaDataToTensors can both call this function. I also added a note to computation_client.h to be wary of the GIL when calling TransferFromServer, since it's the only call to PJRT that is blocking.

I'm open to better names for ReleaseGilAndTransferData.

[yeounoh]

Should we require the caller to manage the ownership & lifetime of TensorSource*, i.e., const TensorSource* instead, or it's necessary to ensure that the memory is held during the client ops and in the client?

[will-cromar]

Good question. PJRT lets us tie the lifetime of an object to an operation by capturing it in a callback. std::functions have to be copyable, so shared_ptr is our best choice here. TensorSource itself may be expensive or impossible to copy.

The caller of TransferToServer will be much shorter-lived than the actual transfer, so ownership should pass down. We could tighten up the interface here and consume a unique_ptr since we only need copyability within the implementation of TransferToServer. What do you think?

[yeounoh]

LGTM

[JackCaoG]

I will try to take a look today