Load a Llama or Code Llama checkpoint and return a new generator for this model.
(
ckpt_dir: str,
gen_args: GenArgs,
device: Union[torch.device, str],
tokenizer_path: Optional[str] = None,
num_layers: int = 13,
use_full_vocab: bool = False,
)
| 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, |