Gradient checkpointing

Gradient checkpointing, also called activation checkpointing, activation recomputation, or rematerialization, is a memory-saving technique for training deep neural networks that trades extra compute for much lower peak memory: instead of storing every intermediate activation from the forward pass, it stores only a small subset (the "checkpoints") and recomputes the rest on demand during the backward pass.^[1][2] The technique was introduced by Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin in the 2016 paper "Training Deep Nets with Sublinear Memory Cost" (arXiv:1604.06174), which showed that a network of n layers can be trained with $O(\sqrt{n})$ activation memory at the cost of a single extra forward pass per mini-batch.^[1] In their experiments this cut the memory of a 1,000-layer residual network from roughly 48 GB to about 7 GB for around 30 percent additional running time.^[1] As the paper states, "we propose a novel method to trade computation for memory."^[1]

Gradient checkpointing is implemented in every major training framework, including pytorch (torch.utils.checkpoint), jax (jax.checkpoint, aliased to jax.remat), tensorflow (tf.recompute_grad), deepspeed, and fsdp (apply_activation_checkpointing).^[2][3][4][5] Higher-level libraries expose it as a single flag: Hugging Face transformers enables it with model.gradient_checkpointing_enable() or with gradient_checkpointing=True on TrainingArguments.^[16] It is a standard component of large language model training, used in llama, GPT-NeoX-20B, and the megatron lm family among many others.^[6][7][13]

What is gradient checkpointing? (background)

Training a neural network with backpropagation requires the gradient of the loss with respect to every parameter. Computing those gradients via the chain rule needs intermediate values from the forward pass, called activations, at every layer where the gradient is propagated through a non-linear or otherwise non-invertible operation.^[8] In the most direct implementation of reverse-mode automatic differentiation, all activations are stored during the forward pass and read back during the backward pass. Memory consumption therefore grows linearly with the depth of the network and with the batch size, sequence length, and hidden dimension.^[8]

For a transformer block, the activations that must be stored include the input to attention, the queries, keys, values, the attention probability matrix (of size sequence_length squared), the attention output, the input to the feed-forward layer, the intermediate feed-forward activation (typically four times the hidden dimension), and the layer-norm statistics. Korthikanti and colleagues derive an approximate per-layer formula and show that, without recomputation, per-layer activation memory is about $sbh(34 + 5as/h)$ bytes: linear in batch size and hidden dimension, plus an attention term that grows with the number of heads and quadratically with sequence length.^[6] At a context length of 2,048, 96 attention heads and hidden size of 12,288 (the GPT-3 configuration), the attention term $5as/h$ equals 80, more than twice the factor of 34 for the rest of the layer, so the attention activations make up most of each layer's activation memory.^[6]

As neural networks grew in depth and width during the 2010s, this linear scaling became a bottleneck. Chen and colleagues observed in 2016 that on a 1,000-layer deep residual network the activations occupied roughly 48 GB, exceeding the memory of available GPUs.^[1] The same constraint reappeared at much larger scale with the transformer architecture: when training a gpt-3 style 530-billion-parameter model, activation memory alone exceeded what a single NVIDIA A100 GPU with 80 GB of memory could hold even after applying tensor and pipeline parallelism.^[6] Backpropagation thus requires explicit strategies to control activation memory, of which gradient checkpointing is the most widely adopted.^[1][6]

How is it different from a model checkpoint?

Gradient checkpointing is often confused with model checkpointing, but the two are unrelated. A model checkpoint periodically saves model weights (and optimizer state) to disk for fault tolerance and resumption during long training runs. Gradient checkpointing, by contrast, operates entirely within a single backward pass and concerns intermediate activations held in GPU memory, not parameter snapshots written to disk. One is about saving progress; the other is about saving memory during the gradient computation.

Predecessor ideas

The core idea (trade compute for memory by recomputing intermediates) predates the 2016 paper by decades. The automatic differentiation literature describes "checkpointing" or "revolve" schemes for adjoint computation in the context of ODE solvers and PDE-constrained optimisation, including Andreas Griewank's work in the 1990s on recursive binomial checkpointing. Chen et al. cite this earlier literature and adapt it to deep network training graphs, but their contribution is the specific $O(\sqrt{n})$ recipe and an automated, framework-level implementation.^[1] The reframing for deep learning, and the practical demonstration on a 1,000-layer ResNet, opened the technique to widespread adoption in mainstream training frameworks.^[1]

