|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +# pyre-unsafe |
| 8 | + |
| 9 | +import json |
| 10 | +import os |
| 11 | +from typing import Any, Dict |
| 12 | + |
| 13 | +import torch |
| 14 | + |
| 15 | +from executorch.examples.models.checkpoint import get_checkpoint_dtype |
| 16 | +from executorch.examples.models.llama.model_args import ModelArgs |
| 17 | +from executorch.examples.models.llama.rope import Rope, RotaryEmbedding |
| 18 | +from executorch.examples.models.model_base import EagerModelBase |
| 19 | +from executorch.extension.llm.modules.attention import ( |
| 20 | + replace_mha_with_inference_mha, |
| 21 | + replace_rope_with_inference_rope, |
| 22 | +) |
| 23 | + |
| 24 | +from torchtune.models import convert_weights |
| 25 | + |
| 26 | +from torchtune.models.llama3_1._position_embeddings import Llama3ScaledRoPE |
| 27 | + |
| 28 | +from torchtune.models.llama3_2._component_builders import lora_llama3_2 |
| 29 | + |
| 30 | + |
| 31 | +class Llama3_2_Lora(EagerModelBase): |
| 32 | + def __init__(self, **kwargs): |
| 33 | + # Set member vars from kwargs. |
| 34 | + self.max_seq_len = kwargs.get( |
| 35 | + "max_seq_len", 8192 |
| 36 | + ) # Trained to be a lot larger, but this value is kept small because of static kv cache at the moment. |
| 37 | + # self.encoder_max_seq_len = kwargs.get( |
| 38 | + # "encoder_max_seq_len", int(4 * (448 / 14) ** 2 + 1) |
| 39 | + # ) # Same as above. |
| 40 | + self.generate_full_logits = kwargs.get("generate_full_logits", False) |
| 41 | + self.enable_dynamic_shape = kwargs.get("enable_dynamic_shape", True) |
| 42 | + self.output_prune_map_path = kwargs.get("output_prune_map_path", None) |
| 43 | + self.use_kv_cache = kwargs.get("use_kv_cache", False) |
| 44 | + self.verbose = kwargs.get("verbose", False) |
| 45 | + self.args = kwargs.get("args", None) |
| 46 | + self.dtype = kwargs.get("dtype", torch.float16) |
| 47 | + self.use_checkpoint = False |
| 48 | + self.max_context_len = kwargs.get("max_context_len", 8192) |
| 49 | + |
| 50 | + # Single checkpoint file. |
| 51 | + checkpoint_path = kwargs.get("checkpoint") |
| 52 | + |
| 53 | + if os.path.isfile(checkpoint_path): |
| 54 | + self.use_checkpoint = True |
| 55 | + |
| 56 | + params_path = kwargs.get("params") |
| 57 | + adapter_path = kwargs.get("adapter") |
| 58 | + |
| 59 | + # self.input_pos = torch.arange(self.max_seq_len, dtype=torch.int64) |
| 60 | + # Load checkpoint and params. |
| 61 | + device = "cpu" |
| 62 | + if self.use_checkpoint: |
| 63 | + checkpoint = torch.load( |
| 64 | + checkpoint_path, map_location=device, weights_only=False, mmap=True |
| 65 | + ) |
| 66 | + checkpoint = convert_weights.meta_to_tune(checkpoint) |
| 67 | + self.dtype = get_checkpoint_dtype(checkpoint) |
| 68 | + |
| 69 | + adapter = torch.load( |
| 70 | + adapter_path, map_location="cpu", mmap=True, weights_only=False |
| 71 | + ) |
| 72 | + |
| 73 | + checkpoint.update(adapter) |
| 74 | + |
| 75 | + with open(params_path, "r") as f: |
| 76 | + params = json.loads(f.read()) |
| 77 | + |
| 78 | + # Load model. |
| 79 | + # Cannot use "with torch.device("meta"):" because it causes some exceptions during export, |
| 80 | + # i.e. the model isn't fully initialized or something. |
| 81 | + self.model_ = lora_llama3_2( |
| 82 | + lora_attn_modules=[ |
| 83 | + "q_proj", |
| 84 | + ], |
| 85 | + apply_lora_to_mlp=False, |
| 86 | + apply_lora_to_output=False, |
| 87 | + # llama3_2 args |
| 88 | + vocab_size=params["vocab_size"], |
| 89 | + num_layers=params["n_layers"], |
| 90 | + num_heads=params["n_heads"], |
| 91 | + num_kv_heads=params["n_kv_heads"], |
| 92 | + embed_dim=params["dim"], |
| 93 | + max_seq_len=self.max_seq_len, # 131072 |
| 94 | + # intermediate_dim=params["intermediate_dim"], # 8192, calc is 4096 |
| 95 | + # LoRA args. TODO take in the adapter config. |
| 96 | + lora_rank=8, |
| 97 | + lora_alpha=16, |
| 98 | + ) |
| 99 | + self.model_.requires_grad_(False) |
| 100 | + for param_name, param_val in params.items(): |
| 101 | + setattr(self.model_, param_name, param_val) |
| 102 | + |
| 103 | + setattr(self.model_, "enable_dynamic_shape", self.enable_dynamic_shape) |
| 104 | + # Source transformation for MultiHeadAttention |
| 105 | + self.model_ = replace_mha_with_inference_mha(self.model_) |
| 106 | + |
| 107 | + model_args: ModelArgs = ModelArgs( |
| 108 | + max_seq_len=self.max_seq_len, |
| 109 | + max_context_len=self.max_context_len, |
| 110 | + use_kv_cache=self.use_kv_cache, |
| 111 | + generate_full_logits=self.generate_full_logits, |
| 112 | + enable_dynamic_shape=self.enable_dynamic_shape, |
| 113 | + **params, |
| 114 | + ) |
| 115 | + # Source transformation for RoPE |
| 116 | + # self.model_ = replace_rope_with_inference_rope(self.model_, model_args) |
| 117 | + |
| 118 | + setattr(self.model_, "checkpoint_dtype", self.dtype) |
| 119 | + if self.use_checkpoint: |
| 120 | + # Load checkpoint. |
| 121 | + missing, unexpected = self.model_.load_state_dict( |
| 122 | + checkpoint, |
| 123 | + strict=False, |
| 124 | + assign=True, |
| 125 | + ) |
| 126 | + if kwargs.get("verbose", False): |
| 127 | + print("============= missing keys ================") |
| 128 | + print(missing) |
| 129 | + print("============= /missing ================") |
| 130 | + print("============= unexpected keys ================") |
| 131 | + print(unexpected) |
| 132 | + print("============= /unexpected ================") |
| 133 | + |
| 134 | + self.model_.to(self.dtype) |
| 135 | + # breakpoint() # 2, OK. |
| 136 | + |
| 137 | + def get_eager_model(self) -> torch.nn.Module: |
| 138 | + return self.model_ |
| 139 | + |
| 140 | + def get_example_inputs(self): |
| 141 | + return (torch.tensor([[2, 3, 4]], dtype=torch.int64),) |
| 142 | + # return ( |
| 143 | + # torch.tensor([[2, 3, 4]], dtype=torch.long), |
| 144 | + # {"input_pos": torch.tensor([0], dtype=torch.long)}, |
| 145 | + # ) |
| 146 | + # return (torch.ones(1, self.n_tokens, dtype=torch.int64),) |
| 147 | + |
| 148 | + # eg=torch.tensor([[2, 3, 4]], dtype=torch.int64) |
| 149 | + # ip=torch.tensor([[0, 1, 2]], dtype=torch.long) |
| 150 | + def get_example_kwarg_inputs(self): |
| 151 | + return {"input_pos": torch.tensor([[0, 1, 2]], dtype=torch.long)} |
| 152 | + |
| 153 | + def get_dynamic_shapes(self): |
| 154 | + dim = torch.export.Dim("token_dim", min=1, max=self.max_seq_len - 1) |
| 155 | + return ({1: dim}, {1: dim}) |
0 commit comments