Split sharded Llama dataset exporting and loading in export scripts (#327)

Separate the 2 steps. We need exported irpa files for the IREE module
anyway.

Author: sogartar

sharktank/sharktank/examples/export_paged_llm_v1.py

@@ -59,31 +59,26 @@ def main():
         default="decomposed",
         choices=["decomposed", "torch_sdpa"],
     )
-    parser.add_argument(
-        "--tensor-parallelism-size",
-        type=int,
-        default=1,
-        help="How many devices are involved for tensor parallel sharding.",
-    )
 
     args = cli.parse(parser)
     dataset_type = cli.get_input_data_files(args)
     dataset_type = "irpa" if "irpa" in dataset_type else "gguf"
     dataset = cli.get_input_dataset(args)
 
     hp = configs.LlamaHParams.from_gguf_props(dataset.properties)
-    llama_config = LlamaModelConfig(hp)
-    if args.tensor_parallelism_size > 1:
-        dataset.root_theta = shard_theta(dataset.root_theta, llama_config)
-    llama_config.use_hf = False
-    llama_config.static_tables = False  # Rely on the compiler for hoisting tables.
-    llama_config.kv_cache_type = "direct" if args.bs == [1] else "paged"
-    llama_config.attention_kernel = args.attention_kernel
-
-    # This is a bit gross and should be changed in the future. Best Idea I had so far.
-    attn_q_weight = dataset.root_theta.tensor("blk")["0"]["attn_q"]["weight"]
-    if isinstance(attn_q_weight, SplitPrimitiveTensor):
-        llama_config.tensor_parallelism_size = attn_q_weight.shard_count
+    tensor_parallelism_size = (
+        dataset.properties["tensor_parallelism_size"]
+        if "tensor_parallelism_size" in dataset.properties
+        else 1
+    )
+    llama_config = LlamaModelConfig(
+        hp,
+        tensor_parallelism_size=tensor_parallelism_size,
+        use_hf=False,
+        static_tables=False,  # Rely on the compiler for hoisting tables.
+        kv_cache_type="direct" if args.bs == [1] else "paged",
+        attention_kernel=args.attention_kernel,
+    )
 
     if llama_config.hp.expert_count:
         if llama_config.hp.model_arch == "grok":

sharktank/sharktank/examples/sharding/shard_llm_dataset.py

@@ -10,7 +10,8 @@
 weights of an LLM by converting the RHS of all eligible layers to a sharded
 form.
 """
-from ...transforms.dataset import MmtRHSShardingTransform
+from ...models.llama.sharding import shard_theta
+from ...layers import LlamaHParams, LlamaModelConfig
 from ...types import *
 
 
@@ -21,16 +22,22 @@ def main(raw_args=None):
     cli.add_input_dataset_options(parser)
     cli.add_output_dataset_options(parser)
     parser.add_argument(
-        "--num-shards", type=int, required=True, help="Number of shards to split"
+        "--tensor-parallelism-size",
+        type=int,
+        required=True,
+        help="Number of shards to split",
     )
     args = cli.parse(parser, args=raw_args)
     dataset = cli.get_input_dataset(args)
 
-    tr = MmtRHSShardingTransform(
-        r"^blk\.[0-9]+\.(attn_k|attn_q|attn_v|ffn_gate|ffn_up|ffn_down)\.weight$",
-        num_shards=8,
+    hp = LlamaHParams.from_gguf_props(dataset.properties)
+    llama_config = LlamaModelConfig(
+        hp, tensor_parallelism_size=args.tensor_parallelism_size
     )
-    dataset.transform(tr)
+    sharded_theta = shard_theta(dataset.root_theta, llama_config)
+    sharded_theta.rename_tensors_to_paths()
+    dataset.root_theta = sharded_theta
+    dataset.properties["tensor_parallelism_size"] = args.tensor_parallelism_size
     dataset.save(args.output_irpa_file, io_report_callback=print)
 
 

sharktank/tests/transforms/dataset_transforms_test.py

@@ -19,18 +19,28 @@
 from sharktank.utils.testing import MainRunnerTestBase
 
 
-class MmtRHSShardingTransformTest(MainRunnerTestBase):
-    def testPrimitive(self):
+class DatasetShardingTransformTest(MainRunnerTestBase):
+    def testShardLlmDataset(self):
         orig_pts = [
             DefaultPrimitiveTensor(
                 name="blk.1.attn_k.weight", data=torch.randn([32, 128])
             ),
             DefaultPrimitiveTensor(
                 name="blk.2.attn_q.weight", data=torch.randn([48, 64])
             ),
-            DefaultPrimitiveTensor(name="other", data=torch.randn([2, 2])),
         ]
-        ds_orig = Dataset({}, Theta(orig_pts))
+        ds_orig = Dataset(
+            {
+                "general.architecture": "llm",
+                "llm.attention.head_count": 1,
+                "llm.context_length": 2,
+                "llm.embedding_length": 3,
+                "llm.block_count": 4,
+                "llm.feed_forward_length": 5,
+                "llm.attention.layer_norm_rms_epsilon": 0.1,
+            },
+            Theta(orig_pts),
+        )
         input_path = self.save_dataset(ds_orig, "input")
         output_path = self.get_irpa_path("output")
         from sharktank.examples.sharding import shard_llm_dataset
@@ -41,38 +51,38 @@ def testPrimitive(self):
             input_path,
             "--output-irpa-file",
             output_path,
-            "--num-shards",
+            "--tensor-parallelism-size",
             8,
         )
         ds_tran = Dataset.load(output_path, mmap=False)
 
+        ds_tran.properties["tensor_parallelism_size"] = 8
+
         # Verify.
         flat_sts = ds_tran.root_theta.flatten()
-        self.assertEqual(3, len(flat_sts))
+        self.assertEqual(2, len(flat_sts))
         st_1 = flat_sts["blk.1.attn_k.weight"]
         st_2 = flat_sts["blk.2.attn_q.weight"]
-        pt_3 = flat_sts["other"]
         self.assertIsInstance(st_1, SplitPrimitiveTensor)
         self.assertIsInstance(st_2, SplitPrimitiveTensor)
-        self.assertIsInstance(pt_3, DefaultPrimitiveTensor)
         self.assertListEqual(st_1.shape, [32, 128])
         self.assertListEqual(st_2.shape, [48, 64])
 
         # Verify component shapes for st_1.
         self.assertEqual(8, len(st_1.shards))
-        self.assertTrue(all(pt.shape == [32, 16] for pt in st_1.shards))
+        self.assertTrue(all(pt.shape == [4, 128] for pt in st_1.shards))
         self.assertTrue(
-            all(list(pt.as_torch().shape) == [32, 16] for pt in st_1.shards)
+            all(list(pt.as_torch().shape) == [4, 128] for pt in st_1.shards)
         )
 
         # Verify component shapes for st_2.
         self.assertEqual(8, len(st_2.shards))
-        self.assertTrue(all(pt.shape == [48, 8] for pt in st_2.shards))
-        self.assertTrue(all(list(pt.as_torch().shape) == [48, 8] for pt in st_2.shards))
+        self.assertTrue(all(pt.shape == [6, 64] for pt in st_2.shards))
+        self.assertTrue(all(list(pt.as_torch().shape) == [6, 64] for pt in st_2.shards))
 
         # Verify contents for one shard for sanity.
         new_t = st_1.shards[0].as_torch()
-        torch.testing.assert_close(new_t, orig_pts[0].as_torch().split(16, dim=1)[0])
+        torch.testing.assert_close(new_t, orig_pts[0].as_torch().split(4, dim=0)[0])
 
 
 if __name__ == "__main__":