How does gradient checkpointing work? (algorithm)

The algorithm proposed by Chen et al. partitions an n-layer feed-forward network into segments and stores only the activation at the boundary of each segment.^[1] During the backward pass, when gradients are needed for a layer inside a segment, the forward pass is re-executed from the most recent stored checkpoint up to that layer, regenerating the missing activations on demand.^[1] Once the gradient for the layer has been computed, the recomputed activations are discarded. In the PyTorch implementation, as the documentation puts it, "the unsaved tensors are recomputed by re-invoking function in the backward pass as needed for gradient computation."^[2]

Formally, consider a network expressed as a composition $f = f_L \circ \cdots \circ f_2 \circ f_1$, where each $f_i$ takes the activation $a_{i-1}$ and produces $a_i$. Standard backpropagation stores all $a_i$ during the forward pass. With checkpointing, the layers are grouped into segments; only the activations at segment boundaries are persisted. To compute the gradient $da_i/d\theta_i$ inside a segment, the activations $a_{j+1}, \ldots, a_{i-1}$ (where $a_j$ is the previous checkpoint) are recomputed by re-running the forward pass on the segment.^[1] The gradient is then accumulated along the chain rule exactly as in standard backpropagation; the only difference is when and how the activations are materialised.^[1]

Chen et al. observed that placing approximately $\sqrt{n}$ checkpoints in an n-layer network is the optimum of the simple memory-vs-compute tradeoff: each segment then holds at most $\sqrt{n}$ activations in memory at any one time, and the total recomputed work is one extra forward pass.^[1] Recursive application of the same strategy reduces memory further to $O(\log n)$ at the cost of $O(n \log n)$ extra forward computation, useful only in extreme regimes.^[1] In an implementation released alongside the paper, the authors reduced memory of a 1,000-layer residual network from 48 GB to 7 GB with 30 percent additional running time on the ImageNet problem.^[1]

A widely cited illustrative reimplementation by Tim Salimans and Yaroslav Bulatov demonstrated that a TensorFlow training script using their automatic graph rewrite could fit "more than 10x larger models" onto a single GPU at roughly 20 percent additional compute time, popularising the technique under the name "gradient checkpointing".^[9]

How much memory does it save? (mathematical analysis)

Let n be the number of layers and c be the number of checkpoints. Activation memory after checkpointing scales as $O(n/c + c)$, since at any instant the implementation holds c checkpoints plus the activations recomputed for the current segment, which contains $n/c$ layers. Minimising $n/c + c$ over $c$ yields $c = \sqrt{n}$ and a memory cost of $O(\sqrt{n})$.^[1] Compute, in contrast, scales as $O(n)$ per backward pass plus $O(n)$ for the additional segment-by-segment recomputation, so the forward computation is effectively performed twice. In practice, because the backward pass already does about twice as much work as the forward pass under standard backpropagation, the total wall-clock cost of training with checkpointing is roughly 33 percent higher than without it, in line with the 30 percent number reported by Chen et al.^[1] The paper frames the result precisely: it gives "an algorithm that costs $O(\sqrt{n})$ memory to train a n layer network, with only the computational cost of an extra forward pass per mini-batch."^[1]

A worked numerical example clarifies the magnitudes. Consider a model shaped like LLaMA-65B (80 layers, hidden size 8,192, 64 attention heads) at sequence length 4,096 and a micro-batch of 1. Korthikanti et al. estimate per-layer activation memory without recomputation, in bytes with 16-bit storage and no model parallelism, as $sbh(34 + 5as/h)$, while full recomputation at layer boundaries keeps only each layer's $2sbh$-byte input.^[6] By that formula, the example needs about 6.5 GB per layer, or roughly 520 GB for the 80-layer stack, without checkpointing, against about 67 MB per layer (about 5.4 GB in total) with full checkpointing, plus the activations of the one layer being recomputed; tensor and sequence parallelism divide these figures further.^[6] Korthikanti et al. depict this graphically across the 22B, 175B, 530B, and 1T model configurations, showing total baseline memory (parameters, optimizer state and activations) above the 80 GB A100 ceiling in every case, while with the "present work" (sequence parallelism plus selective recomputation) every configuration fits, with activations taking under about 30 GB.^[6]

