Remove shard_llama, add checks to shard_llm_dataset

Signed-off-by: aviator19941 <[email protected]>
nod-ai · Oct 31, 2024 · 87c4bfd · 87c4bfd
1 parent e30f0b6
commit 87c4bfd
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 56 deletions.
diff --git a/sharktank/sharktank/examples/sharding/shard_llm_dataset.py b/sharktank/sharktank/examples/sharding/shard_llm_dataset.py
@@ -30,6 +30,12 @@ def main(raw_args=None):
     args = cli.parse(parser, args=raw_args)
     dataset = cli.get_input_dataset(args)
 
+    if args.output_file is None:
+        raise RuntimeError(f"Need file destination for IRPA file")
+
+    if args.tensor_parallelism_size < 2:
+        raise RuntimeError(f"Expect sharding greater than 1 found {args.shard_count}")
+
     hp = LlamaHParams.from_gguf_props(dataset.properties)
     llama_config = LlamaModelConfig(
         hp, tensor_parallelism_size=args.tensor_parallelism_size

diff --git a/sharktank/sharktank/models/llama/tools/shard_llama.py b/sharktank/sharktank/models/llama/tools/shard_llama.py
diff --git a/sharktank/tests/models/llama/benchmark_amdgpu_test.py b/sharktank/tests/models/llama/benchmark_amdgpu_test.py
@@ -56,7 +56,7 @@ class BenchmarkLlama3_1_8B(BaseBenchmarkTest):
     def setUp(self):
         super().setUp()
         # TODO: add numpy files to Azure and download from it
-        self.artifacts_dir = Path("/data/llama-3.1/8b")
+        self.artifacts_dir = Path("/data/extra/models/llama3.1_8B")
         self.gguf_path = self.artifacts_dir / "llama8b_f16.gguf"
         self.irpa_path = self.artifacts_dir / "llama8b_f16.irpa"
         self.irpa_path_fp8 = self.artifacts_dir / "llama8b_fp8.irpa"