pytorch/torch/distributed
Brian Hirsh e68f5087d8 update _unsafe_set_version_counter to accept lists of tensors (#137921)
See the comment [here](https://github.com/pytorch/pytorch/issues/132014#issuecomment-2379547400) (cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @XilunWu @rec) - this PR updates `_unsafe_set_version_counter` to accept a list of tensors, for overhead-sensitive users (e.g. distributed) who need to hide VC bumps from autograd on a large list of tensors without wanting to suffer the overhead of going from python->C++ separately for every tensor in the list.

I left the binding in pybind, and used a `std::vector`. if we **really** need to optimize overhead even further, we could write a manual cpython binding.

I use this updated API in the next PR to fix FSDP2, so that it properly hides the VC of all `all_gather_buffer` tensors in its call to `split_with_sizes_copy.out(all_gather_buffers)`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/137921
Approved by: https://github.com/awgu, https://github.com/albanD
2025-02-04 04:51:11 +00:00
..
_composable PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_shard PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_sharded_tensor
_sharding_spec
_symmetric_memory [async-TP] Fix scheduling in matmul+reduce-scatter for 2 ranks (#145846) 2025-01-30 18:26:34 +00:00
_tensor
_tools Turn on mypy for _dynamo/variables/builtin.py (#145552) 2025-01-30 22:21:32 +00:00
algorithms PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
autograd
benchmarks
checkpoint [DCP] Remove all-gather of state dict keys (#145998) 2025-02-04 03:16:13 +00:00
elastic Make torchelastic etcd rendezvous publicly importable (#145396) 2025-01-23 23:56:45 +00:00
examples
fsdp update _unsafe_set_version_counter to accept lists of tensors (#137921) 2025-02-04 04:51:11 +00:00
launcher PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
nn PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
optim PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
pipelining [NFC] Fix some minor typos. (#145599) 2025-01-24 18:58:59 +00:00
rpc PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
tensor DeepSpeed github repo move sync (#146320) 2025-02-03 23:20:49 +00:00
__init__.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_checkpointable.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_composable_state.py
_functional_collectives.py [2/N] Enable ruff F841 on distributed tests (#146132) 2025-02-02 03:44:48 +00:00
_functional_collectives_impl.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
_state_dict_utils.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
argparse_util.py
c10d_logger.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
collective_utils.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
constants.py
CONTRIBUTING.md
device_mesh.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
distributed_c10d.py [NFC] Fix some minor typos. (#145599) 2025-01-24 18:58:59 +00:00
launch.py
logging_handlers.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
remote_device.py
rendezvous.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00
run.py Improve torchrun documentation (#144354) 2025-01-24 20:40:05 +00:00
utils.py PEP585 update - torch/distributed (#145164) 2025-01-21 04:23:29 +00:00