saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-15 21:00:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
pytorch/test/distributed/_tensor
History
wz337 fed536dbcf [DTensor][Optim] Add support for fused_adam and fused_adamw when lr is a tensor (#126750) 
Fixes #126670

In this PR, we update the following:
1. lr is an kwarg. Add support to automatically turn on implict replication for kwarg. We only did this for arg previously.
2. add associated tensor_lr ops in pointwises.py
3. add associated unit test in test_optimizers.py

Pull Request resolved: https://github.com/pytorch/pytorch/pull/126750
Approved by: https://github.com/wanchaol, https://github.com/msaroufim
2024-05-21 21:38:05 +00:00
..
debug [comm_mode] adding some initial c10d ops to CommDebugMode (#125475) 2024-05-04 04:20:46 +00:00
experimental [RELAND] Switch default behavoir of export IR to be predispatch (#125860) 2024-05-10 17:36:53 +00:00
__init__.py
README.md
test_api.py Revert "try to fix the warning in distribute_tensor (#125476)" 2024-05-04 22:25:32 +00:00
test_attention.py distributed: templated ring attention (#124215) 2024-04-19 00:57:08 +00:00
test_common_rules.py [dtensor] refactor schema suggestions in output sharding (#122929) 2024-04-01 17:39:39 +00:00
test_convolution_ops.py
test_dtensor.py [dtensor] early return for _split_tensor (#125810) 2024-05-14 04:59:27 +00:00
test_dtensor_compile.py AsyncCollectiveTensor: prevent wait_tensor() calls on graph inputs from getting DCEd (#125677) 2024-05-08 15:54:01 +00:00
test_dtensor_ops.py Fix index_reduce sampler filter when op_info.variant_test_name is specified (#123375) 2024-04-17 15:31:28 +00:00
test_embedding_ops.py
test_experimental_ops.py
test_init.py
test_math_ops.py
test_matrix_ops.py [dtensor] add op support for memory efficient attention (#122996) 2024-05-08 17:08:27 +00:00
test_op_strategy.py
test_optimizers.py [DTensor][Optim] Add support for fused_adam and fused_adamw when lr is a tensor (#126750) 2024-05-21 21:38:05 +00:00
test_pointwise_ops.py [dtensor] use str for reduce_op (#125172) 2024-04-29 23:30:24 +00:00
test_random_ops.py
test_redistribute.py [dtensor] implement shard dim change with alltoall (#124872) 2024-04-30 18:30:34 +00:00
test_tensor_ops.py [dtensor] improve new factory strategy (#122995) 2024-05-08 17:05:07 +00:00
test_utils.py [comm_mode] adding some initial c10d ops to CommDebugMode (#125475) 2024-05-04 04:20:46 +00:00
test_view_ops.py [dtensor] refactor view ops to use OpStrategy (#126011) 2024-05-17 05:39:21 +00:00
test_xla_integration.py

README.md

Run distributed tensor tests:

from root, run (either CPU or GPU)

pytest test/spmd/tensor/test_tensor.py

pytest test/spmd/tensor/test_ddp.py

run specific test case and print stdout/stderr:

pytest test/spmd/tensor/test_tensor.py -s -k test_tensor_from_local