fix zero3 integration (#1897)

* fix zero3 integration * bump transformers and accelerate too
2024-09-05 10:58:50 -04:00
parent ab461d83c4
commit 6e354682e3
2 changed files with 4 additions and 3 deletions
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -425,7 +425,8 @@ def setup_deepspeed_env(cfg, stage=None):
        os.environ["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(stage)
        if stage == 3:
            os.environ["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = "true"
-    HfTrainerDeepSpeedConfig(cfg.deepspeed)
+    # If we don't assign this, it doesn't actually get set in the accelerate weakref
+    _ = HfTrainerDeepSpeedConfig(cfg.deepspeed)


 def setup_fsdp_envs(cfg):