Custom generate

Author: jlamypoirier

src/main.py

@@ -26,16 +26,19 @@ def get_arg_parser() -> ArgumentParser:
     parser.add_argument("config_args", nargs="*")
 
     # Runtime
+    parser.add_argument("-c", "--custom_generate", action="store_true")
     parser.add_argument("--pipeline_class", default="HF_Pipeline")
     parser.add_argument("--device", default="cuda", type=torch.device)
     parser.add_argument("--dtype", default="float16", type=lambda x: getattr(torch, x))
     parser.add_argument("--local_rank", type=int)
-    parser.add_argument("--no_fast_init", dest="fast_init", action="store_false")
+    parser.add_argument("--no_fast_init","--nf", dest="fast_init", action="store_false")
+    parser.add_argument("--no_cache","--nc", dest="use_cache", action="store_false")
+    parser.add_argument("--no_prefill","--np", dest="do_prefill", action="store_false")
 
     # Input and output
-    parser.add_argument("--batch_size", default=1, type=int)
-    parser.add_argument("--max_input_length", default=-1, type=int)
-    parser.add_argument("--max_new_tokens", default=100, type=int)
+    parser.add_argument("--batch_size","-b", default=1, type=int)
+    parser.add_argument("--max_input_length","-i", default=-1, type=int)
+    parser.add_argument("--max_new_tokens","-g", default=100, type=int)
 
     # Cleanup
     parser.add_argument("--clear_every_run", action="store_true")
@@ -47,10 +50,11 @@ def get_arg_parser() -> ArgumentParser:
 
     # Profiling and logging
     parser.add_argument("--max_log_outputs", type=int)
-    parser.add_argument("--profile", action="store_true")
-    parser.add_argument("--profile_cycles", type=int)
-    parser.add_argument("--full_trace", action="store_true")
-    parser.add_argument("--show_op_names", action="store_true")
+    parser.add_argument("--breakdown_latency","--bl", action="store_true")
+    parser.add_argument("--profile","-p", action="store_true")
+    parser.add_argument("--profile_cycles","--pc", type=int)
+    parser.add_argument("--full_trace","--pt", action="store_true")
+    parser.add_argument("--show_op_names","--pn", action="store_true")
     parser.add_argument("--save", type=Path)
 
     return parser
@@ -61,7 +65,6 @@ def main(argv: Optional[List[str]] = None) -> None:
     parser = get_arg_parser()
     args = parser.parse_args(argv)
     config_args = parse_config_args(args.config_args)
-    generate_kwargs = {"max_new_tokens": args.max_new_tokens, "do_sample": False}
     inputs = get_dummy_batch(args.batch_size, args.max_input_length)
     separate_profile = args.profile and args.profile_cycles is not None
     warmup = args.profile if args.warmup is None else args.warmup
@@ -88,6 +91,10 @@ def main(argv: Optional[List[str]] = None) -> None:
         dtype=args.dtype,
         fast_init=args.fast_init,
         trust_remote_code=args.trust_remote_code,
+        custom_generate=args.custom_generate,
+        use_cache=args.use_cache,
+        do_prefill=args.do_prefill,
+        breakdown_latency=args.breakdown_latency,
     )
 
     all_metrics = []
