Method build_model

lmdeploy/vl/model/molmo.py:26–50

Build the vision part of a VLM model when backend is turbomind, or load the whole VLM model when `self.with_llm==True`

(self, trust_remote_code: bool = False)

                                                       device_map='auto')

    def build_model(self, trust_remote_code: bool = False):
        """Build the vision part of a VLM model when backend is turbomind, or
        load the whole VLM model when `self.with_llm==True`"""
        from accelerate import init_empty_weights, load_checkpoint_and_dispatch
        with init_empty_weights():
            model = AutoModelForCausalLM.from_config(self.hf_config, trust_remote_code=trust_remote_code)

            self.vl_model = model
            if not self.with_llm:
                # Remove nn modules other than embedding from the LLM model
                for key in ['emb_drop', 'ln_f', 'blocks', 'ff_out']:
                    del model.model.transformer[key]
            self.token_embedding = model.model.transformer.wte

        with disable_logging():
            load_checkpoint_and_dispatch(model=model,
                                         checkpoint=self.model_path,
                                         device_map='auto' if not self.with_llm else {'': 'cpu'},
                                         max_memory=self.max_memory,
                                         no_split_module_classes=['ResidualAttentionBlock', 'Embedding'],
                                         dtype=torch.half)

        # We need eval mode to freeze the weights in model, thus,
        # avoid randomness in inference.
        self.model = model.eval()

    def preprocess(self, messages: list[dict]) -> list[dict]:
        """Refer to the `super.preprocess() for spec."""