I wonder why we want to implement mark_sharding as a custom dynamo op

I wonder why we want to implement mark_sharding as a custom dynamo op

This is a follow-up PR to #6161, which originally aimed to support Dynamo + SPMD activation sharding (i.e. allow mark_sharding in a torch.compile call). In order to allow this, we need to make mark_sharding traceable and making the entire mark_sharding a python custom op was the easiest way to make it traceable.

To reviewers, I have some refactoring I want to do here based on my offline sync with Jack and Jiewen. Please don't review this yet, I'll tag you here once it's ready for review (ETA tomorrow). Thanks!

Based on our conversation yesterday, we ideally want to keep the same existing xs.mark_sharding API so customers' UX stays the same. However, one blocker I saw while trying to do this is that if we want to enable mark_sharding inside torch.compile (i.e. support activation sharding), we cannot keep the existing xs.mark_sharding because it takes in Mesh as an input. And Dynamo does not know how to trace a custom type like Mesh. So it seems like we have to introduce a new API to make it traceable.

cc @JackCaoG

And Dynamo does not know how to trace a custom type like Mesh Can you post the error message here? Dynamo should be able to understand basic custom python types.

By adding xs.Mesh into the forward function that is being torch.compile'ed like such:

         # # TODO(yeounoh) remove this after activation sharding support is enabled.
         # num_devices = xr.global_runtime_device_count()
         # device_ids = torch.arange(num_devices)
-        # data_model_mesh = xs.Mesh(device_ids, (4, 1, 2))
+        data_model_mesh = xs.Mesh(device_ids, (4, 1, 2))
         # xs.mark_sharding(output, data_model_mesh, (0, 1, 2), use_dynamo_custom_op=True)

The error we get is:

  File "/home/wonjoo/miniconda3/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: undefined LOAD_FAST

from user code:
   File "/home/wonjoo/llama/llama/generation.py", line 183, in _generate_one_token
    logits = self.model(input_tokens, input_pos_tensor, output_pos_tensor)
  File "/home/wonjoo/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wonjoo/miniconda3/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/wonjoo/llama/llama/model.py", line 374, in forward
    h = layer(h, freqs_cis, mask, input_indexes)
  File "/home/wonjoo/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/wonjoo/llama/llama/model.py", line 310, in forward
    h = x + self.attention.forward(
  File "/home/wonjoo/llama/llama/model.py", line 210, in forward
    data_model_mesh = xs.Mesh(device_ids, (4, 1, 2))

Full error logs is at https://gist.github.com/wonjoolee95/937575e2f5498a2179554974e0a52264.

This error is similar we saw in our prior attempt to make mark_sharding traceable, similar errors would be thrown for dynamo complaining to be unable to trace code like os.environ.

  File "/home/wonjoo/miniconda3/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 839, in LOAD_FAST
    unimplemented("undefined LOAD_FAST")

        if name not in self.symbolic_locals:
            unimplemented("undefined LOAD_FAST")

can you print the name here? Mesh is a relativelly easy python types, I wonder if there is an easy way to make dynamo understand it.

Posting an update after offline sync with Jack: so the issue with xs.Mesh with an issue with my user code (llama/model.py).

The current issue under discussion is from torch_xla.distributed.spmd.mark_sharding -> torch_xla.experimental.dynamo_mark_sharding. It does not seem to like the input type (full error logs):

torch._dynamo.exc.TorchRuntimeError: Failed running call_function xla.dynamo_mark_sharding(*(FakeTensor(..., device='xla:0', size=(1, 128, 512)), [u0], '(4, 1, 2)', 'None', '(0, 1, 2)'), **{}):
xla::dynamo_mark_sharding() Expected a value of type 'List[int]' for argument 'device_ids' but instead found type 'immutable_list'.
Position: 1
Value: [u0]
Declaration: xla::dynamo_mark_sharding(Tensor t, int[] device_ids, str mesh_shape, str axis_names, str partition_spec) -> Tensor
Cast error details: Unable to cast Python instance of type <class 'torch.fx.immutable_collections.immutable_list'> to C++ type 'std::vector<long, std::allocator<long> >'

This error is happening at exactly this line in the new xs.mark_sharding:

return torch.ops.xla.dynamo_mark_sharding(t, device_ids_tensor_list, mesh_shape_str, axis_names_str, partition_spec_str)

Due to Dynamo capturing a python native list type as an torch.fx.immutable_collections.immutable_list as described in #6524 (comment), using the existing xs.mark_sharding to directly call the custom op does not seem possible for now.

I can follow-up with PyTorch to see if we can do something to get this resolved, but for now, we can stick to having a separate custom op to enable dynamo mark sharding.

With these changes, I can confirm that LLaMa 2 Dynamo + SPMD with activation sharding is passing as following: pytorch-tpu/llama#53.

@JackCaoG, this should be ready to be reviewed now. Thanks!

test_parallel_cow_materialize_error_xla_bfloat16 test failing on CPU/GPU tests, seemingly unrelated. Re-running the CI.

Thanks for the review, Yeounoh. I'll go ahead and merge this for now and open a follow-up PR soon.