@@ -104,7 +111,7 @@ def main(argv: Optional[List[str]] = None) -> None:
         profiler = contextlib.nullcontext()
 
     benchmark_metrics = {
-        **generate_kwargs,
+        "max_new_tokens": args.max_new_tokens,
         "Model parameters": pipeline.get_num_parameters(),
         "Cycles (warmup)": args.skip + warmup,
         "Cycles (benchmark)": args.cycles,
@@ -124,7 +131,7 @@ def main(argv: Optional[List[str]] = None) -> None:
             if step == args.skip + warmup:
                 t2 = time.perf_counter()
                 benchmark_metrics[Metrics.RUNTIME_WARMUP] = t2 - t1
-            generated_text, metrics = pipeline(inputs, **generate_kwargs)
+            generated_text, metrics = pipeline(inputs, args.max_new_tokens)
             if args.profile:
                 p.step()
 

src/metrics.py

@@ -17,14 +17,20 @@ def format_ms(t: float) -> str:
     return f"{1000 * t:.2f} ms"
 
 
+def format_ms_dict(t_dict: Dict[str,float]) -> Dict[str,str]:
+    return {key:format_ms(value) for key, value in t_dict.items()}
+
+
 def format_mib(m: float) -> str:
     return f"{m/2**20:.0f} MiB"
 
 
 class Metrics:
     LATENCY_E2E = "Latency (end to end)"
     LATENCY_TOKEN = "Latency (tokenization)"
-    LATENCY_MODEL = "Latency (model)"
+    LATENCY_MODEL = "Latency (generate)"
+    LATENCY_GENERATE_START = "Latency (prepare for generation)"
+    LATENCY_GENERATE_BREAKDOWN = "Latency (generate breakdown)"
     LATENCY_DECODE = "Latency (decode)"
     LATENCY_MAX = "Latency (max)"
     LATENCY_MIN = "Latency (min)"
@@ -59,6 +65,8 @@ class Metrics:
         LATENCY_E2E: format_ms,
         LATENCY_TOKEN: format_ms,
         LATENCY_MODEL: format_ms,
+        LATENCY_GENERATE_START: format_ms,
+        LATENCY_GENERATE_BREAKDOWN: format_ms_dict,
         LATENCY_DECODE: format_ms,
         LATENCY_MAX: format_ms,
         LATENCY_MIN: format_ms,

src/pipeline.py

@@ -18,6 +18,7 @@
     AutoTokenizer,
     PretrainedConfig,
     PreTrainedModel,
+    GPTBigCodeConfig,GPTBigCodeForCausalLM
 )
 
 
@@ -37,25 +38,41 @@ def __init__(
         dtype: torch.dtype,
         fast_init: bool = True,
         trust_remote_code: bool = False,
+        custom_generate:bool=False,
+        use_cache: bool = True,
+        do_prefill: bool = True,
+        breakdown_latency=False,
     ):
         self.global_metrics = {}
         log_rank_n("*** Setting up tokenizer", logger.info)
-        t0 = time.perf_counter()
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+        t0 = self._get_time()
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, padding_side="left")
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token=self.tokenizer.eos_token
 
-        self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-        t1 = time.perf_counter()
+        t1 = self._get_time()
 
         self.device = device
+        if self.device==torch.device("cuda"):
+            self.device=torch.device("cuda:0")
+
         self.dtype = dtype
         self.is_int8 = self.dtype == torch.int8
         self.fast_init = fast_init
         self.trust_remote_code = trust_remote_code
-        if self.is_int8 and self.device != torch.device("cuda"):
+        self.use_cache = use_cache
+        self.do_prefill = do_prefill
+        if not self.do_prefill:
+            assert custom_generate
+            assert self.use_cache
+        self.breakdown_latency=breakdown_latency
+        if self.is_int8 and self.device != torch.device("cuda:0"):
             raise ValueError(f"Model quantization not supported on device {self.device}")
 
+        self._generate=self._generate_custom if custom_generate else self._generate_hf
+
         self.config = self._get_config(model_type, pretrained_config or pretrained_model, config_args)
-        t2 = time.perf_counter()
+        t2 = self._get_time()
 
         logger.info(f"Model configuration: {self.config}")
 
@@ -67,27 +84,27 @@ def __init__(
             self.model = self._load_pretrained(pretrained_model)
 
         self.model.eval()
-        t3 = time.perf_counter()
+        t3 = self._get_time()
         self.global_metrics[Metrics.INIT_TOKEN] = t1 - t0
         self.global_metrics[Metrics.INIT_CONFIG] = t2 - t1
         self.global_metrics[Metrics.INIT_TOTAL] = t3 - t0
 
     def _create_model(self) -> PreTrainedModel:
-        t0 = time.perf_counter()
+        t0 = self._get_time()
         log_rank_n("*** Creating model", logger.info)
         with fast_init(self.device) if self.fast_init else contextlib.nullcontext():
             torch_dtype = torch.float16 if self.is_int8 else self.dtype
             model = AutoModelForCausalLM.from_config(
                 config=self.config, torch_dtype=torch_dtype, trust_remote_code=self.trust_remote_code
             )
-        t1 = time.perf_counter()
+        t1 = self._get_time()
         log_rank_n("*** Moving to device", logger.info)
         model.to(self.device)
-        t2 = time.perf_counter()
+        t2 = self._get_time()
         log_rank_n("*** Initializing weights", logger.info)
         # Initialization is ~1000x faster on GPU.
         model.init_weights()
-        t3 = time.perf_counter()
+        t3 = self._get_time()
         self.global_metrics[Metrics.INIT_CREATE] = t1 - t0
         self.global_metrics[Metrics.INIT_DEVICE] = t2 - t1
         self.global_metrics[Metrics.INIT_WEIGHTS] = t3 - t2
@@ -101,14 +118,14 @@ def _reload_model(self):
         self.model = self._load_pretrained("tmp")
 
     def _save_pretrained(self, pretrained_model: str):
-        t0 = time.perf_counter()
+        t0 = self._get_time()
         log_rank_n(f"*** Saving model to {pretrained_model}", logger.info)
-        t1 = time.perf_counter()
+        t1 = self._get_time()
         self.global_metrics[Metrics.INIT_SAVE] = t1 - t0
         self.model.save_pretrained(pretrained_model)
 
     def _load_pretrained(self, pretrained_model: str) -> PreTrainedModel:
-        t0 = time.perf_counter()
+        t0 = self._get_time()
         log_rank_n(f"*** Loading model from {pretrained_model}", logger.info)
         kwargs = {"load_in_8bit": True, "device_map": "auto"} if self.is_int8 else {"torch_dtype": self.dtype}
         with fast_init(self.device) if self.fast_init else contextlib.nullcontext():
@@ -120,12 +137,12 @@ def _load_pretrained(self, pretrained_model: str) -> PreTrainedModel:
                 trust_remote_code=self.trust_remote_code,
                 **kwargs,
             )
