Repository: BlackSamorez/tensor_parallel (language: python)

Issue:
py3.7/8/9 support, assorted minor fixes/refactor
- removed match/case for py3.7/8/9 compatibility
- gather_op dim=-1 is no longer hard-coded into AllGather
- output_rules ALL is now 0, rationale: less things to worry about
- config "reduce" is now "sum"
- ran affected files through black / isort



Discussion / hints:

Fix patch:
diff --git a/petals_local_parallel/__init__.py b/petals_local_parallel/__init__.py
index 5b3db59..3ad9c14 100644
--- a/petals_local_parallel/__init__.py
+++ b/petals_local_parallel/__init__.py
@@ -1,1 +1,6 @@
-# from slicer_wrapper_interface import SlicingConfig, tensor_parallel
+from .slicer_wrapper_interface import SlicingConfig, tensor_parallel
+
+__all__ = [
+    "SlicingConfig",
+    "tensor_parallel"
+]
diff --git a/petals_local_parallel/all_to_all_communication_primitives.py b/petals_local_parallel/all_to_all_communication_primitives.py
index 7907b91..d154906 100644
--- a/petals_local_parallel/all_to_all_communication_primitives.py
+++ b/petals_local_parallel/all_to_all_communication_primitives.py
@@ -1,35 +1,40 @@
 import threading
+from typing import List, Optional
 
 import torch
 
 
-class AllReduce:
+class CollectiveOpetationBase:
+    def __call__(self, x: torch.Tensor, rank: int):
+        raise NotImplementedError()
+
+
+class AllReduce(CollectiveOpetationBase):
     def __init__(self, world_size: int, reduce_op: callable = sum):
         self.scatter_reduce = ScatterReduce(world_size, reduce_op)
         self.all_gather = AllGather(world_size, gather_op=torch.cat)
-    
+
     def __call__(self, x: torch.Tensor, rank: int):
         reduced_part = self.scatter_reduce(x, rank)
-        return self.all_gather(reduced_part, rank).reshape_as(x)
+        return self.all_gather(reduced_part, rank).view_as(x)
 
 
-class ScatterReduce:
+class ScatterReduce(CollectiveOpetationBase):
     def __init__(self, world_size: int, reduce_op: callable = sum):
         self.world_size = world_size
         self.tensor_parts = [[] for _ in range(world_size)]
         self.parts_ready = [threading.Event() for _ in range(world_size)]
         self.reduce_op = reduce_op
-        
+
     def __call__(self, x: torch.Tensor, rank: int):
         try:
             for i, part in enumerate(x.flatten().tensor_split(self.world_size)):
-                self.tensor_parts[i].append(part) # append is thread-safe. thanks, GIL!
+                self.tensor_parts[i].append(part)  # append is thread-safe. thanks, GIL!
                 if len(self.tensor_parts[i]) == self.world_size:
-                    self.parts_ready[i].set() # can be called more than once; we dont care
+                    self.parts_ready[i].set()  # can be called more than once; we dont care
 
             self.parts_ready[rank].wait()
-            parts_to_reduce = [part.to(x.device, non_blocking=True)
-                               for part in self.tensor_parts[rank]]
+            parts_to_reduce = [part.to(x.device, non_blocking=True) for part in self.tensor_parts[rank]]
             reduced_part = self.reduce_op(parts_to_reduce)
             return reduced_part
         finally:
@@ -38,29 +43,30 @@ def __call__(self, x: torch.Tensor, rank: int):
             self.parts_ready[rank].clear()
 
 
-class AllGather:
-    def __init__(self, world_size: int, gather_op = torch.cat, ordered: bool = True):
+class AllGather(CollectiveOpetationBase):
+    def __init__(self, world_size: int, gather_op: callable = torch.cat):
         self.world_size = world_size
-        self.parts = []
+        self.parts: List[Optional[torch.Tensor]] = [None for _ in range(world_size)]
+        self.ranks_updated = []
         self.parts_ready = threading.Event()
         self.gather_op = gather_op
-        self.ordered = ordered
-        
+
     def __call__(self, x: torch.Tensor, rank: int):