For more general computation graphs, the optimisation problem becomes a graph partition: pick a subset S of edges (the "cut") such that every node has an ancestor in S, minimise the maximum size of a connected component of (graph minus S), and minimise total recomputation cost. Chen et al. handle general graphs with a heuristic plan search under a memory budget (their Algorithm 3), and approaches close to the original sqrt-n scheme work well for the layer-stacked computation graphs typical of feed-forward networks and transformer models.^[1] More sophisticated dynamic programming solutions exist for chain-structured graphs and are sometimes called Treeverse or revolve algorithms in the automatic differentiation literature.

A second important consideration is the size of each saved activation. Two checkpoints of equal "count" can have very different memory footprints if the network's hidden dimension changes across layers (e.g., U-Nets or models with bottlenecks). Implementations therefore use byte-size-weighted variants of the cut problem in practice. The original Chen et al. analysis assumes uniform layer cost; this matches deep ResNets and isotropic transformers well but is less appropriate for encoder-decoder architectures with varying feature-map sizes.^[1]

What variants of gradient checkpointing exist?

Selective activation recomputation

Vijay Korthikanti and collaborators at NVIDIA introduced selective activation recomputation in the May 2022 paper "Reducing Activation Recomputation in Large Transformer Models".^[6] Their observation: in a transformer block, certain operations (notably the $QK^\top$ matrix multiply, softmax, softmax dropout and attention over V inside the attention block) consume a disproportionate amount of activation memory relative to the compute they require, while other operations (large matrix multiplications) are the opposite.^[6] The selective strategy stores activations for the compute-expensive operations and recomputes only the memory-heavy, compute-light operations.^[6]

Combined with sequence parallelism, this technique reduces activation memory by 5x while reducing the execution-time overhead of activation recomputation by over 90 percent.^[6] When training a 530B GPT-3 style model on 2,240 A100 GPUs, the authors achieved 54.2 percent Model FLOPs Utilisation, a 29 percent improvement over the 42.1 percent obtained with full activation recomputation.^[6] The implementation is part of the open-source megatron lm codebase.^[6]

Asynchronous offload

A second variant moves checkpointed activations from GPU memory to either CPU host memory or even disk between the forward and backward passes, prefetching them back asynchronously when needed.^[4] This adds a third axis to the tradeoff: GPU memory, CPU memory, and compute. deepspeed supports CPU offloading via the checkpoint_in_cpu parameter inside its activation checkpointing configuration, which the configure() docstring limits to activations partitioned across model parallelism ranks; the current configuration documentation (JSON key cpu_checkpointing) says that without partitioning it instead uses an asynchronous pinned side-stream copy that overlaps the transfer with compute.^[4][10] PyTorch's distributed checkpoint_wrapper.py similarly provides a separate offload_wrapper that offloads a wrapped module's saved activations to pinned CPU memory and can be composed with recomputation-based checkpointing.^[5]

Partitioned activation checkpointing

In tensor parallelism setups (as used in megatron lm), each GPU holds only a shard of the model. DeepSpeed's partition_activations option distributes the activation checkpoints themselves across the tensor parallel ranks, dividing activation memory per GPU by the tensor parallel size.^[4] When combined with contiguous_memory_optimization, the partitioned checkpoints are copied into a single contiguous buffer, reducing memory fragmentation.^[4]

Compiler-managed rematerialisation

In jax, the jax.checkpoint decorator (alias jax.remat) is implemented as a transformation that interacts with the XLA compiler.^[3] Recent versions add a policy argument that controls which intermediate values are saved versus recomputed, with built-in policies such as dots_with_no_batch_dims_saveable, save_only_these_names, and offload_dot_with_no_batch_dims.^[11] The JAX documentation notes that when functions are compiled with jax.jit, XLA may apply rematerialisation transforms automatically, so an explicit jax.checkpoint is most useful around staged control flow such as jax.lax.scan, which is common in transformer implementations.^[11]

