Update app.py

app.py

@@ -1,91 +1,109 @@
+import os
 import gradio as gr
 import torch
+import torch.distributed as dist
 import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
 import warnings
-from accelerate import Accelerator, DistributedType
-import os
 
 # disable some warnings
 transformers.logging.set_verbosity_error()
 transformers.logging.disable_progress_bar()
 warnings.filterwarnings('ignore')
 
-# Initialize Accelerator
-accelerator = Accelerator()
-
-model_name = 'cognitivecomputations/dolphin-vision-72b'
-
-# Determine the number of GPUs available
-num_gpus = torch.cuda.device_count()
-print(f"Number of GPUs available: {num_gpus}")
-
-# Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True
-)
-
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    trust_remote_code=True
-)
-
-# Prepare model
-model = accelerator.prepare(model)
-
-def inference(prompt, image, temperature, beam_size):
-    messages = [
-        {"role": "user", "content": f'<image>\n{prompt}'}
-    ]
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-
-    text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
-    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
-
-    image_tensor = model.process_images([image], model.config)
-
-    # Move tensors to the appropriate device
-    input_ids = input_ids.to(accelerator.device)
-    image_tensor = image_tensor.to(accelerator.device)
-
-    # generate
-    with torch.cuda.amp.autocast():
-        output_ids = accelerator.unwrap_model(model).generate(
-            input_ids,
-            images=image_tensor,
-            max_new_tokens=1024,
-            temperature=temperature,
-            num_beams=beam_size,
-            use_cache=True
-        )[0]
-
-    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
-
-# Create Gradio interface
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail")
-            image_input = gr.Image(label="Image", type="pil")
-            temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
-            beam_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Beam Size")
-            submit_button = gr.Button("Submit")
-        with gr.Column():
-            output_text = gr.Textbox(label="Output")
-
-    submit_button.click(
-        fn=inference, 
-        inputs=[prompt_input, image_input, temperature_input, beam_size_input], 
-        outputs=output_text
+def setup(rank, world_size):
+    os.environ['MASTER_ADDR'] = 'localhost'
+    os.environ['MASTER_PORT'] = '12355'
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+
+def cleanup():
+    dist.destroy_process_group()
+
+def load_model_on_gpus(model_name, num_gpus):
+    # Calculate number of layers to assign to each GPU
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, trust_remote_code=True)
+    num_layers = len(model.model.layers)
+    layers_per_gpu = num_layers // num_gpus
+
+    # Assign layers to GPUs
+    device_map = {}
+    for i in range(num_layers):
+        device_map[f'model.layers.{i}'] = i // layers_per_gpu
+    
+    # Assign other components
+    device_map['model.embed_tokens'] = 0
+    device_map['model.norm'] = num_gpus - 1
+    device_map['lm_head'] = num_gpus - 1
+
+    return AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map=device_map,
+        torch_dtype=torch.float16,
+        trust_remote_code=True
     )
 
-# Launch the app
-demo.launch()
+def run_distributed(rank, world_size, model_name):
+    setup(rank, world_size)
+    
+    if rank == 0:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    
+    model = load_model_on_gpus(model_name, world_size)
+    
+    def inference(prompt, image, temperature, beam_size):
+        if rank == 0:
+            messages = [{"role": "user", "content": f'<image>\n{prompt}'}]
+            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
+            input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(rank)
+            image_tensor = model.process_images([image], model.config).to(rank)
+        else:
+            input_ids = torch.zeros(1, 1, dtype=torch.long).to(rank)
+            image_tensor = torch.zeros(1, 3, 224, 224).to(rank)
+
+        dist.broadcast(input_ids, src=0)
+        dist.broadcast(image_tensor, src=0)
+
+        with torch.cuda.amp.autocast():
+            output_ids = model.generate(
+                input_ids,
+                images=image_tensor,
+                max_new_tokens=1024,
+                temperature=temperature,
+                num_beams=beam_size,
+                use_cache=True
+            )[0]
+
+        if rank == 0:
+            return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
+        else:
+            return ""
+
+    if rank == 0:
+        with gr.Blocks() as demo:
+            with gr.Row():
+                with gr.Column():
+                    prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail")
+                    image_input = gr.Image(label="Image", type="pil")
+                    temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
+                    beam_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Beam Size")
+                    submit_button = gr.Button("Submit")
+                with gr.Column():
+                    output_text = gr.Textbox(label="Output")
+
+            submit_button.click(
+                fn=inference, 
+                inputs=[prompt_input, image_input, temperature_input, beam_size_input], 
+                outputs=output_text
+            )
+
+        demo.launch(share=True)
+
+    cleanup()
+
+if __name__ == "__main__":
+    model_name = 'cognitivecomputations/dolphin-vision-72b'
+    world_size = torch.cuda.device_count()
+    print(f"Running on {world_size} GPUs")
+    torch.multiprocessing.spawn(run_distributed, args=(world_size, model_name), nprocs=world_size, join=True)