+        parts, ranks_updated, parts_ready = self.parts, self.ranks_updated, self.parts_ready
+        # ^-- note: we copy properties to locals so that the "finally" clause is thread-safe
         try:
-            parts, parts_ready = self.parts, self.parts_ready
-            # ^-- note: we copy properties to locals so that the "finally" clause is thread-safe
-            parts.append((rank, x))  # append is thread-safe. thanks, GIL!
-            if len(parts) == self.world_size:
-                parts_ready.set() # can be called more than once; we dont care
+            parts[rank] = x  # no race b/c each rank writes to a separate location
+            ranks_updated.append(rank)  # append is thread-safe. thanks, GIL!
+            if len(ranks_updated) == self.world_size:
+                parts_ready.set()  # can be called more than once; we dont care
             parts_ready.wait()
-            if self.ordered:
-                parts = sorted(parts)  # sorted by rank
-            parts = [part.to(x.device, non_blocking=True) for r, part in parts]
+            parts = [part.to(x.device, non_blocking=True) for part in parts]
             # note: for one of the parts with r == rank, part.to(device) is a no-op
             return self.gather_op(parts)
         finally:
-            self.parts = []
-            self.parts_ready = threading.Event()
-            # note: we can safely clear now because all ranks have copied 
-            # self.parts_* to locals before passing parts_ready.wait
+            if ranks_updated[-1] == rank:
+                self.parts = [None for _ in range(self.world_size)]
+                self.ranks_updated = []
+                self.parts_ready = threading.Event()
+            # note: we can safely clear now because all ranks have
+            # copied self.parts_* to locals before passing parts_ready.wait
diff --git a/petals_local_parallel/autoconfig.py b/petals_local_parallel/autoconfig.py
new file mode 100644
index 0000000..1bbdd5a
--- /dev/null
+++ b/petals_local_parallel/autoconfig.py
@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+
+from .slicing_config import SlicingConfig
+
+def build_default_slicing_config(model: nn.Module) -> SlicingConfig:
+    slicing_config = SlicingConfig({}, {})
+
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            slicing_config.tensor_rules[name + ".(weight|bias)"] = "vertical"
+            slicing_config.module_rules[name] = {"input": {}, "output": {0: "gather"}, "attributes": {}}
+        elif isinstance(module, (nn.Embedding, nn.EmbeddingBag)):
+            slicing_config.tensor_rules[name + ".weight"] = "horizontal"
+            slicing_config.module_rules[name] = {"input": {}, "output": {0: "gather"}, "attributes": {}}
+    
+    return slicing_config
diff --git a/petals_local_parallel/communications.py b/petals_local_parallel/communications.py
index 6611e07..e8779a2 100644
--- a/petals_local_parallel/communications.py
+++ b/petals_local_parallel/communications.py
@@ -6,6 +6,7 @@
 from all_to_all_communication_primitives import AllGather, AllReduce
 from torch import Tensor
 
+from .all_to_all_communication_primitives import AllReduce, AllGather
 
 class Communicator(ABC):
     @abstractmethod
diff --git a/petals_local_parallel/slicer_wrapper.py b/petals_local_parallel/slicer_wrapper.py
index cc936e4..e8c24db 100644
--- a/petals_local_parallel/slicer_wrapper.py
+++ b/petals_local_parallel/slicer_wrapper.py
@@ -1,115 +1,88 @@
 import re
+from typing import Callable, Dict, Iterator, Tuple, Union
 
-import communications
 import torch
-import torch.distributed as dist
 import torch.nn as nn
+from .autoconfig import build_default_slicing_config
+from .all_to_all_communication_primitives import AllReduce, AllGather
+from .slicing_config import SlicingConfig
 from torch.nn.parallel import parallel_apply
+from transformers import PreTrainedModel, PretrainedConfig
 
 
-class SlicingConfig():
-    def __init__(self, tensor_rules: dict, module_rules: dict):
-        self.tensor_rules = tensor_rules
-        self.module_rules = module_rules
-                
+Pattern, Arg = str, Union[int, str]
 
