Upload folder using huggingface_hub

README.md

@@ -166,6 +166,32 @@ def load_image(image_file, input_size=448, max_num=6):
     return pixel_values
 
 
+def split_model(model_name):
+    device_map = {}
+    world_size = torch.cuda.device_count()
+    num_layers = {'InternVL2-8B': 32, 'InternVL2-26B': 48, 'InternVL-Chat-V1-5': 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/InternVL-Chat-V1-5'
 # If you have an 80G A100 GPU, you can put the entire model on a single GPU.
 model = AutoModel.from_pretrained(
@@ -173,15 +199,15 @@ model = AutoModel.from_pretrained(
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=True,
     trust_remote_code=True).eval().cuda()
-# Otherwise, you need to set device_map='auto' to use multiple GPUs for inference.
-# import os
-# os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+# Otherwise, you need to set device_map to use multiple GPUs for inference.
+# device_map = split_model('InternVL-Chat-V1-5')
+# print(device_map)
 # model = AutoModel.from_pretrained(
 #     path,
 #     torch_dtype=torch.bfloat16,
 #     low_cpu_mem_usage=True,
 #     trust_remote_code=True,
-#     device_map='auto').eval()
+#     device_map=device_map).eval()
 
 tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
 # set the max number of tiles in `max_num`