| 17 | |
| 18 | |
| 19 | class MetaModel(nn.Module): |
| 20 | def __init__( |
| 21 | self, llama_type: str, llama_config: str|List[str], tokenizer_path: str, |
| 22 | with_visual: bool = False, max_seq_len: int = 4096 |
| 23 | ) -> None: |
| 24 | super().__init__() |
| 25 | |
| 26 | self.llama_type = llama_type |
| 27 | self.with_visual = with_visual |
| 28 | |
| 29 | model_module = importlib.import_module(f"accessory.model.LLM.{llama_type}") |
| 30 | ModelArgs = model_module.ModelArgs |
| 31 | Transformer = model_module.Transformer |
| 32 | |
| 33 | llama_args = {} |
| 34 | if isinstance(llama_config, str): |
| 35 | llama_config = [llama_config] |
| 36 | for _ in llama_config: |
| 37 | with open(_, "r") as f: |
| 38 | llama_args.update(json.loads(f.read())) |
| 39 | llama_args['max_seq_len'] = max_seq_len |
| 40 | llama_args['max_batch_size'] = 32 |
| 41 | |
| 42 | tokenizer = Tokenizer(model_path=tokenizer_path) |
| 43 | llama_args['vocab_size'] = tokenizer.n_words |
| 44 | |
| 45 | llama_args: ModelArgs = ModelArgs(**llama_args) |
| 46 | |
| 47 | if "tokenizer" in inspect.signature(Transformer.__init__).parameters: |
| 48 | # generally it means the inner llm modify change the tokenizer |
| 49 | model = Transformer(llama_args, tokenizer, with_visual=with_visual) |
| 50 | assert hasattr(model, "tokenizer") |
| 51 | self.tokenizer = model.tokenizer |
| 52 | else: |
| 53 | model = Transformer(llama_args, with_visual=with_visual) |
| 54 | self.tokenizer = tokenizer |
| 55 | |
| 56 | print("Model Args:\n", model.args) |
| 57 | |
| 58 | self.llma = model |
| 59 | |
| 60 | self.criterion = torch.nn.CrossEntropyLoss(ignore_index=0) |
| 61 | |
| 62 | self._set_default_trainability() |
| 63 | |
| 64 | self.is_peft = getattr(model, "is_peft", False) |
| 65 | print(f"Model is Peft: {self.is_peft}") |
| 66 | |
| 67 | misc.mark_mp_params(self) |
| 68 | |
| 69 | param_count_local, param_count_all = 0, 0 |
| 70 | for name, param in self.named_parameters(): |
| 71 | is_model_parallel = getattr(param, "is_model_parallel", False) |
| 72 | if param.requires_grad: |
| 73 | if is_model_parallel: |
| 74 | param_count_all += param.numel() * fs_init.get_model_parallel_world_size() |
| 75 | else: |
| 76 | param_count_all += param.numel() |