Advanced Training Techniques: Gradient Checkpointing, Mixed Precision, ZeRO Offload

L7-02 covered the high-level parallelism. This piece: the workhorse techniques that make training memory-feasible.

Without these, training a 70B model would need a $1M+ GPU cluster. With them, a single RTX 4090 can fine-tune one.

Memory Budget Recap

Training memory = weights + gradients + optimizer state + activations.

For a 7B model in FP16:

• Weights: 14GB
• Gradients: 14GB
• Adam (m, v): 28GB
• Activations: ~30GB (depends on batch/seq)
• Total: ~86GB

H100 has 80GB. Doesn’t quite fit. Need tricks.

Technique 1: Gradient Checkpointing

The biggest memory saver in deep nets.

Normal forward pass

Save activations at every layer (for backward):

forward: x → h1 → h2 → ... → h_L → loss
                  ↑   ↑         ↑
          all activations stored

Memory: O(L) — grows linearly with depth.

Gradient checkpointing

Save only every K-th activation. Recompute the rest during backward:

forward: save h_1, h_5, h_10, ...

backward: when computing gradient for h_3:
          recompute from h_1 → h_2 → h_3
          → use h_3 for gradient
          → discard

Memory: O(L / K) Compute: 1.3× more (the recomputation)

For a 7B model — activations drop from 30GB to ~5GB. Huge win.

In PyTorch:

from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(nn.Module):
    def forward(self, x):
        return checkpoint(self._forward, x)

    def _forward(self, x):
        # Actual block computation
        return ...

Technique 2: Mixed Precision

Compute in FP16 / BF16, master weights in FP32.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for x, y in loader:
    with autocast(dtype=torch.bfloat16):
        loss = model(x, y)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Effect:

• 2× faster (Tensor Cores love FP16)
• ~50% memory savings on activations
• Same final quality

BF16 vs FP16:

• FP16: better precision, smaller range — risk of overflow
• BF16: same range as FP32, less precision — more stable for training

Modern frontier training (LLaMA, DeepSeek) all use BF16.

Technique 3: ZeRO Offload

ZeRO Stage 3 (FSDP) shards weights / gradients / optimizer across GPUs. ZeRO Offload goes further: move some to CPU RAM.

from deepspeed import DeepSpeedConfig

config = DeepSpeedConfig({
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu"},  # Adam state on CPU
        "offload_param": {"device": "cpu"},      # Some weights on CPU
    }
})

Effect: train 7B model on 24GB GPU (4090) by offloading Adam state to CPU RAM.

Cost: ~3-5× slower (PCIe bandwidth bottleneck). Worth it: if you don’t have datacenter GPUs.

Technique 4: 8-bit Adam

Adam state is 8× larger than weights (m + v + FP32 master).

8-bit Adam (Tim Dettmers 2022) quantizes optimizer state to INT8:

import bitsandbytes as bnb

optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=1e-4)

Effect: optimizer memory 4× smaller.

Used in many “fit big model on small GPU” setups.

Technique 5: PEFT (Parameter-Efficient Fine-Tuning)

The biggest trick: only train a tiny fraction of weights.

LoRA

Already covered in L4-05. Train low-rank deltas:

ΔW = A @ B  where A, B are tiny

Effect:

• Train 0.01-0.5% of parameters
• Optimizer state: 100-1000× smaller
• Full fine-tune quality on most tasks

Prefix tuning

Prepend learned vectors to each layer’s input:

# Standard input: [token_1, token_2, ...]
# Prefix-tuned:    [learned_1, learned_2, ..., token_1, token_2, ...]

Even less params than LoRA. Slightly worse quality.

IA³

Element-wise scaling per layer:

output = layer(x) * learnable_vector  (per layer)

Tiny param count. Used for very-low-data scenarios.

Technique 6: Gradient Accumulation

If you want batch size 32 but only fit batch size 4:

for batch in loader:
    loss = model(batch)
    loss = loss / 8  # Scale to compensate
    loss.backward()

    if (step + 1) % 8 == 0:
        optimizer.step()
        optimizer.zero_grad()

8 mini-batches accumulate → optimizer sees gradient of effective batch size 32.

No memory increase, just slower per-step.

Technique 7: Sequence Packing

Variable-length sequences = wasted compute (padding).

Pack multiple short sequences into one long one:

Naive:
  [Hello world      <pad><pad><pad>]
  [The quick brown fox jumps over   ]
  [A    <pad><pad><pad><pad><pad><pad>]

Packed:
  [Hello world </s> The quick brown fox jumps over </s> A </s> <pad>]

With attention masking to prevent cross-sequence interference, you train on 100% useful tokens.

~2× throughput improvement when input lengths vary.

Technique 8: FlashAttention (in training too)

Same trick as inference (L7-01) but for training:

• O(N²) → O(N) memory
• Faster long-context training

Default in modern PyTorch via F.scaled_dot_product_attention.

Technique 9: Activation Recomputation

A variant of gradient checkpointing — for specific layers.

Common: don’t store attention map (large), recompute during backward.

# In FlashAttention's training mode
# Forward: don't materialize attention matrix
# Backward: recompute it as needed

Saves significant memory on long-context training.

Technique 10: CPU / Disk Offload

Last resort: move things to slow storage.

# DeepSpeed NVMe offload
config["zero_optimization"]["offload_optimizer"]["device"] = "nvme"
config["zero_optimization"]["offload_optimizer"]["nvme_path"] = "/mnt/nvme"

Effect: optimizer state on SSD instead of GPU/CPU. Cost: 10-30× slower than GPU-only. Worth it: research / experimentation on small budgets.

A Practical Training Recipe

Fine-tune Llama 3 8B on a single 24GB GPU:

# 1. Load 4-bit model
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B",
    quantization_config=bnb_config,
)

# 2. Add LoRA adapter
from peft import LoraConfig, get_peft_model
peft_config = LoraConfig(
    r=16, lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
)
model = get_peft_model(model, peft_config)

# 3. Use gradient checkpointing
model.gradient_checkpointing_enable()

# 4. 8-bit Adam
import bitsandbytes as bnb
optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=1e-4)

# 5. BF16 mixed precision
trainer_args = TrainingArguments(
    bf16=True,
    gradient_checkpointing=True,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    ...
)

Memory used: ~20GB out of 24GB. Speed: ~1000 tokens/sec. Quality: 95%+ of full fine-tuning.

That’s QLoRA — the modern standard for low-resource fine-tuning.

When These Techniques Hurt

Each technique has costs:

Technique	Cost
Gradient checkpointing	1.3× slower
Mixed precision	Some loss / stability issues if not careful
ZeRO Offload	3-5× slower
8-bit Adam	Small quality risk
PEFT	Limited capacity for hard tasks
Gradient accumulation	Slower convergence
CPU offload	10-30× slower

Combine techniques wisely — don’t apply all at once if you don’t need to.

Training Curve

What real frontier training looks like:

Step 0: model is random
Step 100k: model speaks coherent English
Step 1M: model writes good code, does math
Step 10M: model approaches frontier quality

Each "step" includes:
- forward pass
- backward pass
- optimizer step
- (with our techniques) recomputation, offloading, sync

For a 70B model, 1 step ≈ 1-5 seconds on a good cluster. 1M steps = 2-5 weeks straight.

This is why training is expensive — not just compute, but time.

Next recommended: L7-07 Monitoring or L7-02 Distributed Training.