Wei Feng cfb8aec1a4 [FSDP2] idempotent reset_sharded_param: no-op if _local_tensor is already padded (#163130) ...

resolves https://github.com/pytorch/torchtitan/issues/1136

torchtitan use cached state dict for ft. reset_sharded_param should be idempotent if model.parameters() are padded already

```
# pad DTensor._local_tensor
fully_shard(model)
sd = fsdp_model.state_dict()
# reset_sharded_param should be a no-op in lazy_init
loss = fsdp_model(inp).sum()
```

this PR make `reset_sharded_param` idempotent by checking storage data ptr and return early

unit test
```
pytest -s test/distributed/_composable/fsdp/test_fully_shard_state_dict.py -k test_cached_state_dict
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/163130
Approved by: https://github.com/tianyu-l

2025-09-18 09:20:37 +00:00

..

_fully_shard

[FSDP2] idempotent reset_sharded_param: no-op if _local_tensor is already padded (#163130)

2025-09-18 09:20:37 +00:00

__init__.py

[FSDP2] Move to public torch.distributed.fsdp (#141868)

2024-12-07 01:24:28 +00:00

_common_utils.py

[BE][Ez]: Use itertools.chain.from_iterable when possible (#148190)

2025-03-06 20:37:06 +00:00

_debug_utils.py

[BE][5/16] fix typos in torch/ (torch/distributed/) (#156315)

2025-06-23 02:57:28 +00:00

_dynamo_utils.py

PEP585 update - torch/distributed/fsdp (#145162)

2025-01-19 20:04:05 +00:00

_exec_order_utils.py

[BE][PYFMT] migrate PYFMT for torch.{distributed,distributions} to ruff format (#144547)

2025-02-28 07:35:56 +00:00

_flat_param.py

Fix: Ensure writeback handles NO_SHARD correctly by flattening tensors before copying (#154369)

2025-07-06 09:20:31 +00:00

_fsdp_extensions.py

PEP585 update - torch/distributed/fsdp (#145162)

2025-01-19 20:04:05 +00:00

_init_utils.py

[BE][5/16] fix typos in torch/ (torch/distributed/) (#156315)

2025-06-23 02:57:28 +00:00

_limiter_utils.py

PEP585 update - torch/distributed/fsdp (#145162)

2025-01-19 20:04:05 +00:00

_optim_utils.py

[doc]: Small typos (#162982)

2025-09-16 17:42:19 +00:00

_runtime_utils.py

[BE][5/16] fix typos in torch/ (torch/distributed/) (#156315)

2025-06-23 02:57:28 +00:00

_shard_utils.py

[BE][PYFMT] migrate PYFMT for torch.{distributed,distributions} to ruff format (#144547)

2025-02-28 07:35:56 +00:00

_state_dict_utils.py

[BE] add noqa for flake8 rule B036: found except BaseException without re-raising (#159043)

2025-07-25 02:56:34 +00:00

_trace_utils.py

[BE][PYFMT] migrate PYFMT for torch.{distributed,distributions} to ruff format (#144547)

2025-02-28 07:35:56 +00:00

_traversal_utils.py

[BE][5/16] fix typos in torch/ (torch/distributed/) (#156315)

2025-06-23 02:57:28 +00:00

_unshard_param_utils.py

[BE][PYFMT] migrate PYFMT for torch.{distributed,distributions} to ruff format (#144547)

2025-02-28 07:35:56 +00:00

_wrap_utils.py

PEP585 update - torch/distributed/fsdp (#145162)

2025-01-19 20:04:05 +00:00

api.py

[BE][PYFMT] migrate PYFMT for torch.{distributed,distributions} to ruff format (#144547)

2025-02-28 07:35:56 +00:00

fully_sharded_data_parallel.py

[doc] remove FSDP1 developer note (#158991)

2025-07-24 08:21:54 +00:00

sharded_grad_scaler.py

[BE][PYFMT] migrate PYFMT for torch/[a-c]*/ to ruff format (#144554)

2025-07-03 18:56:07 +00:00

wrap.py

Use Python 3.9 typing (#148157)

2025-03-04 03:09:55 +00:00