Merge pull request #47 from ai-forever/kirillova/lita_multigpu_fix

fix: change device choice in model loader function

Author: boomb0om

DPF/filters/videos/lita_filter.py

@@ -1,4 +1,5 @@
 import os
+import warnings
 from io import BytesIO
 from typing import Any, Optional
 
@@ -7,17 +8,20 @@
 from lita.constants import (
     DEFAULT_IM_END_TOKEN,
     DEFAULT_IM_START_TOKEN,
+    DEFAULT_IMAGE_PATCH_TOKEN,
     DEFAULT_IMAGE_TOKEN,
     IMAGE_TOKEN_INDEX,
+    TIME_TOKEN_TEMPLATE,
 )
-from lita.model.builder import load_pretrained_model
+from lita.model.language_model.lita_llama import LitaLlamaForCausalLM
 from lita.utils import load_video
 from llava.conversation import SeparatorStyle, conv_templates
 from llava.mm_utils import (
     KeywordsStoppingCriteria,
     get_model_name_from_path,
     tokenizer_image_token,
 )
+from transformers import AutoConfig, AutoTokenizer, BitsAndBytesConfig
 
 from DPF.types import ModalityToDataMapping
 
@@ -29,6 +33,98 @@
     from torch.utils.data import default_collate
 
 
+def load_pretrained_model(model_path: str,
+                          model_base: str,
+                          model_name: str,
+                          load_8bit: bool = False,
+                          load_4bit: bool = False,
+                          device_map: str = "auto",
+                          device: str = "cuda"):
+    kwargs = {"device_map": device_map}
+
+    if device != "cuda":
+        kwargs['device_map'] = {"": device}  # type: ignore
+
+    if load_8bit:
+        kwargs['load_in_8bit'] = True  # type: ignore
+    elif load_4bit:
+        kwargs['load_in_4bit'] = True  # type: ignore
+        kwargs['quantization_config'] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+    else:
+        kwargs['torch_dtype'] = torch.float16  # type: ignore
+
+    if 'lita' not in model_name.lower():
+        warnings.warn("this function is for loading LITA models", stacklevel=2)
+    if 'lora' in model_name.lower():
+        warnings.warn("lora is currently not supported for LITA", stacklevel=2)
+    if 'mpt' in model_name.lower():
+        warnings.warn("mpt is currently not supported for LITA", stacklevel=2)
+
+    if model_base is not None:
+        print('Loading LITA from base model...')
+        tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+        cfg_pretrained = AutoConfig.from_pretrained(model_path)
+        model = LitaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+
+        mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
+        mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items() if 'mm_projector' in k}
+        model.load_state_dict(mm_projector_weights, strict=False)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+        model = LitaLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+
+    mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+    mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", False)
+    if mm_use_im_patch_token:
+        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+    if mm_use_im_start_end:
+        tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+    model.resize_token_embeddings(len(tokenizer))
+
+    vision_tower = model.get_vision_tower()
+    if not vision_tower.is_loaded:
+        vision_tower.load_model()
+    vision_tower.to(device=device, dtype=torch.float16)
+    image_processor = vision_tower.image_processor
+
+    # time tokens and embeddings
+    num_time_tokens = getattr(model.config, "num_time_tokens", 0)
+    if num_time_tokens > 0:
+        time_tokens = [TIME_TOKEN_TEMPLATE.format(t=x) for x in range(num_time_tokens)]
+        num_new_tokens = tokenizer.add_tokens(time_tokens)
+
+        if model_base is None:
+            assert num_new_tokens == 0, "time tokens should already be in the tokenizer for full finetune model"
+
+        if num_new_tokens > 0:
+            warnings.warn("looking for weights in mm_projector.bin", stacklevel=2)
+            assert num_new_tokens == num_time_tokens
+            model.resize_token_embeddings(len(tokenizer))
+            input_embeddings = model.get_input_embeddings().weight.data
+            output_embeddings = model.get_output_embeddings().weight.data
+            weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
+            assert 'model.embed_tokens.weight' in weights and 'lm_head.weight' in weights
+
+            dtype = input_embeddings.dtype
+            device = input_embeddings.device
+
+            tokenizer_time_token_ids = tokenizer.convert_tokens_to_ids(time_tokens)
+            time_token_ids = getattr(model.config, 'time_token_ids', tokenizer_time_token_ids)
+            input_embeddings[tokenizer_time_token_ids] = weights['model.embed_tokens.weight'][time_token_ids].to(dtype).to(device)
+            output_embeddings[tokenizer_time_token_ids] = weights['lm_head.weight'][time_token_ids].to(dtype).to(device)
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+    return tokenizer, model, image_processor, context_len
+
+
 def disable_torch_init() -> None:
     """
     Disable the redundant torch default initialization to accelerate model creation.
@@ -79,7 +175,7 @@ def __init__(
 
         disable_torch_init()
 
-        pretrainers = load_pretrained_model(weights_path, model_base, self.model_name, load_8bit, load_4bit)
+        pretrainers = load_pretrained_model(weights_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device)  # type: ignore
         self.tokenizer, self.model, self.processor, self.context_len = pretrainers
 
         self.model_num_frames = self.model.config.num_frames

pyproject.toml

@@ -78,7 +78,7 @@ disable_error_code = ["import-not-found"]
 
 [[tool.mypy.overrides]]
 module = "DPF.filters.videos.lita_filter"
-disable_error_code = ["import-not-found", "call-overload"]
+disable_error_code = ["import-not-found", "call-overload", "no-untyped-def"]
 
 [[tool.mypy.overrides]]
 module = "DPF.filters.images.llava_captioning_filter"