How is gradient checkpointing implemented in PyTorch, JAX, and TensorFlow?

PyTorch

pytorch exposes the technique through torch.utils.checkpoint.^[2] The PyTorch documentation summarises the idea in one line: "Activation checkpointing is a technique that trades compute for memory."^[2] The primary entry point is torch.utils.checkpoint.checkpoint(function, *args, use_reentrant=None, ...), which runs function(*args) without saving its intermediate activations and arranges for them to be recomputed during the backward pass.^[2] A checkpoint_sequential(functions, segments, input) helper applies the technique to a nn.Sequential model by splitting it into a chosen number of segments.^[2]

The library supports two modes selected by the use_reentrant flag. The original, reentrant implementation invokes a nested backward pass through the autograd engine; it has limitations around inputs that require gradients and around hooks. The non-reentrant implementation (use_reentrant=False) was added later, supports more functionality, and is the recommended choice; PyTorch documentation states that the parameter must be explicitly specified and that future versions will raise an exception if it is omitted.^[2] Checkpointed functions preserve the random number generator state by default so that operations such as dropout produce the same values on the recomputation pass; this can be disabled with preserve_rng_state=False at the cost of correctness for stochastic operations.^[2]

For distributed training, PyTorch ships apply_activation_checkpointing in torch.distributed.algorithms._checkpoint.checkpoint_wrapper, which walks a model and wraps submodules selected by a user-supplied check_fn or auto_wrap_policy.^[5] For fsdp training, the recommended pattern is to wrap each transformer block (attention plus feed-forward) and to use non-reentrant checkpointing, which is the wrapper's default (CheckpointImpl.NO_REENTRANT).^[5] For CPU offload the same module provides a separate offload_wrapper, which can be passed to apply_activation_checkpointing as the checkpoint_wrapper_fn.^[5]

The Hugging Face transformers library exposes the same PyTorch mechanism through a single switch: calling model.gradient_checkpointing_enable(), or setting gradient_checkpointing=True on the TrainingArguments object, wraps the model's layers with torch.utils.checkpoint, which is why the flag is the most common way practitioners turn the technique on for fine-tuning.^[16]

JAX

jax provides jax.checkpoint, aliased to jax.remat, with the signature jax.checkpoint(fun, *, prevent_cse=True, policy=None, static_argnums=(), static_argnames=()).^[3] Applying it as a decorator marks the wrapped function so that during reverse-mode autodiff its intermediates are recomputed rather than stored.^[3] As noted above, JAX exposes several policies for fine-grained control over what is checkpointed.^[11]

TensorFlow

tensorflow exposes the technique through tf.recompute_grad(f), a decorator that wraps a function so that its forward computation is repeated during the backward pass instead of caching intermediate tensors.^[12] The early popular implementation of the technique on TensorFlow was the cybertronai/gradient-checkpointing library by Tim Salimans and Yaroslav Bulatov, which automatically rewrote a TensorFlow computation graph to insert checkpoints; the project's README lists its status as maintenance (bug fixes and minor updates), and TensorFlow's built-in tf.recompute_grad API now provides the same capability.^[9][12]

DeepSpeed

deepspeed provides a dedicated activation checkpointing API in the deepspeed.checkpointing module. Key functions include deepspeed.checkpointing.configure() to set options, deepspeed.checkpointing.checkpoint() to checkpoint a block, deepspeed.checkpointing.is_configured(), deepspeed.checkpointing.reset() to clear buffers between forward passes, and deepspeed.checkpointing.model_parallel_cuda_manual_seed() to initialise RNG state for model-parallel runs.^[4] The configuration options include partition_activations (split checkpoints across tensor-parallel ranks), contiguous_memory_optimization (copy checkpoints into a contiguous buffer, requires partitioning to be enabled), checkpoint_in_cpu (offload checkpoints to CPU memory; cpu_checkpointing in the JSON config, which as of November 2026 also works without partitioning via an asynchronous copy), and profile (log forward and backward time per checkpoint invocation).^[4][10]

