diff --git a/train.py b/train.py index 2117664..a2f9370 100644 --- a/train.py +++ b/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}") - - # -------------------------------------------------------------- - # 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}") - - # -------------------------------------------------------------- - # 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}") - - # ---------------------------------------------------------- - # 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}") - - # ------------------------------------------------------ - # 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}" - ) + # 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}") + + # Use accelerate to create device_map + from accelerate import infer_auto_device_map + print(" Creating device_map with accelerate...") + + # 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, + ) + + # 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") + + # 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(