mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
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 |
||
|---|---|---|
| .. | ||
| debug | ||
| experimental | ||
| parallel | ||
| __init__.py | ||
| README.md | ||
| test_api.py | ||
| test_attention.py | ||
| test_common_rules.py | ||
| test_convolution_ops.py | ||
| test_dtensor_compile.py | ||
| test_dtensor_ops.py | ||
| test_dtensor.py | ||
| test_embedding_ops.py | ||
| test_experimental_ops.py | ||
| test_init.py | ||
| test_math_ops.py | ||
| test_matrix_ops.py | ||
| test_op_strategy.py | ||
| test_optimizers.py | ||
| test_pointwise_ops.py | ||
| test_random_ops.py | ||
| test_redistribute.py | ||
| test_tensor_ops.py | ||
| test_utils.py | ||
| test_view_ops.py | ||
| test_xla_integration.py | ||
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