pytorch/torch/distributed/algorithms
Rohan Varma a197f3fe52 [FSDP/Checkpoint] Activation offload support in checkpoint_wrapper (#70165)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/70165

Implements activation offload support in checkpoint_wrapper API via
save_on_cpu hooks. We avoid modifying the torch.utils.checkpoint implementation
and instead compose offload + checkpoint using the save_on_cpu hook for the
former.
ghstack-source-id: 146078900

Test Plan: CI

Reviewed By: zhaojuanmao

Differential Revision: D33228820

fbshipit-source-id: 98b4da0828462c41c381689ee07360ad014e808a
2021-12-21 10:08:18 -08:00
..
_checkpoint [FSDP/Checkpoint] Activation offload support in checkpoint_wrapper (#70165) 2021-12-21 10:08:18 -08:00
ddp_comm_hooks [DDP Comm Hook] Create a noop hook for performance debugging (#64344) 2021-09-01 17:36:22 -07:00
model_averaging [Reland] [Model Averaging] Simplify PostLocalSGD Optimizer API (#65197) 2021-09-17 10:31:58 -07:00
quantization [BE] minor improvement to dist quantization (#67401) 2021-11-21 23:31:59 -08:00
__init__.py Make _Join, _Joinable, _JoinHook public (#62605) 2021-08-03 12:20:11 -07:00
join.py Make _Join, _Joinable, _JoinHook public (#62605) 2021-08-03 12:20:11 -07:00