Transfer data directly to the device by will-cromar · Pull Request #5752 · pytorch/xla

will-cromar · 2023-10-31T23:21:25Z

Remove populate_fn entirely. We had to implement this weird indirection before to prevent our TensorFlow code from taking a dependency on PyTorch code.
Copy data directly from PyTorch Tensors into runtime buffers without having to stage them in an xla::Literal first.
Re-implement TensorSource as an interface that can directly wrap and own a data source.
- Ownership makes it easier to implement async data transfers. Otherwise, we risk the source tensor being deleted before the transfer is done.
~~Rename TransferToServer since there's no Server.~~ Reverted to make review easier. Will rebase commit into follow-up PR.
Re-enable and update PJRT client unit test, since it's the best unit test to cover this change.

This reverts commit 6384516.

will-cromar · 2023-11-01T20:45:00Z

This very slightly improves performance on ResNet50 with fake data on TPU v4 by removing synchronous overhead. Key metrics after epoch 1:

Before:

Metric: DeviceLockWait
  TotalSamples: 2582
  Accumulator: 01s292ms557.601us
  ValueRate: 007ms301.380us / second
  Rate: 13.7858 / second
  Percentiles: 1%=000.880us; 5%=001.049us; 10%=001.240us; 20%=001.550us; 50%=005.870us; 80%=008.320us; 90%=002ms368.870us; 95%=005ms894.689us; 99%=006ms147.740us
...
Metric: TransferToServerTime
  TotalSamples: 1951
  Accumulator: 20s711ms722.341us
  ValueRate: 109ms777.976us / second
  Rate: 7.31885 / second
  Percentiles: 1%=065.640us; 5%=070.460us; 10%=073.750us; 20%=080.750us; 50%=094.750us; 80%=061ms422.564us; 90%=066ms168.024us; 95%=068ms557.643us; 99%=070ms677.393us

After:

Metric: DeviceLockWait
  TotalSamples: 2582
  Accumulator: 794ms078.846us
  ValueRate: 008ms619.516us / second
  Rate: 13.8462 / second
  Percentiles: 1%=001.020us; 5%=001.180us; 10%=001.310us; 20%=001.449us; 50%=005.920us; 80%=007.830us; 90%=002ms445.669us; 95%=005ms409.450us; 99%=006ms294.320us
...
Metric: TransferToDeviceTime
  TotalSamples: 1951
  Accumulator: 193ms746.667us
  ValueRate: 001ms032.600us / second
  Rate: 7.3311 / second
  Percentiles: 1%=047.810us; 5%=051.320us; 10%=057.360us; 20%=061.490us; 50%=074.340us; 80%=107.620us; 90%=123.420us; 95%=142.660us; 99%=197.790us

Note the drop from >20 seconds to ~200 ms in TransferToServerTime. In practice, this shaves off about a half second of waiting since ResNet50 is not actually transfer-bound.

JackCaoG

This is great! Thanks @will-cromar

This reverts commit 5ca36cf.

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

yeounoh · 2023-11-09T07:16:30Z

Just learning about this change, this is great feat -- thanks @will-cromar !

