fix zero3 integration (#1897)

* fix zero3 integration * bump transformers and accelerate too
axolotl-ai-cloud · Sep 5, 2024 · 6e35468 · 6e35468
1 parent ab461d8
commit 6e35468
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,10 +1,10 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 packaging==23.2
 peft==0.12.0
-transformers==4.44.0
+transformers==4.44.2
 tokenizers>=0.19.1
 bitsandbytes==0.43.3
-accelerate==0.33.0
+accelerate==0.34.0
 datasets==2.20.0
 deepspeed==0.14.4
 pydantic==2.6.3

diff --git 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):