From 7134a1e73fce7bb9386bd7d6c823efc80fec3439 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 6 Feb 2025 17:19:57 +0000 Subject: [PATCH] Revert "Before the first training step, the model has no optimizer: fix ds3" This reverts commit bf6e7edea54f2e34b2f6802468ee3224c4aa8030. --- trl/models/utils.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/trl/models/utils.py b/trl/models/utils.py index dce9d60228..22a30c0afb 100644 --- a/trl/models/utils.py +++ b/trl/models/utils.py @@ -137,8 +137,6 @@ def setup_chat_format( def remove_hooks(model: "DeepSpeedEngine") -> None: """Removes the optimizer hooks from a DeepSpeed ZeRO-3 model.""" - if not hasattr(model, "optimizer"): # before the first training step, the model has no optimizer - return if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): optimizer_offload = model.optimizer.parameter_offload elif model.optimizer is not None: @@ -166,8 +164,6 @@ def iter_params(module, recurse=False): def add_hooks(model: "DeepSpeedEngine") -> None: """Adds the optimizer hooks from a DeepSpeed ZeRO-3 model.""" - if not hasattr(model, "optimizer"): # before the first training step, the model has no optimizer - return if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): optimizer_offload = model.optimizer.parameter_offload elif model.optimizer is not None: