def _remote_worker_process(self, ddp_mode):
gLogger.info("The remote worker is running.")
dist.init_process_group(
backend="gloo",
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
world_size=self.world_size,
rank=self.rank,
)
if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE):
# new_group needs to be called on ranks.
dist.new_group(TRAINER_RANKS)
global shutdown_signal
with shutdown_signal:
shutdown_signal.wait()
gLogger.info("Exiting remote worker.")
dist.destroy_process_group()
def _master_process(self, ddp_mode: DdpMode, simulate_uneven_inputs: bool):
gLogger.info("Running the master process...")
dist.init_process_group(
backend="gloo",
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
world_size=self.world_size,
rank=self.rank,
)
remote_em_rref = rpc.remote(self.remote_worker_name(),
RemoteEM,
args=(NUM_EM_ROW, D_SPARSE))
remote_net_rref = rpc.remote(self.remote_worker_name(),
RemoteNet,
args=(D_DENSE + D_SPARSE, D_HID))
gLogger.info("Created remote rrefs on master")
self.do_test_on_master(ddp_mode, simulate_uneven_inputs,
remote_em_rref, remote_net_rref)
def _trainer_process(self, rank: int):
gLogger.info(f"Running the trainer #{rank}...")
gLogger.info(
f"Initing trainer process group by trainer #{rank} with ranks {TRAINER_RANKS}"
)
dist.init_process_group(
backend="gloo",
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
world_size=self.world_size,
rank=self.rank,
)
gLogger.info(f"Waiting for shutdown signal on trainer #{rank}...")
global shutdown_signal
with shutdown_signal:
shutdown_signal.wait()
gLogger.info(f"Exiting the trainer #{rank}...")
dist.destroy_process_group()
def test_ddp_dist_autograd_local_vs_remote(self):
# Each trainer uses a different random seed. Otherwise, they are going
# to have exactly the same initial model parameters, input, and
# therefore grads. That means the grads will be the same before and
# after DDP's all-reduce.
torch.manual_seed(self.rank)
dist.init_process_group(
backend="gloo",
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
world_size=self.world_size,
rank=self.rank,
)
remote_layer1 = RemoteModule("worker0",
device="cpu",
module_cls=nn.Linear,
args=(10, 5, False))
layer1 = nn.Linear(10, 5, False)
# Start with the same parameters for remote and local
layer1.weight = remote_layer1.module_rref.to_here().weight
# Run local case.
layer2 = nn.Linear(5, 1)
inputs = torch.rand((10, 10))
ddp_model = DistributedDataParallel(layer2)
loss = ddp_model(layer1(inputs)).sum()
loss.backward()
# Run remote case.
with dist_autograd.context() as context_id:
loss = ddp_model(remote_layer1(inputs)).sum()
dist_autograd.backward(context_id, [loss])
grads_dict = dist_autograd.get_gradients(context_id)
dist.barrier()
self.assertEqual(layer2.weight.grad, grads_dict[layer2.weight])
self.assertEqual(
layer1.weight.grad,
rpc.rpc_sync(
"worker0",
DdpComparisonTest.get_remote_grads,
args=(remote_layer1.module_rref, context_id),
),
)
def _run_basic_test(self,
backend,
checkpoint,
find_unused_parameters=False):
dist.init_process_group(
backend="nccl",
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
world_size=self.world_size,
rank=self.rank,
)
# Use 4 GPUs, two replicas of a pipe across GPU 0 and 1 and another
# pipe between GPU 2 and 3. Both replicas are replicated via DDP.
fc1 = nn.Linear(16, 8, bias=False).cuda(2 * self.rank)
class MyModule(nn.Module):
def __init__(self, device):
super(MyModule, self).__init__()
self.fc2 = nn.Linear(8, 4, bias=False).cuda(device)
self.fc3 = nn.Linear(4, 2, bias=False).cuda(device)
def forward(self, inp):
if find_unused_parameters:
return self.fc2(inp)
else:
return self.fc3(self.fc2(inp))
layer2 = MyModule(2 * self.rank + 1)
model = nn.Sequential(fc1, layer2)
model = Pipe(model, chunks=2, checkpoint=checkpoint)
model = DistributedDataParallel(
model, find_unused_parameters=find_unused_parameters)
out = model(torch.rand(16, 16).cuda(2 * self.rank)).local_value()
out.sum().backward()
# Run forward again for find_unused_parameters to trigger any potential errors.
if find_unused_parameters:
model(torch.rand(16, 16).cuda(2 * self.rank))
# Check grads
output = [
torch.empty_like(fc1.weight.grad),
torch.empty_like(fc1.weight.grad)
]
dist.all_gather(output, fc1.weight.grad)
self.assertEqual(output[0], output[1])
output = [
torch.empty_like(layer2.fc2.weight.grad),
torch.empty_like(layer2.fc2.weight.grad)
]
dist.all_gather(output, layer2.fc2.weight.grad)
self.assertEqual(output[0], output[1])
if not find_unused_parameters:
output = [
torch.empty_like(layer2.fc3.weight.grad),
torch.empty_like(layer2.fc3.weight.grad)
]
dist.all_gather(output, layer2.fc3.weight.grad)
self.assertEqual(output[0], output[1])
def _run_basic_test(self,
backend,
checkpoint,
find_unused_parameters=False,
static_graph=False):
dist.init_process_group(
backend=backend,
init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name),
world_size=self.world_size,
rank=self.rank,
)
# Use 4 GPUs, two replicas of a pipe across GPU 0 and 1 and another
# pipe between GPU 2 and 3. Both replicas are replicated via DDP.
fc1 = nn.Linear(16, 8, bias=False).cuda(2 * self.rank)
class MyModule(nn.Module):
def __init__(self, device):
super(MyModule, self).__init__()
self.fc2 = nn.Linear(8, 4, bias=False).cuda(device)
self.fc3 = nn.Linear(4, 2, bias=False).cuda(device)
def forward(self, inp):
if find_unused_parameters:
return self.fc2(inp)
else:
return self.fc3(self.fc2(inp))
layer2 = MyModule(2 * self.rank + 1)
model = nn.Sequential(fc1, layer2)
model = Pipe(model, chunks=2, checkpoint=checkpoint)
model = DistributedDataParallel(
model,
find_unused_parameters=find_unused_parameters,
static_graph=static_graph,
)
# Ensure inputs are different across ranks to verify that gradient
# sync indeed occurs.
model_input = torch.rand(16, 16).cuda(2 * self.rank) * (self.rank + 1)
out = model(model_input).local_value()
out.sum().backward()
# Run forward again for find_unused_parameters to trigger any potential errors.
if find_unused_parameters:
# Ensure inputs are different across ranks to verify that gradient
# sync indeed occurs.
unused_param_input = torch.rand(16, 16).cuda(
2 * self.rank) * (self.rank + 1)
model(unused_param_input).local_value().sum().backward()
# Run a few more iterations of fwd + bwd to ensure gradient synchronization
# occurs properly across iterations via delay_all_reduce/bucketized allreduce.
for _ in range(3):
model_input = torch.rand(16, 16).cuda(
2 * self.rank) * (self.rank + 1)
out = model(model_input).local_value()
out.sum().backward()
# Check grads
output = [
torch.empty_like(fc1.weight.grad),
torch.empty_like(fc1.weight.grad)
]
dist.all_gather(output, fc1.weight.grad)
self.assertEqual(output[0], output[1])
output = [
torch.empty_like(layer2.fc2.weight.grad),
torch.empty_like(layer2.fc2.weight.grad)
]
dist.all_gather(output, layer2.fc2.weight.grad)
self.assertEqual(output[0], output[1])
if not find_unused_parameters:
output = [
torch.empty_like(layer2.fc3.weight.grad),
torch.empty_like(layer2.fc3.weight.grad)
]
dist.all_gather(output, layer2.fc3.weight.grad)
self.assertEqual(output[0], output[1])