|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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() |
|
|
|
|
|
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 |
|
|
|
|
|
|
