| 35 | |
| 36 | |
| 37 | class FastGen: |
| 38 | GRAPH_WARMUPS: int = 1 |
| 39 | tokenizer: Tokenizer |
| 40 | |
| 41 | @staticmethod |
| 42 | def build( |
| 43 | ckpt_dir: str, |
| 44 | gen_args: GenArgs, |
| 45 | device: Union[torch.device, str], |
| 46 | tokenizer_path: Optional[str] = None, |
| 47 | num_layers: int = 13, |
| 48 | use_full_vocab: bool = False, |
| 49 | ) -> "FastGen": |
| 50 | """ |
| 51 | Load a Llama or Code Llama checkpoint and return a new |
| 52 | generator for this model. |
| 53 | """ |
| 54 | start_time = time.time() |
| 55 | |
| 56 | model_args_prefill = fast.ModelArgs(use_kernel=False) |
| 57 | model_args_decode = fast.ModelArgs(use_kernel=True) |
| 58 | tokenizer = Tokenizer("./tokenizer.model") |
| 59 | |
| 60 | torch.set_default_device(device) |
| 61 | torch.set_default_dtype(torch.bfloat16) |
| 62 | |
| 63 | prefill_model = fast.Transformer(model_args_prefill) |
| 64 | decode_model = fast.Transformer(model_args_decode) |
| 65 | |
| 66 | fp16_ckpt_path = str(Path(ckpt_dir) / "model_state_fp16.pt") |
| 67 | fp16_checkpoint = torch.load(fp16_ckpt_path, map_location="cpu", weights_only=True) |
| 68 | int2_ckpt_path = str(Path(ckpt_dir) / "model_state_int2.pt") |
| 69 | int2_checkpoint = torch.load(int2_ckpt_path, map_location="cpu", weights_only=True) |
| 70 | prefill_model.load_state_dict(fp16_checkpoint, strict=True) |
| 71 | decode_model.load_state_dict(int2_checkpoint, strict=True) |
| 72 | |
| 73 | torch.cuda.synchronize() |
| 74 | print(f"loaded model in {time.time() - start_time:.2f} seconds") |
| 75 | start_time = time.time() |
| 76 | |
| 77 | return FastGen(gen_args, model_args_prefill, prefill_model, decode_model, tokenizer) |
| 78 | |
| 79 | def __init__( |
| 80 | self, |
| 81 | args: GenArgs, |
| 82 | model_args: fast.ModelArgs, |
| 83 | prefill_model: fast.Transformer, |
| 84 | decode_model: fast.Transformer, |
| 85 | tokenizer: Tokenizer, |
| 86 | ): |
| 87 | self.gen_args = args |
| 88 | self.max_seq_length = args.prompt_length + args.gen_length |
| 89 | self.model_args = model_args |
| 90 | # self.model = model |
| 91 | self.prefill_model = prefill_model |
| 92 | self.decode_model = decode_model |
| 93 | self.tokenizer = tokenizer |
| 94 | self._prefill_cuda_graph, self._prefill_compile_model, self._prefill_inputs, self._prefill_logits = None, None, None, None |