Inference (meta-pytorch#619)

Author: kartikayk

docs/source/api_ref_utilities.rst

@@ -67,16 +67,6 @@ Data
 
 .. _gen_label:
 
-Generation
-----------
-
-.. autosummary::
-    :toctree: generated/
-    :nosignatures:
-
-    generation.GenerationUtils
-    generation.generate_from_prompt
-
 
 Miscellaneous
 -------------

recipes/alpaca_generate.py

@@ -0,0 +1,31 @@
+
+# Model arguments
+model:
+  _component_: torchtune.models.llama2.llama2_13b
+
+checkpointer:
+  _component_: torchtune.utils.FullModelHFCheckpointer
+  checkpoint_dir: /tmp/Llama-2-13b-hf/
+  checkpoint_files: [
+    pytorch_model-00001-of-00003.bin,
+    pytorch_model-00002-of-00003.bin,
+    pytorch_model-00003-of-00003.bin
+  ]
+  output_dir: /tmp/Llama-2-13b-hf/
+  model_type: LLAMA2
+
+device: cuda
+dtype: bf16
+
+seed: 1234
+
+# Tokenizer arguments
+tokenizer:
+  _component_: torchtune.models.llama2.llama2_tokenizer
+  path: /tmp/Llama-2-13b-hf/tokenizer.model
+
+# Generation arguments; defaults taken from gpt-fast
+prompt: "Hello, my name is"
+max_new_tokens: 300
+temperature: 0.8
+top_k: 300

recipes/configs/alpaca_generate.yaml

@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+import sys
+import time
+from typing import Any, Dict
+
+import torch
+from omegaconf import DictConfig
+
+from torch import nn
+
+from torchtune import config, utils
+
+logger = utils.get_logger("DEBUG")
+
+
+class InferenceRecipe:
+    """
+    Recipe for generating tokens from a dense Transformer-based LLM.
+
+    Currently this recipe support single-GPU generation only. Speculative
+    decoding is not supported.
+    """
+
+    def __init__(self, cfg: DictConfig) -> None:
+        self._device = utils.get_device(device=cfg.device)
+        self._dtype = utils.get_dtype(dtype=cfg.dtype)
+
+        utils.set_seed(seed=cfg.seed)
+
+    def load_checkpoint(self, checkpointer_cfg: DictConfig) -> Dict[str, Any]:
+        checkpointer = config.instantiate(checkpointer_cfg)
+        checkpoint_dict = checkpointer.load_checkpoint()
+        return checkpoint_dict
+
+    def setup(self, cfg: DictConfig) -> None:
+        ckpt_dict = self.load_checkpoint(cfg.checkpointer)
+        self._model = self._setup_model(
+            model_cfg=cfg.model,
+            model_state_dict=ckpt_dict[utils.MODEL_KEY],
+        )
+        self._tokenizer = config.instantiate(cfg.tokenizer)
+
+    def _setup_model(
+        self,
+        model_cfg: DictConfig,
+        model_state_dict: Dict[str, Any],
+    ) -> nn.Module:
+        with utils.set_default_dtype(self._dtype), self._device:
+            model = config.instantiate(model_cfg)
+
+        model.load_state_dict(model_state_dict)
+
+        # Validate model was loaded in with the expected dtype.
+        utils.validate_expected_param_dtype(model.named_parameters(), dtype=self._dtype)
+        logger.info(f"Model is initialized with precision {self._dtype}.")
+
+        # Ensure the cache is setup on the right device
+        with self._device:
+            model.setup_caches(max_batch_size=1, dtype=self._dtype)
+
+        return model
+
+    @torch.no_grad()
+    def generate(self, cfg: DictConfig):
+        tokens = self._tokenizer.encode(cfg.prompt, add_bos=True, add_eos=False)
+        prompt = torch.tensor(tokens, dtype=torch.int, device=self._device)
+
+        t0 = time.perf_counter()
+        generated_tokens = utils.generate(
+            model=self._model,
+            prompt=prompt,
+            max_generated_tokens=cfg.max_new_tokens,
+            temperature=cfg.temperature,
+            top_k=cfg.top_k,
+            eos_id=self._tokenizer.eos_id,
+        )
+        t = time.perf_counter() - t0
+
+        logger.info(self._tokenizer.decode(generated_tokens))
+
+        tokens_generated = len(generated_tokens) - prompt.size(0)
+        tokens_sec = tokens_generated / t
+        logger.info(
+            f"Time for inference: {t:.02f} sec total, {tokens_sec:.02f} tokens/sec"
+        )
+        logger.info(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")
+
+
+@config.parse
+def main(cfg: DictConfig) -> None:
+    recipe = InferenceRecipe(cfg=cfg)
+    recipe.setup(cfg=cfg)
+    recipe.generate(cfg=cfg)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

recipes/configs/generate.yaml

@@ -48,7 +48,6 @@ def llama2_test_config(max_batch_size: Optional[int] = None) -> List[str]:
         "model.max_seq_len=2048",
         "model.norm_eps=1e-5",
         "model.num_kv_heads=8",
-        f"model.max_batch_size={max_batch_size if max_batch_size else 'null'}",
     ]
 
 
@@ -75,7 +74,6 @@ def lora_llama2_test_config(
         "model.max_seq_len=2048",
         "model.norm_eps=1e-5",
         "model.num_kv_heads=8",
-        f"model.max_batch_size={max_batch_size if max_batch_size else 'null'}",
         f"model.lora_rank={lora_rank}",
         f"model.lora_alpha={lora_alpha}",
         "model.lora_dropout=0.0",