[dtensor] add op support for aten.gather.default

[tianyu-l]

Stack from ghstack (oldest at bottom):

• [dtensor] add op support for nll_loss_forward #118917
• -> [dtensor] add op support for aten.gather.default #118513

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @wconstab @yf225

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/118513

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 80826b4 with merge base 0f7e636 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_huggingface, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
  AlbertForQuestionAnswering

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[XilunWu]

The CI report says "test_dtensor_op_db_take_along_dim_cpu_float32 in test_dtensor_ops.py has unexpected success". This means your change has made it working correctly. You can remove "xfail(take_along_dim)" from the file which will mark this test as a passing test instead of an expected failing test.

Another thing you can do in future is to locally run pytest test/distributed/_tensor/test_dtensor_ops.py to see if the test result has changed with your PR.

[XilunWu]

left 2 suggestions. I'm not very clear about the main body of gather_strategy and the use of _MaskPartial though.

[XilunWu]

I dont think we need to instantiate another comm_mode but re-use the one instantiated above. We can instantiate one comm_mode at the beginning of the test and reuse it everywhere in the test,

[XilunWu]

We also need to check if the input_shape and output_shape is eligible to perform torch.gather:
https://pytorch.org/docs/stable/generated/torch.gather.html

[tianyu-l]

If we don't check in dtensor sharding prop, these errors would pop up from local tensor ops, which seems still is the expected behavior?

[wanchaol]

lgtm, some minor suggestions inlined.

[wanchaol]

I think the main reason here is that the embedding output would produce an additional dimension compare to the input, hence the output masking logic become different? maybe have more clarification in the comment

[wanchaol]

given that we are reusing the logic, let's put this logic as a common method in MaskBuffer.

[tianyu-l]

The CI report says "test_dtensor_op_db_take_along_dim_cpu_float32 in test_dtensor_ops.py has unexpected success". This means your change has made it working correctly. You can remove "xfail(take_along_dim)" from the file which will mark this test as a passing test instead of an expected failing test.

Another thing you can do in future is to locally run pytest test/distributed/_tensor/test_dtensor_ops.py to see if the test result has changed with your PR.

This PR doesn't touch take_along_dim. Not sure how it could make the test pass... Should I still remove it?

[XilunWu]

The CI report says "test_dtensor_op_db_take_along_dim_cpu_float32 in test_dtensor_ops.py has unexpected success". This means your change has made it working correctly. You can remove "xfail(take_along_dim)" from the file which will mark this test as a passing test instead of an expected failing test.
Another thing you can do in future is to locally run pytest test/distributed/_tensor/test_dtensor_ops.py to see if the test result has changed with your PR.

This PR doesn't touch take_along_dim. Not sure how it could make the test pass... Should I still remove it?

@tianyu-l check

pytorch/torch/_refs/__init__.py

Lines 4470 to 4504 in fb8ffba

	@out_wrapper()
	def take_along_dim(
	a: torch.Tensor, indices: torch.Tensor, dim: Optional[int] = None
	) -> torch.Tensor:
	torch._check(
	a.ndim == indices.ndim,
	lambda: (
	"torch.take_along_dim(): input and indices should have the same "
	f"number of dimensions, but got {a.ndim} dimensions for input, and "
	f"{indices.ndim} dimensions for indices"
	),
	)
	torch._check(
	utils.is_integer_dtype(indices.dtype),
	lambda: (
	"torch.take_along_dim(): dtype of indices should be int but got "
	f"{indices.dtype} instead"
	),
	)
	if dim is None:
	return torch.gather(a.view(-1), 0, indices.view(-1))
	else:
	self_sizes = list(a.shape)
	self_sizes[dim] = indices.size(dim)
	broadcast_shape = utils.infer_size_shapes(self_sizes, indices.size())
	indices_broadcast = broadcast_to(indices, broadcast_shape)
	indices_sizes = list(indices.shape)
	indices_sizes[dim] = a.size(dim)
	broadcast_shape = utils.infer_size_shapes(indices_sizes, a.size())
	self_broadcast = broadcast_to(a, broadcast_shape)
	return torch.gather(self_broadcast, dim, indices_broadcast)

[XilunWu]

LGTM! Thx for the work!

[tianyu-l]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here