feat: make cross_entropy_chunk_size tunable, with an auto memory budget (#2190) [3/3]

litgpt/args.py

@@ -2,6 +2,7 @@
 import math
 import warnings
 from dataclasses import dataclass
+from typing import Literal
 
 
 @dataclass
@@ -31,6 +32,20 @@ class TrainArgs:
     """Limits the number of seconds to train for"""
     max_seq_length: int | None = None
     """Limits the length of samples"""
+    cross_entropy_chunk_size: int | Literal["auto"] = 128
+    """Chunk size used when computing cross-entropy loss during training, to reduce the memory spike
+    during the backward pass at the cost of extra compute (see `litgpt.utils.chunked_cross_entropy`).
+    Set to `0` to disable chunking and compute exact cross-entropy in one shot (uses more peak memory,
+    especially with large vocab sizes / sequence lengths). Set to `"auto"` to instead derive the chunk
+    size from `cross_entropy_memory_budget_bytes` and the model's vocab size, via
+    `litgpt.utils.auto_cross_entropy_chunk_size` — a real memory budget rather than a fixed guess."""
+    cross_entropy_memory_budget_bytes: int = 32 * 1024 * 1024
+    """Only used when `cross_entropy_chunk_size="auto"`. Target peak memory (bytes) for a single
+    cross-entropy chunk's forward+backward intermediates. Derived from `torch.profiler` measurements
+    on real GPU hardware; see `litgpt.utils.auto_cross_entropy_chunk_size` and
+    docs/profiling/op_table_gpu.md for the byte-per-element estimate this is based on. The 32MiB
+    default keeps a single chunk's log_softmax intermediates well under typical GPU memory pressure
+    regardless of vocab size."""
     tie_embeddings: bool | None = None
     """Whether to tie the embedding weights with the language modeling head weights"""
 

litgpt/finetune/adapter.py

@@ -304,7 +304,12 @@ def fit(
             logits = model(input_ids, lm_head_chunk_size=128)
             # shift the targets such that output n predicts token n+1
             logits[-1] = logits[-1][..., :-1, :]
-            loss = chunked_cross_entropy(logits, targets[..., 1:])
+            loss = chunked_cross_entropy(
+                logits,
+                targets[..., 1:],
+                chunk_size=train.cross_entropy_chunk_size,
+                memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+            )
             fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
 
         running_loss.update(loss.detach())

litgpt/finetune/adapter_v2.py

@@ -331,7 +331,12 @@ def fit(
             logits = model(input_ids, lm_head_chunk_size=128)
             # shift the targets such that output n predicts token n+1
             logits[-1] = logits[-1][..., :-1, :]
-            loss = chunked_cross_entropy(logits, targets[..., 1:])
+            loss = chunked_cross_entropy(
+                logits,
+                targets[..., 1:],
+                chunk_size=train.cross_entropy_chunk_size,
+                memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+            )
             fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
 
         running_loss.update(loss.detach())

litgpt/finetune/full.py

@@ -281,7 +281,12 @@ def fit(
         with fabric.no_backward_sync(model, enabled=is_accumulating):
             logits = model(input_ids)
             # shift the targets such that output n predicts token n+1
-            loss = chunked_cross_entropy(logits[..., :-1, :], targets[..., 1:])
+            loss = chunked_cross_entropy(
+                logits[..., :-1, :],
+                targets[..., 1:],
+                chunk_size=train.cross_entropy_chunk_size,
+                memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+            )
             fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
 
         running_loss.update(loss.detach())

litgpt/finetune/lora.py

@@ -355,7 +355,12 @@ def fit(
             logits = model(input_ids, lm_head_chunk_size=128)
             # shift the targets such that output n predicts token n+1
             logits[-1] = logits[-1][..., :-1, :]
-            loss = chunked_cross_entropy(logits, targets[..., 1:])
+            loss = chunked_cross_entropy(
+                logits,
+                targets[..., 1:],
+                chunk_size=train.cross_entropy_chunk_size,
+                memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+            )
             fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
 
         running_loss.update(loss.detach())

litgpt/finetune/lora_legacy.py

@@ -338,7 +338,12 @@ def fit(
             logits = model(input_ids, lm_head_chunk_size=128)
             # shift the targets such that output n predicts token n+1
             logits[-1] = logits[-1][..., :-1, :]
-            loss = chunked_cross_entropy(logits, targets[..., 1:])
+            loss = chunked_cross_entropy(
+                logits,
+                targets[..., 1:],
+                chunk_size=train.cross_entropy_chunk_size,
+                memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+            )
             fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
 
         running_loss.update(loss.detach())

litgpt/pretrain.py

@@ -300,11 +300,11 @@ def fit(
     optimizer = state["optimizer"]
 
     if eval.initial_validation:
-        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
+        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters, train=train)
         val_loss = f"{val_loss:.3f}"
     else:
         fabric.print("Verifying settings ...")
-        validate(fabric, model, val_dataloader, max_iters=2, verbose=False)  # sanity check
+        validate(fabric, model, val_dataloader, max_iters=2, train=train, verbose=False)  # sanity check
         val_loss = "n/a"
 
     throughput = ThroughputMonitor(fabric, window_size=5)
@@ -351,7 +351,12 @@ def fit(
         is_accumulating = state["iter_num"] % train.gradient_accumulation_iters(devices, num_nodes) != 0
         with fabric.no_backward_sync(model, enabled=is_accumulating):
             logits = model(input_ids)
-            loss = chunked_cross_entropy(logits, targets)
+            loss = chunked_cross_entropy(
+                logits,
+                targets,
+                chunk_size=train.cross_entropy_chunk_size,
+                memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+            )
             fabric.backward(loss / train.gradient_accumulation_iters(devices, num_nodes))
 
         running_loss.update(loss.detach())
@@ -402,7 +407,7 @@ def fit(
 
         if val_dataloader is not None and not is_accumulating and state["step_count"] % eval.interval == 0:
             t0 = time.perf_counter()
-            val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
+            val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters, train=train)
             val_loss = val_loss.item()
             td = time.perf_counter() - t0
 
@@ -416,15 +421,20 @@ def fit(
 
     # Final validation
     if eval.final_validation:
-        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
+        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters, train=train)
         metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
         fabric.log_dict(metrics, step=state["iter_num"])
         fabric.print(f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}")
 
 
 @torch.no_grad()
 def validate(
-    fabric: L.Fabric, model: nn.Module, val_dataloader: DataLoader, max_iters: int, verbose: bool = True
+    fabric: L.Fabric,
+    model: nn.Module,
+    val_dataloader: DataLoader,
+    max_iters: int,
+    train: TrainArgs,
+    verbose: bool = True,
 ) -> torch.Tensor:
     fabric.barrier()
     if verbose:
@@ -438,7 +448,12 @@ def validate(
         input_ids = batch[:, 0 : model.max_seq_length].contiguous().long()
         targets = batch[:, 1 : (model.max_seq_length + 1)].contiguous().long()
         logits = model(input_ids)
-        loss = chunked_cross_entropy(logits, targets)
+        loss = chunked_cross_entropy(
+            logits,
+            targets,
+            chunk_size=train.cross_entropy_chunk_size,
+            memory_budget_bytes=train.cross_entropy_memory_budget_bytes,
+        )
         losses.append(loss)
 
     val_loss = torch.stack(losses).mean()

litgpt/utils.py

@@ -298,18 +298,51 @@ def __exit__(self, type, value, traceback):
 
 T = TypeVar("T")
 
+# bytes of intermediate memory (forward log_softmax output, its backward gradient, and the
+# nll_loss/allocator overhead around them) held per row of a cross-entropy chunk, per unit of
+# vocab_size and per byte of dtype itemsize. Calibrated on an NVIDIA T4 by measuring actual
+# torch.cuda.max_memory_allocated() peak (not just one op's self-CUDA-mem from a profiler table) for
+# `torch.nn.functional.cross_entropy` on a (chunk_size, vocab_size) slice, at chunk sizes chosen by
+# this exact formula, swept across vocab_size = 8k..152k (GPT-2 to Llama-3/Qwen2.5 scale). The
+# measured peak came out flat across every vocab_size tested (as intended -- see
+# docs/profiling/budget_formula_sweep.png) at a consistent 1.5x the target `memory_budget_bytes`,
+# which is where the factor of 3 below comes from (an earlier factor of 2, based only on a
+# profiler-table self-CUDA-mem estimate for one op, undershot the real allocator peak by that same
+# 1.5x -- see docs/profiling/ for both measurements).
+_CROSS_ENTROPY_BYTES_PER_CHUNK_ELEMENT = 3
+
+
+def auto_cross_entropy_chunk_size(
+    vocab_size: int, dtype: torch.dtype, memory_budget_bytes: int, min_chunk_size: int = 1
+) -> int:
+    """Computes a `chunked_cross_entropy` `chunk_size` that keeps a single chunk's log_softmax
+    forward+backward intermediates within `memory_budget_bytes`, given the model's `vocab_size` and
+    logits `dtype`. This is the actual memory-budget knob behind `TrainArgs.cross_entropy_chunk_size
+    ="auto"` — see the module-level comment above for where the underlying byte-per-element estimate
+    comes from.
+    """
+    itemsize = torch.tensor([], dtype=dtype).element_size()
+    bytes_per_row = vocab_size * itemsize * _CROSS_ENTROPY_BYTES_PER_CHUNK_ELEMENT
+    return max(min_chunk_size, memory_budget_bytes // bytes_per_row)
+
 
 def chunked_cross_entropy(
     logits: torch.Tensor | list[torch.Tensor],
     targets: torch.Tensor,
-    chunk_size: int = 128,
+    chunk_size: int | Literal["auto"] = 128,
     ignore_index: int = -100,
+    memory_budget_bytes: int = 32 * 1024 * 1024,
 ) -> torch.Tensor:
     # with large max_sequence_lengths, the beginning of `backward` allocates a large memory chunk which can dominate
     # the memory usage in fine-tuning settings with low number of parameters.
     # as a workaround hack, the cross entropy computation is chunked to force it to deallocate on the go, reducing
     # the memory spike's magnitude
 
+    if chunk_size == "auto":
+        vocab_size = logits[0].size(-1) if isinstance(logits, list) else logits.size(-1)
+        dtype = logits[0].dtype if isinstance(logits, list) else logits.dtype
+        chunk_size = auto_cross_entropy_chunk_size(vocab_size, dtype, memory_budget_bytes)
+
     # lm_head was chunked (we are fine-tuning)
     if isinstance(logits, list):
         # don't want to chunk cross entropy

tests/test_utils.py

@@ -30,6 +30,7 @@
     CLI,
     CycleIterator,
     _RunIf,
+    auto_cross_entropy_chunk_size,
     capture_hparams,
     check_file_size_on_cpu_and_warn,
     check_nvlink_connectivity,
@@ -162,6 +163,102 @@ def test_chunked_cross_entropy(ignore_index, B):
     torch.testing.assert_close(chunked_loss, baseline_loss)
 
 
+def test_chunked_cross_entropy_equivalence_at_scale():
+    # exercises `chunked_cross_entropy` at a realistic sequence length / vocab size (issue #2190,
+    # "test with full context lengths and realistic batch sizes"), not just the tiny shapes above.
+    B, T, V = 2, 2048, 32000
+    logits = torch.randn(B, T, V)
+    targets = torch.randint(0, V, (B, T))
+
+    unchunked_loss = chunked_cross_entropy(logits, targets, chunk_size=0)
+    chunked_loss = chunked_cross_entropy(logits, targets, chunk_size=128)
+    torch.testing.assert_close(chunked_loss, unchunked_loss)
+
+
+@_RunIf(min_cuda_gpus=1)
+def test_chunked_cross_entropy_peak_memory_decreases_with_smaller_chunks():
+    # confirms the memory-saving claim behind the `chunked_cross_entropy` "workaround hack"
+    # (litgpt/utils.py) actually holds at a realistic scale, on the CUDA allocator the issue is about.
+    B, T, V = 4, 4096, 32000
+    device = torch.device("cuda")
+
+    def peak_memory_for(chunk_size):
+        logits = torch.randn(B, T, V, device=device, requires_grad=True)
+        targets = torch.randint(0, V, (B, T), device=device)
+        torch.cuda.reset_peak_memory_stats(device)
+        loss = chunked_cross_entropy(logits, targets, chunk_size=chunk_size)
+        loss.backward()
+        return torch.cuda.max_memory_allocated(device)
+
+    peak_unchunked = peak_memory_for(chunk_size=0)
+    peak_chunked = peak_memory_for(chunk_size=128)
+    assert peak_chunked < peak_unchunked
+
+
+def test_auto_cross_entropy_chunk_size():
+    # the byte-per-row factor was calibrated on an NVIDIA T4 by measuring real
+    # torch.cuda.max_memory_allocated() peaks (not just one profiler op's self-CUDA-mem) across a
+    # vocab_size sweep from 8k to 152k, at chunk sizes chosen by this exact formula -- see
+    # docs/profiling/budget_formula_sweep.png. The measured peak came out flat across every
+    # vocab_size, at ~1.5x the target budget, which is what the module-level comment's factor of 3
+    # is calibrated against; this test only checks the formula's own arithmetic (the byte accounting
+    # this function does, not the real allocator peak, which needs a GPU to measure).
+    budget = 32 * 1024 * 1024
+    chunk_size = auto_cross_entropy_chunk_size(vocab_size=32000, dtype=torch.float32, memory_budget_bytes=budget)
+    assert chunk_size > 0
+    predicted_bytes = chunk_size * 32000 * 4
+    assert predicted_bytes <= budget
+
+    # smaller budget -> smaller chunk size (the actual "budget system" behavior issue #2190 asked for)
+    smaller_chunk_size = auto_cross_entropy_chunk_size(
+        vocab_size=32000, dtype=torch.float32, memory_budget_bytes=budget // 4
+    )
+    assert smaller_chunk_size < chunk_size
+
+    # never returns 0 (which would silently disable chunking, defeating the memory budget)
+    tiny_chunk_size = auto_cross_entropy_chunk_size(vocab_size=200000, dtype=torch.float32, memory_budget_bytes=1)
+    assert tiny_chunk_size == 1
+
+
+def test_chunked_cross_entropy_auto_matches_manual_chunk_size():
+    # chunk_size="auto" must resolve to the same numerical result as an equivalent manual chunk_size
+    # (chunking doesn't change the math, only how much memory is held at once)
+    B, T, V = 2, 64, 500
+    logits = torch.randn(B, T, V)
+    targets = torch.randint(0, V, (B, T))
+
+    budget = 32 * 1024 * 1024
+    resolved_chunk_size = auto_cross_entropy_chunk_size(vocab_size=V, dtype=logits.dtype, memory_budget_bytes=budget)
+    auto_loss = chunked_cross_entropy(logits, targets, chunk_size="auto", memory_budget_bytes=budget)
+    manual_loss = chunked_cross_entropy(logits, targets, chunk_size=resolved_chunk_size)
+    torch.testing.assert_close(auto_loss, manual_loss)
+
+
+@_RunIf(min_cuda_gpus=1)
+def test_chunked_cross_entropy_auto_reduces_peak_memory_like_manual_chunking():
+    # the actual "memory budget system" issue #2190 asked for: chunk_size="auto", derived from a
+    # memory_budget_bytes rather than a hardcoded chunk_size, should give the same peak-memory win as
+    # picking a good chunk_size by hand (test_..._peak_memory_decreases_with_smaller_chunks above).
+    # Note: peak CUDA memory here also includes the unavoidable full-size logits.grad tensor, so this
+    # is a real-hardware sanity check on the *relative* saving, not a literal check that peak memory
+    # stays under `memory_budget_bytes` (that bound only applies to the chunk's own intermediates,
+    # which chunked_cross_entropy doesn't isolate from the rest of the backward pass).
+    B, T, V = 4, 4096, 32000
+    device = torch.device("cuda")
+    targets = torch.randint(0, V, (B, T), device=device)
+
+    def peak_memory_for(chunk_size, **kwargs):
+        logits = torch.randn(B, T, V, device=device, requires_grad=True)
+        torch.cuda.reset_peak_memory_stats(device)
+        loss = chunked_cross_entropy(logits, targets, chunk_size=chunk_size, **kwargs)
+        loss.backward()
+        return torch.cuda.max_memory_allocated(device)
+
+    peak_unchunked = peak_memory_for(chunk_size=0)
+    peak_auto = peak_memory_for(chunk_size="auto", memory_budget_bytes=32 * 1024 * 1024)
+    assert peak_auto < peak_unchunked
+
+
 def test_num_parameters():
     model = torch.nn.Linear(2, 2)
     assert num_parameters(model) == 6