From 8e836666b61ef9bf27678c2d689dcc6296194850 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Sat, 17 Jan 2026 18:39:43 -0800 Subject: [PATCH] Fix issue with BF16 optimizer selection Signed-off-by: Masahiro Tanaka --- deepspeed/runtime/base_optimizer.py | 4 +++- deepspeed/runtime/bf16_optimizer.py | 5 +++++ deepspeed/runtime/constants.py | 4 +++- tests/unit/runtime/zero/test_zero_tensor_fragment.py | 5 +++++ 4 files changed, 16 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/base_optimizer.py b/deepspeed/runtime/base_optimizer.py index f2ca7fd887ac..4163f719c1a8 100644 --- a/deepspeed/runtime/base_optimizer.py +++ b/deepspeed/runtime/base_optimizer.py @@ -355,7 +355,9 @@ def scale_if_loss(self, value: Any) -> Any: return self.external_loss_scale * value if self.torch_autocast_gradscaler: return self.torch_autocast_gradscaler.scale(value) - return self.loss_scaler.scale_loss(value) + # Only call loss_scaler if it exists (not present in BF16_Optimizer) + if hasattr(self, 'loss_scaler') and self.loss_scaler is not None: + return self.loss_scaler.scale_loss(value) return value diff --git a/deepspeed/runtime/bf16_optimizer.py b/deepspeed/runtime/bf16_optimizer.py index 3225dfbcd104..cc5f9959f57a 100644 --- a/deepspeed/runtime/bf16_optimizer.py +++ b/deepspeed/runtime/bf16_optimizer.py @@ -59,6 +59,11 @@ def __init__(self, ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}" self.grad_acc_dtype = grad_acc_dtype + # BF16 doesn't use loss scaling, but these attributes are needed for API compatibility + self.custom_loss_scaler = False + self.external_loss_scale = None + self.torch_autocast_gradscaler = None + self.immediate_grad_update = bfloat16_config.immediate_grad_update self.clip_grad = clip_grad diff --git a/deepspeed/runtime/constants.py b/deepspeed/runtime/constants.py index 428496a03c42..9e73bad73376 100755 --- a/deepspeed/runtime/constants.py +++ b/deepspeed/runtime/constants.py @@ -144,7 +144,9 @@ BFLOAT16_OPTIMIZER_STATES_DEFAULT = False # DDP variant of BFLOAT16 -DDP_BFLOAT16 = "bf16" +# DDP variant: bf16 model with bf16 grad accumulation (uses FP16_Optimizer in bf16 mode) +# Must be different from BFLOAT16 to allow proper optimizer selection +DDP_BFLOAT16 = "ddp_bf16" ######################################### # FP16 support diff --git a/tests/unit/runtime/zero/test_zero_tensor_fragment.py b/tests/unit/runtime/zero/test_zero_tensor_fragment.py index 977920bd23c2..31d8b0990bf3 100644 --- a/tests/unit/runtime/zero/test_zero_tensor_fragment.py +++ b/tests/unit/runtime/zero/test_zero_tensor_fragment.py @@ -173,6 +173,11 @@ def test_bf16_optimizer_fragments(self, frozen_weights): "bf16": { "enabled": True }, + # Use fp32 gradient accumulation to ensure BF16_Optimizer is used + # (bf16 model + bf16 grad_accum uses FP16_Optimizer which doesn't support tensor fragment APIs) + "data_types": { + "grad_accum_dtype": "fp32" + }, "zero_optimization": { "stage": 0, }