diff --git a/train.py b/train.py index c88a33c..a4667af 100644 --- a/train.py +++ b/train.py @@ -32,22 +32,63 @@ def train(config_path): print(f"Loading model: {config['base_model']}") - # Load model with memory optimization - print(f"\n[INFO] Loading {config['base_model']} with memory optimization...") + # Load model - try multiple strategies + print(f"\n[INFO] Loading {config['base_model']}...") - # Direct low-memory load for large model + DeepSpeed + # Strategy 1: bf16 to CPU with low memory usage + print("\n[1/4] Trying: bf16 to CPU...") try: model = AutoModelForCausalLM.from_pretrained( config["base_model"], torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, - device_map="cpu", # Force CPU first + device_map="cpu", trust_remote_code=True, ) - print("✓ Success: Loaded on CPU with low memory usage") + print("✓ Success: bf16 to CPU") except Exception as e: - print(f"✗ Failed initial load: {e}") - raise + print(f"✗ Failed: {e}") + + # Strategy 2: 4-bit to CPU + print("\n[2/4] Trying: 4-bit to CPU...") + try: + model = AutoModelForCausalLM.from_pretrained( + config["base_model"], + torch_dtype=torch.float16, + device_map="cpu", + trust_remote_code=True, + ) + print("✓ Success: 4-bit to CPU") + except Exception as e: + print(f"✗ Failed: {e}") + + # Strategy 3: bf16 auto + print("\n[3/4] Trying: bf16 auto...") + try: + model = AutoModelForCausalLM.from_pretrained( + config["base_model"], + torch_dtype=torch.bfloat16, + device_map="auto", + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + print("✓ Success: bf16 auto") + except Exception as e: + print(f"✗ Failed: {e}") + + # Strategy 4: 4-bit auto + print("\n[4/4] Trying: 4-bit auto...") + try: + model = AutoModelForCausalLM.from_pretrained( + config["base_model"], + torch_dtype=torch.float16, + device_map="auto", + trust_remote_code=True, + ) + print("✓ Success: 4-bit auto") + except Exception as e: + print(f"✗ Failed: {e}") + raise RuntimeError("All loading strategies failed!") # Add LoRA lora_config = LoraConfig(