def chunks_less_than_1(): model = nn.Sequential(nn.Linear(1, 1)) with pytest.raises(ValueError): Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=0) with pytest.raises(ValueError): Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=-1)
def checkpoint_mode_when_chunks_1(): model = nn.Sequential(nn.Linear(1, 1)) # All checkpoint modes are fine. Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=1, checkpoint="except_last", ) Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=1, checkpoint="always") Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=1, checkpoint="never")
def balance_less_than_1(): a = nn.Linear(1, 1) b = nn.Linear(1, 1) model = nn.Sequential(a, b) with pytest.raises(ValueError): Pipe(model, balance=[0, 2], style=Pipe.MultiProcess, worker_map=get_worker_map()) with pytest.raises(ValueError): Pipe(model, balance=[-1, 3], style=Pipe.MultiProcess, worker_map=get_worker_map())
def balance_wrong_length(): a = nn.Linear(1, 1) b = nn.Linear(1, 1) model = nn.Sequential(a, b) with pytest.raises(ValueError): Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map()) with pytest.raises(ValueError): Pipe(model, balance=[3], style=Pipe.MultiProcess, worker_map=get_worker_map())
def checkpoint_mode(): def count_grad_fn(grad_fn, name, visited=set()): if grad_fn in visited: return 0 visited.add(grad_fn) if grad_fn is None: return 0 if grad_fn.__class__.__name__ == name: return 1 counter = 0 for next_grad_fn, _ in grad_fn.next_functions: counter += count_grad_fn(next_grad_fn, name, visited=visited) return counter model = nn.Sequential(nn.Linear(1, 1)) input = torch.rand(2, 1) always = Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=2, checkpoint="always", pipelined_backward=False, ) except_last = Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=2, checkpoint="except_last", pipelined_backward=False, ) never = Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=2, checkpoint="never", pipelined_backward=False, ) always_output = always(input) except_last_output = except_last(input) never_output = never(input) assert count_grad_fn(always_output.grad_fn, "CheckpointBackward") == 2 assert count_grad_fn(except_last_output.grad_fn, "CheckpointBackward") == 1 assert count_grad_fn(never_output.grad_fn, "CheckpointBackward") == 0
def construct_only_rank_zero(): model = [nn.Linear(10, 10), nn.ReLU()] if torch.distributed.get_rank() == 0: PipeRPCWrapper(model, [1, 1], worker_map=get_worker_map()) rpc.shutdown() else: # Must enter rpc loop to complte PipeRPCWrapper constructor above rpc.shutdown() with pytest.raises(AssertionError): PipeRPCWrapper(model, [1, 1], worker_map=get_worker_map())
def pipelined_backward(): model = nn.Sequential(nn.ReLU(), nn.ReLU()) destroy_model_parallel() initialize_model_parallel(1, 4) pipe = Pipe(model, [1, 1], style=Pipe.MultiProcess, worker_map=get_worker_map()) assert pipe.pipelined_backward is False destroy_model_parallel() initialize_model_parallel(2, 2) pipe = Pipe(model, [1, 1], style=Pipe.MultiProcess, worker_map=get_worker_map()) assert pipe.pipelined_backward is True
def exception_no_hang(pipeline_style): # In v0.0.2, once a failed partition receives a normal message # (non-closing) for the next micro-batch, a hang occured. The reason was # that a failed partition didn't call in_queue.task_done() on a normal # message. So the former partition was blocked at out_queue.join() for the # next of next micro-batch. class ExpectedException(Exception): pass class Pass(nn.Module): def forward(self, x): return x class Raise(nn.Module): def forward(self, x): raise ExpectedException() model = nn.Sequential(Pass(), Pass(), Raise()) model = Pipe(model, [1, 1, 1], style=pipeline_style, worker_map=get_worker_map(), chunks=3) model.eval() if model.group.rank() == 2: with pytest.raises(ExpectedException): model(torch.rand(3)) else: model(torch.rand(3)) torch.distributed.barrier()
def no_grad(pipeline_style): model = nn.Sequential(nn.Linear(1, 1)) model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2) input = torch.rand(2, 1) latent = None def hook(module, input, output): _ = module _ = input nonlocal latent latent = output partition = model.mp_partitions[0] partition.module.register_forward_hook(hook) with torch.no_grad(): model(input) assert latent.grad_fn is None
def input_singleton(pipeline_style): class One(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(1, 1) def forward(self, only_a): (a, ) = only_a return (self.fc(a), ) model = nn.Sequential(One()) model = Pipe( model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False, ) a = torch.rand(10, 1, requires_grad=True) (a_out, ) = model((a, )) loss = a_out.mean() loss.backward() assert all(p.grad is not None for p in model.parameters()) assert a.grad is not None
def sequential_like(balance): a = nn.Linear(1, 1) b = nn.Linear(1, 1) model = nn.Sequential(a, b) model = Pipe(model, balance, style=Pipe.MultiProcess, worker_map=get_worker_map()) if balance == [2]: if torch.distributed.get_rank() == 0: assert len(model) == 2 assert list(model) == [a, b] assert model[0] is a assert model[1] is b with pytest.raises(IndexError): _ = model[2] assert model[-1] is b assert model[-2] is a else: assert len(model) == 0 assert list(model) == [] else: assert len(model) == 1 if torch.distributed.get_rank() == 0: assert list(model) == [a] assert model[0] is a assert model[-1] is a else: assert list(model) == [b] assert model[0] is b assert model[-1] is b with pytest.raises(IndexError): _ = model[1]
def public_attrs(): class MyString: def __init__(self, value): self.value = value def __str__(self): return self.value model = nn.Sequential(nn.Linear(1, 1)) pipe = Pipe( model, balance=(1,), style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=42.000, checkpoint=MyString("always"), ) print(f"balance = {pipe.devices}") assert pipe.balance == [1] assert pipe.devices is None assert pipe.chunks == 42 assert isinstance(pipe.chunks, int) assert pipe.checkpoint == "always" assert isinstance(pipe.checkpoint, str)
def lazy_skippable_error(pipeline_style): """Using skippable layers in combination with lazy construction is currently not supported, check that it raises an Exception""" @skippable(stash=["1to3"]) class Layer1(nn.Linear): pass @skippable(pop=["1to3"]) class Layer3(nn.Linear): pass model = [ LazyModule(lambda: Layer1(10, 10)), LazyModule(lambda: nn.Linear(10, 10)), LazyModule(lambda: Layer3(10, 10)), ] with pytest.raises( ValueError, match= "Can't use Skippable layers with multi-process pipe and lazy construction" ): Pipe( model, [2, 1], style=pipeline_style, worker_map=get_worker_map(), )
def parallel_randoms(pipeline_style): class Dropouts(nn.Module): def forward(self, x): for _ in range(100): x = F.dropout(x, p=0.001) return x model = nn.Sequential(Dropouts(), Dropouts()) x = torch.rand(10, 10, requires_grad=True).cuda() x.retain_grad() model = Pipe( model, [1, 1], style=pipeline_style, input_device=torch.cuda.current_device(), worker_map=get_worker_map(), chunks=10, checkpoint="always", ).cuda() y = model(x) tensor_list = [torch.empty_like(x) for _ in range(2)] if model.group.rank() == 1: y.norm().backward() torch.distributed.barrier() tensor_list[model.group.rank()] = y torch.distributed.all_gather(tensor_list, y, group=model.group) assert tensor_list[0].to(torch.bool).tolist() == tensor_list[1].to( torch.bool).tolist() else: model.back_helper(y) torch.distributed.barrier() tensor_list[model.group.rank()] = x.grad torch.distributed.all_gather(tensor_list, x.grad, group=model.group)
def run_mp_worker(args, available_workers): new_data = True blob = make_model_and_data(args, None, new_data=new_data) model = blob["model"] balance = generate_balance_weighted(get_pipeline_parallel_group().size(), len(model), 0.8) p = pipe.Pipe( model, balance, style=Pipe.AsyncSchedule, chunks=args.chunks, worker_map=get_worker_map(), input_device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"), pipelined_backward=args.pipelined_backward, checkpoint=args.checkpoint, # loss_fn=blob["criterion"], ) if torch.cuda.is_available(): p = p.cuda() if args.all_at_once and p.pipeline: print(f"running all at once") p.pipeline.all_at_once = True if new_data: train(blob["data"], p, blob["criterion"], blob["optimizer"], blob["vocab_size"], args) else: ntokens, train_data, val_data, test_data = blob["data"] benchmark_language_model(train_data, val_data, test_data, p, criterion, optimizer, ntokens, args)
def basic_rpc(): init_rpc() if torch.distributed.get_rank() != 0: rpc.shutdown() torch.distributed.barrier() return model = [nn.Linear(10, 10), nn.ReLU()] pipe = PipeRPCWrapper(model, [1, 1], input_device=torch.cuda.current_device(), worker_map=get_worker_map()) pipe.foreach_worker(register_optimizer, include_self=True) inputs = torch.rand(10).cuda() output = pipe(inputs) loss = output.mean() loss.backward() pipe.foreach_worker(step_optimizer, include_self=True) pipe.eval() rpc.shutdown() torch.distributed.barrier()
def python_autograd_function(): # A Python autograd function might fail with this error: # # RuntimeError: Returning Variables sharing storage with other Variables # that require grad is not supported in Python functions. Please submit a # feature request if you hit this error. # # It doesn't look like an essential restriction. But it happens on the # current PyTorch version. To avoid it, we should detach the tensor before # returning by identity autograd functions, such as Wait, Fork, and Join. torch.manual_seed(0) class Identity(torch.autograd.Function): @staticmethod def forward(ctx, input): return input @staticmethod def backward(ctx, grad): return grad class M(nn.Module): def forward(self, input): return Identity.apply(input) model = nn.Sequential(M(), M()) model = Pipe(model, [1, 1], style=Pipe.MultiProcess, worker_map=get_worker_map(), checkpoint="always").cuda() x = torch.rand(42) y = model(x) if model.group.rank() == 1: assert torch.allclose(x, y) torch.distributed.barrier()
def lazy_construction(): init_count = 0 class Custom(nn.Module): def __init__(self): super(Custom, self).__init__() nonlocal init_count init_count += 1 def forward(self, x): return x model = [ lambda: Custom(), lambda: Custom(), lambda: Custom(), lambda: Custom(), ] pipe = Pipe(model, balance=[2, 2], style=Pipe.MultiProcess, worker_map=get_worker_map()) assert isinstance(pipe[0], Custom) assert isinstance(pipe[1], Custom) assert len(pipe) == 2 assert init_count == 2
def deferred_batch_norm_params(checkpoint, lazy): bn = nn.BatchNorm2d(3) pipe_bn = deepcopy(bn) pipe_fn = lambda: pipe_bn # noqa: E731 if lazy: model = [pipe_fn] else: model = nn.Sequential(pipe_bn) pipe = Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=1, checkpoint=checkpoint, deferred_batch_norm=True, ) x = torch.rand(4, 3, 10, 10) pipe(x).mean().backward() bn(x).mean().backward() assert pipe[0].weight.grad is not None assert pipe[0].bias.grad is not None assert torch.allclose(pipe[0].weight.grad, bn.weight.grad, atol=1e-4) assert torch.allclose(pipe[0].bias.grad, bn.bias.grad, atol=1e-4)
def inplace_incorrect_grad(): class M(nn.Module): def forward(self, foo_bar): # 'foo' requires grad but 'bar' does not. In-place operation on # 'bar' won't cause a RuntimeError. foo, bar = foo_bar # add_(1) is not idempotent, in contrast to relu_(). If it is # executed multiple times, it will accumulates each difference onto # 'bar'. bar.add_(1) # 'bar' is still captured by checkpointing. 'foo' will get # incorrect grad. return foo * bar model = nn.Sequential(M()) model = Pipe(model, [1], style=Pipe.MultiProcess, worker_map=get_worker_map(), checkpoint="always") foo = torch.tensor([1.0], requires_grad=True) bar = torch.tensor([1.0]) output = model((foo, bar)) del model output.backward() # The gradient of 'foo' should be 2, but it is 3 actually because # bar.add_(1) was executed twice due to checkpointing. assert foo.grad.item() == 2.0
def input_pair(): class Two(nn.Module): def __init__(self): super().__init__() self.fc_a = nn.Linear(1, 1) self.fc_b = nn.Linear(1, 1) def forward(self, a_and_b): a, b = a_and_b return (self.fc_a(a), self.fc_b(b)) model = nn.Sequential(Two()) model = Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=2, pipelined_backward=False, ) a = torch.rand(10, 1, requires_grad=True) b = torch.rand(10, 1, requires_grad=True) a_out, b_out = model((a, b)) loss = (a_out + b_out).mean() loss.backward() assert a.grad is not None assert b.grad is not None
def run_mp_worker(args, available_workers): new_data = True blob = make_model_and_data(args, None, new_data=new_data) model = blob["model"] balance = generate_balance(min(available_workers, 8), len(model)) p = pipe.Pipe( model, balance, style=Pipe.MultiProcess, chunks=args.chunks, worker_map=get_worker_map(), input_device=torch.cuda.current_device(), pipelined_backward=args.pipelined_backward, checkpoint=args.checkpoint, ).cuda() if args.all_at_once and p.pipeline: print(f"running all at once") p.pipeline.all_at_once = True if new_data: train(blob["data"], p, blob["criterion"], blob["optimizer"], blob["vocab_size"], args) else: ntokens, train_data, val_data, test_data = blob["data"] benchmark_language_model(train_data, val_data, test_data, p, criterion, optimizer, ntokens, args)
def exception_early_stop_asap(): """Even the first partitions have finished to process, the partition before the failed partition hould be killed as soon as possible. """ class ExpectedExceptio(Exception): pass class Pass(nn.Module): def forward(self, x): return x counter = 0 class Counter(nn.Module): def forward(self, x): time.sleep(0.1) nonlocal counter counter += 1 return x class Raise(nn.Module): def forward(self, x): raise ExpectedException() model = nn.Sequential(Pass(), Pass(), Counter(), Raise()) model = Pipe(model, [1, 1, 1, 1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=3) with pytest.raises(ExpectedException): model(torch.rand(3)) # If the early stop doesn't work, it would be 3 instead. assert counter == 2
def checkpoint_non_float_input(): class ForkNonFloat(nn.Module): def forward(self, input): return (input * 2, torch.tensor([False])) class JoinNonFloat(nn.Module): def forward(self, input): return input[0] * 2 model = nn.Sequential(ForkNonFloat(), JoinNonFloat()) model = Pipe( model, balance=[1, 1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=1, checkpoint="always", pipelined_backward=False, ) input = torch.rand(1, requires_grad=True) output = model(input) if model.group.rank() == 1: # with torch.autograd.detect_anomaly(): output.backward() else: model.back_helper(output)
def parameters(): model = nn.Sequential(nn.Linear(1, 1)) pipe = Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=1) if torch.distributed.get_rank() == 0: assert list(pipe.parameters()) != [] else: assert list(pipe.parameters()) == []
def checkpoint_eval(): model = nn.Sequential(nn.Linear(1, 1)) model = Pipe( model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=2, pipelined_backward=False, ) input = torch.rand(2, 1) def find_grad_fn(grad_fn, name): if grad_fn is None: return False if grad_fn.__class__.__name__ == name: return True for next_grad_fn, _ in grad_fn.next_functions: if find_grad_fn(next_grad_fn, name): return True return False model.train() train_output = model(input) assert find_grad_fn(train_output.grad_fn, "CheckpointBackward") assert find_grad_fn(train_output.grad_fn, "RecomputeBackward") model.eval() eval_output = model(input) assert not find_grad_fn(eval_output.grad_fn, "CheckpointBackward") assert not find_grad_fn(eval_output.grad_fn, "RecomputeBackward")
def batch_size_small(): model = nn.Sequential(nn.Linear(1, 1)) model = Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=4) with pytest.warns(None) as record: model(torch.rand(2, 1)) # Batch size smaller than chunks is legal. assert not record
def batch_size_indivisible(): model = nn.Sequential(nn.Linear(1, 1)) model = Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map(), chunks=4) with pytest.warns(None) as record: model(torch.rand(7, 1)) # Indivisible batch size is legal. assert not record
def too_few_devices(pipeline_style): model = nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1)) with pytest.raises(IndexError): # len(balance) > len(group.size()) model = Pipe(model, balance=[1, 1, 1, 1], style=pipeline_style, worker_map=get_worker_map())
def input_varargs(): model = nn.Sequential(nn.Linear(1, 1)) model = Pipe(model, balance=[1], style=Pipe.MultiProcess, worker_map=get_worker_map()) a = torch.rand(1) b = torch.rand(1) # TypeError: forward() takes 2 positional arguments but 3 were given with pytest.raises(TypeError): model(a, b)