Support bf16 optimizer states with CPU offload

deepspeed/ops/adam/cpu_adam.py

@@ -59,7 +59,9 @@ def __init__(self,
                 (default: False) NOT SUPPORTED in DeepSpeed CPUAdam!
             adamw_mode: select between Adam and AdamW implementations (default: AdamW)
             fp32_optimizer_states: creates momentum and variance in full precision regardless of
-                        the precision of the parameters (default: True)
+                        the precision of the parameters. Set to False to keep optimizer states
+                        in the parameter dtype (e.g. bf16), which reduces the optimizer-state
+                        memory footprint at the cost of lower state precision. (default: True)
         """
 
         default_args = dict(lr=lr,

deepspeed/runtime/base_optimizer.py

@@ -448,11 +448,14 @@ def _configure_master_weights(self,
                                   fp16_master_weights_and_gradients=False,
                                   bf16_master_weights_and_gradients=False,
                                   bf16_optimizer_states=False,
+                                  offload_enabled=False,
                                   fp16_offload_validator=None,
-                                  bf16_fp32_offload_validator=None):
+                                  bf16_offload_validator=None):
         """
         Common validation and dtype selection for ZeRO optimizer master-weight settings.
         Optionally accepts callables that enforce backend-specific offload requirements.
+        ``offload_enabled`` tells this method whether optimizer-state offload is configured,
+        so the offload requirement is also enforced for the bf16-optimizer-states + offload case.
         """
         self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients
         self.bf16_master_weights_and_gradients = bf16_master_weights_and_gradients
@@ -464,9 +467,18 @@ def _configure_master_weights(self,
             assert self.bf16_master_weights_and_gradients, \
                 "bf16_optimizer_states requires bf16_master_weights_and_gradients."
 
-        if (self.bf16_master_weights_and_gradients and not self.bf16_optimizer_states
-                and bf16_fp32_offload_validator is not None):
-            bf16_fp32_offload_validator()
+        # bf16 master weights require ZeRO-Offload + DeepSpeedCPUAdam whenever the optimizer states
+        # cannot stay on the GPU: either because they remain fp32 (bf16_optimizer_states disabled),
+        # or because CPU offload is explicitly requested alongside bf16 optimizer states.
+        if (self.bf16_master_weights_and_gradients and bf16_offload_validator is not None
+                and (not self.bf16_optimizer_states or offload_enabled)):
+            bf16_offload_validator()
+            # Offloaded bf16 optimizer states need the CPU optimizer to store moments in the
+            # parameter (bf16) precision; otherwise they would silently expand back to fp32.
+            if self.bf16_optimizer_states:
+                assert not getattr(self.optimizer, 'fp32_optimizer_states', True), \
+                    "bf16_optimizer_states with ZeRO-Offload requires DeepSpeedCPUAdam constructed " \
+                    "with fp32_optimizer_states=False so optimizer moments are stored in bf16."
 
         if self.fp16_master_weights_and_gradients and fp16_offload_validator is not None:
             fp16_offload_validator()

deepspeed/runtime/engine.py

@@ -1647,9 +1647,9 @@ def _configure_optimizer(self, client_optimizer, model_parameters):
         self.quantizer = self._configure_quantization()
 
     def _configure_basic_optimizer(self, model_parameters):
-        optimizer_parameters = self.optimizer_params()
-        if optimizer_parameters is None:
-            optimizer_parameters = {}
+        # Copy so the pop() calls below (torch_adam, adam_w_mode, fp32_optimizer_states) do not
+        # mutate the shared config dict returned by optimizer_params().
+        optimizer_parameters = dict(self.optimizer_params() or {})
         # print(optimizer_parameters.keys())
         if "max_grad_norm" in optimizer_parameters.keys():
             raise ValueError(
@@ -1674,9 +1674,24 @@ def _configure_basic_optimizer(self, model_parameters):
                     CPUAdam = ZenFlowCPUAdam if self.zenflow else DeepSpeedCPUAdam
 
                     zenflow_kwargs = {'overlap_step': self.overlap_step} if self.zenflow else {}
+                    # Pop so a user-supplied value does not collide with the keyword built below.
+                    # None means the user did not set it, so no override warning is needed.
+                    user_fp32_optimizer_states = optimizer_parameters.pop('fp32_optimizer_states', None)
+                    if self.bf16_optimizer_states():
+                        # bf16 moments are required so the offloaded state matches the bf16 master weights.
+                        if user_fp32_optimizer_states:
+                            logger.warning("bf16_optimizer_states is enabled; overriding fp32_optimizer_states "
+                                           "to False so CPU Adam moments are stored in bf16.")
+                        fp32_optimizer_states = False
+                    elif user_fp32_optimizer_states is None:
+                        # Default preserves the pre-existing fp32 optimizer-state behavior.
+                        fp32_optimizer_states = True
+                    else:
+                        fp32_optimizer_states = user_fp32_optimizer_states
                     optimizer = CPUAdam(model_parameters,
                                         **optimizer_parameters,
                                         adamw_mode=effective_adam_w_mode,
+                                        fp32_optimizer_states=fp32_optimizer_states,
                                         **zenflow_kwargs)
                 else:
                     from deepspeed.ops.adam import FusedAdam

deepspeed/runtime/zero/stage3.py

@@ -287,8 +287,9 @@ def _enforce_optimizer_offload():
             fp16_master_weights_and_gradients=fp16_master_weights_and_gradients,
             bf16_master_weights_and_gradients=bf16_master_weights_and_gradients,
             bf16_optimizer_states=bf16_optimizer_states,
+            offload_enabled=self.offload_optimizer,
             fp16_offload_validator=_enforce_optimizer_offload,
-            bf16_fp32_offload_validator=_enforce_optimizer_offload)
+            bf16_offload_validator=_enforce_optimizer_offload)
 
         # backup fused_adam optimizer init
         if self.offload_optimizer and self.partial_offload != 1.0:

deepspeed/runtime/zero/stage_1_and_2.py

@@ -280,8 +280,9 @@ def _enforce_cpu_offload():
             fp16_master_weights_and_gradients=fp16_master_weights_and_gradients,
             bf16_master_weights_and_gradients=bf16_master_weights_and_gradients,
             bf16_optimizer_states=bf16_optimizer_states,
+            offload_enabled=self.cpu_offload,
             fp16_offload_validator=_enforce_cpu_offload,
-            bf16_fp32_offload_validator=_enforce_cpu_offload)
+            bf16_offload_validator=_enforce_cpu_offload)
 
         self.low_precision_master_weights_and_grads = self.master_weights_and_grads_dtype != torch.float32
 

docs/_pages/config-json.md

@@ -362,14 +362,14 @@ Example of <i>**scheduler**</i>
 
 | Description | Default |
 | ----------- | ------- |
-| Keep optimizer states in bf16 as well. Requires `bf16_master_weights_and_grads=true`. Enabling this removes the offload requirement because optimizer states no longer stay fp32. | `false` |
+| Keep optimizer states in bf16 as well. Requires `bf16_master_weights_and_grads=true`. Offload is optional: without `offload_optimizer` the bf16 states stay on the GPU; with `offload_optimizer` (`DeepSpeedCPUAdam`) they are offloaded to CPU memory in bf16. The offloaded state (bf16 master weights plus the two bf16 Adam moments) is then ~6 bytes/param, versus ~10 bytes/param when the moments are kept in fp32. | `false` |
 
 **Support matrix (bf16 master weights/gradients)**
 
 | ZeRO stage | bf16_optimizer_states=False | bf16_optimizer_states=True |
 | ---------- | --------------------------- | -------------------------- |
 | 0 | Not supported | Not supported |
-| 1/2/3 | Requires ZeRO-Offload + `DeepSpeedCPUAdam` (optimizer states stay fp32 on CPU) | Supported without offload; optimizer states kept in bf16 |
+| 1/2/3 | Requires ZeRO-Offload + `DeepSpeedCPUAdam` (optimizer states stay fp32 on CPU) | On GPU without offload, or on CPU with `offload_optimizer` + `DeepSpeedCPUAdam`; optimizer states kept in bf16 either way |
 
 ### Automatic mixed precision (AMP) training options
 

tests/unit/ops/adam/test_cpu_adam.py

@@ -120,6 +120,61 @@ def test_torch_adamw_equal(self, dtype, model_size):
                             optimizer2=ref_optimizer)
 
 
+class TestCPUAdamBf16OptimizerStates(DistributedTest):
+    world_size = 1
+    reuse_dist_env = True
+    requires_cuda_env = False
+    if not get_accelerator().is_available():
+        init_distributed = False
+        set_dist_env = False
+
+    @pytest.mark.parametrize('model_size', [64, 1024])
+    def test_bf16_optimizer_states_dtype(self, model_size):
+        """fp32_optimizer_states=False keeps the Adam moments in the bf16 parameter precision."""
+        from deepspeed.ops.adam import DeepSpeedCPUAdam
+
+        param = torch.nn.Parameter(torch.randn(model_size, device='cpu', dtype=torch.bfloat16))
+        optimizer = DeepSpeedCPUAdam([param], fp32_optimizer_states=False)
+        param.grad = torch.randn(model_size, device='cpu', dtype=torch.bfloat16)
+        optimizer.step()
+
+        state = optimizer.state[param]
+        assert state['exp_avg'].dtype == torch.bfloat16
+        assert state['exp_avg_sq'].dtype == torch.bfloat16
+        assert state['exp_avg'].device == torch.device('cpu')
+        assert state['exp_avg_sq'].device == torch.device('cpu')
+
+    @pytest.mark.parametrize('model_size', [64, 1024])
+    def test_bf16_optimizer_states_match_fp32(self, model_size):
+        """bf16 moments should track fp32 moments within bf16 tolerance over several steps."""
+        from deepspeed.ops.adam import DeepSpeedCPUAdam
+
+        torch.manual_seed(0)
+        base = torch.randn(model_size, device='cpu', dtype=torch.float32).to(torch.bfloat16)
+        param_fp32_states = torch.nn.Parameter(base.clone())
+        param_bf16_states = torch.nn.Parameter(base.clone())
+
+        opt_fp32_states = DeepSpeedCPUAdam([param_fp32_states], fp32_optimizer_states=True)
+        opt_bf16_states = DeepSpeedCPUAdam([param_bf16_states], fp32_optimizer_states=False)
+
+        for _ in range(10):
+            grad = torch.randn(model_size, device='cpu', dtype=torch.bfloat16)
+            param_fp32_states.grad = grad.clone()
+            param_bf16_states.grad = grad.clone()
+            opt_fp32_states.step()
+            opt_bf16_states.step()
+
+        assert opt_fp32_states.state[param_fp32_states]['exp_avg'].dtype == torch.float32
+        assert opt_bf16_states.state[param_bf16_states]['exp_avg'].dtype == torch.bfloat16
+
+        # bf16 moments round every Adam update to an 8-bit mantissa, so over 10 steps they
+        # diverge from fp32 moments more than the same-precision comparison in _compare_optimizers
+        # (1e-2). A wider 5% band keeps this stable while still catching gross errors; the dtype
+        # assertions above guard the precision itself. Norm comparison follows _compare_optimizers.
+        tolerance = param_fp32_states.float().norm().detach().numpy() * 5e-2
+        check_equal(param_fp32_states.float().norm(), param_bf16_states.float().norm(), atol=tolerance)
+
+
 class TestCPUAdamGPUError(DistributedTest):
 
     def test_cpu_adam_gpu_error(self):

tests/unit/v1/half_precision/test_bf16.py

@@ -352,10 +352,13 @@ def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False)
 @pytest.mark.parametrize("bf16_optimizer_states,use_cpu_offload,zero_stage", [
     pytest.param(False, True, 1, id="zero_stage_1_cpu_offload"),
     pytest.param(True, False, 1, id="zero_stage_1_bf16_opt_states_True"),
+    pytest.param(True, True, 1, id="zero_stage_1_bf16_opt_states_cpu_offload"),
     pytest.param(False, True, 2, id="zero_stage_2_cpu_offload"),
     pytest.param(True, False, 2, id="zero_stage_2_bf16_opt_states_True"),
+    pytest.param(True, True, 2, id="zero_stage_2_bf16_opt_states_cpu_offload"),
     pytest.param(False, True, 3, id="zero_stage_3_cpu_offload"),
     pytest.param(True, False, 3, id="zero_stage_3_bf16_opt_states_True"),
+    pytest.param(True, True, 3, id="zero_stage_3_bf16_opt_states_cpu_offload"),
 ])
 class TestBF16MasterWeightsGradients(DistributedTest):
     world_size = 2
@@ -442,4 +445,64 @@ def test_gradients_match_ddp(self, bf16_optimizer_states, use_cpu_offload, zero_
         optimizer_ddp.zero_grad()
         engine.step()
         engine.zero_grad()
+
+        if bf16_optimizer_states and use_cpu_offload:
+            # With CPU offload the Adam moments must be allocated in bf16 on the host so the
+            # offloaded optimizer-state footprint is smaller than with fp32 moments.
+            cpu_adam_state = engine.optimizer.optimizer.state
+            moment_tensors = []
+            for param_state in cpu_adam_state.values():
+                for moment_key in ("exp_avg", "exp_avg_sq"):
+                    if moment_key in param_state:
+                        moment_tensors.append(param_state[moment_key])
+            assert moment_tensors, "expected Adam moment tensors to be allocated after a step"
+            for moment in moment_tensors:
+                assert moment.dtype == torch.bfloat16, f"expected bf16 moment, got {moment.dtype}"
+                assert moment.device.type == "cpu", f"expected moment on cpu, got {moment.device}"
+
         engine.destroy()
+
+
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+class TestBF16OptimizerStatesOffloadValidation(DistributedTest):
+    world_size = 1
+
+    def test_user_cpu_adam_must_enable_bf16_states(self, zero_stage):
+        """A user-provided DeepSpeedCPUAdam must be built with fp32_optimizer_states=False
+        to combine bf16_optimizer_states with ZeRO-Offload, otherwise the moments would
+        silently stay fp32 and the memory benefit would be lost."""
+        if not bf16_required_version_check():
+            pytest.skip(
+                " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+        if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        from deepspeed.ops.adam import DeepSpeedCPUAdam
+
+        hidden_dim = 6
+        model = SimpleModel(hidden_dim, nlayers=2)
+        # fp32_optimizer_states defaults to True, which keeps fp32 moments and is
+        # incompatible with bf16_optimizer_states under ZeRO-Offload.
+        optimizer = DeepSpeedCPUAdam(model.parameters())
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 2,
+            "steps_per_print": 1,
+            "bf16": {
+                "enabled": True,
+                "bf16_master_weights_and_grads": True,
+                "bf16_optimizer_states": True,
+            },
+            # offload_optimizer is the current config key for ZeRO optimizer offload
+            # (TestBF16MasterWeightsGradients above still uses the legacy cpu_offload alias).
+            "zero_optimization": {
+                "stage": zero_stage,
+                "offload_optimizer": {
+                    "device": "cpu"
+                },
+            },
+        }
+
+        with pytest.raises(AssertionError, match="fp32_optimizer_states=False"):
+            deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)