Skip to content

GitLab

Unverified Commit eaee5976 authored by msbaines's avatar msbaines Committed by GitHub
Browse files

[refactor] make AsyncPipe its own class (#341)

parent 51625eda
Hide whitespace changes
Inline Side-by-side
Showing with 200 additions and 286 deletions
benchmarks/experimental_ampnet.py
View file @ eaee5976
......@@ -21,7 +21,7 @@ from torchtext.data.utils import get_tokenizer
from experimental.nn.ampnet_pipe import pipe
from fairscale.nn.model_parallel import initialize_model_parallel
from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
from fairscale.nn.pipe import LazyModule, MultiProcessPipe
from fairscale.nn.pipe import LazyModule
from fairscale.optim import GradScaler
from fairscale.utils.testing import dist_init, get_worker_map
......@@ -420,7 +420,6 @@ def run_mp_worker(args, available_workers):
p = pipe.AMPnetPipe(
module=model,
balance=balance,
style=MultiProcessPipe.AsyncSchedule,
chunks=args.chunks,
worker_map=get_worker_map(),
input_device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"),
......
benchmarks/pipe.py
View file @ eaee5976
......@@ -499,7 +499,6 @@ def run_mp_worker(args, available_workers):
pipe_model = MultiProcessPipe(
model,
balance,
style=MultiProcessPipe.AsyncSchedule,
chunks=args.chunks,
worker_map=get_worker_map(),
input_device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"),
......
experimental/nn/ampnet_pipe/ampnet.py
View file @ eaee5976
......@@ -37,7 +37,6 @@ def create_task_without_skip_trackers(
checkpoint_stop: int, i: int, j: int, batch: Batch, partition: nn.Sequential,
) -> Task:
# Determine whether checkpointing or not.
# style is guaranteed to be PipelineStyle.AsyncSchedule
if i < checkpoint_stop:
def function(
......
experimental/nn/ampnet_pipe/pipe.py
View file @ eaee5976
......@@ -11,15 +11,14 @@ from torch import nn
from torch.optim.optimizer import Optimizer
from torch.utils.data import DataLoader
from fairscale.nn.pipe import MultiProcessPipe
from fairscale.nn.pipe.types import PipelineStyle
from fairscale.nn.pipe import AsyncPipe
from .ampnet import AsyncAMPnetEventLoop
__all__ = ["AMPnetPipe"]
class AMPnetPipe(MultiProcessPipe):
class AMPnetPipe(AsyncPipe):
"""
AMPnetPipe is the asynchronous version of the MultiProcessPipe implementation
which avoids the bubble issue, by using stale weights and gradients.
......@@ -44,7 +43,6 @@ class AMPnetPipe(MultiProcessPipe):
# AMPnet implementation doesn't handle skip_trackers!
assert self.pipeline.style is PipelineStyle.AsyncSchedule # type: ignore
assert self.group
rank = self.group.rank()
......
experimental/tests/nn/ampnet_pipe_process/test_ampnet_pipe.py
View file @ eaee5976
......@@ -23,7 +23,6 @@ from torch.optim.optimizer import Optimizer
from torch.utils.data import DataLoader, Dataset
from experimental.nn.ampnet_pipe.pipe import AMPnetPipe
from fairscale.nn.pipe import MultiProcessPipe
from fairscale.utils.testing import get_worker_map, torch_spawn
......@@ -84,14 +83,7 @@ class FakeDataset(Dataset):
@torch_spawn([2])
def async_event_loop_interleave_simple():
model = nn.Sequential(nn.Linear(10, 10), nn.ReLU(inplace=False), nn.Linear(10, 10), nn.ReLU(inplace=False))
pipe = AMPnetPipe(
module=model,
balance=[2, 2],
style=MultiProcessPipe.AsyncSchedule,
worker_map=get_worker_map(),
chunks=10,
checkpoint="never",
)
pipe = AMPnetPipe(module=model, balance=[2, 2], worker_map=get_worker_map(), chunks=10, checkpoint="never",)
fake_dataset = FakeDataset()
fake_dataloader = DataLoader(fake_dataset, batch_size=4, shuffle=True, num_workers=0)
loss = nn.MSELoss()
......@@ -102,14 +94,7 @@ def async_event_loop_interleave_simple():
@torch_spawn([4])
def async_event_loop_interleave_hard():
model = nn.Sequential(nn.Linear(10, 10), nn.Linear(10, 10), nn.Linear(10, 10), nn.Linear(10, 10))
pipe = AMPnetPipe(
module=model,
balance=[1, 1, 1, 1],
style=MultiProcessPipe.AsyncSchedule,
worker_map=get_worker_map(),
chunks=10,
checkpoint="never",
)
pipe = AMPnetPipe(module=model, balance=[1, 1, 1, 1], worker_map=get_worker_map(), chunks=10, checkpoint="never",)
fake_dataset = FakeDataset()
fake_dataloader = DataLoader(fake_dataset, batch_size=4, shuffle=True, num_workers=0)
loss = nn.MSELoss()
......
fairscale/nn/pipe/__init__.py
View file @ eaee5976
......@@ -18,6 +18,7 @@
# limitations under the License.
"""A Pipe implementation in PyTorch."""
from .async_pipe import AsyncPipe
from .checkpoint import is_checkpointing, is_recomputing
from .multiprocess_pipe import LazyModule, MultiProcessPipe
from .pipe import Pipe
......
fairscale/nn/pipe/async_pipe.py 0 → 100644
View file @ eaee5976
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.
from .multiprocess_pipe import MultiProcessPipe
from .types import PipelineStyle
class AsyncPipe(MultiProcessPipe):
def __init__(self, *args, **kwargs) -> None: # type: ignore
super().__init__(*args, style=PipelineStyle.AsyncSchedule, **kwargs)
fairscale/nn/pipe/multiprocess_pipe.py
View file @ eaee5976
......@@ -386,9 +386,6 @@ class MultiProcessPipe(Module):
"""
MultiProcess: PipelineStyle = PipelineStyle.MultiProcess
AsyncSchedule: PipelineStyle = PipelineStyle.AsyncSchedule
#: The number of layers in each partition.
balance: List[int] = []
# ^^
......
fairscale/nn/pipe/rpc.py
View file @ eaee5976
......@@ -13,6 +13,7 @@ from torch.distributed.distributed_c10d import _get_global_rank
from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
from .async_pipe import AsyncPipe
from .multiprocess_pipe import MultiProcessPipe
from .types import EVENT_LOOP_QUEUE, PipeMessage, TensorOrTensors
......@@ -105,10 +106,9 @@ class PipeRPCWrapper(nn.Module):
else:
kwargs["group"] = self.group
kwargs["style"] = MultiProcessPipe.AsyncSchedule
kwargs["input_device"] = torch.device("cuda", torch.cuda.current_device())
self.model = MultiProcessPipe(*args, **kwargs)
self.model = AsyncPipe(*args, **kwargs)
self.worker_map = kwargs["worker_map"]
self._foreach_worker(self._register_remote_model, args=(args, kwargs))
self.model.cuda()
......
fairscale/nn/pipe/types.py
View file @ eaee5976
......@@ -35,7 +35,6 @@ class LazyModule:
class PipelineStyle(Enum):
SingleProcess = auto()
MultiProcess = auto()
AsyncSchedule = auto()
......
tests/nn/model_parallel/test_layers.py
View file @ eaee5976
......@@ -431,7 +431,6 @@ def run_test_pipe(rank, world_size, filename, filename_rpc, skip_dist_init=False
model[2].weight.data = saved_weight_2
worker_map = {i: f"Test{i}" for i in range(torch.distributed.get_world_size())}
style = MultiProcessPipe.MultiProcess # MultiProcessPipe.AsyncSchedule
if pipe_world_size == 2:
print(f"actually doing pipe stuff now")
......@@ -440,7 +439,6 @@ def run_test_pipe(rank, world_size, filename, filename_rpc, skip_dist_init=False
pipe_model = MultiProcessPipe(
model,
[2, 1],
style=style,
group=pipeline_devices,
worker_map=worker_map,
input_device=torch.cuda.current_device(),
......@@ -507,8 +505,7 @@ def run_test_pipe(rank, world_size, filename, filename_rpc, skip_dist_init=False
failed = False
with torch.autograd.profiler.profile() as prof:
try:
if style == MultiProcessPipe.MultiProcess:
pipe_model.back_helper(pipe_output)
pipe_model.back_helper(pipe_output)
except Exception as e:
failed = True
print(f"got {e} while doing backward, deadlock?")
......
tests/nn/pipe_process/skip/test_gpipe.py
View file @ eaee5976
......@@ -23,7 +23,7 @@ import pytest
import torch
from torch import nn
from fairscale.nn.pipe import LazyModule, MultiProcessPipe
from fairscale.nn.pipe import AsyncPipe, LazyModule, MultiProcessPipe
from fairscale.nn.pipe.skip import pop, skippable, stash
from fairscale.nn.pipe.skip.portal import PortalBlue, PortalCopy, PortalOrange
from fairscale.utils.testing import get_worker_map, torch_spawn
......@@ -33,14 +33,14 @@ from fairscale.utils.testing import get_worker_map, torch_spawn
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("balance", [[3], [1, 2], [2, 1], [1, 1, 1]], ids=["3", "1:2", "2:1", "1:1:1"])
@pytest.mark.parametrize("checkpoint", ["never", "always", "except_last"])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" in os.environ, reason="broken on mpi")
def x1to3(balance, checkpoint, pipeline_style):
def x1to3(balance, checkpoint, pipe_class):
torch.manual_seed(0)
if pipeline_style == MultiProcessPipe.AsyncSchedule and len(balance) > 1:
if pipe_class == AsyncPipe and len(balance) > 1:
print(f"skipping yarg")
pytest.skip("Skip tensors NYI for AsyncSchedule")
pytest.skip("Skip tensors NYI for AsyncPipe")
@skippable(stash=["1to3"])
class Layer1(nn.Module):
......@@ -74,13 +74,12 @@ def x1to3(balance, checkpoint, pipeline_style):
return output
model = nn.Sequential(Layer1(), Layer2(), Layer3())
model = MultiProcessPipe(
model = pipe_class(
model,
balance,
chunks=3,
checkpoint=checkpoint,
input_device=torch.cuda.current_device(),
style=pipeline_style,
worker_map=get_worker_map(),
pipelined_backward=False,
).cuda()
......@@ -106,11 +105,11 @@ def x1to3(balance, checkpoint, pipeline_style):
@torch_spawn([2])
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" in os.environ, reason="broken on mpi")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
@pytest.mark.skip(reason="flaky test")
def none_skip(pipeline_style):
if pipeline_style == MultiProcessPipe.AsyncSchedule:
pytest.skip("Skip tensors NYI for AsyncSchedule")
def none_skip(pipe_class):
if pipe_class == AsyncPipe:
pytest.skip("Skip tensors NYI for AsyncPipe")
@skippable(stash=["none"])
class Stash(nn.Module):
......@@ -126,13 +125,8 @@ def none_skip(pipeline_style):
return input
model = nn.Sequential(Stash(), Pop())
model = MultiProcessPipe(
model,
[1, 1],
style=pipeline_style,
worker_map=get_worker_map(),
input_device=torch.cuda.current_device(),
chunks=5,
model = pipe_class(
model, [1, 1], worker_map=get_worker_map(), input_device=torch.cuda.current_device(), chunks=5,
).cuda()
input = torch.rand(10, requires_grad=True).cuda()
......@@ -161,8 +155,8 @@ def none_skip(pipeline_style):
@torch_spawn([2])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def lazy_skippable_error(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def lazy_skippable_error(pipe_class):
"""Using skippable layers in combination with lazy construction is currently
not supported, check that it raises an Exception"""
......@@ -181,6 +175,6 @@ def lazy_skippable_error(pipeline_style):
]
with pytest.raises(ValueError, match="Can't use Skippable layers with multi-process pipe and lazy construction"):
MultiProcessPipe(
model, [2, 1], style=pipeline_style, worker_map=get_worker_map(),
pipe_class(
model, [2, 1], worker_map=get_worker_map(),
)
tests/nn/pipe_process/skip/test_leak.py
View file @ eaee5976
......@@ -23,7 +23,7 @@ import pytest
import torch
from torch import nn
from fairscale.nn.pipe import MultiProcessPipe, is_checkpointing, is_recomputing
from fairscale.nn.pipe import AsyncPipe, MultiProcessPipe, is_checkpointing, is_recomputing
from fairscale.nn.pipe.skip import pop, skippable, stash
from fairscale.nn.pipe.skip.tracker import current_skip_tracker
from fairscale.utils.testing import get_worker_map, torch_spawn
......@@ -46,10 +46,10 @@ class Pop(nn.Module):
@torch_spawn([2])
@pytest.mark.parametrize("train", [True, False], ids=["train", "eval"])
@pytest.mark.parametrize("checkpoint", ["always", "except_last", "never"])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" in os.environ, reason="broken on mpi")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
def delete_portal_tensor(train, checkpoint, pipeline_style):
def delete_portal_tensor(train, checkpoint, pipe_class):
# Without checkpointing:
# +- Stash --+ +--- Pop ----+ - - - layers
# | 2,blue,1 |--| 1,orange,0 | - - - tensor_life and portal function
......@@ -60,8 +60,8 @@ def delete_portal_tensor(train, checkpoint, pipeline_style):
# | 3,blue,2 |--| 2,orange,1 |--| 1,orange,0 |--| 1,blue,0 |
# +----------+ +------------+ +------------+ +----------+
if pipeline_style == MultiProcessPipe.AsyncSchedule:
pytest.skip("Skip tensors NYI for AsyncSchedule")
if pipe_class == AsyncPipe:
pytest.skip("Skip tensors NYI for AsyncPipe")
def portal_tensor_life_is(tensor_life, skip_tracker=None):
if skip_tracker is None:
......@@ -114,9 +114,7 @@ def delete_portal_tensor(train, checkpoint, pipeline_style):
return self.F.apply(input)
model = nn.Sequential(NoPortalTensorAtBackward(), stash_, pop_)
model = MultiProcessPipe(
model, balance=[2, 1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, checkpoint=checkpoint,
)
model = pipe_class(model, balance=[2, 1], worker_map=get_worker_map(), chunks=2, checkpoint=checkpoint,)
input = torch.rand(10, requires_grad=True)
......
tests/nn/pipe_process/test_bugs.py
View file @ eaee5976
......@@ -22,15 +22,15 @@ import torch
from torch import nn
import torch.nn.functional as F
from fairscale.nn.pipe import MultiProcessPipe
from fairscale.nn.pipe import AsyncPipe, MultiProcessPipe
from fairscale.utils.testing import get_worker_map, torch_spawn
@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def python_autograd_function(pipeline_style):
# FIXME deadlock with MultiProcessPipe.AsyncSchedule?
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def python_autograd_function(pipe_class):
# FIXME deadlock with AsyncPipe?
# A Python autograd function might fail with this error:
#
# RuntimeError: Returning Variables sharing storage with other Variables
......@@ -57,9 +57,7 @@ def python_autograd_function(pipeline_style):
return Identity.apply(input)
model = nn.Sequential(M(), M())
model = MultiProcessPipe(
model, [1, 1], style=pipeline_style, worker_map=get_worker_map(), checkpoint="always"
).cuda()
model = pipe_class(model, [1, 1], worker_map=get_worker_map(), checkpoint="always").cuda()
model.eval()
x = torch.rand(42)
......@@ -73,8 +71,8 @@ def python_autograd_function(pipeline_style):
@torch_spawn([3])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def exception_no_hang(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def exception_no_hang(pipe_class):
# 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
......@@ -92,7 +90,7 @@ def exception_no_hang(pipeline_style):
raise ExpectedException()
model = nn.Sequential(Pass(), Pass(), Raise())
model = MultiProcessPipe(model, [1, 1, 1], style=pipeline_style, worker_map=get_worker_map(), chunks=3)
model = pipe_class(model, [1, 1, 1], worker_map=get_worker_map(), chunks=3)
model.eval()
if model.group.rank() == 2:
......@@ -106,8 +104,8 @@ def exception_no_hang(pipeline_style):
@torch_spawn([2])
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="2 cuda devices required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def tuple_wait(cuda_sleep, pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def tuple_wait(cuda_sleep, pipe_class):
# In v0.0.3, Wait is applied to only the first tensor on a micro-batch.
# Under this behavior, if checkpointing was disabled, there's a possibility
# that gradient accumulations on other tensors are not synchronized
......@@ -135,10 +133,9 @@ def tuple_wait(cuda_sleep, pipeline_style):
return a + b + c
model = nn.Sequential(Layer1(), Layer2())
model = MultiProcessPipe(
model = pipe_class(
model,
[1, 1],
style=pipeline_style,
worker_map=get_worker_map(),
input_device=torch.cuda.current_device(),
chunks=32,
......@@ -160,8 +157,8 @@ def tuple_wait(cuda_sleep, pipeline_style):
@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def parallel_randoms(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def parallel_randoms(pipe_class):
class Dropouts(nn.Module):
def forward(self, x):
for _ in range(100):
......@@ -172,10 +169,9 @@ def parallel_randoms(pipeline_style):
x = torch.rand(10, 10, requires_grad=True).cuda()
x.retain_grad()
model = MultiProcessPipe(
model = pipe_class(
model,
[1, 1],
style=pipeline_style,
input_device=torch.cuda.current_device(),
worker_map=get_worker_map(),
chunks=10,
......
tests/nn/pipe_process/test_inplace.py
View file @ eaee5976
......@@ -21,21 +21,21 @@ import pytest
import torch
from torch import nn
from fairscale.nn.pipe import MultiProcessPipe
from fairscale.nn.pipe import AsyncPipe, MultiProcessPipe
from fairscale.utils.testing import get_worker_map, torch_spawn
@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def inplace_on_requires_grad(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def inplace_on_requires_grad(pipe_class):
model = nn.Sequential(nn.Linear(1, 1), nn.ReLU(inplace=True))
model = MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map(), checkpoint="always")
model = pipe_class(model, [1, 1], worker_map=get_worker_map(), checkpoint="always")
x = torch.rand(1)
if pipeline_style == MultiProcessPipe.AsyncSchedule and model.group.rank() == 0:
# With AsyncSchedule, model will wait forever for gradients if not eval
if pipe_class == AsyncPipe and model.group.rank() == 0:
# With AsyncPipe, model will wait forever for gradients if not eval
model.eval()
y = model(x)
......@@ -50,12 +50,12 @@ def inplace_on_requires_grad(pipeline_style):
@torch_spawn([1])
@pytest.mark.xfail(strict=True)
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def inplace_on_not_requires_grad(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def inplace_on_not_requires_grad(pipe_class):
# In-place operation on a tensor not requiring grad doesn't cause a
# RuntimeError. Currently, we cannot detect this case.
model = nn.Sequential(nn.ReLU(inplace=True))
model = MultiProcessPipe(model, [1], style=pipeline_style, worker_map=get_worker_map(), checkpoint="always")
model = pipe_class(model, [1], worker_map=get_worker_map(), checkpoint="always")
x = torch.rand(1)
y = model(x)
......@@ -70,8 +70,8 @@ def inplace_on_not_requires_grad(pipeline_style):
@torch_spawn([1])
@pytest.mark.xfail(strict=True)
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def inplace_incorrect_grad(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def inplace_incorrect_grad(pipe_class):
class M(nn.Module):
def forward(self, foo_bar):
# 'foo' requires grad but 'bar' does not. In-place operation on
......@@ -88,7 +88,7 @@ def inplace_incorrect_grad(pipeline_style):
return foo * bar
model = nn.Sequential(M())
model = MultiProcessPipe(model, [1], style=pipeline_style, worker_map=get_worker_map(), checkpoint="always")
model = pipe_class(model, [1], worker_map=get_worker_map(), checkpoint="always")
foo = torch.tensor([1.0], requires_grad=True)
bar = torch.tensor([1.0])
......
tests/nn/pipe_process/test_pipe.py
View file @ eaee5976
......@@ -31,15 +31,16 @@ from fairscale.nn.model_parallel.initialize import (
get_pipeline_parallel_group,
initialize_model_parallel,
)
from fairscale.nn.pipe import LazyModule, MultiProcessPipe
from fairscale.nn.pipe import AsyncPipe, LazyModule, MultiProcessPipe
from fairscale.nn.pipe.types import PipelineStyle
from fairscale.utils.testing import get_worker_map, set_random_seed, torch_spawn, torch_version
@torch_spawn([2])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def parameters(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def parameters(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
pipe = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1)
pipe = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=1)
if torch.distributed.get_rank() == 0:
assert list(pipe.parameters()) != []
else:
......@@ -107,8 +108,8 @@ def mpi():
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def public_attrs(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def public_attrs(pipe_class):
class MyString:
def __init__(self, value):
self.value = value
......@@ -118,14 +119,7 @@ def public_attrs(pipeline_style):
model = nn.Sequential(nn.Linear(1, 1))
pipe = MultiProcessPipe(
model,
balance=(1,),
style=pipeline_style,
worker_map=get_worker_map(),
chunks=42.000,
checkpoint=MyString("always"),
)
pipe = pipe_class(model, balance=(1,), worker_map=get_worker_map(), chunks=42.000, checkpoint=MyString("always"),)
assert pipe.balance == [1]
assert pipe.chunks == 42
......@@ -136,13 +130,13 @@ def public_attrs(pipeline_style):
@torch_spawn([2])
@pytest.mark.parametrize("balance", [[2], [1, 1]])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def sequential_like(balance, pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def sequential_like(balance, pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
model = nn.Sequential(a, b)
model = MultiProcessPipe(model, balance, style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, balance, worker_map=get_worker_map())
if balance == [2]:
if torch.distributed.get_rank() == 0:
......@@ -175,62 +169,62 @@ def sequential_like(balance, pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def balance_wrong_length(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def balance_wrong_length(pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
model = nn.Sequential(a, b)
with pytest.raises(ValueError):
MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
pipe_class(model, balance=[1], worker_map=get_worker_map())
with pytest.raises(ValueError):
MultiProcessPipe(model, balance=[3], style=pipeline_style, worker_map=get_worker_map())
pipe_class(model, balance=[3], worker_map=get_worker_map())
@torch_spawn([2])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def balance_less_than_1(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def balance_less_than_1(pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
model = nn.Sequential(a, b)
with pytest.raises(ValueError):
MultiProcessPipe(model, balance=[0, 2], style=pipeline_style, worker_map=get_worker_map())
pipe_class(model, balance=[0, 2], worker_map=get_worker_map())
with pytest.raises(ValueError):
MultiProcessPipe(model, balance=[-1, 3], style=pipeline_style, worker_map=get_worker_map())
pipe_class(model, balance=[-1, 3], worker_map=get_worker_map())
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def chunks_less_than_1(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def chunks_less_than_1(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
with pytest.raises(ValueError):
MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=0)
pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=0)
with pytest.raises(ValueError):
MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=-1)
pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=-1)
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def too_few_devices(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def too_few_devices(pipe_class):
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 = MultiProcessPipe(model, balance=[1, 1, 1, 1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, balance=[1, 1, 1, 1], worker_map=get_worker_map())
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def batch_size_indivisible(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def batch_size_indivisible(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=4)
model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=4)
with pytest.warns(None) as record:
model(torch.rand(7, 1))
......@@ -240,10 +234,10 @@ def batch_size_indivisible(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def batch_size_small(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def batch_size_small(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=4)
model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=4)
with pytest.warns(None) as record:
model(torch.rand(2, 1))
......@@ -253,8 +247,8 @@ def batch_size_small(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def checkpoint_mode(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def checkpoint_mode(pipe_class):
def count_grad_fn(grad_fn, name, visited=set()):
if grad_fn in visited:
return 0
......@@ -273,32 +267,14 @@ def checkpoint_mode(pipeline_style):
model = nn.Sequential(nn.Linear(1, 1))
input = torch.rand(2, 1)
always = MultiProcessPipe(
model,
balance=[1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=2,
checkpoint="always",
pipelined_backward=False,
always = pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=2, checkpoint="always", pipelined_backward=False,
)
except_last = MultiProcessPipe(
model,
balance=[1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=2,
checkpoint="except_last",
pipelined_backward=False,
except_last = pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=2, checkpoint="except_last", pipelined_backward=False,
)
never = MultiProcessPipe(
model,
balance=[1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=2,
checkpoint="never",
pipelined_backward=False,
never = pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=2, checkpoint="never", pipelined_backward=False,
)
always_output = always(input)
......@@ -311,45 +287,34 @@ def checkpoint_mode(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def checkpoint_mode_invalid(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def checkpoint_mode_invalid(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
with pytest.raises(ValueError, match="checkpoint is not one of 'always', 'except_last', or 'never'"):
MultiProcessPipe(
model,
balance=[1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=2,
checkpoint="INVALID_CHECKPOINT",
pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=2, checkpoint="INVALID_CHECKPOINT",
)
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def checkpoint_mode_when_chunks_1(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def checkpoint_mode_when_chunks_1(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
# All checkpoint modes are fine.
MultiProcessPipe(
model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1, checkpoint="except_last",
)
MultiProcessPipe(
model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1, checkpoint="always"
)
MultiProcessPipe(
model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1, checkpoint="never"
pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=1, checkpoint="except_last",
)
pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=1, checkpoint="always")
pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=1, checkpoint="never")
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def checkpoint_eval(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def checkpoint_eval(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
model = MultiProcessPipe(
model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False,
)
model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=2, pipelined_backward=False,)
input = torch.rand(2, 1)
def find_grad_fn(grad_fn, name):
......@@ -375,8 +340,8 @@ def checkpoint_eval(pipeline_style):
@torch_spawn([2])
@pytest.mark.xfail(torch_version() < (1, 6, 0), reason="Doesn't work on torch < 1.6.0", strict=True)
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def checkpoint_non_float_input(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def checkpoint_non_float_input(pipe_class):
class ForkNonFloat(nn.Module):
def forward(self, input):
return (input * 2, torch.tensor([False]))
......@@ -386,14 +351,8 @@ def checkpoint_non_float_input(pipeline_style):
return input[0] * 2
model = nn.Sequential(ForkNonFloat(), JoinNonFloat())
model = MultiProcessPipe(
model,
balance=[1, 1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=1,
checkpoint="always",
pipelined_backward=False,
model = pipe_class(
model, balance=[1, 1], worker_map=get_worker_map(), chunks=1, checkpoint="always", pipelined_backward=False,
)
input = torch.rand(1, requires_grad=True)
......@@ -401,17 +360,17 @@ def checkpoint_non_float_input(pipeline_style):
if model.group.rank() == 1:
# with torch.autograd.detect_anomaly():
output.backward()
elif pipeline_style == MultiProcessPipe.MultiProcess:
elif pipe_class == MultiProcessPipe:
model.back_helper(output)
torch.distributed.barrier()
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def no_grad(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def no_grad(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2)
model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=2)
input = torch.rand(2, 1)
latent = None
......@@ -433,8 +392,8 @@ def no_grad(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def exception(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def exception(pipe_class):
class ExpectedException(Exception):
pass
......@@ -443,7 +402,7 @@ def exception(pipeline_style):
raise ExpectedException()
model = nn.Sequential(Raise())
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1)
model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=1)
with pytest.raises(ExpectedException):
model(torch.rand(1))
......@@ -453,8 +412,8 @@ def exception(pipeline_style):
@torch_spawn([4])
@pytest.mark.skipif(torch.cuda.is_available() and torch.cuda.device_count() < 4, reason="Not enough GPUs")
@pytest.mark.xfail(strict=True)
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def exception_early_stop_asap(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def exception_early_stop_asap(pipe_class):
"""Even the first partitions have finished to process, the partition before
the failed partition hould be killed as soon as possible.
"""
......@@ -482,7 +441,7 @@ def exception_early_stop_asap(pipeline_style):
raise ExpectedException()
model = nn.Sequential(Pass(), Pass(), Counter(), Raise())
model = MultiProcessPipe(model, [1, 1, 1, 1], style=pipeline_style, worker_map=get_worker_map(), chunks=3)
model = pipe_class(model, [1, 1, 1, 1], worker_map=get_worker_map(), chunks=3)
with pytest.raises(ExpectedException):
model(torch.rand(3))
......@@ -492,8 +451,8 @@ def exception_early_stop_asap(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def input_pair(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def input_pair(pipe_class):
class Two(nn.Module):
def __init__(self):
super().__init__()
......@@ -505,9 +464,7 @@ def input_pair(pipeline_style):
return (self.fc_a(a), self.fc_b(b))
model = nn.Sequential(Two())
model = MultiProcessPipe(
model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False,
)
model = pipe_class(model, balance=[1], 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)
......@@ -521,8 +478,8 @@ def input_pair(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def input_singleton(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def input_singleton(pipe_class):
class One(nn.Module):
def __init__(self):
super().__init__()
......@@ -533,9 +490,7 @@ def input_singleton(pipeline_style):
return (self.fc(a),)
model = nn.Sequential(One())
model = MultiProcessPipe(
model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False,
)
model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=2, pipelined_backward=False,)
a = torch.rand(10, 1, requires_grad=True)
......@@ -548,10 +503,10 @@ def input_singleton(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def input_varargs(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def input_varargs(pipe_class):
model = nn.Sequential(nn.Linear(1, 1))
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, balance=[1], worker_map=get_worker_map())
a = torch.rand(1)
b = torch.rand(1)
......@@ -562,14 +517,14 @@ def input_varargs(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def non_tensor(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def non_tensor(pipe_class):
class NonTensor(nn.Module):
def forward(self, _):
return "hello"
model = nn.Sequential(NonTensor())
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, balance=[1], worker_map=get_worker_map())
x = torch.rand(1)
# TypeError: expected Tensor as element 0 in argument 0, but got str
......@@ -582,14 +537,14 @@ def non_tensor(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def non_tensor_tuple(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def non_tensor_tuple(pipe_class):
class NonTensorTuple(nn.Module):
def forward(self, x):
return (x, "hello")
model = nn.Sequential(NonTensorTuple())
model = MultiProcessPipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, balance=[1], worker_map=get_worker_map())
x = torch.rand(1)
# TypeError: CheckpointBackward.forward: expected Variable (got str) for return value 1
......@@ -604,8 +559,8 @@ def non_tensor_tuple(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("checkpoint", ["never", "always", "except_last"])
@pytest.mark.parametrize("lazy", [True, False])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def deferred_batch_norm(checkpoint, lazy, pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def deferred_batch_norm(checkpoint, lazy, pipe_class):
bn = nn.BatchNorm2d(3)
pipe_bn = deepcopy(bn)
pipe_fn = lambda: pipe_bn # noqa: E731
......@@ -613,14 +568,8 @@ def deferred_batch_norm(checkpoint, lazy, pipeline_style):
model = [LazyModule(pipe_fn)]
else:
model = nn.Sequential(pipe_bn)
pipe = MultiProcessPipe(
model,
balance=[1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=2,
checkpoint=checkpoint,
deferred_batch_norm=True,
pipe = pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=2, checkpoint=checkpoint, deferred_batch_norm=True,
)
x = torch.rand(4, 3, 10, 10)
......@@ -634,8 +583,8 @@ def deferred_batch_norm(checkpoint, lazy, pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("checkpoint", ["never", "always"])
@pytest.mark.parametrize("lazy", [True, False])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def deferred_batch_norm_params(checkpoint, lazy, pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def deferred_batch_norm_params(checkpoint, lazy, pipe_class):
bn = nn.BatchNorm2d(3)
pipe_bn = deepcopy(bn)
pipe_fn = lambda: pipe_bn # noqa: E731
......@@ -643,14 +592,8 @@ def deferred_batch_norm_params(checkpoint, lazy, pipeline_style):
model = [LazyModule(pipe_fn)]
else:
model = nn.Sequential(pipe_bn)
pipe = MultiProcessPipe(
model,
balance=[1],
style=pipeline_style,
worker_map=get_worker_map(),
chunks=1,
checkpoint=checkpoint,
deferred_batch_norm=True,
pipe = pipe_class(
model, balance=[1], worker_map=get_worker_map(), chunks=1, checkpoint=checkpoint, deferred_batch_norm=True,
)
x = torch.rand(4, 3, 10, 10)
......@@ -665,15 +608,15 @@ def deferred_batch_norm_params(checkpoint, lazy, pipeline_style):
@torch_spawn([4])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def devices(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def devices(pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
c = nn.Linear(1, 1)
# There are extra two ranks.
model = nn.Sequential(a, b, c)
model = MultiProcessPipe(model, [1, 1, 1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, [1, 1, 1], worker_map=get_worker_map())
# Extra devices must be discarded.
if model.group.rank() == 3:
......@@ -681,13 +624,13 @@ def devices(pipeline_style):
@torch_spawn([2])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def partitions(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def partitions(pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
model = nn.Sequential(a, b)
model = MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, [1, 1], worker_map=get_worker_map())
assert isinstance(model.partitions, list)
assert len(model) == 1
......@@ -701,13 +644,13 @@ def partitions(pipeline_style):
@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def deny_moving(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def deny_moving(pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
model = nn.Sequential(a, b)
model = MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, [1, 1], worker_map=get_worker_map())
model.cuda()
model.cpu()
......@@ -725,11 +668,11 @@ def deny_moving(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def empty_module(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def empty_module(pipe_class):
# Empty sequential module is not illegal.
model = nn.Sequential()
model = MultiProcessPipe(model, [], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, [], worker_map=get_worker_map())
assert model(torch.tensor([42])) == torch.tensor([42])
assert model((torch.tensor([42]),)) == (torch.tensor([42]),)
......@@ -741,13 +684,13 @@ def empty_module(pipeline_style):
@torch_spawn([2])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def named_children(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def named_children(pipe_class):
a = nn.Linear(1, 1)
b = nn.Linear(1, 1)
model = nn.Sequential(OrderedDict([("a", a), ("b", b)]))
model = MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
model = pipe_class(model, [1, 1], worker_map=get_worker_map())
names = set(n for n, _ in model.named_modules())
if model.group.rank() == 0:
......@@ -762,24 +705,24 @@ def named_children(pipeline_style):
@torch_spawn([1])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def recommend_auto_balance(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def recommend_auto_balance(pipe_class):
with pytest.raises(ValueError, match="fairscale.nn.pipe.balance"):
# balance is required
MultiProcessPipe(nn.Sequential())
pipe_class(nn.Sequential())
with pytest.raises(ValueError, match="fairscale.nn.pipe.balance"):
# module and sum of balance have differen length (module: 0, sum of balance: 1)
MultiProcessPipe(nn.Sequential(), [1])
pipe_class(nn.Sequential(), [1])
with pytest.raises(ValueError, match="fairscale.nn.pipe.balance"):
# module and sum of balance have different length (module: 2, sum of balance: 1)
MultiProcessPipe(nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1)), [1])
pipe_class(nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1)), [1])
@torch_spawn([2])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def lazy_construction(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def lazy_construction(pipe_class):
init_count = 0
class Custom(nn.Module):
......@@ -798,7 +741,7 @@ def lazy_construction(pipeline_style):
LazyModule(lambda: Custom()),
]
pipe = MultiProcessPipe(model, balance=[2, 2], style=pipeline_style, worker_map=get_worker_map())
pipe = pipe_class(model, balance=[2, 2], worker_map=get_worker_map())
assert isinstance(pipe[0], Custom)
assert isinstance(pipe[1], Custom)
......@@ -808,18 +751,18 @@ def lazy_construction(pipeline_style):
@torch_spawn([2])
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" in os.environ, reason="doesn't apply to mpi")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def missing_worker_map(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def missing_worker_map(pipe_class):
model = nn.Sequential(nn.ReLU(), nn.ReLU())
with pytest.raises(ValueError, match="'RpcTransport' requires 'worker_map' to be set"):
MultiProcessPipe(model, [1, 1], style=pipeline_style)
pipe_class(model, [1, 1])
@torch_spawn([2])
@pytest.mark.skip(reason="currently broken")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def verify_module_duplicate_parameters_on_distinct_partitions(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def verify_module_duplicate_parameters_on_distinct_partitions(pipe_class):
class Surrogate(nn.Module):
def __init__(self, module):
super().__init__()
......@@ -830,23 +773,23 @@ def verify_module_duplicate_parameters_on_distinct_partitions(pipeline_style):
# FIXME(tom) can't have duplicate params with separate processes
with pytest.raises(ValueError, match="module with duplicate parameters on distinct devices is not supported"):
MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
pipe_class(model, [1, 1], worker_map=get_worker_map())
@torch_spawn([4])
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def pipelined_backward(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def pipelined_backward(pipe_class):
model = nn.Sequential(nn.ReLU(), nn.ReLU())
destroy_model_parallel()
initialize_model_parallel(1, 4)
pipe = MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
pipe = pipe_class(model, [1, 1], worker_map=get_worker_map())
assert pipe.pipelined_backward is False
destroy_model_parallel()
initialize_model_parallel(2, 2)
pipe = MultiProcessPipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
pipe = pipe_class(model, [1, 1], worker_map=get_worker_map())
assert pipe.pipelined_backward is True
......@@ -855,9 +798,7 @@ def pipelined_backward(pipeline_style):
def async_event_loop():
model = nn.Sequential(nn.Linear(10, 10), nn.ReLU(), nn.Linear(10, 10), nn.ReLU())
pipe = MultiProcessPipe(
model, [1, 1, 1, 1], style=MultiProcessPipe.AsyncSchedule, worker_map=get_worker_map(), chunks=10
)
pipe = AsyncPipe(model, [1, 1, 1, 1], worker_map=get_worker_map(), chunks=10)
inputs = torch.rand(100, 10)
......@@ -873,7 +814,7 @@ def reuse_lazy():
reused = LazyModule(lambda: nn.Linear(10, 10))
model = [reused, nn.Linear(10, 10), nn.ReLU(), reused, nn.ReLU(), reused, nn.ReLU()]
# model = [reused, reused, nn.Linear(10, 10), nn.ReLU(), reused, reused, nn.ReLU(), reused, reused, nn.ReLU()]
pipe = MultiProcessPipe(model, [3, 1, 1], style=MultiProcessPipe.AsyncSchedule, worker_map=get_worker_map())
pipe = AsyncPipe(model, [3, 1, 1], worker_map=get_worker_map())
pipe.eval()
output = pipe(torch.rand(10))
......@@ -891,7 +832,7 @@ def reuse_lazy():
# ensure identical weights but no sharing between model and pipe
reused = nn.Linear(10, 10)
layers = [reused, nn.Linear(10, 10), nn.ReLU(), reused, nn.ReLU(), reused, nn.ReLU()]
pipe = MultiProcessPipe(layers, [3, 1, 1], style=MultiProcessPipe.AsyncSchedule, worker_map=get_worker_map())
pipe = AsyncPipe(layers, [3, 1, 1], worker_map=get_worker_map())
pipe.eval()
model_optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
pipe_optimizer = torch.optim.SGD(pipe.parameters(), lr=0.01, momentum=0.9) if len(list(pipe.parameters())) else None
......@@ -964,7 +905,7 @@ def test_instantiate_partition():
# instantiated model
for rank in range(len(balance)):
instantiated = instantiate_partition(
model, balance, FakeGroup(rank, len(balance)), MultiProcessPipe.AsyncSchedule
model, balance, FakeGroup(rank, len(balance)), PipelineStyle.AsyncSchedule
)
for part in instantiated:
assert isinstance(part.module, nn.Sequential)
......
tests/nn/pipe_process/test_transparency.py
View file @ eaee5976
......@@ -21,14 +21,14 @@ import pytest
import torch
from torch import nn
from fairscale.nn.pipe import MultiProcessPipe
from fairscale.nn.pipe import AsyncPipe, MultiProcessPipe
from fairscale.utils.testing import get_worker_map, set_random_seed, torch_spawn
@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@pytest.mark.parametrize("pipeline_style", [MultiProcessPipe.MultiProcess, MultiProcessPipe.AsyncSchedule])
def simple_linears(pipeline_style):
@pytest.mark.parametrize("pipe_class", [MultiProcessPipe, AsyncPipe])
def simple_linears(pipe_class):
def sum_grad(parameters):
return sum([p.grad.sum() for p in parameters if p.grad is not None])
......@@ -54,8 +54,7 @@ def simple_linears(pipeline_style):
zero_grad(model.parameters())
# With MultiProcessPipe
model = MultiProcessPipe(model, [2, 2], style=pipeline_style, worker_map=get_worker_map(), chunks=4)
model = pipe_class(model, [2, 2], worker_map=get_worker_map(), chunks=4)
outputs = model(inputs)
if model.group.rank() == 1:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment