Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,931 Bytes
841bef5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
# --------------------------------------------------------
# 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
|