Reject unsupported BF16 norm fast path

Fail fast when ZeRO-1 selects the dedicated BF16 optimizer with FP32 gradient accumulation, where disabling norm computation is not implemented.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: yh0903 <helloyu0903@gmail.com>

Author: yh0903

deepspeed/runtime/engine.py

@@ -2068,6 +2068,9 @@ def _do_optimizer_sanity_check(self, basic_optimizer):
                     logger.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****")
             if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage(
             ) == 1 and not self.zero_cpu_offload():
+                if not self.zero_compute_grad_norm():
+                    raise ValueError("zero_optimization.compute_grad_norm=false does not support ZeRO Stage 1 with "
+                                     "BF16 parameters and FP32 gradient accumulation")
                 return BFLOAT16
             return ZERO_OPTIMIZATION
         elif amp_enabled:

docs/_pages/config-json.md

@@ -516,7 +516,7 @@ Enabling and configuring ZeRO memory optimizations
 
 | Description                                                                                                                                                                                                                     | Default |
 | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
-| Compute and retain the global gradient norm during ZeRO Stage 1/2 optimizer steps. Set to `false` only with a GPU optimizer, without ZenFlow or gradient clipping, and when callers do not use `get_global_grad_norm()`; finite/overflow checking is unchanged. | `true`  |
+| Compute and retain the global gradient norm during ZeRO Stage 1/2 optimizer steps. Set to `false` only with a GPU optimizer, without ZenFlow, gradient clipping, or ZeRO Stage 1 BF16 parameters with FP32 gradient accumulation, and when callers do not use `get_global_grad_norm()`; finite/overflow checking is unchanged. | `true`  |
 
 <i>**overlap_comm**</i>: [boolean]
 

tests/unit/runtime/zero/test_zero1_optimizer_fastpath.py

@@ -85,14 +85,13 @@ def test_fast_path_matches_default_update(self, stage):
         torch.manual_seed(123)
         baseline_model = SimpleModel(hidden_dim=4)
         fast_model = copy.deepcopy(baseline_model)
-        baseline_engine, baseline_optimizer, _, _ = deepspeed.initialize(
-            model=baseline_model,
-            model_parameters=baseline_model.parameters(),
-            config=_config(compute_grad_norm=True, stage=stage))
-        fast_engine, fast_optimizer, _, _ = deepspeed.initialize(
-            model=fast_model,
-            model_parameters=fast_model.parameters(),
-            config=_config(compute_grad_norm=False, stage=stage))
+        baseline_engine, baseline_optimizer, _, _ = deepspeed.initialize(model=baseline_model,
+                                                                         model_parameters=baseline_model.parameters(),
+                                                                         config=_config(compute_grad_norm=True,
+                                                                                        stage=stage))
+        fast_engine, fast_optimizer, _, _ = deepspeed.initialize(model=fast_model,
+                                                                 model_parameters=fast_model.parameters(),
+                                                                 config=_config(compute_grad_norm=False, stage=stage))
         inputs = torch.randn(1, 4, device=baseline_engine.device, dtype=torch.bfloat16)
         targets = torch.randn(1, 4, device=baseline_engine.device, dtype=torch.bfloat16)
 
@@ -167,3 +166,11 @@ def test_zenflow_rejects_disabled_norm(self):
             deepspeed.initialize(model=model,
                                  model_parameters=model.parameters(),
                                  config=_config(compute_grad_norm=False, zenflow={}))
+
+    def test_fp32_gradient_accumulation_rejects_disabled_norm(self):
+        model = SimpleModel(hidden_dim=4)
+        config = _config(compute_grad_norm=False)
+        config["data_types"] = {"grad_accum_dtype": "fp32"}
+
+        with pytest.raises(ValueError, match="BF16 parameters and FP32 gradient accumulation"):
+            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)