Initial pipeline parallelism support

sharktank/sharktank/examples/pipeline/export_ppffn_net.py

@@ -0,0 +1,121 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""Example program to export a sharded FFN network like what is found in
+a typical transformer layer. This is used for developing and testing various
+tooling flows with a scaled down example.
+
+Generate MLIR and a random inited IRPA file with:
+
+    python -m sharktank.examples.sharding.export_ffn_net \
+        --output-irpa-file=/tmp/ffn.irpa /tmp/ffn.mlir
+"""
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+from ...layers import *
+from ... import ops
+from ...types import *
+
+from iree.turbine.aot import DeviceAffinity, DeviceTensorTrait, export
+
+def create_theta(
+    dim: int, shard_count: int, num_layers: int, save_path
+):
+    split_size = dim // shard_count
+
+    weights = []
+    for layer in range(num_layers):
+        _weight = torch.rand(dim, dim, dtype=torch.float16) / math.sqrt(dim)
+        weights.append(
+            SplitPrimitiveTensor(
+                name=f"w.{layer}", 
+                shard_dim=1,
+                ts=_weight.split(split_size, dim=1),
+                devices = [_d for _d in range(shard_count)]
+            )
+        )
+
+
+    ds = Dataset({}, Theta(weights))
+    ds.save(save_path)
+
+
+def pipeline_parallelize_theta(
+        theta: Theta,
+        pp_count: int
+) -> Theta:
+    num_layers = len(theta.tensor("w"))
+    shard_count = theta.tensor("w", '0').shard_count
+
+    for layer, weight in ((int(l), w) for l, w in theta.tensor("w").items()):
+        pp_group =  int(layer * pp_count / num_layers)
+        zero_4_group = shard_count * pp_group
+        weight.devices = tuple(i + zero_4_group for i in range(shard_count))
+        weight.devices_pinned = True
+        for i, shard in enumerate(weight.shards):
+            DeviceTensorTrait(weight.devices[i]).set(shard._data)
+
+    return theta
+
+
+class PPFFN(ThetaLayer):
+    def forward(self, x: torch.Tensor):
+        num_layers  = len(self.theta.tensor("w"))
+        shard_count = self.theta.tensor("w", '0').shard_count
+
+        x = ReplicatedTensor(ts=x, shard_count=shard_count, devices=tuple(_d for _d in range(shard_count)))
+        for layer in range(num_layers):
+            weight: SplitPrimitiveTensor = self.theta.tensor("w", str(layer))
+            x:      ReplicatedTensor     = ops.all_reduce(ops.linear(x, weight))
+
+        return x
+
+
+def main(raw_args=None):
+    from ...utils import cli
+
+    parser = cli.create_parser()
+    parser.add_argument(
+        "output_file",
+        type=str,
+        nargs="?",
+        default="-",
+        help="Output file to save MLIR to",
+    )
+    cli.add_output_dataset_options(parser)
+    args = cli.parse(parser, args=raw_args)
+
+    bs = 16
+    sl = 128
+    primary_dim = 128 * 2**5
+    shard_count = 2
+    num_layers  = 40
+    create_theta(primary_dim, shard_count, num_layers, save_path=args.output_irpa_file)
+
+    pp_count = 4
+    ds = Dataset.load(args.output_irpa_file)
+    root_theta = pipeline_parallelize_theta(ds.root_theta, pp_count)
+
+    mdl = PPFFN(root_theta)
+
+    example_arg = torch.empty(bs, sl, primary_dim, dtype=torch.float16)
+    ep = torch.export.export(mdl, (example_arg,))#, strict=False)
+    cm = export(ep, arg_device={0: DeviceAffinity(0)})
+
+    if args.output_file == "-":
+        print(cm.mlir_module)
+    else:
+        with open(args.output_file, "wt") as f:
+            f.write(str(cm.mlir_module))
+
+
+if __name__ == "__main__":
+    main()

sharktank/sharktank/examples/sharding/export_ffn_net.py

@@ -30,16 +30,16 @@ def create_theta(
     gate_weight = torch.rand(hidden_dim, primary_dim, dtype=torch.float16)
     up_weight = torch.rand(hidden_dim, primary_dim, dtype=torch.float16)
     sharded_gate_weight = SplitPrimitiveTensor(
-        name="ffn_gate.weight", shard_dim=0, ts=gate_weight.split(split_size, dim=0)
+        name="ffn_gate.weight", shard_dim=0, ts=gate_weight.split(split_size, dim=0), devices=tuple(range(shard_count)), devices_pinned=True
     )
     sharded_up_weight = SplitPrimitiveTensor(
-        name="ffn_up.weight", shard_dim=0, ts=up_weight.split(split_size, dim=0)
+        name="ffn_up.weight", shard_dim=0, ts=up_weight.split(split_size, dim=0), devices=tuple(range(shard_count)), devices_pinned=True
     )
 
     # Rowwise (transposed) sharding of down weight.
     down_weight = torch.rand(primary_dim, hidden_dim, dtype=torch.float16)
     sharded_down_weight = SplitPrimitiveTensor(
-        name="ffn_down.weight", shard_dim=1, ts=down_weight.split(split_size, dim=1)
+        name="ffn_down.weight", shard_dim=1, ts=down_weight.split(split_size, dim=1), devices=tuple(range(shard_count)), devices_pinned=True
     )
 
     return Theta([sharded_gate_weight, sharded_up_weight, sharded_down_weight])

sharktank/sharktank/layers/kv_cache.py

@@ -98,7 +98,7 @@ def unflatten_page_table(
             shards = [
                 shard.unflatten(1, self.sub_page_dims) for shard in page_slab.shards
             ]
-            return SplitPrimitiveTensor(ts=shards, shard_dim=4)
+            return SplitPrimitiveTensor(ts=shards, shard_dim=4, devices=tuple(range(self.shard_count)), devices_pinned=True)  # TODO
 
     def shard_state(
         self, state: List[torch.Tensor]
@@ -127,7 +127,7 @@ def shard_state(
         shards = [
             ops.flatten(shard, start_dim=1) for shard in sharded_page_table.shards
         ]
-        flat_sharded_page_table = SplitPrimitiveTensor(ts=shards, shard_dim=1)
+        flat_sharded_page_table = SplitPrimitiveTensor(ts=shards, shard_dim=1, devices=tuple(range(self.shard_count)), devices_pinned=True)
         return [flat_sharded_page_table]
 
     @property
@@ -152,7 +152,7 @@ def allocate(
         if self.shard_count == 1:
             return shards
 
-        return [SplitPrimitiveTensor(ts=shards, shard_dim=1)]
+        return [SplitPrimitiveTensor(ts=shards, shard_dim=1, devices=tuple(range(self.shard_count)), devices_pinned=True)]
 
     def read(
         self,
@@ -264,8 +264,8 @@ def write_timestep(
                     for _ in range(seq_positions.shard_count)
                 ]
 
-                partitions = ReplicatedTensor(ts=partitions)
-                transformer_block = ReplicatedTensor(ts=transformer_block)
+                partitions = ReplicatedTensor(ts=partitions, devices=tuple(range(self.shard_count)), devices_pinned=True)  # TODO: Should these be pinned?
+                transformer_block = ReplicatedTensor(ts=transformer_block, devices=tuple(range(self.shard_count)), devices_pinned=True)
             else:
                 partitions = torch.tensor(idx).unsqueeze(0)
                 transformer_block = torch.full(

sharktank/sharktank/layers/rotary_embedding.py

@@ -71,7 +71,7 @@ def forward(
             )
             for xt_shard, rotary_shard in zip(xt.shards, rotary_shards)
         ]
-        xt = SplitPrimitiveTensor(ts=xt_shards, shard_dim=xt.shard_dim)
+        xt = SplitPrimitiveTensor(ts=xt_shards, shard_dim=xt.shard_dim, devices=xt.devices, devices_pinned=False)
         return xt
 
     def _create_interleaved_tensor(_, dim):
@@ -216,7 +216,7 @@ def apply_batched_mask(
             )
             for xt_shard, mask_shard in zip(xt.shards, mask.shards)
         ]
-        xt = SplitPrimitiveTensor(ts=xt_shards, shard_dim=xt.shard_dim)
+        xt = SplitPrimitiveTensor(ts=xt_shards, shard_dim=xt.shard_dim, devices=xt.devices, devices_pinned=False)
         return xt
 
     def apply_batched_mask_unsharded(self, *, xt: torch.Tensor, mask: torch.Tensor):

sharktank/sharktank/ops/__init__.py

@@ -33,3 +33,5 @@
 # Comment this out to completely disable optimized quantized implementations.
 from . import qconv_impls
 from . import qlinear_impls
+
+from .sharded_impls import transfer_if_needed  # TODO: Hack just to get tests running, figure out properly later