| """Gradient averaging across ranks (gloo, CPU) for the multi-GPU trainer.""" |
|
|
| from __future__ import annotations |
|
|
| import os |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
|
|
|
|
| def _worker(rank: int, world: int, port: int) -> None: |
| import torch.distributed as dist |
|
|
| from deg.train import _allreduce_grads |
|
|
| os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port)) |
| dist.init_process_group("gloo", rank=rank, world_size=world) |
| params = [torch.nn.Parameter(torch.zeros(n)) for n in (3, 5, 7)] |
| params[0].grad = torch.full((3,), float(rank + 1)) |
| params[1].grad = torch.full((5,), float(10 * (rank + 1))) |
| |
| _allreduce_grads(params, world, bucket=4) |
| assert torch.equal(params[0].grad, torch.full((3,), 1.5)) |
| assert torch.equal(params[1].grad, torch.full((5,), 15.0)) |
| assert torch.equal(params[2].grad, torch.zeros(7)) |
| dist.destroy_process_group() |
|
|
|
|
| def test_allreduce_grads_averages_across_ranks(): |
| import socket |
|
|
| import torch.multiprocessing as mp |
|
|
| with socket.socket() as s: |
| s.bind(("127.0.0.1", 0)) |
| port = s.getsockname()[1] |
| mp.spawn(_worker, args=(2, port), nprocs=2, join=True) |
|
|
