pytorch

mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00

Files

soulitzer 1e4dfeeb06 Add early_stop kwarg to torch.utils.checkpoint (#160781 )

We already have a context manager "set_checkpoint_early_stop". This PR adds a kwarg that toggles the same setting.

It is also useful to have a kwarg version of the setting in addition to the context manager because is annoying to apply a context manager when the AC is being applied via CheckpointWrapper.

Similar to the "debug" kwarg and the corresponding "set_checkpoint_debug_enabled" context manager, the context manager defaults to None and overrides the local setting when non-None.

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

2025-08-26 22:32:35 +00:00

fsdp

[FSDP2] Fix backward-compatible imports (#142419 )

2024-12-09 23:56:32 +00:00

__init__.py

Remove old FSDP1 fully_shard (#141875 )

2024-12-03 17:00:47 +00:00

checkpoint_activation.py

Add early_stop kwarg to torch.utils.checkpoint (#160781 )

2025-08-26 22:32:35 +00:00

contract.py

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