Spaces:
Running
on
Zero
Running
on
Zero
| # -------------------------------------------------------- | |
| # InternVL | |
| # Copyright (c) 2024 OpenGVLab | |
| # Licensed under The MIT License [see LICENSE for details] | |
| # -------------------------------------------------------- | |
| import math | |
| import torch | |
| from internvl.model.internvl_chat import InternVLChatConfig, InternVLChatModel | |
| from transformers import AutoTokenizer | |
| def split_model(num_layers, vit_alpha=0.5): | |
| device_map = {} | |
| world_size = torch.cuda.device_count() | |
| # Since the first GPU will be used for ViT, treat it as half a GPU. | |
| num_layers_per_gpu = math.ceil(num_layers / (world_size - vit_alpha)) | |
| num_layers_per_gpu = [num_layers_per_gpu] * world_size | |
| num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * (1 - vit_alpha)) | |
| layer_cnt = 0 | |
| for i, num_layer in enumerate(num_layers_per_gpu): | |
| for j in range(num_layer): | |
| device_map[f'language_model.model.layers.{layer_cnt}'] = i | |
| layer_cnt += 1 | |
| device_map['vision_model'] = 0 | |
| device_map['mlp1'] = 0 | |
| device_map['language_model.model.tok_embeddings'] = 0 | |
| device_map['language_model.model.embed_tokens'] = 0 | |
| device_map['language_model.output'] = 0 | |
| device_map['language_model.model.norm'] = 0 | |
| device_map['language_model.lm_head'] = 0 | |
| device_map[f'language_model.model.layers.{num_layers - 1}'] = 0 | |
| return device_map | |
| def load_model_and_tokenizer(args): | |
| if args.auto: | |
| config = InternVLChatConfig.from_pretrained(args.checkpoint) | |
| num_hidden_layers = config.llm_config.num_hidden_layers | |
| device_map = split_model(num_hidden_layers) | |
| kwargs = {'device_map': device_map} if args.auto else {} | |
| tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, use_fast=False) | |
| model = InternVLChatModel.from_pretrained( | |
| args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, | |
| load_in_8bit=args.load_in_8bit, load_in_4bit=args.load_in_4bit, **kwargs).eval() | |
| if not args.load_in_8bit and not args.load_in_4bit and not args.auto: | |
| model = model.cuda() | |
| return model, tokenizer | |
| def load_model_and_tokenizer_customed(args): | |
| if args.auto: | |
| config = InternVLChatConfig.from_pretrained(args.checkpoint) | |
| num_hidden_layers = config.llm_config.num_hidden_layers | |
| device_map = split_model(num_hidden_layers) | |
| kwargs = {'device_map': device_map} if args.auto else {} | |
| tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, use_fast=False) | |
| model = InternVLChatModel.from_pretrained( | |
| args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16, | |
| load_in_8bit=args.load_in_8bit, load_in_4bit=args.load_in_4bit, **kwargs).eval() | |
| if not args.load_in_8bit and not args.load_in_4bit and not args.auto: | |
| del model.language_model.model.layers | |
| del model.language_model.output | |
| return model, tokenizer | |