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)