DeepSpeed ZeRO-3 Small Parameter Partition Bug

DeepSpeed ZeRO-3 crashes with UnboundLocalError in partition_parameters.py when the model has very small parameters (fewer elements than the number of GPUs). The partition math divides a small parameter across more GPUs than it has elements. Denpex detects the small-parameter partition failure and recommends the configuration fix.

DeepSpeed ZeRO-3 Small Parameter Partition Bug is a Distributed Training failure seen during ML training runs. DeepSpeed ZeRO-3 crashes with UnboundLocalError in partition_parameters.py when the model has very small parameters (fewer elements than the number of GPUs). The partition math divides a small parameter across more GPUs than it has elements. Denpex detects the small-parameter partition failure and recommends the configuration fix. Common tags: Deepspeed, Zero3, Partition, Small Parameter.

ZeRO-3 partitions each parameter's elements equally across GPUs. When a parameter has fewer elements than the number of GPUs, the partition calculation produces zero or negative slice sizes. The partition_dim variable is never assigned when the parameter is too small to partition, causing the UnboundLocalError. Small parameters (e.g., a bias vector with 8 elements on 16 GPUs) cannot be meaningfully partitioned. The bug is in DeepSpeed's partition_parameters.py which doesn't handle the edge case of parameters smaller than world_size. Taken together, these mechanisms explain why the failure is reproducible, why it tends to surface on specific workloads or scales, and why generic mitigation attempts often fall short without addressing the underlying cause.

• Training crashes with UnboundLocalError in DeepSpeed's partition_parameters.py
• The error occurs during ZeRO-3 initialization for models with very small parameters (e.g., bias vectors, layer norm parameters)
• The same model works with fewer GPUs but crashes when scaled to more GPUs

• DeepSpeed ZeRO-3 with models containing very small parameters
• Training with 8+ GPUs where some parameters have fewer elements than GPU count
• Models with small embedding layers, bias vectors, or layer norm parameters
• DeepSpeed versions before v0.12.0

Use this checklist to test the hypothesis against a small reproduction. No single line proves the root cause, so preserve the preceding events and compare one variable at a time.

• ✓Reproduce the failure from a clean checkpoint/seed: the symptom must appear without warm-up state from a previous run.
• ✓Verified signal present: UnboundLocalError: local variable 'partition_dim' referenced before assignment in partition_parameters.py
• ✓Verified signal present: Error occurs in deepspeed/runtime/zero/partition_parameters.py around line 862
• ✓Verified signal present: Crash happens during DeepSpeed engine initialization, before training starts
• ✓Verified signal present: Reducing the number of GPUs or switching to ZeRO-2 avoids the crash
• ✓A targeted fix from the "How to fix it" section eliminates or substantially reduces the symptom within one validation pass.

Select the traceback or open the failed terminal, then run Denpex locally to see the initiating rank, collateral failures, exact fix, and verification command without uploading the log.

DeepSpeed changes when parameters, gradients and optimizer state are created, partitioned, gathered and offloaded. The hub separates ZeRO, memory, checkpoint and pipeline failures by lifecycle phase.

• ZeRO-3 partitions each parameter's elements equally across GPUs. When a parameter has fewer elements than the number of GPUs, the partition calculation produces zero or negative slice sizes
• The partition_dim variable is never assigned when the parameter is too small to partition, causing the UnboundLocalError
• Small parameters (e.g., a bias vector with 8 elements on 16 GPUs) cannot be meaningfully partitioned
• The bug is in DeepSpeed's partition_parameters.py which doesn't handle the edge case of parameters smaller than world_size