chore : 升级 torch 2.11.0+cu128,移除自定义 Muon,修复 gloo device_id
- torch 2.7.1-cu126 升级至 2.11.0-cu128,numpy 2.3.2 升级至 2.4.4 - 移除 astrai/trainer/optim.py,改用 torch.optim.Muon - parallel setup: gloo 后端不再传递 device_id,单卡多进程不再报错
This commit is contained in:
@@ -58,9 +58,11 @@ def setup_parallel(
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
os.environ["LOCAL_DEVICE"] = str(device_id)
|
||||
|
||||
dist.init_process_group(
|
||||
rank=rank, world_size=world_size, backend=backend, device_id=device_id
|
||||
)
|
||||
pg_kwargs = dict(rank=rank, world_size=world_size, backend=backend)
|
||||
if backend in ("nccl", "ccl"):
|
||||
pg_kwargs["device_id"] = device_id
|
||||
|
||||
dist.init_process_group(**pg_kwargs)
|
||||
|
||||
try:
|
||||
if backend == "nccl" and torch.cuda.is_available():
|
||||
|
||||
Reference in New Issue
Block a user