add tp_size in config doc

This commit is contained in:
Sunny Liu
2025-02-20 00:01:59 -05:00
committed by Sung Ching Liu
parent 64adbf1a15
commit 984be14147
2 changed files with 3 additions and 7 deletions

View File

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