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

prepare accelerate envs for tp

Browse Source
This commit is contained in:
Sunny Liu
2025-02-18 23:34:05 -05:00
committed by Sung Ching Liu
parent a74efcecbe
commit 438b623031
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -543,10 +543,18 @@ def setup_fsdp_envs(cfg):
] = cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap ] = cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap
def setup_tp_envs():
os.environ["ACCELERATE_USE_TP"] = "true"
def prepare_optim_env(cfg): def prepare_optim_env(cfg):
if not check_cuda_p2p_ib_support(): if not check_cuda_p2p_ib_support():
if os.getenv("NCCL_P2P_DISABLE") is None: if os.getenv("NCCL_P2P_DISABLE") is None:
os.environ["NCCL_P2P_DISABLE"] = "1" os.environ["NCCL_P2P_DISABLE"] = "1"
if cfg.tp_size > 1:
setup_tp_envs()
if cfg.fsdp: if cfg.fsdp:
setup_fsdp_envs(cfg) setup_fsdp_envs(cfg)
elif cfg.deepspeed: elif cfg.deepspeed: