39 lines
1.3 KiB
ReStructuredText
39 lines
1.3 KiB
ReStructuredText
Activation Checkpointing
|
|
========================
|
|
|
|
The activation checkpointing API's in DeepSpeed can be used to enable a range
|
|
of memory optimizations relating to activation checkpointing. These include
|
|
activation partitioning across GPUs when using model parallelism, CPU
|
|
checkpointing, contiguous memory optimizations, etc.
|
|
|
|
Please see the `DeepSpeed JSON config <https://www.deepspeed.ai/docs/config-json/>`_
|
|
for the full set.
|
|
|
|
Here we present the activation checkpointing API. Please see the enabling
|
|
DeepSpeed for `Megatron-LM tutorial <https://www.deepspeed.ai/tutorials/megatron/>`_
|
|
for example usage.
|
|
|
|
Configuring Activation Checkpointing
|
|
------------------------------------
|
|
.. autofunction:: deepspeed.checkpointing.configure
|
|
|
|
.. autofunction:: deepspeed.checkpointing.is_configured
|
|
|
|
|
|
Using Activation Checkpointing
|
|
------------------------------
|
|
.. autofunction:: deepspeed.checkpointing.checkpoint
|
|
|
|
.. autofunction:: deepspeed.checkpointing.reset
|
|
|
|
|
|
Configuring and Checkpointing Random Seeds
|
|
------------------------------------------
|
|
.. autofunction:: deepspeed.checkpointing.get_cuda_rng_tracker
|
|
|
|
.. autofunction:: deepspeed.checkpointing.model_parallel_cuda_manual_seed
|
|
|
|
.. autoclass:: deepspeed.checkpointing.CudaRNGStatesTracker
|
|
|
|
.. autoclass:: deepspeed.checkpointing.CheckpointFunction
|