FSDP

PyTorch's fsdp integrates activation checkpointing through apply_activation_checkpointing from torch.distributed.algorithms._checkpoint.checkpoint_wrapper.^[5] The user supplies a checkpoint_wrapper_fn callable that wraps modules and a check_fn or auto_wrap_policy that decides which submodules to wrap. The typical pattern is to wrap each transformer layer, mirroring the strategy used by megatron lm and DeepSpeed.^[5]

Megatron-LM

The original megatron lm codebase implements activation checkpointing at the granularity of transformer layers via the checkpoint-activations flag and the related checkpoint-num-layers.^[13] Setting checkpoint-num-layers to 1 (checkpoint at every layer boundary) is a common configuration; the GPT-NeoX-20B model was trained with exactly these settings.^[13] Since 2022, Megatron-LM also implements the selective recomputation strategy of Korthikanti et al. as a configurable mode alongside the original full recomputation path.^[6]

Where is gradient checkpointing used in modern training?

Activation checkpointing is effectively universal in training large language models. Public training reports from the major labs document its use as follows.

LLaMA

The original llama paper, "LLaMA: Open and Efficient Foundation Language Models" by Touvron et al., released by Meta in February 2023, states that "To further improve training efficiency, we reduced the amount of activations that are recomputed during the backward pass with checkpointing. More precisely, we save the activations that are expensive to compute, such as the outputs of linear layers."^[7] The paper explicitly cites Korthikanti et al. and reports that, in conjunction with sequence and model parallelism, the resulting training pipeline processed roughly 380 tokens per second per GPU on 2,048 A100 GPUs with 80 GB of memory, training the 65B-parameter model in approximately 21 days over 1.4T tokens.^[7] LLaMA's implementers also note that they "manually implemented the backward function for the transformer layers, instead of relying on the PyTorch autograd", which is how they implemented the selective saving of expensive activations.^[7] This pattern of bypassing the framework's default gradient pipeline to gain finer-grained checkpointing control reappears in numerous large-scale training projects.

GPT-3 and Megatron-Turing NLG

The 175B-parameter GPT-3 was trained using model parallelism across V100 GPUs. The 530B-parameter Megatron-Turing NLG model, trained jointly by NVIDIA and Microsoft on NVIDIA's Selene supercomputer of A100 GPUs, used a 3D parallelism strategy (8-way tensor and 35-way pipeline parallelism plus data parallelism) with the activations of each transformer block checkpointed and recomputed.^[14] Korthikanti et al. measure that for that model the activation memory required for the baseline (no recomputation) approach is well above 80 GB per GPU, exceeding the A100 HBM budget; full activation recomputation brings it below the budget at the cost of 30 to 40 percent execution time overhead, while selective recomputation combined with sequence parallelism achieves a 5x reduction and cuts the recomputation overhead to about 2 percent for that model.^[6]

PaLM and others

Google's PaLM 540B was trained with rematerialisation,^[17] GPT-NeoX-20B with per-layer activation checkpointing,^[13] and many other large LLMs use some form of activation checkpointing. The standard configuration in the open-source megatron lm and DeepSpeed stacks is to checkpoint at every transformer layer with non-reentrant checkpointing and either CPU offload or partitioned activations as needed.^[4][5] GPT-NeoX-20B was trained with checkpoint-activations: True and checkpoint-num-layers: 1, the per-layer setting that maximises memory savings at the cost of more recomputation work.^[13]

Training versus fine-tuning

While the discussion above focuses on pre-training, activation checkpointing is also pervasive in fine-tuning. Hugging Face's transformers library exposes gradient_checkpointing_enable() on supported models and a gradient_checkpointing=True flag on the TrainingArguments object, both of which route the model's layers through torch.utils.checkpoint.^[16] LoRA and full fine-tuning pipelines on consumer hardware (24 GB to 80 GB GPUs) routinely enable it; combined with mixed precision training in BF16 and 8-bit optimisers, it helps make LoRA-style fine-tuning of 7B-parameter models feasible on a single GPU.

How does it compare to other memory-saving techniques?

Activation checkpointing addresses one specific source of training memory: intermediate activations. Other techniques target different sources.

