Multi-GPU & Distributeddeepspeedpytorchacceleratecuda

Why DeepSpeed ZeRO throws config errors at startup, and how offload settings actually trade off

Error
AssertionError: Check batch related parameters. train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size

Also appears as

  • ValueError: Please specify either train_batch_size or train_micro_batch_size_per_gpu in the DeepSpeed config
  • RuntimeError: Tensors must be CUDA and dense when using ZeRO stage 3 with offload

Short answer

DeepSpeed refuses to start when its config's batch-size fields do not agree with each other, since train_batch_size must equal train_micro_batch_size_per_gpu times gradient_accumulation_steps times world_size exactly. A second common failure comes from enabling ZeRO stage 3 with CPU or NVMe offload on hardware that lacks enough system RAM or fast enough storage, which surfaces as tensor or contiguity errors rather than a clear resource message.

Affects: DeepSpeed ZeRO stages 1 through 3, with or without HuggingFace accelerate or Trainer integration, any multi-GPU fine-tuning or pretraining job

Fix the batch-size math before touching ZeRO stage or offload settings

  1. 1Open your DeepSpeed JSON config and set only two of train_batch_size, train_micro_batch_size_per_gpu, and gradient_accumulation_steps, leaving the third as 'auto' so DeepSpeed derives it correctly.
  2. 2Confirm world_size (total GPUs across all nodes) matches what you expect, since train_batch_size implicitly depends on it.
  3. 3If using HuggingFace Trainer or accelerate alongside DeepSpeed, remove batch size settings from the Trainer or accelerate config and let the DeepSpeed config be the single source of truth, or vice versa, but never both.
  4. 4Start with ZeRO stage 2 and no offload to confirm the model and data pipeline work, then move to stage 3 or add offload one change at a time.
  5. 5If enabling offload, verify system RAM (for CPU offload) or NVMe read and write speed (for NVMe offload) is sufficient before assuming a resulting error is a config bug.

How to confirm this is your problem

  • Training fails immediately at DeepSpeed initialization with an assertion about batch size parameters
  • Job runs fine at ZeRO stage 2 but fails or hangs the moment stage 3 or offload is enabled
  • Error messages reference contiguous tensors, dense tensors, or parameter partitioning rather than an obvious out-of-memory message
  • Training starts but is far slower than expected once offload is enabled, with heavy CPU or disk activity

Root causes and fixes

Most common

The DeepSpeed config specifies a train_batch_size that is mathematically inconsistent with the other batch fields

DeepSpeed treats train_batch_size, train_micro_batch_size_per_gpu, and gradient_accumulation_steps as a strict equation, not independent settings, because it needs to know exactly how many gradient accumulation steps to run per optimizer step across the whole cluster. If the numbers do not multiply out correctly, most often because world_size changed without updating the other fields, DeepSpeed refuses to start rather than guess.

Fix: Set train_micro_batch_size_per_gpu and gradient_accumulation_steps explicitly, and set train_batch_size to 'auto' so the values can never drift out of sync when your GPU count changes.

Common

ZeRO stage 3 is enabled but a required setting (contiguous_gradients, overlap_comm, or reduce_bucket_size) is misconfigured

Stage 3 partitions parameters, gradients, and optimizer state across all ranks and reconstructs full parameters just-in-time during forward and backward passes. This reconstruction depends on contiguous memory layout assumptions; a mismatched or overly aggressive bucket size setting for a model's specific layer shapes can produce a tensor contiguity error during backward rather than a clear configuration warning.

Fix: Start from a published ZeRO stage 3 example config for a similarly sized model rather than hand-tuning every field, and change one setting at a time while confirming training still runs.

Common

CPU or NVMe offload is enabled but the host lacks enough system RAM or fast enough storage to support it

Offloading optimizer states or parameters to CPU RAM or NVMe trades GPU VRAM for system memory and I/O bandwidth. If system RAM is smaller than the offloaded data, or NVMe throughput is far below what the offload engine expects, the process can fail with allocation or timeout errors that read like config problems but are actually a hardware capacity mismatch.

Fix: Calculate offloaded memory requirements, roughly proportional to optimizer state size, against actual available system RAM before enabling offload, and benchmark NVMe throughput if using NVMe offload.

Commands
free -h
Occasional

DeepSpeed configuration and HuggingFace Trainer or accelerate configuration both try to control batch size or precision

When DeepSpeed is used through the Trainer or accelerate integration, both layers read overlapping settings from different places. If the two are set inconsistently, whichever layer initializes last can override the other in a way that produces the exact assertion the deeper layer expects to prevent.

Fix: Pick one layer to own batch size and precision settings, usually the DeepSpeed JSON when using 'auto' values, and confirm the other layer is not separately overriding the same fields.

Rare

DeepSpeed's version does not match the installed CUDA or PyTorch build, so its custom CUDA ops fail to compile or load

DeepSpeed ships several custom CUDA kernels compiled against a specific CUDA toolkit and PyTorch ABI. Installing a DeepSpeed version built against a different CUDA version than what is actually present can produce confusing errors that look like configuration mistakes but are actually a build or ABI mismatch.

