"""Reduce peak training memory by running the backward walk inside optimizer.step(), updating and freeing each gradient as the walk reaches it. loss.backward() normally computes every gradient before optimizer.step() applies any update, so between the two calls every gradient in the model coexists. Each parameter's update reads only that parameter's own gradient, so the update can fire the moment the gradient lands and the gradient can be freed immediately after; gradients then never coexist. Wrap once and the training loop keeps its standard lines: model, optimizer = BackProp(model, torch.optim.AdamW(model.parameters(), lr=1e-3, foreach=False)) for inputs, targets in batches: loss = loss_fn(model(inputs), targets) loss.backward() # walks only loss -> output optimizer.step() # the network walk; each gradient updates, then frees Use foreach=False: foreach gathers the per-parameter tensors back into full-size batched temporaries, which is the coexistence this wrapper removes. The wrapped model must return a single tensor, and each optimizer.step() consumes the most recent forward and backward pair. Run this file to measure a 12-layer linear stack. Passing an output directory also writes memory_timelines.json (the file the blog reads) plus the raw per-mode exports and counters: python backprop.py [timeline_output_dir] Based on: - https://lightning.ai/blog/faster-pytorch-training-by-reducing-peak-memory - https://gist.github.com/albanD/18c240bd2e09f9d93f5c4a0c9ccda39e - https://docs.pytorch.org/tutorials/intermediate/optimizer_step_in_backward_tutorial.html """ import gc import inspect import json import pathlib import sys import torch class _BackPropModelHook(torch.nn.Module): """Returns a detached view of the wrapped model's output, so loss.backward() walks only from the loss to that view while the network's graph waits for optimizer.step(). """ def __init__(self, model): super().__init__() self.model = model self._held = [] def forward(self, *args, **kwargs): output = self.model(*args, **kwargs) detached_output = output.detach().requires_grad_(True) self._held.append((output, detached_output)) return detached_output class _BackPropOptimizerHook: """Hosts the backward walk; the updates fire from each parameter's own hook.""" def __init__(self, wrapped_model): self._wrapped_model = wrapped_model def step(self): """Walks the held graph end to beginning; each gradient updates, then frees.""" output, detached_output = self._wrapped_model._held.pop() output.backward(detached_output.grad) def zero_grad(self, set_to_none=True): """Nothing to zero: each parameter's hook frees its gradient after the update.""" class BackProp: """Wires the optimizer's work into the backward walk, one parameter at a time. Splits the given optimizer into one instance per parameter, arms each instance on its parameter's gradient-accumulation hook, and moves the walk itself into optimizer.step(). The given optimizer is only a template: its type and settings are copied onto the per-parameter instances and it takes no further part. Constructing BackProp(model, optimizer) returns a (model, optimizer) pair that drops in for the wrapped pair; the model's forward must return a single tensor. """ def __new__(cls, model, optimizer): accepted = inspect.signature(type(optimizer).__init__).parameters for group in optimizer.param_groups: settings = {name: value for name, value in group.items() if name in accepted and name != "params"} for parameter in group["params"]: cls._arm_parameter(parameter, type(optimizer), settings) wrapped_model = _BackPropModelHook(model) return wrapped_model, _BackPropOptimizerHook(wrapped_model) @staticmethod def _arm_parameter(parameter, optimizer_type, settings): per_parameter_optimizer = optimizer_type([parameter], **settings) def update(_parameter): per_parameter_optimizer.step() per_parameter_optimizer.zero_grad() parameter.register_post_accumulate_grad_hook(update) BATCH_SIZE = 2048 FEATURE_DIM = 8 * 1024 NUM_LAYERS = 12 TIMING_ITERATIONS = 20 def _make_optimizer(parameters): return torch.optim.AdamW(parameters, lr=0.01, betas=(0.1, 0.1), foreach=False) def _make_training_step(mode): torch.manual_seed(0) model = torch.nn.Sequential( *[torch.nn.Linear(FEATURE_DIM, FEATURE_DIM, bias=False) for _ in range(NUM_LAYERS)], torch.nn.Linear(FEATURE_DIM, 1), ).cuda() optimizer = _make_optimizer(model.parameters()) if mode == "backprop": model, optimizer = BackProp(model, optimizer) def training_step(): x = torch.ones((BATCH_SIZE, FEATURE_DIM), device="cuda") loss = model(x).mean() loss.backward() optimizer.step() optimizer.zero_grad() return training_step def _updates_match_default(): """One training step on a small stack lands identical weights either way.""" def weights_after_one_step(mode): torch.manual_seed(0) model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).cuda() optimizer = _make_optimizer(model.parameters()) if mode == "backprop": model, optimizer = BackProp(model, optimizer) x = torch.arange(32, dtype=torch.float32, device="cuda").reshape(4, 8) loss = model(x).square().mean() loss.backward() optimizer.step() optimizer.zero_grad() return [parameter.detach().clone() for parameter in model.parameters()] pairs = zip(weights_after_one_step("default"), weights_after_one_step("backprop")) return all(torch.equal(ours, default) for default, ours in pairs) # Column order the blog's stacked plot reads; temporaries sits last so it draws on top. _TIMELINE_CATEGORIES = [ "PARAMETER", "OPTIMIZER_STATE", "INPUT", "ACTIVATION", "GRADIENT", "AUTOGRAD_DETAIL", "TEMPORARY", ] def _combined_timeline(raw_timeline_path): """Reads one export_memory_timeline JSON and returns {t (ms), gb (per category)}. The profiler leaves some allocations uncategorised, which are the same allocations it labels as temporaries in the other run, so they are folded into temporaries. Columns follow _TIMELINE_CATEGORIES. """ times, sizes = json.loads(pathlib.Path(raw_timeline_path).read_text()) def to_gb(byte_column): return [round(byte_count / 1e9, 4) for byte_count in byte_column] # Rows are [unused, parameter, optimizer state, input, temporary, activation, # gradient, autograd detail, uncategorised] in bytes; drop the unused first column. parameter, optimizer_state, layer_input, temporary, activation, gradient, autograd_detail, uncategorised = ( to_gb(column) for column in list(zip(*sizes))[1:] ) temporary = [held + spare for held, spare in zip(temporary, uncategorised)] return { "t": [round(microseconds / 1000, 3) for microseconds in times], "gb": [parameter, optimizer_state, layer_input, activation, gradient, autograd_detail, temporary], } def main(): timeline_dir = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else None print(f"updates match default: {_updates_match_default()}") allocator_counters = {} combined_timelines = {} for mode in ("default", "backprop"): training_step = _make_training_step(mode) events = [[torch.cuda.Event(enable_timing=True) for _ in range(2)] for _ in range(TIMING_ITERATIONS)] for _ in range(3): training_step() torch.cuda.synchronize() for start, end in events: start.record() training_step() end.record() torch.cuda.synchronize() mean_ms = sum(start.elapsed_time(end) for start, end in events) / len(events) torch.cuda.reset_peak_memory_stats() training_step() torch.cuda.synchronize() peak_gb = torch.cuda.max_memory_allocated() / 1e9 print(f"{mode}: {mean_ms:.1f} ms per training step, peak allocated {peak_gb:.1f} GB") # Freed gradients recycle inside the caching allocator: steady-state # steps should make no cudaMalloc or cudaFree calls at all. stats_before = torch.cuda.memory_stats() for _ in range(TIMING_ITERATIONS): training_step() torch.cuda.synchronize() stats_after = torch.cuda.memory_stats() allocator_counters[mode] = { "cuda_mallocs_per_step": (stats_after["num_device_alloc"] - stats_before["num_device_alloc"]) / TIMING_ITERATIONS, "cuda_frees_per_step": (stats_after["num_device_free"] - stats_before["num_device_free"]) / TIMING_ITERATIONS, "reserved_gb": torch.cuda.memory_reserved() / 1e9, } counters = allocator_counters[mode] print( f"{mode}: {counters['cuda_mallocs_per_step']:.1f} cudaMallocs and " f"{counters['cuda_frees_per_step']:.1f} cudaFrees per step, " f"reserved pool {counters['reserved_gb']:.1f} GB" ) if timeline_dir is not None: gc.collect() raw_timeline_path = timeline_dir.joinpath(f"{mode}.json") with torch.profiler.profile(record_shapes=True, profile_memory=True, with_stack=True) as profile: training_step() gc.collect() profile.export_memory_timeline(str(raw_timeline_path), "cuda:0") combined_timelines[mode] = { **_combined_timeline(raw_timeline_path), **allocator_counters[mode], "step_ms": round(mean_ms, 1), } del training_step gc.collect() torch.cuda.empty_cache() if timeline_dir is not None: timeline_dir.joinpath("counters.json").write_text(json.dumps(allocator_counters, indent=2)) # memory_timelines.json is the file the blog's MemoryTimelines.astro imports. timeline_dir.joinpath("memory_timelines.json").write_text( json.dumps({"categories": _TIMELINE_CATEGORIES, **combined_timelines}) ) if __name__ == "__main__": main()