Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

wonjoolee95 · 2023-10-19T20:21:25Z

Implement mark_sharding as a custom op to support dynamo spmd activation sharding.

The PR is a bit messy as it includes fair bit amount refactoring, here is a quick summary:

In init_python_bindings.cpp
- Move the existing _xla_mark_sharding logic out to a helper function, so it can be called by the new custom op xla_mark_sharding_dynamo_custom_op.
- Move the torch custom op registration TORCH_LIBRARY from aten_autograd_ops.h to init_python_bindings.cpp since torch custom ops can registered only in one location.
Update the existing torch_xla/experimental/xla_sharding.py function to accept an additional boolean flag dynamo_custom_op. When set to true, it calls this new custom op xla_mark_sharding_dynamo_custom_op instead of the existing ``_xla_mark_sharding`.

Test plan:
python test/spmd/test_dynamo_spmd.py DynamoSpmdInferenceTest.test_mark_sharding_inside_compile

yeounoh · 2023-10-25T17:41:00Z

torch_xla/csrc/init_python_bindings.cpp

@@ -1626,6 +1626,11 @@ void InitXlaModuleBindings(py::module m) {
    // Register sharded tensor data.
    XLAGraphExecutor::Get()->RegisterTensor(xtensor->data());
  });
+  m.def("_xla_mark_sharding_custom_op",


nit. for future reference, can we make it explicit by calling _xla_mark_sharding_dynamo_custom_op?

Sg, updatd. Also added a new API in xla_sharding.py named mark_sharding_dynamo_custom_op specifically for interacting with this new custom op.

yeounoh · 2023-10-25T17:42:11Z

test/spmd/test_dynamo_spmd.py

+      y = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]],
+                       dtype=torch.float,
+                       device=xm.xla_device())
+      ys = xs.mark_sharding(y, self._get_mesh((1, self.n_devices)), (0, 1))


Yes, this should test the activation sharding use-case. cc @wonjoolee95

yeounoh · 2023-10-25T21:21:46Z

torch_xla/csrc/xla_lower_util.cpp

+xla::XlaOp BuildCustomMarkSharding(const torch::lazy::BackendDevice& device,
+                                   const xla::XlaOp& input,
+                                   const xla::XlaOp sharding) {
+  return xla::CustomCall(input.builder(), /*call_target_name=*/"MarkSharding",


Could we re-use the existing CustomSharding op? mark_sharding can be translated into the CustomSharding op for an IR, which is the activation node. Unless it's the Dynamo side change to register the custom ops, we can reuse the existing extern const OpKindWrapper xla_custom_sharding, or there is more details/differences here?

The original intention here was to have a separate custom op that includes all the logic for mark_sharding so Dynamo can properly capture mark_sharding as an op. All the logic currently at xla_mark_sharding at init_python_bindings.cpp would be moved to the lowering logic of this new custom op, so Dynamo can trace and recognize it properly.

wonjoolee95 · 2023-10-27T00:17:14Z

Did some tests locally, it seems like we need to move rest of mark_sharding logic that's currently at tensor_methods.cpp:: custom_mark_sharding to the actual lowering logic of the new CustomMarkSharding op for dynamo to properly recognize and trace this new op. Now working on this fix.

…ion sharding

torch_xla/experimental/xla_sharding.py

miladm · 2023-11-07T21:49:29Z

Thanks @wonjoolee95 for this PR.
As a follow up to this PR, once landed, can you please evaluate the perf gain impact of activation sharding on the llama2 model?

test/spmd/test_dynamo_spmd.py

JackCaoG · 2023-11-08T21:18:41Z

test/spmd/test_dynamo_spmd.py

+
+    dynamo_linear = torch.compile(linear, backend="openxla")
+    dynamo_res = dynamo_linear(xla_x)
+    torch.allclose(xla_res.cpu(), dynamo_res.cpu())


can we add a counter check? We want to make sure we are not recompiling across differerent runs. You can either add it here or add a separate test.

JackCaoG

mostly lgtm, minor nits

yeounoh · 2023-11-11T01:21:17Z

torch_xla/experimental/xla_sharding.py

-                                      int(sharding_type))
+
+    # If flatten_opsharding = True, return the flattened version of OpSharding
+    if flatten_opsharding:


What is flattening here? Just returning as a tuple? If so, maybe call it as_tuple. I am still hesitant to override the return type here... you can try this using the accessor methods instead, since you only need the sharding_type value in your work. Let me create a follow-up PR and add you, let's land this for now. Thanks @wonjoolee95

yeounoh · 2023-11-11T01:21:48Z

torch_xla/experimental/xla_sharding.py

@@ -471,6 +479,9 @@ def mark_sharding(
        >> mesh_shape = (4, 2)
        >> partition_spec = (0, None)

+        dynamo_custom_op (bool): if set to True, it calls the dynamo custom op variant of mark_sharding
+          to make itself recognizeable and traceable by dynamo.


nit. "recognizeable and traceable" --> "traceable"

yeounoh · 2023-11-11T01:26:33Z

torch_xla/csrc/init_python_bindings.cpp

+        [](const at::Tensor& input, const py::list& tile_assignment,
+           const py::list& group_assignment, const py::list& replication_groups,
+           int sharding_type) {
+          c10::List<at::IntArrayRef> tile_assignment_list =


Can we move the following data processing logic into xla_mark_sharding_dynamo_custom_op and just call

xla_mark_sharding_dynamo_custom_op( input, tile_assignment_list, group_assignment_list, replication_groups_list, sharding_type); });

similar to what you've done for mark_sharing.

yeounoh

LGTM, have some comments that we can address in a follow-up PR. Thank you @wonjoolee95

…ion sharding (pytorch#5712)

…ion sharding (#5712)

…ion sharding (pytorch#5712)

…ion sharding (#5712)

…ion sharding (pytorch#5712)

…ion sharding (#5712)

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch 3 times, most recently from 63f2673 to d6b5852 Compare October 24, 2023 23:07

yeounoh reviewed Oct 25, 2023

View reviewed changes

yeounoh added SPMD / Distributed dynamo labels Oct 25, 2023

yeounoh reviewed Oct 25, 2023

View reviewed changes

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch 3 times, most recently from ad4cc4f to e88c7d3 Compare October 26, 2023 23:56

wonjoolee95 mentioned this pull request Oct 31, 2023

Register xla_mark_sharding_dynamo_custom_op for xla-dynamo pytorch/pytorch#112483

Closed

wonjoolee95 added 8 commits October 31, 2023 23:38

Implement mark_sharding as a custom op to support dynamo spmd activat…

53e7ec0

…ion sharding

Update to include OpSharding as an input

f0e8a94

Rebase with master and run linter

7891b42

Update unit tests

6aeeecf

Refine custom marking sharding op

dc19b9b

Re-run linter due to wrong version

a20d710

Add new API for custom mark sharding op and update tests

bd169c2

Add torch pin

ae05c9a

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch from 1bceb4a to 3c51d5e Compare October 31, 2023 23:38

Clean up some code

08d6296

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch from 3c51d5e to 08d6296 Compare November 2, 2023 18:32

Update code to transfer pylist to xla::OpSharding

9aaa533

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch from 4b453c7 to 9aaa533 Compare November 7, 2023 05:08

wonjoolee95 changed the title ~~[WIP] Implement mark_sharding as a custom op to support dynamo spmd activation sharding~~ Implement mark_sharding as a custom op to support dynamo spmd activation sharding Nov 7, 2023

wonjoolee95 marked this pull request as ready for review November 7, 2023 05:13

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch 2 times, most recently from 281d49a to fbe5c79 Compare November 7, 2023 06:49

JackCaoG reviewed Nov 7, 2023

View reviewed changes

torch_xla/experimental/xla_sharding.py Outdated Show resolved Hide resolved

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch 3 times, most recently from f468ed6 to 3604c15 Compare November 8, 2023 17:41

Address comments -- fix typos and variable names

a4318a6

wonjoolee95 force-pushed the wonjoo/dynamo-custom-op branch from 3604c15 to a4318a6 Compare November 8, 2023 19:57

wonjoolee95 requested a review from JackCaoG November 8, 2023 19:58

Run linter

e35ca64

JackCaoG reviewed Nov 8, 2023

View reviewed changes

test/spmd/test_dynamo_spmd.py Outdated Show resolved Hide resolved

JackCaoG reviewed Nov 8, 2023

View reviewed changes

Add metric assertions to unit tests

ae721ae

wonjoolee95 requested a review from JackCaoG November 9, 2023 05:34

JackCaoG approved these changes Nov 9, 2023

View reviewed changes

yeounoh reviewed Nov 11, 2023

View reviewed changes

yeounoh approved these changes Nov 11, 2023

View reviewed changes

wonjoolee95 merged commit 367f47f into master Nov 11, 2023
17 checks passed

yeounoh mentioned this pull request Nov 15, 2023

Refactor Dynamo (custom op) integration code #5805

Merged

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

Implement mark_sharding as a custom op to support dynamo spmd activat…

452665b

…ion sharding (pytorch#5712)

mbzomowski mentioned this pull request Nov 16, 2023

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

Merged

zpcore pushed a commit that referenced this pull request Nov 21, 2023

Implement mark_sharding as a custom op to support dynamo spmd activat…

f32e276

…ion sharding (#5712)

lsy323 pushed a commit to lsy323/xla that referenced this pull request Nov 28, 2023

Implement mark_sharding as a custom op to support dynamo spmd activat…

3d891e2

…ion sharding (pytorch#5712)

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

Implement mark_sharding as a custom op to support dynamo spmd activat…

f6c149c

…ion sharding (#5712)

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

Implement mark_sharding as a custom op to support dynamo spmd activat…

70b3392

…ion sharding (#5712)

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

Implement mark_sharding as a custom op to support dynamo spmd activat…

c7a04ac

…ion sharding (pytorch#5712)

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

Implement mark_sharding as a custom op to support dynamo spmd activat…

f804705

…ion sharding (#5712)

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

Implement mark_sharding as a custom op to support dynamo spmd activat…

ef06f47

…ion sharding (#5712)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

wonjoolee95 commented Oct 19, 2023 •

edited

Loading

yeounoh Oct 25, 2023

wonjoolee95 Oct 26, 2023

yeounoh Oct 25, 2023

yeounoh Oct 25, 2023 •

edited

Loading

wonjoolee95 Oct 26, 2023

wonjoolee95 commented Oct 27, 2023

miladm commented Nov 7, 2023

JackCaoG Nov 8, 2023

wonjoolee95 Nov 9, 2023

JackCaoG left a comment

yeounoh Nov 11, 2023

yeounoh Nov 11, 2023

yeounoh Nov 11, 2023 •

edited

Loading

yeounoh left a comment

Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

Conversation

wonjoolee95 commented Oct 19, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yeounoh Oct 25, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wonjoolee95 commented Oct 27, 2023

miladm commented Nov 7, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JackCaoG left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yeounoh Nov 11, 2023 • edited Loading

Choose a reason for hiding this comment

yeounoh left a comment

Choose a reason for hiding this comment

wonjoolee95 commented Oct 19, 2023 •

edited

Loading

yeounoh Oct 25, 2023 •

edited

Loading

yeounoh Nov 11, 2023 •

edited

Loading