[SPMD] suppor DTensor API integration #5776

yeounoh · 2023-11-07T23:20:26Z

This is a follow-up to pytorch/pytorch#92909, DTensor & XLA integration. In this PR, we address

move PyTorch/XLA SPMD API to toch_xla/distributed/spmd
introduce xla DTensor API to torch_xla/distributed/spmd so DTensor implementation can access.

yeounoh · 2023-11-08T22:30:22Z

Tested locally, @JackCaoG @jonb377 @alanwaketan I am moving/grouping the SPMD API under toch_xla/experimental/spmd. Later we can move the spmd package out of experimental, when we graduate Beta.

docs/spmd.md

yeounoh · 2023-11-09T17:36:56Z

Tested locally, @JackCaoG @jonb377 @alanwaketan I am moving/grouping the SPMD API under toch_xla/experimental/spmd. Later we can move the spmd package out of experimental, when we graduate Beta.

Had a chat with @JackCaoG , we will move out of experimental and to distributed from BETA.

JackCaoG · 2023-11-09T21:41:19Z

torch_xla/distributed/spmd/api.py

+
+log = logging.getLogger(__name__)
+
+TORCH_XLA_INITIALIZED = False


I don't quite understand this TORCH_XLA_INITIALIZED flag

Ah, we don't need this anymore, due to the fact that this code now lives in torch_xla.

JackCaoG · 2023-11-09T21:41:53Z

torch_xla/distributed/spmd/api.py

+  ) -> None:
+    if TORCH_XLA_INITIALIZED:
+      # TODO(yeounoh) replace this with xr.use_spmd() when we deprecate the flag.
+      os.environ["XLA_USE_SPMD"] = "1"


why not just call xr.use_spmd() now?

there could be a case where the caller already has initialized some tensors before entering the API. We may want to support both, at least in the implementation. I will leave as a TODO here.

JackCaoG

I didn't read the whole or too careful since I felt like most of them are just moving the code around. Is that true?

jonb377

LGTM, thanks Yeounoh!

jonb377 · 2023-11-09T23:11:18Z

docs/spmd.md

@@ -46,8 +46,8 @@ import numpy as np
 import torch
 import torch_xla.core.xla_model as xm
 import torch_xla.runtime as xr
-import torch_xla.experimental.xla_sharding as xs
-from torch_xla.experimental.xla_sharding import Mesh
+import torch_xla.distributed.spmd.xla_sharding as xs


What do you think about making the import path just import torch_xla.distributed.spmd as xs? Not sure if xs is still the right abbreviation in that case (maybe we can think of it as xla spmd), but it would make the spmd package the single entrypoint in userspace.

Good point, let's do that.

…adow xla DTensor API.

alanwaketan

Moving forward, should we use dtensor directly even within our codebase?

@jonb377 do you want to update HF to reflect this change? We can maybe a couple of days in case this get reverted : ).

JackCaoG · 2023-11-14T18:46:51Z

@alanwaketan for HF my major concern is the backward compatibility. I think We should at least wait until 2.2 release is out before updating.

alanwaketan · 2023-11-14T18:47:41Z

@alanwaketan for HF my major concern is the backward compatibility. I think We should at least wait until 2.2 release is out before updating.

Sounds good to me.

* [SPMD] move SPMD package to torch_xla/experimental/spmd, introduce shadow xla DTensor API. * support backward compatibility of the old imports * Move spmd out of experimental * Update spmd.md for distributed/spmd

In response to the change pytorch/xla#5776 and #92909 Pull Request resolved: #113214 Approved by: https://github.com/wanchaol

* [SPMD] move SPMD package to torch_xla/experimental/spmd, introduce shadow xla DTensor API. * support backward compatibility of the old imports * Move spmd out of experimental * Update spmd.md for distributed/spmd

yeounoh added the SPMD / Distributed label Nov 7, 2023

yeounoh self-assigned this Nov 7, 2023

yeounoh marked this pull request as draft November 7, 2023 23:20

yeounoh added the DO_NOT_MERGE_YET For PRs which cannot be merged, despite tests passing label Nov 7, 2023