* Transfer data directly to the device (#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

…#5765) This reverts commit 5ca36cf.

* Transfer data directly to the device (pytorch#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Transfer data directly to the device (#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Transfer data directly to the device (pytorch#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

This reverts commit 5ca36cf.

* Transfer data directly to the device (#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

This reverts commit 5ca36cf.

* Transfer data directly to the device (#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

…#5765) This reverts commit 5ca36cf.

* Transfer data directly to the device (pytorch#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

This reverts commit 5ca36cf.

* Transfer data directly to the device (#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

* Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516.

This reverts commit 5ca36cf.

* Transfer data directly to the device (#5752) * Remove `populate_fn` from `TensorSource` * Make TensorSource an interface * Re-enable pjrt_computation_client_test * server -> device * add comment * fix outbound data metric * formatting * implement byte_strides in TensorSource * more formatting * remove extra deps * add missing deps * Revert "server -> device" This reverts commit 6384516. * Use `at::Tensor`'s layout for byte strides * Downcast at::Tensor if required * formatting * Simplify AtenSource * fix build * formatting * fix typo that makes us ignore input type * Revert "Simplify AtenSource" This reverts commit 4225deb. * Skip hanging test * fix gil deadlock * formatting

Remove populate_fn from TensorSource

d2dfe24

will-cromar added the runtime label Oct 31, 2023

will-cromar added 7 commits November 1, 2023 19:07

Make TensorSource an interface

fad33dd

Re-enable pjrt_computation_client_test

7b9cfe3

server -> device

6384516

add comment

8d9f2d1

fix outbound data metric

c0cca42

formatting

ff2ffb9

implement byte_strides in TensorSource

ab88d50

will-cromar changed the title ~~Remove populate_fn from TensorSource~~ Transfer data directly to the device Nov 1, 2023

will-cromar added 4 commits November 1, 2023 19:59

more formatting

904eae5

remove extra deps

59e73ef

add missing deps

9b7d481

Revert "server -> device"

2a6433d

This reverts commit 6384516.

will-cromar requested a review from JackCaoG November 1, 2023 20:45

will-cromar marked this pull request as ready for review November 1, 2023 20:45

JackCaoG approved these changes Nov 1, 2023

View reviewed changes

will-cromar merged commit 5ca36cf into master Nov 2, 2023

will-cromar added a commit that referenced this pull request Nov 2, 2023

Revert "Transfer data directly to the device (#5752)"

700b276

This reverts commit 5ca36cf.

will-cromar added a commit that referenced this pull request Nov 2, 2023

Revert "Transfer data directly to the device (#5752)" (#5765)

a93e14b

This reverts commit 5ca36cf.

will-cromar mentioned this pull request Nov 6, 2023

Transfer data directly to the device #5772

Merged

mbzomowski pushed a commit to mbzomowski-test-org/xla that referenced this pull request Nov 16, 2023

Revert "Transfer data directly to the device (pytorch#5752)" (pytorch…

898b36e

…#5765) This reverts commit 5ca36cf.

mbzomowski mentioned this pull request Nov 16, 2023

tpu ci module refactor mbzomowski-test-org/xla#7

Merged

ManfeiBai pushed a commit that referenced this pull request Nov 29, 2023

Revert "Transfer data directly to the device (#5752)" (#5765)

43a2e02

This reverts commit 5ca36cf.

ManfeiBai pushed a commit that referenced this pull request Nov 29, 2023

Revert "Transfer data directly to the device (#5752)" (#5765)

eba3afc

This reverts commit 5ca36cf.

chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023

Revert "Transfer data directly to the device (pytorch#5752)" (pytorch…

ec70430

…#5765) This reverts commit 5ca36cf.

golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024

Revert "Transfer data directly to the device (#5752)" (#5765)

8e2a4c1

This reverts commit 5ca36cf.

bhavya01 pushed a commit that referenced this pull request Apr 22, 2024

Revert "Transfer data directly to the device (#5752)" (#5765)

b54b842

This reverts commit 5ca36cf.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Transfer data directly to the device#5752

Transfer data directly to the device#5752
will-cromar merged 12 commits intomasterfrom
wcromar/remove-populate-fn

will-cromar commented Oct 31, 2023 •

edited

Loading

Uh oh!

will-cromar commented Nov 1, 2023

Uh oh!

JackCaoG left a comment

Uh oh!

yeounoh commented Nov 9, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

will-cromar commented Oct 31, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

will-cromar commented Nov 1, 2023

Uh oh!

JackCaoG left a comment

Choose a reason for hiding this comment

Uh oh!

yeounoh commented Nov 9, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

will-cromar commented Oct 31, 2023 •

edited

Loading