simplify: only train, no quantization logic

train.py

@@ -33,77 +33,15 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Check if model is quantized (has model-XXXXX-of-XXXXX files)
+    print(f"\n[INFO] Loading pre-quantized BnB 4-bit model...")
-    import glob
+    model = AutoModelForCausalLM.from_pretrained(
-    import json
+        config["base_model"],
-    safetensor_files = glob.glob(f"{config['base_model']}/*.safetensors")
+        device_map="cpu",
-    is_quantized = any("of-" in Path(f).name for f in safetensor_files)
+        torch_dtype=torch.float16,
-    
+        trust_remote_code=True,
-    if is_quantized:
+        low_cpu_mem_usage=True,
-        # Create proper index file from safetensors metadata
+    )
-        index_file = Path(config["base_model"]) / "model.safetensors.index.json"
+    print("✓ Model loaded to CPU (BnB 4-bit)")
-        if not index_file.exists():
-            print(f"\n[INFO] Creating model index file from safetensors metadata...")
-            shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
-            
-            # Load metadata from first shard to get weight names
-            from safetensors import safe_open
-            weight_map = {}
-            for shard_name in shards:
-                shard_path = Path(config["base_model"]) / shard_name
-                with safe_open(str(shard_path), framework="pt", device="cpu") as f:
-                    for key in f.keys():
-                        weight_map[key] = shard_name
-            
-            # Create index
-            index = {
-                "metadata": {"total_size": sum((Path(config["base_model"]) / s).stat().st_size for s in shards)},
-                "weight_map": weight_map
-            }
-            with open(index_file, 'w') as f:
-                json.dump(index, f)
-            print(f"✓ Created index: {index_file} ({len(weight_map)} weights mapped)")
-        
-        # Remove quantization_config from config.json to prevent re-quantization
-        import json as json_mod
-        config_json_path = Path(config["base_model"]) / "config.json"
-        if config_json_path.exists():
-            with open(config_json_path, 'r') as f:
-                config_data = json_mod.load(f)
-            if 'quantization_config' in config_data:
-                del config_data['quantization_config']
-                with open(config_json_path, 'w') as f:
-                    json_mod.dump(config_data, f)
-                print(f"✓ Removed quantization_config from config.json")
-        
-        print(f"\n[INFO] Loading pre-quantized BnB 4-bit model...")
-        model = AutoModelForCausalLM.from_pretrained(
-            config["base_model"],
-            device_map="cpu",
-            torch_dtype=torch.float16,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-        )
-        print("✓ Model loaded to CPU (BnB 4-bit)")
-    else:
-        # Load bf16 with BnB 4-bit quantization
-        print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
-        from transformers import BitsAndBytesConfig
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.float16,
-        )
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            config["base_model"],
-            quantization_config=bnb_config,
-            device_map="cpu",
-            torch_dtype=torch.float16,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-        )
-        print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
     
     # Move to GPU
     print("   Moving to GPU 0...")