One GPU, a 13B Model, and the DeepSpeed Trick That Saved My Weekend

Press enter or click to view image in full size

A no-nonsense, slightly unhinged intro to ZeRO — the wizardry quietly holding today’s giant AI models together.

A few weeks ago, my GPU and I were locked in emotional combat over this beauty of an error:

RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB

I wasn’t doing anything reckless. Just innocently trying to fine-tune a 13-billion-parameter model on a single A100 — the exact thing the internet swears is “super easy now, bro.” The internet lied. Again.

What eventually saved my sanity (and possibly my weekend) was a library I’d heard people mumble about for years but never actually understood: DeepSpeed, and its star player, ZeRO. I spent two days reading the original paper, questioning my life choices, and finally bullying a real training run into working. Here’s everything I wish someone had tattooed on my arm on day one.

No equations. No PhD. Just vibes and intuition.

The Trap Nobody Bothers to Mention

Here’s the dirty little secret of big-model training: the model weights aren’t what murder your GPU. It’s all the emotional baggage they bring along.

Train with the Adam optimizer in mixed precision, and every single parameter secretly hauls around:

• The FP16 weight — 2 bytes
• A full-precision FP32 master copy — 4 bytes (the “just in case” copy)
• The gradient — 2 bytes
• Adam’s momentum — 4 bytes
• Adam’s variance — 4 bytes

Total: 16 bytes per parameter. For a 10B model that’s roughly 160 GB gone before a single token even says hello. Your fancy 80GB A100? Already gargling water at the bottom of the pool.

So how does anyone train GPT-3-scale models? The classic move is data parallelism: clone the entire model onto every GPU, hand each one a different batch, then average the gradients. Elegant, simple, and about as efficient as heating your house by leaving the fridge open.

Picture it — 8 GPUs, each clutching its own identical 160 GB copy of the same state. That’s over a terabyte of memory spent storing the same thing eight times. Basically photocopying one textbook eight times because eight people refuse to share.

Someone had to fix this madness. In 2020, a Microsoft team rolled up their sleeves.

The Question That Broke the Game Wide Open

The DeepSpeed crew asked something almost insultingly simple:

“What if each GPU only kept the part of the model it actually owns, and just… asked the others for the rest when needed?”

Revolutionary. Like discovering you can split the dinner bill instead of everyone paying full price. This became ZeRO — the Zero Redundancy Optimizer, and it comes in three increasingly chaotic stages.

Stage 1 — Split the optimizer states. Adam’s momentum and variance are the real gluttons — about 8 bytes per parameter, half the whole buffet. ZeRO-1 slices them across your GPUs. Four GPUs? Each keeps a quarter. Roughly 4× memory saved with almost no extra chit-chat. Basically a free lunch, and those don’t exist, so enjoy it.

Stage 2 — Split the gradients too. This is where my brain short-circuited. “But don’t all GPUs compute the full gradient anyway?” Yeah — for a hot second. The trick: instead of a normal all-reduce (which dumps the full averaged gradient onto everyone), ZeRO uses a reduce-scatter. After it runs, each GPU keeps only its little gradient slice and yeets the rest into the void. Same math, same traffic, half the memory. Beautiful.

Stage 3 — Split the weights themselves. Okay this one’s certified unhinged. ZeRO-3 shards the actual model weights. At rest, no GPU holds the full model. When a layer’s forward pass shows up, the GPUs do an all-gather to briefly Frankenstein that layer’s weights back together, run the math, then immediately shred the copy. Same drama on the backward pass. Imagine a study group where everyone owns exactly one chapter and they frantically pass chapters around before the exam. Chaotic? Yes. Works? Also yes.

Where It Gets Genuinely Ridiculous: Offloading

This is the moment DeepSpeed went from “huh, neat” to “wait, I can actually afford to do this?”

ZeRO-Offload kicks the optimizer states and gradients clean off the GPU and into your humble CPU RAM. The CPU even does the optimizer step itself and mails the updated weights back like a helpful roommate. The result? A lonely little 32GB V100 can fine-tune a 13-billion-parameter model. Not a typo. One. GPU.

ZeRO-Infinity takes it even further by spilling onto NVMe SSDs. With a quick Gen4 drive, you can poke at models in the hundreds of billions of parameters on a setup that costs less than a used hatchback.

