github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: avoid double init_process_group in DeepSpeed pretrain

pull/791/head
baihaochen 2026-05-23 16:21:57 +08:00
parent ed006a254f
commit 4f5cdfe59a
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
trainer/train_pretrain_deepspeed.py
View File

@ -129,11 +129,14 @@ if __name__ == "__main__":
parser.add_argument("--local_rank", type=int, default=-1, help="DeepSpeed 自动注入") parser.add_argument("--local_rank", type=int, default=-1, help="DeepSpeed 自动注入")
args = parser.parse_args() args = parser.parse_args()
# ========== 1. 初始化环境 ========== # ========== 1. 初始化环境（DeepSpeed 全权负责）==========
local_rank = init_distributed_mode() # deepspeed.initialize() 内部会自动 init_process_group + set_device
if dist.is_initialized(): local_rank = int(os.environ.get("LOCAL_RANK", 0))
args.device = f"cuda:{local_rank}" if not dist.is_initialized():
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) dist.init_process_group(backend="nccl")
torch.cuda.set_device(local_rank)
args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
setup_seed(42 + local_rank)
# ========== 2. 模型配置 ========== # ========== 2. 模型配置 ==========
os.makedirs(args.save_dir, exist_ok=True) os.makedirs(args.save_dir, exist_ok=True)