From 392246feaf43788336be3cfaa1274361a5a34d74 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 23 Jul 2026 18:33:10 +0200 Subject: [PATCH] feat(diffusion): add gradient checkpointing for memory optimization (#4127) * feat(diffusion): add gradient checkpointing for memory optimization Add gradient_checkpointing config option to DiffusionPolicy. When enabled, wraps the UNet encoder, mid, and decoder residual blocks with torch.utils.checkpoint.checkpoint to trade compute for memory. Allows training with larger batch sizes or higher-resolution inputs on memory-constrained GPUs. Disabled by default. Usage: --policy.gradient_checkpointing=true Part of the 0.6.0 roadmap item 3.3 (gradient checkpointing for all policies). * test(diffusion): verify gradient checkpointing parity --------- Co-authored-by: Jash Shah --- .../diffusion/configuration_diffusion.py | 3 +++ .../policies/diffusion/modeling_diffusion.py | 24 +++++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/src/lerobot/policies/diffusion/configuration_diffusion.py b/src/lerobot/policies/diffusion/configuration_diffusion.py index ed04ab54d..eba395e93 100644 --- a/src/lerobot/policies/diffusion/configuration_diffusion.py +++ b/src/lerobot/policies/diffusion/configuration_diffusion.py @@ -79,6 +79,8 @@ class DiffusionConfig(PreTrainedConfig): use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning. Bias modulation is used be default, while this parameter indicates whether to also use scale modulation. + gradient_checkpointing: Whether to checkpoint the Unet residual blocks during training. This reduces + activation memory at the cost of recomputing those blocks during the backward pass. noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"]. num_train_timesteps: Number of diffusion steps for the forward diffusion schedule. beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers. @@ -132,6 +134,7 @@ class DiffusionConfig(PreTrainedConfig): n_groups: int = 8 diffusion_step_embed_dim: int = 128 use_film_scale_modulation: bool = True + gradient_checkpointing: bool = False # Noise scheduler. noise_scheduler_type: str = "DDPM" num_train_timesteps: int = 100 diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 8758a7e29..aeee35505 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -31,6 +31,7 @@ import torch import torch.nn.functional as F # noqa: N812 import torchvision from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE from lerobot.utils.import_utils import _diffusers_available, require_package @@ -727,22 +728,35 @@ class DiffusionConditionalUnet1d(nn.Module): else: global_feature = timesteps_embed + use_gc = self.config.gradient_checkpointing and self.training + # Run encoder, keeping track of skip features to pass to the decoder. encoder_skip_features: list[Tensor] = [] for resnet, resnet2, downsample in self.down_modules: - x = resnet(x, global_feature) - x = resnet2(x, global_feature) + if use_gc: + x = checkpoint(resnet, x, global_feature, use_reentrant=False) + x = checkpoint(resnet2, x, global_feature, use_reentrant=False) + else: + x = resnet(x, global_feature) + x = resnet2(x, global_feature) encoder_skip_features.append(x) x = downsample(x) for mid_module in self.mid_modules: - x = mid_module(x, global_feature) + if use_gc: + x = checkpoint(mid_module, x, global_feature, use_reentrant=False) + else: + x = mid_module(x, global_feature) # Run decoder, using the skip features from the encoder. for resnet, resnet2, upsample in self.up_modules: x = torch.cat((x, encoder_skip_features.pop()), dim=1) - x = resnet(x, global_feature) - x = resnet2(x, global_feature) + if use_gc: + x = checkpoint(resnet, x, global_feature, use_reentrant=False) + x = checkpoint(resnet2, x, global_feature, use_reentrant=False) + else: + x = resnet(x, global_feature) + x = resnet2(x, global_feature) x = upsample(x) x = self.final_conv(x)