-        t1 = time.perf_counter()
+        t1 = self._get_time()
         self.global_metrics["load pretrained model"] = t1 - t0
         if not self.is_int8:
             log_rank_n("*** Moving to device", logger.info)
             model = model.to(self.device)
-            t2 = time.perf_counter()
+            t2 = self._get_time()
             self.global_metrics[Metrics.INIT_DEVICE] = t2 - t1
         return model
 
@@ -171,26 +188,103 @@ def _get_config(
 
         return config
 
-    def __call__(self, text: List[str], **generate_kwargs) -> Tuple[List[str], Dict[str, Any]]:
-        t0 = time.perf_counter()
-        inputs = self.tokenizer(text, return_tensors="pt", padding=True)
+    def _get_time(self, synchronize=False):
+        if synchronize:
+            torch.cuda.synchronize()
+        return time.perf_counter()
+
+    def _generate_custom(self, inputs:Dict, max_new_tokens:int):
+        t0 = self._get_time(self.breakdown_latency)
+        batch_size, input_length = inputs["input_ids"].shape
+        output_length = input_length + max_new_tokens
+        input_ids = torch.empty([batch_size, output_length], dtype=torch.int64, device=self.device)
+        input_ids[:, :input_length].copy_(inputs["input_ids"])
 
+        attention_mask = torch.empty([batch_size, output_length], dtype=torch.bool, device=self.device)
+        attention_mask[:, :input_length].copy_(inputs["attention_mask"])
+        attention_mask[:, input_length:].fill_(True)
+
+        position_ids = attention_mask.long().cumsum(-1, dtype=torch.int64) - 1
+        # TODO: Useless?
+        position_ids[:, :input_length].masked_fill_(attention_mask[:, :input_length] == 0, 1)
+
+        if self.do_prefill or input_length<=1:
+            past_key_values=None
+            past_key_length=0
+        else:
+            # Generate mock `past_key_values`
+            past_key_length=input_length-1
+            if isinstance(self.config, GPTBigCodeConfig):
+                if self.config.pre_allocate_kv_cache:
+                    past_key_values=[past_key_length]*self.config.n_layer
+                    for block in self.model.transformer.h:
+                        block.attn.get_kv_cache(batch_size, past_key_length, dtype=self.dtype, device=self.device).normal_()
+                else:
+                    kv_dim=self.config.n_embd // self.config.n_head if self.config.multi_query else self.config.n_embd
+                    past_key_values=[torch.randn([batch_size, past_key_length, 2*kv_dim], dtype=self.dtype, device=self.device) for _ in  range(self.config.n_layer)]
+            else:
+                past_key_values = [
+                    [torch.randn([batch_size, past_key_length, self.config.n_embd], dtype=self.dtype, device=self.device) for _ in range(2)] for _ in
+                    range(self.config.n_layer)]
+
+        t1 = self._get_time(self.breakdown_latency)
+        last_time=t1
+        generate_times={}
+        for key_length in range(input_length, output_length):
+            outputs = self.model(
+                input_ids=input_ids[:, past_key_length:key_length],
+                past_key_values=past_key_values,
+                attention_mask=attention_mask[:, :key_length],
+                position_ids=position_ids[:, past_key_length:key_length],
+                return_dict=True,
+                use_cache=self.use_cache,
+            )
+            if self.use_cache:
+                past_key_values=outputs.past_key_values
+                past_key_length=key_length
+            next_tokens = torch.argmax(outputs.logits[:, -1, :], dim=-1)
+            input_ids[:, key_length] = next_tokens
+            t2 = self._get_time(self.breakdown_latency)
+            generate_times[key_length]=t2-last_time
+            last_time=t2
+
+        metrics={}
+        if self.breakdown_latency:
+            metrics[Metrics.LATENCY_GENERATE_START]=t1-t0
+            metrics[Metrics.LATENCY_GENERATE_BREAKDOWN]=generate_times
+
+        return input_ids, metrics
+
+    def _generate_hf(self, inputs:Dict, max_new_tokens:int):
         inputs = {key: value.to(self.device) if torch.is_tensor(value) else value for key, value in inputs.items()}
+        output = self.model.generate(
+            **inputs,
+            return_dict_in_generate=True,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            pad_token_id=self.tokenizer.pad_token_id,
+            use_cache=self.use_cache,
+        )
+        return output.sequences, {}
 
-        t1 = time.perf_counter()
-        with torch.inference_mode():
-            output = self.model.generate(**inputs, return_dict_in_generate=True, **generate_kwargs)
-        t2 = time.perf_counter()
 
-        output_tokens = output.sequences
+    def __call__(self, text: List[str], max_new_tokens:int) -> Tuple[List[str], Dict[str, Any]]:
+        t0 = self._get_time()
+        inputs = self.tokenizer(text, return_tensors="pt", padding=True)
+
+        t1 = self._get_time()
+        with torch.inference_mode():
+            output_tokens, generate_metrics = self._generate(inputs, max_new_tokens)
+        t2 = self._get_time(True)
 
         batch_size, input_length = inputs["input_ids"].shape
         output_length = output_tokens.size(1)
 
         output_text = self.tokenizer.batch_decode(output_tokens.cpu(), skip_special_tokens=True)
-        t3 = time.perf_counter()
+        t3 = self._get_time()
 
         metrics = {
+            **generate_metrics,
             Metrics.BATCH_SIZE: batch_size,
             Metrics.INPUT_LENGTH: input_length,
             Metrics.OUTPUT_LENGTH: output_length,
@@ -218,14 +312,23 @@ def aggregate_metrics(self, metrics: List[Dict[str, Any]]):
                 Metrics.TOKENS_BATCH,
                 Metrics.LATENCY_TOKEN,
                 Metrics.LATENCY_MODEL,
+                Metrics.LATENCY_GENERATE_START,
+                Metrics.LATENCY_GENERATE_BREAKDOWN,
                 Metrics.LATENCY_DECODE,
                 Metrics.LATENCY_E2E,
             )
         }
