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

automatically set env vars for single gpu deepspeed zero3 (#3118) [skip ci]

Browse Source 
* automatically set env vars for single gpu deepspeed zero3

* use setdefault
This commit is contained in:
Wing Lian
2025-08-29 13:36:47 -04:00
committed by GitHub
parent 5b6ec2820f
commit 7ed40f1d70
2 changed files with 7 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -547,6 +547,13 @@ def setup_deepspeed_env(cfg, stage=None):
if stage == 3:
os.environ["ACCELERATE_DEEPSPEED_ZERO3_INIT"] = "true"
device_count = torch.cuda.device_count()
if device_count == 1:
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")
os.environ.setdefault("MASTER_ADDR", "0.0.0.0") # nosec B104
os.environ.setdefault("MASTER_PORT", "29500")
# NOTE(djsaunde): The distribued state cannot be initialized prior to the
# ACCELERATE_USE_DEEPSPEED assignment, but it must be initialized some time prior
# to model load.