Closed HubHop closed 1 day ago
Could you try this split_model
function:
def split_model(model_name):
device_map = {}
world_size = torch.cuda.device_count()
num_layers = {'InternVL2-8B': 32, 'InternVL2-26B': 48,
'InternVL2-40B': 60, 'InternVL2-Llama3-76B': 80}[model_name]
# 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 - 0.5))
num_layers_per_gpu = [num_layers_per_gpu] * world_size
num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
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
path = 'OpenGVLab/InternVL2-26B'
device_map = split_model('InternVL2-26B')
print(device_map)
model = AutoModel.from_pretrained(
path,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
trust_remote_code=True,
device_map=device_map).eval()
Thanks for the help @czczup, I just tried this new code and unfortunately it still doesn't work. My setting is based on 8 GPUs and I can run other scales (e.g., 1B, 4B, 40B, 76B) successfully.
Hi authors,
Was trying to run InternVL-8B and InternVL-26B on 4 GPUs, but I got this,
Code is exactly the same as in huggingface demo.
My python env:
I guess the error is due to flash attention, but not 100% sure. Any help?
Thanks in advance!