Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
activation-checkpointing.rst 1.28 KB
Newer Older
Jeff Rasley's avatar
ZeRO-2 (#217)  
Jeff Rasley committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38 
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