Fix BF16_Optimizer being used without ZeRO (#7790)

tohtana · web-flow · commit 09885efefd56 · 2026-01-20T23:04:30.000Z
The test_ds_initialize.py::TestOptimizerImplementation expects the
configuration (None, 'bf16', 'fp32') is unimplemented.
However, it is actually supported by DeepSpeed and uses FP16_Optimizer
in bf16 mode.
This PR adds the configuration in the dict of supported combinations.

---------

Signed-off-by: Masahiro Tanaka &lt;mtanaka@anyscale.com&gt;
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -1481,8 +1481,6 @@ def _do_optimizer_sanity_check(self, basic_optimizer):
                 )
                 return BFLOAT16
             return FP16 if model_dtype == torch.float16 else DDP_BFLOAT16
-        elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32:
-            return BFLOAT16
         else:
             raise NotImplementedError(f"unsupported mix of {model_dtype=} and {grad_accum_dtype=}")
 

Original file line number	Diff line number	Diff line change
`@@ -1481,8 +1481,6 @@ def _do_optimizer_sanity_check(self, basic_optimizer):`
`1481`	`1481`	`)`
`1482`	`1482`	`return BFLOAT16`
`1483`	`1483`	`return FP16 if model_dtype == torch.float16 else DDP_BFLOAT16`
`1484`		`- elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32:`
`1485`		`- return BFLOAT16`
`1486`	`1484`	`else:`
`1487`	`1485`	`raise NotImplementedError(f"unsupported mix of {model_dtype=} and {grad_accum_dtype=}")`
`1488`	`1486`