Fix: Run ds_report to check which ops compiled successfully and match against your installed CUDA and PyTorch versions, reinstalling DeepSpeed after confirming CUDA toolkit compatibility if any op shows as not installed.

Commands
ds_report

Diagnostic commands

Validate the DeepSpeed config's op compatibility

ds_report

Lists each DeepSpeed op and whether it compiled successfully against your current CUDA and PyTorch install; any op marked not installed or incompatible explains downstream errors that otherwise look unrelated.

Manually verify the batch-size equation

python -c "print(4 * 8 * 8)"

Substitute your actual train_micro_batch_size_per_gpu, gradient_accumulation_steps, and world_size; the result must equal your intended train_batch_size exactly, or that alone explains the assertion error.

Check available system RAM before enabling CPU offload

free -h

Compare available memory against the estimated offloaded optimizer state size, roughly twice the model parameter count in FP32 for Adam; insufficient headroom here explains offload-related failures or extreme slowdowns.

Stopping it from happening again

  • Always set 'auto' for at least one of the three batch-size fields so DeepSpeed derives it, instead of hardcoding all three.
  • Start every new model or GPU count at ZeRO stage 2 with no offload to confirm the baseline works, then add complexity incrementally.
  • Benchmark actual NVMe or CPU RAM offload throughput on your specific hardware before committing a production training plan to it.
  • Keep DeepSpeed, PyTorch, and CUDA versions pinned together in your container image and re-validate with ds_report after any upgrade.

When this becomes an architecture problem

If a model genuinely does not fit even with ZeRO stage 3 and full offload on your current hardware, or offload throughput makes training impractically slow, that is a capacity and hardware decision (more GPUs, more VRAM per GPU, or faster NVMe and interconnect) rather than a configuration problem to keep tuning around.

Frequently asked questions

Why does DeepSpeed care about train_batch_size if I already set micro batch size and gradient accumulation?

DeepSpeed uses train_batch_size as a consistency check across the whole cluster, since the effective global batch size depends on world_size, which can silently change when you scale the job to more or fewer GPUs. Setting it to 'auto' avoids the check ever going stale.

Is CPU offload always slower than keeping everything on GPU?

Yes, offload always adds overhead compared to fitting everything in GPU memory, because data has to move across the PCIe bus to and from system RAM or NVMe. It is a deliberate tradeoff: offload lets you train models that would not otherwise fit at all, at the cost of throughput, and is best used only for the portion of memory you cannot avoid offloading.

Can I mix ZeRO stage 3 with tensor parallelism?

Yes, and DeepSpeed and other frameworks increasingly support this combination for very large models, but it substantially increases configuration complexity. Get ZeRO working alone first, then add tensor parallelism, changing one dimension of parallelism at a time so you can isolate which layer introduced any new error.

Related problems

NCCL error during multi-GPU training or inference

An NCCL error during multi-GPU training or inference is almost always a symptom of a rank that crashed, a version mismatch across processes, or bad GPU topology, not a bug in NCCL itself. Enable NCCL_DEBUG=INFO first and read the per-rank logs before touching timeouts or retry logic.

Gradient checkpointing errors during fine-tuning

Gradient checkpointing errors during fine-tuning almost always come from three sources: the use_reentrant parameter left unset (it now must be explicit and False is usually correct for transformer models), an attention implementation that isn't fully compatible with checkpointing's re-computation approach, or leaving use_cache=True enabled while checkpointing is on, which conflicts because checkpointing recomputes the forward pass and a live KV cache assumes it won't be recomputed. Set use_reentrant=False and use_cache=False together.

HuggingFace Accelerate config mismatch causes wrong distributed launch

Accelerate errors during launch almost always mean the saved configuration does not match the current machine's actual GPU count, node count, or distributed type, or that the model was loaded with device_map='auto' inference-style sharding and then also handed to accelerate's training-mode preparation, which are two incompatible placement strategies. Regenerating the config for the current machine, or passing explicit CLI overrides, resolves most cases.

CUDA out of memory even though nvidia-smi shows free VRAM

This almost always means memory fragmentation: the allocator has enough total free memory but no single contiguous block large enough for the requested allocation, because the address space is broken into many small free-and-used segments from prior allocations of different sizes. The fix is enabling PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, reducing allocation size variability, or restarting the process to reset the address space.

Guide

LoRA vs QLoRA: Choosing the Right Fine-Tuning Method

LoRA vs QLoRA for enterprise fine-tuning: rank and alpha choices, real VRAM math by model size, and when each method actually wins.

Guide

Fine-Tuning Failure Modes: What Actually Goes Wrong

Fine-tuning failure modes that actually derail enterprise projects: catastrophic forgetting, eval overfitting, data leakage, and how to catch each one.

Guide

Multi-Node LLM Training Infrastructure: Networking and Storage

Multi-node LLM training infrastructure explained: InfiniBand vs RoCE tradeoffs, storage throughput needs, and cluster topology for enterprise fine-tuning.

Still stuck, or tired of fighting your own infrastructure?

Netray deploys and operates on-prem AI for regulated manufacturers and defense suppliers. We have debugged this stack in production, on air-gapped networks, at scale.