xla
[SPMD][Virtual Device]All tensors should be in SPMD:0 C++ device
#5284

Merged

[SPMD][Virtual Device]All tensors should be in SPMD:0 C++ device #5284

JackCaoG merged 11 commits into master from JackCaog/spmd_device_default

Move all tensors to SPMD:0 C++ device under spmd context

5d1683f0

fix load shards

bbcb321e

fix test_mark_sharding_2d by not creating placeholder for virtual device

0aba6a5b

JackCaoG force pushed from 92fa0cc5 to 0aba6a5b 2 years ago

fix the waitdeviceop for spmd case

132fd6e7

Fix test_shard_hashing

e56d25b4

fix spmd device casting issue

e7d460d6

JackCaoG added distributed

remove hacks in test_xla_virtual_device.py

95448436

add test for new virtual device usage

511c3892

JackCaoG changed the title ~~[WIP]Move all tensors to SPMD:0 C++ device under spmd context~~ [SPMD][Virtual Device]All tensors should be in SPMD:0 C++ device 2 years ago

JackCaoG requested a review from

jonb377 2 years ago

JackCaoG requested a review from

alanwaketan 2 years ago

JackCaoG requested a review from

wonjoo-wj 2 years ago

jonb377 commented on 2023-07-10

fix review comments

ad3bda93

fix IsTpuDevice

58363642

alanwaketan commented on 2023-07-10

linter

e4439cf3

JackCaoG requested a review from

jonb377 2 years ago

JackCaoG requested a review from

alanwaketan 2 years ago

alanwaketan commented on 2023-07-11

JackCaoG commented on 2023-07-11

jonb377 approved these changes on 2023-07-11

alanwaketan approved these changes on 2023-07-11

JackCaoG merged e42fffa0 into master 2 years ago

Reviewers

alanwaketan

jonb377

wonjoo-wj

Assignees

No one assigned

Labels

distributed

Milestone

No milestone

xla [SPMD][Virtual Device]All tensors should be in SPMD:0 C++ device #5284 Merged

[SPMD][Virtual Device]All tensors should be in SPMD:0 C++ device #5284

xla
[SPMD][Virtual Device]All tensors should be in SPMD:0 C++ device
#5284

Merged