-def slice_weight_vertical(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
-    assert(tensor.shape[-2] % world_size == 0)
-    slice_size = tensor.shape[-2] // world_size
 
-    return tensor[..., rank * slice_size: (rank + 1) * slice_size, :]
+def slice_weight_vertical(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
+    return tensor.tensor_split(world_size, dim=-2)[rank]
 
 
 def slice_bias_vertical(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
-    assert(tensor.shape[-1] % world_size == 0)
-    slice_size = tensor.shape[-1] // world_size
-
-    return tensor[rank * slice_size: (rank + 1) * slice_size]
+    return tensor.tensor_split(world_size)[rank]
 
 
 def slice_weight_horizontal(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
-    assert(tensor.shape[-1] % world_size == 0)
-    slice_size = tensor.shape[-1] // world_size
-
-    return tensor[..., rank * slice_size: (rank + 1) * slice_size]
+    return tensor.tensor_split(world_size, dim=-1)[rank]
 
 
 def slice_bias_horizontal(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
     return tensor / world_size
 
 
-def slice_tensors(key_parameter_iterator, tensor_rules: dict, rank: int, world_size: int):
+SLICING_RULES = {
+    ("vertical", "weight"): slice_weight_vertical,
+    ("vertical", "bias"): slice_bias_vertical,
+    ("horizontal", "weight"): slice_weight_horizontal,
+    ("horizontal", "bias"): slice_bias_horizontal,
+}
+
+
+def slice_tensors(
+    key_parameter_iterator: Iterator[Tuple[str, nn.Parameter]], tensor_rules: Dict[Arg, str], rank: int, world_size: int
+):
     regular_rules = [(re.compile(key), value) for key, value in tensor_rules.items()]
 
     with torch.no_grad():
         for name, param in key_parameter_iterator:
             for pattern, rule in regular_rules:
-                if not pattern.search(name) is None:
-                    name_ending = name.split('.')[-1]
-                    match (rule, name_ending):
-                        case ("vertical", "weight"):
-                            param.data = slice_weight_vertical(param.data, rank=rank, world_size=world_size)
-                        case ("vertical", "bias"):
-                            param.data = slice_bias_vertical(param.data, rank=rank, world_size=world_size)
-                        case ("horizontal", "weight"):
-                            param.data = slice_weight_horizontal(param.data, rank=rank, world_size=world_size)
-                        case ("horizontal", "bias"):
-                            param.data = slice_bias_horizontal(param.data, rank=rank, world_size=world_size)
-                        case _:
-                            raise Exception("Fuck you tensor!")
-
-
-def process_input(rules, rank, world_size, *args, **kwargs):
-    args = list(args)
+                if pattern.search(name) is not None:
+                    name_ending = name.split(".")[-1]
+                    param.data = SLICING_RULES[rule, name_ending](param.data, rank=rank, world_size=world_size)
+
+
+def process_input(rules: Dict[Arg, str], rank: int, world_size: int, *args, **kwargs):
+    extended_kwargs = dict(kwargs)
+    extended_kwargs.update(enumerate(args))
     for target, action in rules.items():
-        match action.split():
-            case "cut", dim:
-                dim = int(dim)
-                match target:
-                    case int(idx):
-                        slice_size= args[idx].shape[dim] // world_size
-                        args[idx] = args[idx][..., rank * slice_size: (rank + 1) * slice_size, :]
-                    case str(name):
-                        slice_size= kwargs[name].shape[dim] // world_size
-                        kwargs[name] = kwargs[name][rank * slice_size: (rank + 1) * slice_size, ...]
-                    case _:
-                        raise Exception("Fuck you cut input!")
-
-            case "scale":
-                match target:
-                    case int(idx):
-                        args[idx] = args[idx] / world_size
-                    case str(name):
-                        kwargs[name] = kwargs[name]/ world_size
-                    case _:
-                        raise Exception("Fuck you scale input!")
-            case _:
-                raise Exception("Fuck you input action!")
-
-    return args, kwargs
-
-
-def process_output(output, rules, rank: int):
+        if not isinstance(extended_kwargs.get(target), torch.Tensor):
+            continue  # optional parameter is None or False
+        action_type, *opts = action.split()
+        if action_type == "cut":
+            dim = int(opts[0])
+            extended_kwargs[target] = extended_kwargs[target].tensor_split(world_size, dim=dim)[rank]
+        elif action_type == "scale":
+            extended_kwargs[target] = extended_kwargs[target] / world_size
+        else:
+            raise Exception(f"unexpected action {action_type}")
+
+    args = [extended_kwargs.pop(i) for i in range(len(args))]
+    return args, extended_kwargs
+
+
+def process_output(output, rules: Dict[Arg, Callable[[torch.Tensor, int], torch.Tensor]], rank: int):
+    if isinstance(output, torch.Tensor):
+        return process_output([output], rules, rank)[0]
     for target, action in rules.items():
-        match action:
-            case "reduce":
-                match target:
-                    case "ALL":
-                        output = communications.TENSOR_PARALLEL_COMMUNICATOR.all_reduce(output, rank)
-                    case int(idx):
-                        output[idx] = communications.TENSOR_PARALLEL_COMMUNICATOR.all_reduce(output[idx], rank)
-                    case _:
-                        raise Exception("Fuck you output taget!")
-            case _:
-                raise Exception("Fuck you output action!")
+        output[target] = action(output[target], rank)
     return output
 
 
-def process_attr(module, rules, rank, world_size):
+def process_attr(module: nn.Module, rules: Dict[Arg, str], rank: int, world_size: int):
     for attr, action in rules.items():
-            match action:
-                case "scale_int":
-                    setattr(module, attr, getattr(module, attr) // world_size)
+        if action == "scale_int":
+            setattr(module, attr, getattr(module, attr) // world_size)
+        else:
+            raise NotImplementedError(action)
 
 
 class ParallelLayerWrapper(nn.Module):
@@ -124,7 +97,6 @@ def __init__(self, module: nn.Module, module_rules: dict, rank: int, world_size:
         self.rank = rank
         self.world_size = world_size
 
-        
     def forward(self, *args, **kwargs):
         args, kwargs = process_input(self.input_rules, self.rank, self.world_size, *args, **kwargs)
         output = self.module(*args, **kwargs)
@@ -132,6 +104,22 @@ def forward(self, *args, **kwargs):
 
 
 def wrap_submodules(model: nn.Module, module_rules: dict, rank: int, world_size: int):
+    unique_output_transforms = {op for rules in module_rules.values() for op in rules['output'].values()}
+    transform_map = {}
+    for transform in unique_output_transforms:
+        if transform == 'sum':
+            transform_map[transform] = AllReduce(world_size, reduce_op=sum)
+        elif transform == 'gather':
+            transform_map[transform] = AllGather(world_size, gather_op=lambda xs: torch.cat(xs, dim=-1))
+        elif callable(transform):
+            transform_map[transform] = transform  # user-defined transform, no action needed
+        else:
+            raise NotImplementedError(f"Unknown output transform {transform}")
+
+    for rules in module_rules.values():
+        for key, rule in rules['output'].items():
+            rules['output'][key] = transform_map[rule]
+
     unique_wrappers = {}
     with torch.no_grad():
         for name, module in model.named_modules():
@@ -147,36 +135,46 @@ def wrap_submodules(model: nn.Module, module_rules: dict, rank: int, world_size:
 
 def get_tensor_parallel_model_slice(model_cls, slicing_config: SlicingConfig, rank: int, world_size: int):
     class _TensorParallelSlice(model_cls):
-        slicing_config = None
-        rank = None
-        world_size = None
         def __new__(cls, *args, __slicing_config=slicing_config, __rank=rank, __world_size=world_size, **kwargs):
             _TensorParallelSlice.slicing_config = __slicing_config
             _TensorParallelSlice.rank = __rank
             _TensorParallelSlice.world_size = __world_size
 
             model = model_cls(*args, **kwargs)  # Create an instance of vanilla model
-            
+
+            if _TensorParallelSlice.slicing_config is None:
+                _TensorParallelSlice.slicing_config = build_default_slicing_config(model)
+
             # modify untrained parameters/buffers
-            slice_tensors(model.named_parameters(), slicing_config.tensor_rules, rank, world_size)
+            slice_tensors(model.named_parameters(), _TensorParallelSlice.slicing_config.tensor_rules, rank, world_size)
             return model
 
         @classmethod
         def _load_pretrained_model(cls, model: model_cls, state_dict, loaded_keys, *args, **kwargs):
-            slice_tensors(state_dict.items(), slicing_config.tensor_rules, rank, world_size)
+            slice_tensors(
+                state_dict.items(),
+                _TensorParallelSlice.slicing_config.tensor_rules,
+                _TensorParallelSlice.rank,
+                _TensorParallelSlice.world_size,
+            )
             result = super()._load_pretrained_model(model, state_dict, loaded_keys, *args, **kwargs)
-            
-            wrap_submodules(model, slicing_config.module_rules, rank, world_size)
+
+            wrap_submodules(
+                model,
+                _TensorParallelSlice.slicing_config.module_rules,
+                _TensorParallelSlice.rank,
+                _TensorParallelSlice.world_size,
+            )
 
             return result
-        
+
     return _TensorParallelSlice
 
 
-class MultithreadedModule(nn.Module):
+class MultithreadedModule(PreTrainedModel):
     def __init__(self, slice_types, devices) -> None:
-        super().__init__()
-        assert(len(slice_types) == len(devices))
+        super().__init__(PretrainedConfig())  # Temporary empty config. Gets replaced in from_pretrained
+        assert len(slice_types) == len(devices)
         self.slice_types = slice_types
         self.slices = torch.nn.ModuleList()
         self.devices = devices
@@ -201,10 +199,13 @@ def scatter_map(obj):
         return parallel_apply(self.slices, inputs, kwargs_tup=kwargs_tup)[0]
 
     def from_pretrained(self, *args, **kwargs):
-        self.slices = torch.nn.ModuleList([slice_type.from_pretrained(*args, **kwargs) for slice_type in self.slice_types])
+        self.slices = torch.nn.ModuleList(
+            [slice_type.from_pretrained(*args, **kwargs) for slice_type in self.slice_types]
+        )
+        self.config = self.slices[0].config
         return self.scatter()
 
     def scatter(self):
         for slice, device in zip(self.slices, self.devices):
-            slice = slice.to(device)
+            slice.to(device, non_blocking=True)
         return self
diff --git a/petals_local_parallel/slicer_wrapper_interface.py b/petals_local_parallel/slicer_wrapper_interface.py
index 0c1c993..62906eb 100644
--- a/petals_local_parallel/slicer_wrapper_interface.py
+++ b/petals_local_parallel/slicer_wrapper_interface.py
@@ -1,4 +1,8 @@
-import communications
+from .slicing_config import SlicingConfig
+from .slicer_wrapper import get_tensor_parallel_model_slice, MultithreadedModule
+from .slicing_configs import SLICING_CONFIGS
+from . import communications
+
 import torch
 import torch.distributed as dist
 from slicer_wrapper import MultithreadedModule, SlicingConfig, get_tensor_parallel_model_slice
@@ -10,7 +14,7 @@ def tensor_parallel(model_cls, devices, slicing_config: SlicingConfig = None):
         try:
             slicing_config = SLICING_CONFIGS[model_cls.__name__]
         except KeyError:
-            raise NotImplemented(f"Unknown model type {model_cls.__name__} and lazy mode not implemented yet. Must specify config")
+            print("No slicing_config provided. Using nn.Linear slicing fallback")
 
     if dist.is_initialized():
         communications.TENSOR_PARALLEL_COMMUNICATOR = communications.TorchrunCommunicator()
diff --git a/petals_local_parallel/slicing_config.py b/petals_local_parallel/slicing_config.py
new file mode 100644
index 0000000..7d198ae
--- /dev/null
+++ b/petals_local_parallel/slicing_config.py
@@ -0,0 +1,4 @@
+class SlicingConfig():
+    def __init__(self, tensor_rules: dict, module_rules: dict):
+        self.tensor_rules = tensor_rules
+        self.module_rules = module_rules
\ No newline at end of file
diff --git a/petals_local_parallel/slicing_configs.py b/petals_local_parallel/slicing_configs.py
index f0c3983..9c925df 100644
--- a/petals_local_parallel/slicing_configs.py
+++ b/petals_local_parallel/slicing_configs.py
@@ -1,4 +1,4 @@
-from slicer_wrapper import SlicingConfig
+from .slicer_wrapper import SlicingConfig
 
 SLICING_CONFIGS = {
 "BloomModel": SlicingConfig(
@@ -10,8 +10,8 @@
     },
     {
         ".*self_attention$": {"input": {"alibi": "cut 0"}, "output": {}, "attributes": {"num_heads": "scale_int"}},
-        ".*self_attention\.dense$": {"input": {}, "output": {"ALL": "reduce"}, "attributes": {}},
-        ".*mlp\.dense_4h_to_h$": {"input": {}, "output": {"ALL": "reduce"}, "attributes": {}},
+        ".*self_attention\.dense$": {"input": {}, "output": {0: "sum"}, "attributes": {}},
+        ".*mlp\.dense_4h_to_h$": {"input": {}, "output": {0: "sum"}, "attributes": {}},
     },
 ),
 
@@ -26,8 +26,8 @@
     },
     {
         ".*attention\.self$": {"input": {}, "output": {}, "attributes": {"num_attention_heads": "scale_int", "all_head_size": "scale_int"}},
-        ".*attention\.output\.dense$": {"input": {}, "output": {0: "reduce"}, "attributes": {}},
-        ".*[0-9]\.output\.dense$": {"input": {}, "output": {"ALL": "reduce"}, "attributes":{}},
+        ".*attention\.output\.dense$": {"input": {}, "output": {0: "sum"}, "attributes": {}},
+        ".*[0-9]\.output\.dense$": {"input": {}, "output": {0: "sum"}, "attributes":{}},
     },
 ),
 
@@ -44,9 +44,9 @@
         ".*DenseReluDense\.wo\.(weight|bias)": "horizontal",
     },
     {
-        ".*SelfAttention$": {"input": {}, "output": {0: "reduce"}, "attributes": {"n_heads": "scale_int", "inner_dim": "scale_int"}},
+        ".*SelfAttention$": {"input": {}, "output": {0: "sum"}, "attributes": {"n_heads": "scale_int", "inner_dim": "scale_int"}},
         ".*relative_attention_bias$": {"input": {}, "output": {}, "attributes": {"embedding_dim": "scale_int"}},
-        ".*DenseReluDense$": {"input": {}, "output": {"ALL": "reduce"}, "attributes": {}},
+        ".*DenseReluDense$": {"input": {}, "output": {0: "sum"}, "attributes": {}},
     },
 )
 }
\ No newline at end of file


Test patch:
diff --git a/petals_local_parallel/slicer_wrapper_test.py b/petals_local_parallel/slicer_wrapper_test.py
index 5bd30ec..452abe2 100644
--- a/petals_local_parallel/slicer_wrapper_test.py
+++ b/petals_local_parallel/slicer_wrapper_test.py
@@ -13,8 +13,9 @@
 logging.set_verbosity_error()
 
 
+from .slicer_wrapper_interface import tensor_parallel
+
 import torch.distributed as dist
-from slicer_wrapper_interface import tensor_parallel
 
 BACKEND = 'nccl' if torch.cuda.is_available() else 'gloo'
 torch.set_num_threads(1)
diff --git a/petals_local_parallel/slicer_wrapper_test_cpu.py b/petals_local_parallel/slicer_wrapper_test_cpu.py
index 6925338..ef627fc 100644
--- a/petals_local_parallel/slicer_wrapper_test_cpu.py
+++ b/petals_local_parallel/slicer_wrapper_test_cpu.py
@@ -11,7 +11,7 @@
 logging.set_verbosity_error()
 
 
-from slicer_wrapper_interface import tensor_parallel
+from .slicer_wrapper_interface import tensor_parallel
 
 
 def converter_main():
@@ -26,6 +26,7 @@ def converter_main():
     print(f"Loading slices")
     model = tensor_parallel(MODEL_CLS, devices=["cpu", "cpu"]).from_pretrained(NAME)
     
+    print(f"Slices forward")
     sharded_output = model(test_input).last_hidden_state.cpu()
 
     print(f"Asserting allclose")