Is it slower? Oh absolutely — usually 2–3× slower, because PCIe and SSD bandwidth become the new office bottleneck. But “slow” absolutely destroys “impossible” in a fair fight, every single time.

What It Looks Like in Actual Code

Bolting DeepSpeed onto a PyTorch script is honestly just a few lines, which feels illegal for how much it does:

import deepspeed


model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)


# The training loop barely notices anything happened
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

Most of the real magic hides in ds_config.json — that's where you pick your ZeRO stage, offload settings, and precision. The library handles the rest like a quiet genius who never brags. When my run finally started without crashing, I genuinely felt like I'd pulled off a heist.

So what actually lives inside that mysterious ds_config.json? This is where all the real decisions get made. Don't panic though — it looks way scarier than it is. A minimal config looks something like this:

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "overlap_comm": true,
    "contiguous_gradients": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  }
}

That’s the ZeRO-2 config I started with. That one little line — "stage": 2 — basically decides whether your GPU celebrates or cries.

Now say you hit the memory wall even on ZeRO-2 (like I did). Just bump stage up to 3 and add a bit of extra tuning so the parameter-sharding runs smoothly:

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "stage3_gather_16bit_weights_on_model_save": true,
    "stage3_max_live_parameters": 1e9,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6
  }
}

And if you’re stuck with a single lonely GPU staring down a 13B model like I was — this is the config that saved my weekend. It offloads both the optimizer states and the parameters onto the CPU:

{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 8,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

Little bonus: if your CPU RAM taps out too, swap "device": "cpu" for "nvme" and add an nvme_path — congrats, you're now officially in ZeRO-Infinity territory, where your SSD becomes your extra memory. Slower, sure, but "it runs" always beats "it doesn't."

The Trade-Offs, Said Out Loud (Because Someone Has To)

DeepSpeed is not a magic wand, sadly. After a few projects, here’s the stuff the tutorials conveniently skip:

• Your interconnect speed is basically destiny. ZeRO-3 is a chatterbox. Without NVLink between cards or InfiniBand between nodes, you’ll watch your GPUs sit at 30% utilization, sighing, waiting for data. On plain PCIe consumer rigs, the gains evaporate fast.
• Debugging distributed runs is a personality test. NCCL timeouts, silent deadlocks, mysterious hangs that appear only when you’re tired — welcome to your new hobby. The error messages read like ancient prophecy.
• Small models get nothing. Under a billion parameters, boring old DistributedDataParallel is usually faster. All this partitioning only pays off when things get chunky.
• Offloading has a hangover. It makes giant models possible, but throughput drops noticeably. Got multiple GPUs? Scaling out almost always beats offloading. Offloading is the “I have no other choice” option.

Why Any of This Actually Matters

DeepSpeed quietly props up some of the most important models of recent years — BLOOM (176B), Microsoft’s Turing-NLG, MT-NLG (530B). It’s part of the plumbing that dragged “training a frontier model” from a thing only three companies on Earth could do, to something labs, startups, and sleep-deprived individuals can seriously attempt.

And that last part is what actually gets me. Five years ago, fine-tuning a 13B model on your own hardware was a straight-up fever dream. Today, with ZeRO-Offload and one decent GPU, it’s a Saturday afternoon side quest.

The wall keeps getting shorter. And gloriously unglamorous libraries like DeepSpeed are doing most of the quiet heavy lifting to knock it down.

If You Want to Try It Yourself

The official DeepSpeed repo has surprisingly solid starter examples, and Hugging Face’s Trainer plugs into it natively — you can flip ZeRO on with a single config flag, like a cheat code.

Start at ZeRO-2. Only climb to ZeRO-3 when you genuinely hit a wall. And treat offloading like the emergency exit — good to have, not where you start.

When that training run finally boots up without crashing, take a second to just… sit with it. You’re standing on top of a quietly brilliant piece of engineering. Also, you didn’t have to sell a kidney. Win-win.

If this saved you even one CUDA-out-of-memory meltdown, a clap or two genuinely helps it reach more struggling GPUs out there. And if you’ve got your own DeepSpeed war stories, drop them in the comments — misery loves company, and I love a good debugging horror tale.