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

fsdp2 support

Browse Source
This commit is contained in:
Wing Lian
2025-03-22 21:47:54 -04:00
parent a8f38c367c
commit 03dcf1a5ea
2 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -538,6 +538,8 @@ def setup_deepspeed_env(cfg, stage=None):
def setup_fsdp_envs(cfg):
os.environ["ACCELERATE_USE_FSDP"] = "true"
if str(cfg.fsdp_version) == "2":
os.environ["FSDP_VERSION"] = "2"
if cfg.fsdp_config.fsdp_activation_checkpointing:
os.environ["FSDP_ACTIVATION_CHECKPOINTING"] = "true"
if cfg.fsdp_config.fsdp_offload_params:
@@ -556,6 +558,10 @@ def setup_fsdp_envs(cfg):
os.environ["FSDP_TRANSFORMER_CLS_TO_WRAP"] = (
cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap
)
if cfg.fsdp_config.fsdp_reshard_after_forward is not None:
os.environ["FSDP_RESHARD_AFTER_FORWARD"] = (
"true" if cfg.fsdp_config.fsdp_reshard_after_forward else "false"
)
def prepare_optim_env(cfg):