yeounoh mentioned this pull request Nov 7, 2023

[DTensor][XLA] refactor DTensor _xla API pytorch/pytorch#113214

Closed

yeounoh force-pushed the dtensor_integration_support branch 7 times, most recently from 6cef706 to ccb2c4c Compare November 8, 2023 22:21

yeounoh requested review from alanwaketan, wanchaol, jonb377 and JackCaoG November 8, 2023 22:28

yeounoh removed the DO_NOT_MERGE_YET For PRs which cannot be merged, despite tests passing label Nov 8, 2023

yeounoh marked this pull request as ready for review November 8, 2023 22:28

JackCaoG reviewed Nov 8, 2023

View reviewed changes

docs/spmd.md Show resolved Hide resolved

yeounoh force-pushed the dtensor_integration_support branch 5 times, most recently from e60ce71 to 19fd03a Compare November 9, 2023 17:35

JackCaoG reviewed Nov 9, 2023

View reviewed changes

jonb377 approved these changes Nov 9, 2023

View reviewed changes

JackCaoG approved these changes Nov 10, 2023

View reviewed changes

yeounoh force-pushed the dtensor_integration_support branch 5 times, most recently from be48d90 to c0cdd77 Compare November 13, 2023 17:54

yeounoh added 6 commits November 13, 2023 12:24

[SPMD] move SPMD package to torch_xla/experimental/spmd, introduce sh…

a314e59

…adow xla DTensor API.

linter fix

768ddf6

support backward compatibility of the old imports

646b6da

Move spmd out of experimental

0bbd7d2

Update spmd.md for distributed/spmd

25f6dce

Disable test_set functionalization ops test for now.

794cbc7

yeounoh force-pushed the dtensor_integration_support branch from c7b779a to 794cbc7 Compare November 13, 2023 20:24

yeounoh merged commit e60428d into master Nov 13, 2023
1 check failed

alanwaketan reviewed Nov 14, 2023

View reviewed changes

mbzomowski mentioned this pull request Nov 16, 2023

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

Merged

pytorchmergebot pushed a commit to pytorch/pytorch that referenced this pull request Mar 7, 2024

[DTensor][XLA] refactor DTensor _xla API (#113214)

4f9d4e1

In response to the change pytorch/xla#5776 and #92909 Pull Request resolved: #113214 Approved by: https://github.com/wanchaol

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPMD] suppor DTensor API integration #5776

[SPMD] suppor DTensor API integration #5776

yeounoh commented Nov 7, 2023 •

edited

Loading

yeounoh commented Nov 8, 2023

yeounoh commented Nov 9, 2023

JackCaoG Nov 9, 2023

yeounoh Nov 10, 2023

JackCaoG Nov 9, 2023

yeounoh Nov 10, 2023

JackCaoG left a comment

jonb377 left a comment

jonb377 Nov 9, 2023

yeounoh Nov 10, 2023

alanwaketan left a comment

JackCaoG commented Nov 14, 2023

alanwaketan commented Nov 14, 2023


		log = logging.getLogger(__name__)

		TORCH_XLA_INITIALIZED = False

[SPMD] suppor DTensor API integration #5776

[SPMD] suppor DTensor API integration #5776

Conversation

yeounoh commented Nov 7, 2023 • edited Loading

yeounoh commented Nov 8, 2023

yeounoh commented Nov 9, 2023

JackCaoG Nov 9, 2023

Choose a reason for hiding this comment

yeounoh Nov 10, 2023

Choose a reason for hiding this comment

JackCaoG Nov 9, 2023

Choose a reason for hiding this comment

yeounoh Nov 10, 2023

Choose a reason for hiding this comment

JackCaoG left a comment

Choose a reason for hiding this comment

jonb377 left a comment

Choose a reason for hiding this comment

jonb377 Nov 9, 2023

Choose a reason for hiding this comment

yeounoh Nov 10, 2023

Choose a reason for hiding this comment

alanwaketan left a comment

Choose a reason for hiding this comment

JackCaoG commented Nov 14, 2023

alanwaketan commented Nov 14, 2023

yeounoh commented Nov 7, 2023 •

edited

Loading