Fix TransformersModelManager: Prevents crash by dispatching model on all GPUs

lib/TransformersModelManager.py

@@ -8,6 +8,7 @@
 from transformers import LogitsProcessor
 import torch
 import time
+import gc
 
 from transformers import (
     TemperatureLogitsWarper,
@@ -129,19 +130,23 @@ def switch_model(self, model_name: str):
 
         # Clear existing GPU model if present
         self.clear_model()
-        
-        # Move new model to GPU from CPU
+
+        #  MODE: CPU → MULTI-GPU ---
         if model_config.location == "cpu":
-            self.current_gpu_model = self.cpu_models[model_name].to("cuda")
+            self.current_gpu_model = AutoModelForCausalLM.from_pretrained(
+                model_config.model_name,
+                device_map="auto",   # <== use all available gpus
+                torch_dtype=torch.float16,
+                cache_dir=MODELS_FOLDER,
+                **model_config.model_kwargs
+            )
             self.current_gpu_model_name = model_name
-
-            torch.cuda.synchronize()  # Synchronize CUDA operations
 
         # Load model from disk to GPU
         elif model_config.location == "disk":
             self.current_gpu_model = AutoModelForCausalLM.from_pretrained(
                 model_config.model_name,
-                device_map="cuda",
+                device_map="auto",
                 torch_dtype=torch.float16,  # Use float16 for memory efficiency
                 cache_dir=MODELS_FOLDER,
                 # attn_implementation="flash_attention_2",
@@ -156,25 +161,43 @@ def switch_model(self, model_name: str):
         print(f"Time taken to switch model: {time.time() - time_start:.2f}s")
         return self.current_gpu_model
 
+
     def clear_model(self):
         """
-        Remove the current model from GPU if one exists.
+        Free the current model in VRAM.
+        Keep preloaded CPU models in RAM.
         """
         time_start = time.time()
-        if self.current_gpu_model is not None:
-            if self.models_config[self.current_gpu_model_name].location == "cpu":
-                # Move model back to CPU
-                self.cpu_models[self.current_gpu_model_name] = self.current_gpu_model.to("cpu")
-
-            # Clear GPU model references
-            self.current_gpu_model = None
-            self.current_gpu_model_name = None
-
-            # Clear CUDA cache
-            torch.cuda.empty_cache()
 
-            # Ensure all CUDA operations are finished
-            torch.cuda.synchronize()
+        if self.current_gpu_model is None:
+            return
+
+        try:
+            # For multi-GPU models
+            if hasattr(self.current_gpu_model, "hf_device_map"):
+                # Unload each sub module
+                for device in set(self.current_gpu_model.hf_device_map.values()):
+                    if device != "disk" and "cpu" not in str(device):
+                        with torch.cuda.device(device):
+                            torch.cuda.empty_cache()
+
+            # Delete reference to GPU
+            del self.current_gpu_model
+
+        except Exception as e:
+            print(f"[WARN] clear_model: Erreur lors de la libération : {e}")
+
+        # Force garbage collector Python
+        gc.collect()
+
+        # Free memory from all GPUs
+        if torch.cuda.is_available():
+            for i in range(torch.cuda.device_count()):
+                with torch.cuda.device(i):
+                    torch.cuda.empty_cache()
+
+        self.current_gpu_model = None
+        self.current_gpu_model_name = None
 
         print(f"Time taken to clear model for TRANSFORMERS MODEL MANAGER: {time.time() - time_start:.2f}s")
 
@@ -492,4 +515,4 @@ def get_tokenizer(self, model_name: str):
             "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='', is_first_sp=true, is_last_user=false) %}{%- for message in messages %}{%- if message['role'] == 'system' %}{%- if ns.is_first_sp %}{% set ns.system_prompt = ns.system_prompt + message['content'] %}{% set ns.is_first_sp = false %}{%- else %}{% set ns.system_prompt = ns.system_prompt + '\n\n' + message['content'] %}{%- endif %}{%- endif %}{%- endfor %}{{ bos_token }}{{ ns.system_prompt }}{%- for message in messages %}{% set content = message['content'] %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{%- set ns.is_first = false -%}{%- set ns.is_last_user = true -%}{{'<｜User｜>' + content + '<｜Assistant｜>'}} {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}{%- if message['role'] == 'assistant' %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{% endif %}{%- if message['role'] == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] is not none %}{%- set ns.is_last_user = false -%}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{%- endif %}{%- set ns.is_first = false %}{%- set ns.is_tool = false -%}{%- set ns.is_output_first = true %}{%- for tool in message['tool_calls'] %}{%- if not ns.is_first %}{%- if content is none %}{{'<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{%- else %}{{content + '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{%- endif %}{%- set ns.is_first = true -%}{%- else %}{{'\n' + '<｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + '<｜tool▁call▁end｜>'}}{%- endif %}{%- endfor %}{{'<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'}}{%- endif %}{%- if message['role'] == 'assistant' and (message['tool_calls'] is not defined or message['tool_calls'] is none)%}{%- set ns.is_last_user = false -%}{%- if ns.is_tool %}{{'<｜tool▁outputs▁end｜>' + content + '<｜end▁of▁sentence｜>'}}{%- set ns.is_tool = false -%}{%- else %}{{content + '<｜end▁of▁sentence｜>'}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_last_user = false -%}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{'<｜tool▁outputs▁begin｜><｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}}{%- set ns.is_output_first = false %}{%- else %}{{'\n<｜tool▁output▁begin｜>' + content + '<｜tool▁output▁end｜>'}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{'<｜tool▁outputs▁end｜>'}}{% endif %}{% if add_generation_prompt and not ns.is_last_user and not ns.is_tool %}{{'<｜Assistant｜>'}} \n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n {% endif %}"
         }
     ),
-    }
+    }