tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

ensure that the hftrainer deepspeed config is set before the trainer class is ever init'ed (#1850) [skip ci]

Browse Source
This commit is contained in:
Wing Lian
2024-08-22 13:10:40 -04:00
committed by GitHub
parent de4ea2d1f2
commit 2f8037fee6
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/utils/trainer.py
View File

@@ -399,12 +399,15 @@ def setup_torch_compile_env(cfg):
def setup_deepspeed_env(cfg, stage=None): def setup_deepspeed_env(cfg, stage=None):
from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
os.environ["ACCELERATE_USE_DEEPSPEED"] = "true" os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
os.environ["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = cfg.deepspeed os.environ["ACCELERATE_DEEPSPEED_CONFIG_FILE"] = cfg.deepspeed
if stage: if stage:
os.environ["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(stage) os.environ["ACCELERATE_DEEPSPEED_ZERO_STAGE"] = str(stage)
if stage == 3: if stage == 3:
os.environ["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = "true" os.environ["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = "true"
HfTrainerDeepSpeedConfig(cfg.deepspeed)
def setup_fsdp_envs(cfg): def setup_fsdp_envs(cfg):