feat: use accelerate device_map for DISTRIBUTED model loading (Test 7 method)

train.py

@@ -33,100 +33,48 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model with multiple strategies
+    # Load model with accelerate device_map (Test 7 method - DISTRIBUTED)
     print(f"\n[INFO] Loading {config['base_model']}...")
-    errors = []
     
-    # ------------------------------------------------------------------
-    # Strategy 1: Load already-quantized 4-bit model (distributed across GPUs)
-    # ------------------------------------------------------------------
-    print("\n[1/4] Trying: 4-bit model AS-IS (distributed across GPUs)...")
-    try:
-        model = AutoModelForCausalLM.from_pretrained(
-            config["base_model"],
-            device_map="auto",  # Distribute layers across GPUs automatically
-            torch_dtype=torch.float16,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-        )
-        print("✓ Success: 4-bit model distributed across GPUs (no FSDP)")
-    except Exception as e:
-        errors.append(("QLoRA 4-bit", e))
-        print(f"✗ Failed: {e}")
+    # Detect layer names dynamically
+    from transformers import AutoConfig
+    print("   Detecting layer names...")
+    config_model = AutoConfig.from_pretrained(config["base_model"], trust_remote_code=True)
+    layer_names = [f"{config_model.model_type.title().replace('_', '')}DecoderLayer"]
+    if 'moe' in config_model.model_type.lower():
+        layer_names.append(f"{config_model.model_type.title().replace('_', '')}SparseMoeBlock")
+    print(f"   Detected layers: {layer_names}")
     
-        # --------------------------------------------------------------
-        # Strategy 2: 4-bit model with FSDP (load to GPU, FSDP shards)
-        # --------------------------------------------------------------
-        print("\n[2/4] Trying: 4-bit model with FSDP (load to GPU)...")
-        try:
-            model = AutoModelForCausalLM.from_pretrained(
-                config["base_model"],
-                device_map="auto",  # Load to GPU first
-                torch_dtype=torch.float16,
-                trust_remote_code=True,
-                low_cpu_mem_usage=True,
-            )
-            print("✓ Success: 4-bit model loaded to GPU (FSDP will shard)")
-        except Exception as e:
-            errors.append(("QLoRA 4-bit FSDP", e))
-            print(f"✗ Failed: {e}")
+    # Use accelerate to create device_map
+    from accelerate import infer_auto_device_map
+    print("   Creating device_map with accelerate...")
     
-            # --------------------------------------------------------------
-            # Strategy 3: BF16 CPU (model too large for single GPU)
-            # --------------------------------------------------------------
-            print("\n[3/4] Trying: bf16 CPU...")
-            try:
-                model = AutoModelForCausalLM.from_pretrained(
-                    config["base_model"],
-                    torch_dtype=torch.bfloat16,
-                    device_map="cpu",
-                    low_cpu_mem_usage=True,
-                    trust_remote_code=True,
-                )
-                print("✓ Success: bf16 CPU")
-            except Exception as e:
-                errors.append(("bf16 CPU", e))
-                print(f"✗ Failed: {e}")
+    # First load to CPU to get the model structure
+    model_temp = AutoModelForCausalLM.from_pretrained(
+        config["base_model"],
+        device_map="cpu",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
     
-            # ----------------------------------------------------------
-            # Strategy 4: FP16 CPU (fallback)
-            # ----------------------------------------------------------
-            print("\n[4/4] Trying: fp16 CPU...")
-            try:
-                model = AutoModelForCausalLM.from_pretrained(
-                    config["base_model"],
-                    torch_dtype=torch.float16,
-                    device_map="cpu",
-                    trust_remote_code=True,
-                    low_cpu_mem_usage=True,
-                )
-                print("✓ Success: fp16 CPU")
-            except Exception as e:
-                errors.append(("fp16 CPU", e))
-                print(f"✗ Failed: {e}")
+    # Create device_map
+    device_map = infer_auto_device_map(
+        model_temp,
+        max_memory={i: "15GB" for i in range(torch.cuda.device_count())},
+        no_split_module_classes=layer_names,
+    )
+    print(f"   Created device_map with {len(device_map)} entries")
     
-                # ------------------------------------------------------
-                # Strategy 5: 4-bit AS-IS (already quantized)
-                # ----------------------------------------------------------
-                print("\n[5/5] Trying: 4-bit AS-IS...")
-                try:
-                    model = AutoModelForCausalLM.from_pretrained(
-                        config["base_model"],
-                        torch_dtype=torch.float16,
-                        device_map="cpu",
-                        trust_remote_code=True,
-                        low_cpu_mem_usage=True,
-                    )
-                    print("✓ Success: 4-bit AS-IS")
-                except Exception as e:
-                    errors.append(("4-bit AS-IS", e))
-                    print(f"✗ Failed: {e}")
-                    msg = "\n".join(
-                        f"{name}: {err}" for name, err in errors
-                    )
-                    raise RuntimeError(
-                        f"All loading strategies failed:\n\n{msg}"
-                    )
+    # Reload with device_map
+    model = AutoModelForCausalLM.from_pretrained(
+        config["base_model"],
+        device_map=device_map,
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+    print("✓ Success: Model distributed across GPUs via accelerate device_map")
     
     # Add LoRA
     lora_config = LoraConfig(