wz337 0cdc6a8dcd [DSD] Fix distributed state dict full_state_dict option hang during set_state_dict (#135725) ...

Fix https://github.com/pytorch/pytorch/issues/134095
This fix distributed state dict full_state_dict option hang during set_state_dict. We switch `_distribute_tensors` in _state_dict_utils.py to use `DTensor.from_local` instead of `distribute_tensor` to support FSDP2+TP 2D strided sharding use case, as `distribute_tensor` cannot handle strided sharding yet. `distribute_tensor` incurs a scatter behind the scenes, while `DTensor.from_local` takes the local slice from the full tensor on each rank to create the DTensor (no collective).  This means it's the user's responsibility to make sure the full_tensor from the full_state_dict is the same across all ranks.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135725
Approved by: https://github.com/fegin

2024-09-13 03:26:36 +00:00

..

_composable

[FSDP2] better error msg for cpu offloading (#135156)

2024-09-12 00:05:07 +00:00

_shard

[BE][Easy] enable ruff rule PIE790: unnecessary pass statement (#133200)

2024-08-15 15:50:19 +00:00

_sharded_tensor

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

_sharding_spec

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

_symmetric_memory

[micro_pipeline_tp] support all _scaled_mm args (#131984)

2024-08-05 21:44:37 +00:00

_tensor

[reland][dtensor] move DTensor to public namespace (#134203)

2024-09-08 17:08:40 +00:00

_tools

Runtime Estimator for estimating GPU compute time (#134243)

2024-08-28 20:06:54 +00:00

algorithms

[BE][Easy] enable ruff rule PIE790: unnecessary pass statement (#133200)

2024-08-15 15:50:19 +00:00

autograd

[BE][Easy] enable UFMT for torch/distributed/{algorithms,autograd,benchmarks,checkpoint,elastic}/ (#128866)

2024-06-18 13:51:53 +00:00

benchmarks

[BE][Easy] enable ruff rule PIE790: unnecessary pass statement (#133200)

2024-08-15 15:50:19 +00:00

checkpoint

[DCP] Fixes the stateless optimizer issue of distributed state_dict (#135535)

2024-09-10 03:10:00 +00:00

elastic

Adding entry-point based support for out-of-tree rendezvous plugins (#132633)

2024-09-11 03:35:02 +00:00

examples

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

fsdp

[reland][dtensor] move DTensor to public namespace (#134203)

2024-09-08 17:08:40 +00:00

launcher

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

nn

Revert "added persistent option to buffers and namedbuffers (#132994)"

2024-08-09 18:14:53 +00:00

optim

Revert "[BE] typing for decorators - _jit_internal (#131573)"

2024-07-28 03:29:32 +00:00

pipelining

[PP] Fix zero bubble composability with DP (#134052)

2024-09-04 23:46:29 +00:00

rpc

[BE][Easy] enable ruff rule PIE790: unnecessary pass statement (#133200)

2024-08-15 15:50:19 +00:00

tensor

[dtensor][BE] replace compute_local_shape with compute_local_shape_and_global_offset (#135554)

2024-09-12 06:30:09 +00:00

__init__.py

Remove ProcessGroupRoundRobin (#132888)

2024-08-08 01:07:40 +00:00

_checkpointable.py

[torchrec][pt-d][model store] introduce LocalShardsWrapper for DTensor (#129150)

2024-06-21 01:58:51 +00:00

_composable_state.py

Fix docstring errors in _composable_state.py, remote_device.py, value_ranges.py, utils.py, run.py, rendezvous.py, launch.py, argparse_util.py, __init__.py, _cycles.py (#112953)

2023-11-08 01:13:09 +00:00

_functional_collectives_impl.py

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

_functional_collectives.py

[reland][dtensor] move DTensor to public namespace (#134203)

2024-09-08 17:08:40 +00:00

_state_dict_utils.py

[DSD] Fix distributed state dict full_state_dict option hang during set_state_dict (#135725)

2024-09-13 03:26:36 +00:00

argparse_util.py

Flip default value for mypy disallow_untyped_defs [5/11] (#127842)

2024-06-08 18:49:18 +00:00

c10d_logger.py

[DCP] Fix duplicated logging messages when enable both c10d and dcp l… (#130423)

2024-07-11 13:43:39 +00:00

collective_utils.py

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

constants.py

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

CONTRIBUTING.md

Clean up distributed/CONTRIBUTING.md (#128450)

2024-06-22 02:41:22 +00:00

device_mesh.py

[DeviceMesh][Easy] Make RuntimeError a bit more descriptive by including the actual world_size (#135271)

2024-09-06 06:23:20 +00:00

distributed_c10d.py

Revert "[c10d] Remove Option for ProcessGroup and Expose backend Options to reflect the correct code structure (#132931)"

2024-08-30 16:27:40 +00:00

launch.py

Flip default value for mypy disallow_untyped_defs [6/11] (#127843)

2024-06-08 18:49:29 +00:00

logging_handlers.py

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

remote_device.py

[BE][Easy] fix ruff rule needless-bool (SIM103) (#130206)

2024-07-14 08:17:52 +00:00

rendezvous.py

[BE][Easy] enable UFMT for torch/distributed/ (#128870)

2024-06-22 18:53:28 +00:00

run.py

fix torchrun log message (#131652)

2024-07-25 14:50:10 +00:00

utils.py

[FSDP] casting input args with dataclass(frozen=True) (#135067)

2024-09-05 01:19:53 +00:00