From 6e354682e3c1735d3f7fb9e362280c38e922260f Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 5 Sep 2024 10:58:50 -0400 Subject: [PATCH] fix zero3 integration (#1897) * fix zero3 integration * bump transformers and accelerate too --- requirements.txt | 4 ++-- src/axolotl/utils/trainer.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/requirements.txt b/requirements.txt index b8d0a388b..c61216e63 100644 --- 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 index 1029fff13..89ae4e697 100644 --- 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):