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

set env vars trainer needs for FSDP
Some checks failed
pre-commit / pre-commit (push) Has been cancelled
Details
PyTest / test (3.10) (push) Has been cancelled
Details
PyTest / test (3.9) (push) Has been cancelled
Details

Browse Source
This commit is contained in:
Wing Lian
2023-08-11 08:37:33 -04:00
parent 6b5cf8b5ea
commit 64af21bcb2
1 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -353,7 +353,17 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
return total_num_steps return total_num_steps
def setup_fsdp_envs(cfg):
os.environ["ACCELERATE_USE_FSDP"] = "true"
if cfg.fsdp_config.fsdp_sync_module_states:
os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
if cfg.fsdp_config.fsdp_state_dict_type:
os.environ["FSDP_STATE_DICT_TYPE"] = cfg.fsdp_config.fsdp_state_dict_type
def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps): def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
if cfg.fsdp:
setup_fsdp_envs(cfg)
warmup_steps = ( warmup_steps = (
cfg.warmup_steps cfg.warmup_steps
if cfg.warmup_steps is not None if cfg.warmup_steps is not None