[sharktank] Refactor sharktank llm modeling (#1207)

- Refactor various sharktank llms into a single `llm.py` script
- Remove all references to direct cache scripts and tests
- Make `--attention-kernel="torch"` default for all llms and tests,
except grok as softcap+sdpa isn't supported yet.

Author: archana-ramalingam

docs/model_cookbook.md

@@ -269,14 +269,3 @@ python -m sharktank.examples.paged_llm_v1 \
   --dump-decode-steps=1 \
   --dump-path='/tmp'
 ```
-
-## Generating data for llama models
-
-```bash
-set TURBINE_DEBUG=log_level=info
-python -m sharktank.models.llama.tools.generate_data \
-  --tokenizer=openlm-research/open_llama_3b_v2 \
-  --config=/tmp/open_llama_3b_v2/open-llama-3b-v2-f16.json \
-  --output-dir=/tmp/open_llama_3b_v2/inputs \
-  --prompt="What is the meaning of life?"
-```

sharktank/sharktank/evaluate/perplexity_iree.py

@@ -21,11 +21,8 @@
 from torch.nn import CrossEntropyLoss
 import iree.runtime
 
-from sharktank.models.llama.llama import *
-from sharktank.models.mixtral.mixtral import *
-from sharktank.models.grok.grok import *
-
-from ..models.llama.sharding import shard_theta
+from sharktank.models.llm import *
+from sharktank.models.llama.sharding import shard_theta
 
 from sharktank.layers import *
 from sharktank.types import *
@@ -187,13 +184,7 @@ def load_model(self, weight_path, tokenizer):
 
         theta = weight_path.root_theta
 
-        if self.config.hp.expert_count:
-            if self.config.hp.model_arch == "grok":
-                model = PagedGrokModelV1(theta, self.config)
-            else:
-                model = PagedMixtralModelV1(theta, self.config)
-        else:
-            model = PagedLlamaModelV1(theta, self.config)
+        model = PagedLlmModelV1(theta, self.config)
 
         self.generator = TorchGenerator(model, tokenizer)
 

sharktank/sharktank/evaluate/perplexity_prefill.py

@@ -18,11 +18,8 @@
 from sharktank.layers import *
 from sharktank.types import *
 
-from sharktank.models.llama.llama import *
-from sharktank.models.mixtral.mixtral import *
-from sharktank.models.grok.grok import *
-
-from ..models.llama.sharding import shard_theta
+from sharktank.models.llm import *
+from sharktank.models.llama.sharding import shard_theta
 
 from sharktank.utils import cli
 from sharktank.utils.load_llm import *
@@ -106,13 +103,7 @@ def load_model(self, dataset, tokenizer, tensor_parallelism_size, attention_kern
 
         theta = dataset.root_theta
 
-        if self.config.hp.expert_count:
-            if self.config.hp.model_arch == "grok":
-                model = PagedGrokModelV1(theta, self.config)
-            else:
-                model = PagedMixtralModelV1(theta, self.config)
-        else:
-            model = PagedLlamaModelV1(theta, self.config)
+        model = PagedLlmModelV1(theta, self.config)
 
         self.generator = TorchGenerator(model, tokenizer)