These approaches are complementary; the standard recipe for training a multi-billion-parameter LLM combines activation checkpointing, mixed precision, ZeRO or FSDP sharding, and FlashAttention.^[4][5][6][7] Korthikanti et al. note that ZeRO-style optimisations and offload techniques "are complementary to the techniques presented here and could be additionally employed for even greater memory savings."^[6]

The choice between full and selective activation checkpointing depends on architecture and scale. For pre-transformer architectures (ResNets, large CNNs), the original Chen et al. sqrt-n strategy or its per-block analogue suffices. For transformers above ~1B parameters, selective recomputation or flash attention is usually preferable because the attention softmax and dropout layers are exactly the memory-heavy, compute-light operations the selective strategy targets.^[6]

What are the limitations of gradient checkpointing?

Activation checkpointing's main cost is wall-clock time. Korthikanti et al. measured 30 to 40 percent execution time overhead from full activation recomputation on their training runs.^[6] An ApxML distributed-training course gives a similar 20 to 30 percent range as a rule of thumb.^[15] Cybertronai's TensorFlow implementation reported around 20 percent.^[9] The exact overhead depends on the architecture (the ratio of compute to activation memory in each block), the chosen checkpoint granularity, and whether selective recomputation is in use.^[6]

A subtler limitation concerns randomness. Operations such as dropout produce different outputs on each forward pass; if the second (recomputation) pass uses a different RNG state from the first, the gradient is incorrect. PyTorch's torch.utils.checkpoint preserves the RNG state by default but emits warnings that operations involving randomness across multiple devices may still produce incorrect gradients.^[2] Disabling RNG-state preservation via preserve_rng_state=False is an option only when the wrapped function contains no stochastic operations.^[2]

The reentrant implementation in PyTorch, the original variant, places additional restrictions: at least one input must have requires_grad=True, certain backward hooks can interfere with the recursive autograd call, and it does not record the autograd graph during the forward pass, supports only torch.autograd.backward() without its inputs argument, and does not support detached tensors inside the checkpointed region.^[2] These issues motivated the non-reentrant implementation, which the PyTorch documentation strongly recommends.^[2]

Activation checkpointing also has no effect on memory consumed by parameters, gradients, or optimiser state. For these, sharding techniques such as ZeRO and fsdp are required. In practice, parameter and optimiser memory often dominate at the smaller end of the LLM range; activation memory dominates as sequence length and batch size grow, which is why checkpointing becomes essential at long-context training.^[6]

Finally, the technique can interact awkwardly with certain compilation paths. For example, the PyTorch documentation notes that under torch.compile the preserve_rng_state flag has no effect and RNG state is always preserved, and it recommends the non-reentrant variant generally.^[2]

ELI5: gradient checkpointing in plain terms

Imagine solving a long math problem and writing down every single intermediate step so you can check your work later. That uses a lot of paper. Gradient checkpointing is like keeping only a few key intermediate answers and, when you need a step you erased, quickly redoing the small piece of arithmetic that produces it. You spend a little more time recomputing, but you need far less paper (memory). For a deep network, the trick is to keep about $\sqrt{n}$ of the n intermediate results, which is why a model that would not fit in a GPU's memory suddenly does, in exchange for roughly one extra forward pass of compute.^[1]

Continue reading

• CCheckpointIn machine learning, a checkpoint is a saved snapshot of a model's state captured at a specific point during the training process, used to resume training, recover from failures, or deploy and fine-tune the…Mentioned 48 times above
• TTrainingTraining in machine learning is the process of fitting a model's parameters to data so that the model can make accurate predictions or generate useful outputs.Mentioned 30 times above
• JAXJAX is an open-source Python library for accelerator-oriented array computation and program transformation.Mentioned 17 times above
• PyTorchPyTorch is an open-source software library for tensor computation and machine learning.Mentioned 16 times above
• DeepSpeedDeepSpeed is an open-source deep learning optimization library, originally developed by Microsoft, that makes distributed training and inference of large models efficient, easy to use, and cost-effective.Mentioned 13 times above
• TensorFlowTensorFlow is an open-source software platform for numerical computation and machine learning.Mentioned 9 times above

