pytorch/test/distributed/tensor
Ruisi Zhang f74842d665 [DTensor] enable SimpleFSDP's composability with Tensor Parallel (#152286)
This PR adds support for SimpleFSDP's composability with Tensor Parallel + torch.compile.

`_StridedShard` is used in SimpleFSDP/FSDP2 to support correct distributed checkpointing when FSDP+TP is applied. Previously, `_StridedShard` is not guarded by torch.compile. This PR adds `_StridedShard` as an additional placement type to be guarded by torch.compile.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152286
Approved by: https://github.com/bdhirsh
2025-05-23 01:40:38 +00:00
..
debug [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
experimental [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
parallel [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
__init__.py [dtensor] move all tests to distribute/tensor folder (#144166) 2025-01-08 00:32:33 +00:00
README.md [dtensor] move all tests to distribute/tensor folder (#144166) 2025-01-08 00:32:33 +00:00
test_api.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_attention.py [CP] Use TorchFunctionMode to dispatch SDPA for CP (#147902) 2025-04-25 23:33:48 +00:00
test_common_rules.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_convolution_ops.py [DTensor][conv] add DTensor convolution_backward op support for case where the input Tensor has requires_grad=False (#142278) 2025-02-10 07:06:40 +00:00
test_dtensor_compile.py [DTensor] enable SimpleFSDP's composability with Tensor Parallel (#152286) 2025-05-23 01:40:38 +00:00
test_dtensor_ops.py torch.tensordot: performance improvements when contracting to a scalar. (#145936) 2025-05-13 10:57:30 +00:00
test_dtensor.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_embedding_ops.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_experimental_ops.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_init.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_math_ops.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_matrix_ops.py [dtensor] add op support for torch._grouped_mm (#151072) 2025-04-12 07:07:44 +00:00
test_op_strategy.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_optimizers.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_pointwise_ops.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_random_ops.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_redistribute.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_tensor_ops.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00
test_utils.py Implement util function compute_global_tensor_shape for 1D device mesh (#152751) 2025-05-05 02:44:31 +00:00
test_view_ops.py [DTensor] Error on illegal view op during sharding prop (#149764) 2025-04-28 18:21:49 +00:00
test_xla_integration.py [BE][DTensor] move torch.distributed._tensor import to torch.distributed.tensor in test files (#153225) 2025-05-09 20:40:54 +00:00

Run distributed tensor tests:

from root, run (either CPU or GPU)

pytest test/distributed/tensor/test_dtensor.py

run specific test cases and print stdout/stderr:

pytest test/distributed/tensor/test_dtensor.py -s -k test_from_local