+
+        breakdown=all_metrics.pop(Metrics.LATENCY_GENERATE_BREAKDOWN, [])
+
         mean_metrics = {key: np.mean(value).item() for key, value in all_metrics.items() if len(value) > 0}
         throughput = mean_metrics[Metrics.TOKENS_BATCH] / mean_metrics[Metrics.LATENCY_E2E]
         model_throughput = mean_metrics[Metrics.TOKENS_BATCH] / mean_metrics[Metrics.LATENCY_MODEL]
 
+        if len(breakdown) > 0:
+            mean_metrics[Metrics.LATENCY_GENERATE_BREAKDOWN] = {
+                str(key): np.mean([values[key] for values in breakdown]).item() for key in breakdown[0]}
+
         return {
             **self.global_metrics,
             **mean_metrics,

src/utils.py

@@ -82,9 +82,13 @@ def log_rank_n(msg: str, logger: Callable = logging.info, rank: int = 0):
             logger(line)
 
 
-def log_dict(data: dict, logger: Callable = logging.info, rank: int = 0):
+def log_dict(data: dict, logger: Callable = logging.info, rank: int = 0, _prefix=""):
     for key, value in data.items():
-        log_rank_n(f"{key}: {value}", logger, rank)
+        if isinstance(value, dict):
+            log_rank_n(f"{_prefix}{key}:", logger, rank)
+            log_dict(value, logger, rank, _prefix+"  ")
+        else:
+            log_rank_n(f"{_prefix}{key}: {value}", logger, rank)
 
 
 dummy_input_sentences = [