More from See also

• Backpropagation
• Automatic Differentiation
• Mixed-Precision Training
• Flash Attention
• Fully Sharded Data Parallel (FSDP)
• Megatron-LM
• Data Parallelism
• Tensor Parallelism
• Pipeline Parallelism
• Model Parallelism
• Deep Learning

Topics:Deep LearningTraining & Optimization

References

1. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9 ^10 ^11 ^12 ^13 ^14 ^15 ^16 ^17 ^18 ^19 ^20 ^21Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin, "Training Deep Nets with Sublinear Memory Cost", arXiv, 2016-04-21. arxiv.org/...1604.06174. Accessed 2026-06-27.
2. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9 ^10 ^11 ^12 ^13 ^14PyTorch documentation, "torch.utils.checkpoint", PyTorch Foundation, 2026. docs.pytorch.org/...checkpoint. Accessed 2026-06-27.
3. ^1 ^2 ^3 ^4JAX documentation, "jax.checkpoint", JAX project, 2026. docs.jax.dev/...jax.checkpoint. Accessed 2026-06-27.
4. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9DeepSpeed documentation, "Activation Checkpointing", DeepSpeed project, 2026. deepspeed.readthedocs.io/...activation-checkpointing. Accessed 2026-06-27.
5. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9PyTorch source, "checkpoint_wrapper.py", PyTorch GitHub, 2026. github.com/...checkpoint_wrapper.py. Accessed 2026-06-27.
6. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9 ^10 ^11 ^12 ^13 ^14 ^15 ^16 ^17 ^18 ^19 ^20 ^21 ^22Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, Bryan Catanzaro, "Reducing Activation Recomputation in Large Transformer Models", arXiv, 2022-05-10. arxiv.org/...2205.05198. Accessed 2026-06-27.
7. ^1 ^2 ^3 ^4 ^5Hugo Touvron et al., "LLaMA: Open and Efficient Foundation Language Models", arXiv, 2023-02-27. arxiv.org/...2302.13971. Accessed 2026-06-27.
8. ^1 ^2David E. Rumelhart, Geoffrey E. Hinton, Ronald J. Williams, "Learning representations by back-propagating errors", Nature, 1986-10-09. nature.com/...323533a0. Accessed 2026-06-27.
9. ^1 ^2 ^3 ^4Tim Salimans, Yaroslav Bulatov, "cybertronai/gradient-checkpointing: Make huge neural nets fit in memory", GitHub, 2018. github.com/...gradient-checkpointing. Accessed 2026-06-27.
10. ^1 ^2DeepSpeed documentation, "DeepSpeed Configuration JSON", DeepSpeed project, 2026. deepspeed.ai/...config-json. Accessed 2026-06-27.
11. ^1 ^2 ^3JAX documentation, "Gradient checkpointing with jax.checkpoint (aka jax.remat)", JAX project, 2026. docs.jax.dev/...gradient-checkpointing. Accessed 2026-06-27.
12. ^1 ^2TensorFlow documentation, "tf.recompute_grad", Google, 2026. tensorflow.org/...recompute_grad. Accessed 2026-06-27.
13. ^1 ^2 ^3 ^4 ^5Sid Black et al., "GPT-NeoX-20B: An Open-Source Autoregressive Language Model", arXiv, 2022-04-14. arxiv.org/...2204.06745. Accessed 2026-06-27.
14. ^Shaden Smith et al., "Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model", arXiv, 2022-01-28. arxiv.org/...2201.11990. Accessed 2026-06-27.
15. ^1 ^2ApxML, "Activation Checkpointing Mechanics", apxml.com distributed training course, 2025. apxml.com/...activation-checkpointing-mechanics. Accessed 2026-06-27.
16. ^1 ^2 ^3Hugging Face, "Models" (PreTrainedModel.gradient_checkpointing_enable), Transformers documentation, 2026. huggingface.co/...model. Accessed 2026-11-03.
17. ^Aakanksha Chowdhery et al., "PaLM: Scaling Language Modeling with Pathways", arXiv, 2022-04-05. arxiv.org/...2204.02311. Accessed 2026-11-03.

Show all 17 references