feat: add proper BnB quantization script for 2x RTX 5090

.gitignore

@@ -23,5 +23,5 @@ Thumbs.db
 # Logs
 *.log
 
-# Large dataset (optional, use git-lfs if needed)
+# Large dataset (use git-lfs or exclude)
-# dataset/*.jsonl
+dataset/*.jsonl

check_model_size.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+"""Test loading bf16 model with BnB 4-bit to CPU, then report size."""
+
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+model_path = "/data/models/Ornith-1.0-35B"
+
+print(f"Loading model from: {model_path}")
+print(f"\nApplying BnB 4-bit quantization...")
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    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")
+
+# Count parameters
+total_params = sum(p.numel() for p in model.parameters())
+print(f"\nTotal parameters: {total_params / 1e9:.2f}B")
+
+# Check for quantized parameters
+bnb_params = 0
+bf16_params = 0
+for name, p in model.named_parameters():
+    if hasattr(p, 'quant_state') and p.quant_state is not None:
+        bnb_params += p.numel()
+    else:
+        bf16_params += p.numel()
+
+print(f"BnB 4-bit parameters: {bnb_params / 1e9:.2f}B")
+print(f"BF16 parameters: {bf16_params / 1e9:.2f}B")
+print(f"Estimated size: {(bnb_params * 0.5 + bf16_params * 2) / 1e9:.2f} GB")
+
+if bnb_params / total_params > 0.9:
+    print("\n✓ SUCCESS: Model is properly quantized to 4-bit!")
+else:
+    print(f"\n✗ FAILED: Only {bnb_params/total_params*100:.1f}% of parameters are 4-bit")
+    print("   Expected: ~100%, Got: this percentage")
+
+del model
+import gc
+gc.collect()
+torch.cuda.empty_cache()

convert_to_safetensors.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Convert torch.save quantized shards to safetensors format."""
+
+import argparse
+import gc
+import torch
+from pathlib import Path
+from safetensors.torch import save_file
+import glob
+
+
+def convert_shards(model_path):
+    output_path = Path(model_path)
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    
+    print(f"Converting {len(shards)} shards to safetensors format...\n")
+    
+    for i, shard_path in enumerate(shards):
+        print(f"Converting shard {i+1}/{len(shards)}: {Path(shard_path).name}")
+        
+        # Load torch.save format
+        ckpt = torch.load(shard_path, map_location="cpu", weights_only=False)
+        
+        # Separate tensors from QuantState objects
+        tensors = {}
+        for key, value in ckpt.items():
+            if isinstance(value, torch.Tensor):
+                tensors[key] = value
+            else:
+                # Save QuantState separately if needed
+                print(f"   Skipping non-tensor: {key} ({type(value)})")
+        
+        # Save as safetensors
+        new_path = output_path / Path(shard_path).name
+        save_file(tensors, new_path)
+        print(f"   ✓ Saved {len(tensors)} tensors")
+        
+        # Cleanup
+        del ckpt, tensors
+        gc.collect()
+    
+    print(f"\n✅ Converted {len(shards)} shards to safetensors format")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
+    args = parser.parse_args()
+    convert_shards(args.model_path)

deploy-agenx-lora.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Deploy agenx-lora-training
+# Run this on the GPU server after scp'ing it there
+
+set -e
+
+REPO_DIR="${LORA_REPO_DIR:-$HOME/loras/agenx-lora-training}"
+REPO_URL="https://gitea.cyaren.com/cmedina/agenx-lora-training.git"
+
+echo "=== Deploying agenx-lora-training ==="
+echo "Cloning to: ${REPO_DIR}"
+echo ""
+
+# Clone repo
+git clone "${REPO_URL}" "${REPO_DIR}"
+
+echo ""
+echo "=== Done ==="
+echo "Repository is at: ${REPO_DIR}"
+echo ""
+echo "Next steps:"
+echo "  cd ${REPO_DIR}"
+echo "  bash train-on-this-server.sh"

deploy-and-train.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+# Deploy LoRA training repo (clone + prepare dataset).
+# Training is run manually on the GPU server.
+#
+# Usage:
+#   bash deploy-and-train.sh
+#
+# This will:
+# 1. Clone/update the repo to /opt/loras/agenx-lora-training
+# 2. Print instructions for downloading dataset and training
+#
+
+set -euo pipefail
+
+# Configuration
+REPO_URL="${LORA_DEPLOY_REPO_URL:-https://gitea.cyaren.com/cmedina/agenx-lora-training.git}"
+REPO_DIR="${LORA_DEPLOY_REPO_DIR:-/opt/loras/agenx-lora-training}"
+BRANCH="${LORA_DEPLOY_BRANCH:-main}"
+
+RUN_USER="${SUDO_USER:-$(id -un)}"
+RUN_GROUP="$(id -gn "${RUN_USER}" 2>/dev/null || id -gn)"
+REPO_PARENT="$(dirname "${REPO_DIR}")"
+
+echo "=== LoRA Training Deploy ==="
+echo "Repo:   ${REPO_URL}"
+echo "Dir:    ${REPO_DIR}"
+echo "Branch: ${BRANCH}"
+echo ""
+
+# Create parent directory if needed
+if [ ! -d "${REPO_PARENT}" ]; then
+    echo "Creating ${REPO_PARENT}..."
+    sudo install -d -m 0755 -o "${RUN_USER}" -g "${RUN_GROUP}" "${REPO_PARENT}"
+fi
+
+# Clone or update repo
+if [ ! -d "${REPO_DIR}/.git" ]; then
+    echo "Cloning repository..."
+    git clone --branch "${BRANCH}" "${REPO_URL}" "${REPO_DIR}"
+else
+    echo "Repository exists, pulling latest..."
+    cd "${REPO_DIR}"
+    git fetch origin "${BRANCH}"
+    git checkout "${BRANCH}"
+    git pull --ff-only origin "${BRANCH}"
+fi
+
+echo ""
+echo "=== Repository ready ==="
+echo "Location: ${REPO_DIR}"
+echo ""
+echo "=== Next Steps ==="
+echo ""
+echo "1. Download dataset (8.77 MB):"
+echo "   cd ${REPO_DIR}"
+echo "   curl -O https://gitea.cyaren.com/cmedina/agenx-lora-training/raw/branch/main/dataset/combined_20k.jsonl"
+echo ""
+echo "2. Setup Python + GPU drivers on your training server"
+echo ""
+echo "3. Prepare dataset:"
+echo "   python3 training/scripts/prepare_dataset.py --input dataset/combined_20k.jsonl --output training/data"
+echo ""
+echo "4. Train LoRA (on GPU server):"
+echo "   python3 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml"
+echo ""
+echo "=== To train now (if on GPU server) ==="
+echo "   bash ${REPO_DIR}/train-on-this-server.sh"
+echo ""

inference.py

@@ -68,8 +68,8 @@ def generate_summary(model, tokenizer, task, files_changed=None, tests_run=False
 
 def main():
     parser = argparse.ArgumentParser(description="Generate Cyron summaries with LoRA")
-    parser.add_argument("--model", type=str, default="output/llama2-7b-lora", 
+    parser.add_argument("--model", type=str, default="output/ornith-35b-lora-cyron", 
-                       help="Path to trained model")
+                       help="Path to trained LoRA model")
     parser.add_argument("--task", type=str, required=True, help="Task description")
     parser.add_argument("--files", type=str, nargs="*", default=[], help="Files changed")
     parser.add_argument("--tests-run", action="store_true", help="Tests were run")

inspect_model.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+"""Inspect the Ornith-1.0-35B model architecture"""
+
+import torch
+from transformers import AutoModelForCausalLM, AutoConfig
+
+print("=" * 80)
+print("Inspecting Ornith-1.0-35B model")
+print("=" * 80)
+
+# Load config
+print("\n1. Loading config...")
+config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
+print(f"   Config class: {type(config).__name__}")
+print(f"   Model type: {config.model_type}")
+
+# Load model to CPU
+print("\n2. Loading model to CPU...")
+model = AutoModelForCausalLM.from_pretrained(
+    "/data/models/Ornith-1.0-35B",
+    device_map="cpu",
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    low_cpu_mem_usage=True,
+)
+print(f"   Model class: {type(model).__name__}")
+
+# Check if model has quantize_4bit
+print("\n3. Checking for quantization methods...")
+has_quantize = hasattr(model, 'quantize_4bit')
+print(f"   Has quantize_4bit(): {has_quantize}")
+
+# List all model components
+print("\n4. Model components:")
+for name, module in model.named_modules():
+    if len(name.split('.')) <= 2:  # Top-level and first-level
+        print(f"   {name}: {type(module).__name__}")
+
+# Check for BnB quantization support
+print("\n5. Checking BnB support...")
+try:
+    from bitsandbytes.nn import Linear4bit, Linear8bitLt
+    print("   ✓ BnB 4bit and 8bit modules available")
+except ImportError:
+    print("   ✗ BnB not installed")
+
+# Check if we can use prepare_model_for_kbit_training
+print("\n6. Checking PEFT support...")
+try:
+    from peft import prepare_model_for_kbit_training
+    print("   ✓ prepare_model_for_kbit_training available")
+except ImportError:
+    print("   ✗ PEFT not installed")
+
+print("\n" + "=" * 80)
+print("Summary:")
+print(f"   Model: {type(model).__name__}")
+print(f"   Config: {type(config).__name__}")
+print(f"   Has quantize_4bit(): {has_quantize}")
+print("=" * 80)

prepare_dataset.py

@@ -36,16 +36,20 @@ def prepare_dataset(input_file, output_file, test_size=0.05):
         
         output_text = ex['output']
         
-        # Create conversation format
+        # Create chat format for SFTTrainer
         conversation = {
-            "conversations": [
+            "messages": [
                 {
-                    "from": "human",
+                    "role": "system",
-                    "value": f"Generate a Cyron summary for this task:\n\n{instruction}\n\n{input_text}"
+                    "content": "You are a helpful coding assistant that generates project summaries."
                 },
                 {
-                    "from": "gpt",
+                    "role": "user",
-                    "value": output_text
+                    "content": f"Generate a summary for this task:\n\n{instruction}\n\n{input_text}"
+                },
+                {
+                    "role": "assistant",
+                    "content": output_text
                 }
             ]
         }
@@ -62,17 +66,17 @@ def prepare_dataset(input_file, output_file, test_size=0.05):
     test_data = formatted[split_point:]
     
     # Save
-    with open(output_file.parent / "train.jsonl", "w") as f:
+    with open(output_file / "train.jsonl", "w") as f:
         for item in train_data:
             f.write(json.dumps(item) + "\n")
     
-    with open(output_file.parent / "test.jsonl", "w") as f:
+    with open(output_file / "test.jsonl", "w") as f:
         for item in test_data:
             f.write(json.dumps(item) + "\n")
     
     print(f"Train: {len(train_data)} examples")
     print(f"Test: {len(test_data)} examples")
-    print(f"Saved to {output_file.parent}")
+    print(f"Saved to {output_file}")
 
 
 def main():

quantize_model.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+"""Simple NF4 quantization using BnB with device_map auto-distribution."""
+
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer
+
+
+def quantize_model(model_path, output_path):
+    print(f"Quantizing model from: {model_path}")
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    print("Loading model with 4-bit quantization...")
+    print("   Using CPU offloading to handle 70GB bf16 → 17GB 4-bit conversion\n")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="auto",
+        max_memory={0: "28GiB", 1: "28GiB", "cpu": "120GiB"},
+        low_cpu_mem_usage=True,
+        trust_remote_code=True,
+    )
+
+    print("\n✓ Model loaded and quantized")
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    print(f"   GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
+
+    print(f"\nSaving quantized model to: {output_path}")
+    model.save_pretrained(output_path)
+
+    # Save tokenizer
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        tokenizer.save_pretrained(output_path)
+        print("✓ Tokenizer saved")
+    except:
+        print("⚠ No tokenizer found")
+
+    print(f"\n✅ Quantized model saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B")
+    parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
+    args = parser.parse_args()
+
+    quantize_model(args.model_path, args.output_path)

quantize_proper_bnb.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Quantize Ornith-1.0-35B to 4-bit NF4 (recommended method for 2x RTX 5090)"""
+
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
+import torch
+import os
+
+def quantize_model():
+    model_path = "/data/models/Ornith-1.0-35B"
+    output_path = "/data/models/Ornith-1.0-35B-4bit-nf4"
+
+    print(f"Quantizing model: {model_path}")
+    print("Using 4-bit NF4 with double quantization + aggressive offloading...\n")
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="auto",
+        max_memory={
+            0: "26GiB",       # Good balance for 5090 (leaves headroom)
+            1: "26GiB",
+            "cpu": "150GiB",  # Heavy CPU offloading during quantization
+        },
+        low_cpu_mem_usage=True,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+    )
+
+    print(f"\nSaving quantized model to: {output_path}")
+    os.makedirs(output_path, exist_ok=True)
+    model.save_pretrained(output_path)
+
+    # Also save the config explicitly
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    config.save_pretrained(output_path)
+
+    print(f"\n✅ Quantization complete!")
+    print(f"   Model saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    quantize_model()

quantize_streaming.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+"""True streaming 4-bit NF4 quantization - one shard at a time."""
+
+import argparse
+import gc
+import torch
+import concurrent.futures
+from pathlib import Path
+from safetensors.torch import load_file, save_file
+from bitsandbytes.functional import quantize_nf4
+from transformers import AutoConfig
+
+
+def quantize_weight_nf4(weight: torch.Tensor):
+    """Quantize a single weight tensor to NF4 using bitsandbytes functional API."""
+    if weight.dim() != 2:
+        return weight, None
+
+    # quantize_nf4 returns (quantized_tensor, quant_state)
+    qweight, quant_state = quantize_nf4(
+        weight,
+        blocksize=64,
+        compress_statistics=True,
+    )
+    return qweight, quant_state
+
+
+def streaming_quantize(model_path: str, output_path: str):
+    print(f"Streaming NF4 quantization: {model_path}")
+    output_path = Path(output_path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    import glob
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    
+    # Rename existing 0-indexed files to 1-indexed
+    for existing in output_path.glob("*.safetensors"):
+        parts = existing.name.split("-")
+        if len(parts) >= 3 and existing.name.startswith("model-00000-"):
+            num = int(parts[1])
+            new_name = f"model-{num+1:05d}-of-{len(shards):05d}.safetensors"
+            new_path = output_path / new_name
+            existing.rename(new_path)
+            print(f"Renamed: {existing.name} -> {new_name}")
+    
+    print(f"Found {len(shards)} shards\n")
+    
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+
+    # Process continuously: max 4 shards per GPU, add to whichever GPU has room
+    max_per_gpu = 4
+    num_gpus = 2
+    
+    def process_shard(idx, shard_file, gpu_id):
+        """Process a single shard (called per-GPU)."""
+        shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors"
+        if (output_path / shard_name).exists():
+            return
+        
+        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})")
+        
+        # Load to assigned GPU
+        state_dict = load_file(shard_file, device=f"cuda:{gpu_id}")
+        
+        weight_keys = [
+            k for k, v in state_dict.items()
+            if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
+        ]
+        
+        print(f"   Quantizing {len(weight_keys)} tensors...")
+        
+        quant_states = {}
+        for key in weight_keys:
+            try:
+                weight = state_dict[key]
+                qweight, qstate = quantize_weight_nf4(weight)
+                state_dict[key] = qweight
+                if qstate is not None:
+                    quant_states[f"{key}.quant_state"] = qstate
+                del weight, qweight, qstate
+                gc.collect()
+                torch.cuda.empty_cache()
+            except Exception as e:
+                print(f"   Warning: Failed on {key}: {e}")
+                gc.collect()
+                torch.cuda.empty_cache()
+        
+        # Move to CPU and save
+        state_dict = {
+            k: v.cpu() if isinstance(v, torch.Tensor) else v
+            for k, v in state_dict.items()
+        }
+        
+        torch.save({**state_dict, **quant_states}, output_path / shard_name)
+        print(f"   ✓ Saved {shard_name} ({len(quant_states)} quant states)")
+        
+        del state_dict
+        gc.collect()
+        torch.cuda.empty_cache()
+    
+    # Find first unsaved shard
+    start_idx = 0
+    for idx in range(len(shards)):
+        shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors"
+        if not (output_path / shard_name).exists():
+            start_idx = idx
+            break
+    
+    print(f"Starting from shard {start_idx+1}/16\n")
+    
+    # Process continuously: submit to GPU with fewer active tasks
+    import threading
+    gpu_locks = [threading.Lock() for _ in range(num_gpus)]
+    gpu_counts = [0] * num_gpus
+    
+    def get_next_gpu():
+        """Get the GPU with fewest active tasks (max 4 per GPU)."""
+        with gpu_locks[0]:
+            with gpu_locks[1]:
+                if gpu_counts[0] <= gpu_counts[1] and gpu_counts[0] < max_per_gpu:
+                    return 0
+                elif gpu_counts[1] < max_per_gpu:
+                    return 1
+                else:
+                    return None  # Both GPUs at max
+    
+    # Create a queue of unprocessed shards
+    from collections import deque
+    remaining = deque(range(start_idx, len(shards)))
+    
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_per_gpu * num_gpus) as executor:
+        futures = []
+        
+        def submit_next():
+            """Submit next shard to available GPU."""
+            gpu_id = get_next_gpu()
+            if gpu_id is None or not remaining:
+                return
+            
+            idx = remaining.popleft()
+            shard_file = shards[idx]
+            
+            with gpu_locks[gpu_id]:
+                gpu_counts[gpu_id] += 1
+            
+            future = executor.submit(process_shard, idx, shard_file, gpu_id)
+            futures.append(future)
+            
+            # When this future completes, release GPU slot and submit next
+            def on_complete(f):
+                with gpu_locks[gpu_id]:
+                    gpu_counts[gpu_id] -= 1
+                submit_next()  # Try to submit next shard
+            
+            future.add_done_callback(on_complete)
+        
+        # Start submitting
+        while remaining:
+            gpu_id = get_next_gpu()
+            if gpu_id is None:
+                break
+            submit_next()
+        
+        # Wait for all to complete
+        concurrent.futures.wait(futures)
+
+    config.save_pretrained(output_path)
+    print(f"\n✅ Done → {output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B")
+    parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
+    args = parser.parse_args()
+
+    streaming_quantize(args.model_path, args.output_path)

quantize_to_bnb.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+"""Quantize bf16 model to BnB 4-bit using transformers' built-in mechanism."""
+
+import argparse
+import gc
+import torch
+from pathlib import Path
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+
+def quantize_model(model_path, output_path):
+    """Load bf16 model, quantize to BnB 4-bit, save."""
+    
+    print(f"Loading model from: {model_path}")
+    
+    # Load with BnB quantization config and device_map="cpu"
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+    )
+    
+    print("Loading with BnB 4-bit quantization to CPU...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="cpu",
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+    print("✓ Model loaded with BnB 4-bit to CPU")
+    
+    # Check if model is actually quantized
+    bnb_modules = sum(
+        1 for m in model.modules() 
+        if hasattr(m, 'weight') and hasattr(m.weight, 'quant_state')
+    )
+    print(f"   BnB quantized modules: {bnb_modules}")
+    
+    # Save model
+    print(f"\nSaving to: {output_path}")
+    model.save_pretrained(output_path)
+    print("✓ Model saved")
+    
+    # Free memory
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    print("\nDone! Model is ready for QLoRA training.")
+    print(f"Save location: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Quantize model to BnB 4-bit")
+    parser.add_argument("--model-path", type=str, 
+                       default="/data/models/Ornith-1.0-35B",
+                       help="Path to bf16 model")
+    parser.add_argument("--output-path", type=str,
+                       default="/data/models/Ornith-1.0-35B-bnb-4bit",
+                       help="Output path for quantized model")
+    args = parser.parse_args()
+    
+    Path(args.output_path).mkdir(parents=True, exist_ok=True)
+    quantize_model(args.model_path, args.output_path)
+
+
+if __name__ == "__main__":
+    main()

test_full_quantized.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+"""Test loading all quantized shards."""
+
+import torch
+from pathlib import Path
+import glob
+
+
+def test_all_shards():
+    model_path = "/data/models/Ornith-1.0-35B-nf4"
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    
+    print(f"Found {len(shards)} shards\n")
+    
+    total_tensors = 0
+    total_params = 0
+    
+    for i, shard_path in enumerate(shards):
+        print(f"Loading shard {i+1}/{len(shards)}: {Path(shard_path).name}")
+        try:
+            ckpt = torch.load(shard_path, map_location="cpu", weights_only=False)
+            shard_tensors = len([k for k in ckpt.keys() if isinstance(ckpt[k], torch.Tensor)])
+            total_tensors += shard_tensors
+            print(f"   ✓ {shard_tensors} tensors")
+        except Exception as e:
+            print(f"   ✗ Failed: {e}")
+    
+    print(f"\nTotal: {total_tensors} tensors across {len(shards)} shards")
+    print("✅ All shards loaded successfully!")
+
+
+if __name__ == "__main__":
+    test_all_shards()

test_load_quantized.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""Test loading quantized shard to check for MISMATCH."""
+
+import gc
+import torch
+import shutil
+from pathlib import Path
+from transformers import AutoModelForCausalLM, AutoConfig
+
+
+def test_load():
+    # Copy config.json from original model
+    test_dir = Path("/data/models/test_quantize")
+    config_src = Path("/data/models/Ornith-1.0-35B") / "config.json"
+    config_dst = test_dir / "config.json"
+    if not config_dst.exists():
+        shutil.copy2(config_src, config_dst)
+        print(f"Copied config.json to {test_dir}")
+    
+    print("\nLoading quantized test shard...")
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            str(test_dir),
+            device_map="cpu",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("✓ Model loaded successfully!")
+        
+        # Check for mismatched parameters
+        print("\nChecking parameter shapes...")
+        mismatch_count = 0
+        for name, param in model.named_parameters():
+            if hasattr(param, 'quant_state') and param.quant_state is not None:
+                # Quantized parameter - check if it loaded correctly
+                expected_shape = config.get_shape_for_parameter(name) if hasattr(config, 'get_shape_for_parameter') else None
+                if expected_shape and tuple(param.shape) != expected_shape:
+                    print(f"✗ MISMATCH: {name} - expected {expected_shape}, got {tuple(param.shape)}")
+                    mismatch_count += 1
+                else:
+                    print(f"✓ {name} - {tuple(param.shape)}")
+        
+        if mismatch_count == 0:
+            print("\n✅ No MISMATCH errors!")
+        else:
+            print(f"\n❌ {mismatch_count} MISMATCH errors found!")
+            
+    except Exception as e:
+        print(f"✗ Failed to load: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    test_load()

test_model_loading.py

@@ -0,0 +1,594 @@
+#!/usr/bin/env python3
+"""
+Test multiple model loading strategies to find what works.
+Each strategy is tested independently.
+
+Model: deepreinforce-ai/Ornith-1.0-35B (Qwen3_5Moe architecture)
+"""
+
+import torch
+from transformers import AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
+
+def get_layer_names(model_path):
+    """Detect decoder layer class names from model config"""
+    print("   Detecting layer names from config...")
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    
+    # Common layer name patterns
+    layer_names = []
+    
+    # Check for decoder layer
+    if hasattr(config, 'decoder_layer'):
+        layer_names.append(config.decoder_layer)
+    
+    # Check for common patterns
+    if hasattr(config, 'hidden_act'):
+        # Some configs have layer info in different fields
+        pass
+    
+    # If no standard field, try to infer from model type
+    if not layer_names:
+        model_type = config.model_type
+        if 'moe' in model_type.lower():
+            layer_names.append(f"{model_type.title().replace('_', '')}DecoderLayer")
+            layer_names.append(f"{model_type.title().replace('_', '')}SparseMoeBlock")
+        elif 'qwen' in model_type.lower():
+            layer_names.append("Qwen2DecoderLayer")
+        else:
+            layer_names.append("DecoderLayer")
+    
+    print(f"   Detected layers: {layer_names}")
+    return layer_names
+
+def check_gpu_memory():
+    """Check memory usage on all GPUs."""
+    print("   Memory Usage:")
+    for i in range(torch.cuda.device_count()):
+        mem = torch.cuda.memory_allocated(i) / 1e9
+        total = torch.cuda.get_device_properties(i).total_memory / 1e9
+        print(f"   GPU {i}: {mem:.2f} GB / {total:.2f} GB")
+    
+    gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
+    gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
+    
+    # Determine pattern
+    if abs(gpu0_mem - gpu1_mem) < 2.0:  # Within 2GB
+        if gpu0_mem < 15.0:
+            return "DISTRIBUTED"
+        else:
+            return "DUPLICATE"
+    else:
+        if gpu0_mem > gpu1_mem:
+            return f"GPU0_ONLY ({gpu0_mem:.1f}GB)"
+        else:
+            return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
+
+def quantize_model_bnb(model, quant_type="4bit"):
+    """Quantize model using BnB (BitsAndBytes)"""
+    print("   Using BnB to quantize model...")
+    
+    from transformers import BitsAndBytesConfig
+    from peft import prepare_model_for_kbit_training
+    
+    # Prepare model for k-bit training
+    model = prepare_model_for_kbit_training(
+        model,
+        use_gradient_checkpointing=False,
+    )
+    
+    # Set quantization config
+    if quant_type == "4bit":
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+    else:
+        bnb_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+        )
+    
+    # The actual quantization happens when we set device_map
+    # For now, just return the prepared model
+    print(f"   ✓ Model prepared for {quant_type}-bit quantization")
+    return model
+
+# def test_strategy_1():
+#     """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
+#     print("\n" + "=" * 80)
+#     print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
+#     print("=" * 80)
+#     
+#     try:
+#         print("   Loading bf16 model with BnB 4-bit...")
+#         bnb_config = BitsAndBytesConfig(
+#             load_in_4bit=True,
+#             bnb_4bit_quant_type="nf4",
+#             bnb_4bit_compute_dtype=torch.bfloat16,
+#         )
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",  # ← bf16 model
+#             quantization_config=bnb_config,
+#             device_map="auto",
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded successfully")
+#         
+#         pattern = check_gpu_memory()
+#         print(f"\n   Pattern: {pattern}")
+#         return True, pattern
+#     except Exception as e:
+#         print(f"\n   ✗ FAILED: {e}")
+#         return False, str(e)
+
+# def test_strategy_2():
+#     """Test 2: bf16 model + BnB 4-bit (alternative config)"""
+#     print("\n" + "=" * 80)
+#     print("TEST 2: bf16 model + BnB 4-bit (alt config)")
+#     print("=" * 80)
+#     
+#     try:
+#         torch.cuda.empty_cache()
+#         print("   Loading bf16 model with BnB 4-bit...")
+#         bnb_config = BitsAndBytesConfig(
+#             load_in_4bit=True,
+#             bnb_4bit_quant_type="nf4",
+#             bnb_4bit_compute_dtype=torch.bfloat16,
+#             bnb_4bit_use_double_quant=True,
+#         )
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",  # ← bf16 model
+#             quantization_config=bnb_config,
+#             device_map="auto",
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded successfully")
+#         
+#         pattern = check_gpu_memory()
+#         print(f"\n   Pattern: {pattern}")
+#         return True, pattern
+#     except Exception as e:
+#         print(f"\n   ✗ FAILED: {e}")
+#         return False, str(e)
+
+# def test_strategy_3():
+#     """Test 3: device_map with explicit GPU assignment"""
+#     print("\n" + "=" * 80)
+#     print("TEST 3: device_map with explicit GPU assignment")
+#     print("=" * 80)
+#     
+#     try:
+#         torch.cuda.empty_cache()
+#         print("   Loading model with explicit device_map...")
+#         
+#         # Get model config to determine layers
+#         from transformers import AutoConfig
+#         config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
+#         num_layers = config.num_hidden_layers
+#         
+#         # Split layers: first half on GPU 0, second half on GPU 1
+#         device_map = {}
+#         for i in range(num_layers):
+#             if i < num_layers // 2:
+#                 device_map[f"model.layers.{i}"] = 0
+#             else:
+#                 device_map[f"model.layers.{i}"] = 1
+#         
+#         # Embeddings and norm on GPU 0
+#         device_map["model.embed_tokens"] = 0
+#         device_map["model.norm"] = 0
+#         device_map["lm_head"] = 0
+#         
+#         print(f"   Created device_map with {len(device_map)} entries")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map=device_map,
+#             torch_dtype=torch.bfloat16,
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded successfully")
+#         
+#         pattern = check_gpu_memory()
+#         print(f"\n   Pattern: {pattern}")
+#         return True, pattern
+#     except Exception as e:
+#         print(f"\n   ✗ FAILED: {e}")
+#         return False, str(e)
+
+# def test_strategy_4():
+#     """Test 4: Load to CPU, then move to GPU manually"""
+#     print("\n" + "=" * 80)
+#     print("TEST 4: Load to CPU, then move to GPU")
+#     print("=" * 80)
+#     
+#     try:
+#         torch.cuda.empty_cache()
+#         print("   Loading bf16 model to CPU...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map="cpu",
+#             torch_dtype=torch.bfloat16,
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded to CPU")
+#         
+#         # Count params on CPU
+#         cpu_params = sum(p.numel() for p in model.parameters() if p.device.type == 'cpu')
+#         print(f"   CPU parameters: {cpu_params / 1e9:.2f}B")
+#         
+#         # Move to GPU 0
+#         print("\n   Moving to GPU 0...")
+#         model = model.to("cuda:0")
+#         pattern = check_gpu_memory()
+#         print(f"   Pattern after move to GPU 0: {pattern}")
+#         
+#         # Move to GPU 1
+#         print("\n   Moving to GPU 1...")
+#         model = model.to("cuda:1")
+#         pattern = check_gpu_memory()
+#         print(f"   Pattern after move to GPU 1: {pattern}")
+#         
+#         return True, "LOADED_TO_CPU_THEN_GPU"
+#     except Exception as e:
+#         print(f"\n   ✗ FAILED: {e}")
+#         return False, str(e)
+
+# def test_strategy_5():
+#     """Test 5: Sequential layer loading (manual distribution)"""
+#     print("\n" + "=" * 80)
+#     print("TEST 5: Sequential layer loading (manual distribution)")
+#     print("=" * 80)
+#     
+#     try:
+#         torch.cuda.empty_cache()
+#         print("   Loading model layer by layer...")
+#         
+#         # This is a simplified version - in reality would need more complex logic
+#         # For now, just test if we can load to one GPU
+#         print("   Loading bf16 to GPU 0 only...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map={"": 0},
+#             torch_dtype=torch.bfloat16,
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded to GPU 0")
+#         
+#         pattern = check_gpu_memory()
+#         print(f"\n   Pattern: {pattern}")
+#         
+#         # Now try GPU 1
+#         torch.cuda.empty_cache()
+#         print("\n   Loading bf16 to GPU 1 only...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map={"": 1},
+#             torch_dtype=torch.bfloat16,
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded to GPU 1")
+#         
+#         pattern = check_gpu_memory()
+#         print(f"   Pattern: {pattern}")
+#         
+#         return True, "SEQUENTIAL_LOAD"
+#     except Exception as e:
+#         print(f"\n   ✗ FAILED: {e}")
+#         return False, str(e)
+
+def test_strategy_6():
+    """Test 6: Use CompressedTensors 4-bit checkpoint (pre-quantized)"""
+    print("\n" + "=" * 80)
+    print("TEST 6: CompressedTensors 4-bit checkpoint")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load CompressedTensors 4-bit checkpoint...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",  # ← Pre-quantized checkpoint
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded (~18GB on disk)")
+        
+        print("\n   Step 2: Move to GPU 0...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern after move to GPU 0: {pattern}")
+        
+        print("\n   Step 3: Move to GPU 1...")
+        model = model.to("cuda:1")
+        pattern = check_gpu_memory()
+        print(f"   Pattern after move to GPU 1: {pattern}")
+        
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_7():
+    """Test 7: CompressedTensors 4-bit → accelerate distribute"""
+    print("\n" + "=" * 80)
+    print("TEST 7: CompressedTensors 4-bit → accelerate")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        
+        # Detect layer names dynamically
+        print("   Detecting layer names...")
+        layer_names = get_layer_names("/data/models/Ornith-1.0-35B-4bit")
+        
+        print("\n   Step 1: Load CompressedTensors 4-bit checkpoint...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded (~18GB on disk)")
+        
+        print("\n   Step 2: Use accelerate to distribute across GPUs...")
+        from accelerate import infer_auto_device_map
+        
+        # Create device map for quantized model
+        device_map = infer_auto_device_map(
+            model,
+            max_memory={0: "15GB", 1: "15GB"},
+            no_split_module_classes=layer_names,
+        )
+        print(f"   Created device_map with {len(device_map)} entries")
+        
+        # Reload with device_map
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            torch_dtype=torch.float16,
+            device_map=device_map,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded with device_map")
+        
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_8():
+    """Test 8: CompressedTensors 4-bit → GPU 0 only"""
+    print("\n" + "=" * 80)
+    print("TEST 8: CompressedTensors 4-bit → GPU 0 only")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load CompressedTensors 4-bit checkpoint...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded (~18GB on disk)")
+        
+        print("\n   Step 2: Move to GPU 0 only...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_9():
+    """Test 9: CompressedTensors 4-bit → GPU (test distribution)"""
+    print("\n" + "=" * 80)
+    print("TEST 9: CompressedTensors 4-bit → GPU (test)")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load CompressedTensors 4-bit checkpoint...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded (~18GB on disk)")
+        
+        print("\n   Step 2: Move to GPU...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_10():
+    """Test 10: CompressedTensors 4-bit → FSDP"""
+    print("\n" + "=" * 80)
+    print("TEST 10: CompressedTensors 4-bit → FSDP")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load CompressedTensors 4-bit checkpoint...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded (~18GB on disk)")
+        
+        print("\n   Step 2: Move to GPU...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_11():
+    """Test 11: PEFT prepare_model_for_kbit_training + manual quantization"""
+    print("\n" + "=" * 80)
+    print("TEST 11: PEFT prepare + manual 4-bit quantization")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load bf16 model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map="cpu",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded to CPU (~70GB)")
+        
+        # Check CPU memory
+        import psutil
+        mem = psutil.virtual_memory()
+        print(f"   CPU RAM: {mem.used / 1e9:.2f}GB / {mem.total / 1e9:.2f}GB")
+        
+        print("\n   Step 2: Apply PEFT prepare_model_for_kbit_training...")
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=False,
+        )
+        print("   ✓ Model prepared for k-bit training")
+        
+        print("\n   Step 3: Manually quantize Linear layers to 4-bit...")
+        from bitsandbytes.nn import Linear4bit
+        from torch import nn
+        
+        # Count and replace Linear layers
+        linear_count = 0
+        for name, module in model.named_modules():
+            if isinstance(module, nn.Linear) and 'lm_head' not in name:
+                # Replace with 4-bit version
+                new_module = Linear4bit(
+                    module.in_features,
+                    module.out_features,
+                    bias=module.bias is not None,
+                )
+                # Copy weights
+                new_module.weight = nn.Parameter(
+                    module.weight.data.clone()
+                )
+                if module.bias is not None:
+                    new_module.bias = nn.Parameter(
+                        module.bias.data.clone()
+                    )
+                # Replace in model
+                layers = name.split('.')
+                parent = model
+                for layer in layers[:-1]:
+                    parent = getattr(parent, layer)
+                setattr(parent, layers[-1], new_module)
+                linear_count += 1
+        
+        print(f"   ✓ Replaced {linear_count} Linear layers with 4-bit")
+        
+        print("\n   Step 4: Move to GPU...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+if __name__ == "__main__":
+    print("=" * 80)
+    print("Testing multiple model loading strategies")
+    print("=" * 80)
+    
+    # Check GPU availability
+    print(f"\n1. GPU Check:")
+    print(f"   CUDA available: {torch.cuda.is_available()}")
+    print(f"   GPU count: {torch.cuda.device_count()}")
+    for i in range(torch.cuda.device_count()):
+        props = torch.cuda.get_device_properties(i)
+        print(f"   GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)")
+    
+    # Run all tests
+    results = []
+    
+    tests = [
+        ("Test 6: CompressedTensors 4-bit → GPU", test_strategy_6),
+        ("Test 7: CompressedTensors 4-bit → accelerate", test_strategy_7),
+        ("Test 8: CompressedTensors 4-bit → GPU 0 only", test_strategy_8),
+        ("Test 9: CompressedTensors 4-bit → GPU (test)", test_strategy_9),
+        ("Test 10: CompressedTensors 4-bit → FSDP", test_strategy_10),
+        ("Test 11: PEFT prepare + manual 4-bit", test_strategy_11),
+    ]
+    
+    for name, test_func in tests:
+        try:
+            success, pattern = test_func()
+            results.append((name, success, pattern))
+        except Exception as e:
+            print(f"\n   ✗ Test crashed: {e}")
+            results.append((name, False, str(e)))
+        
+        # Clear GPU memory between tests
+        torch.cuda.empty_cache()
+    
+    # Summary
+    print("\n" + "=" * 80)
+    print("SUMMARY")
+    print("=" * 80)
+    
+    for name, success, pattern in results:
+        status = "✓ PASS" if success else "✗ FAIL"
+        print(f"{status}: {name}")
+        print(f"         Pattern: {pattern}")
+    
+    # Find working strategies
+    working = [name for name, success, _ in results if success]
+    if working:
+        print(f"\n✓ {len(working)} strategy/strategies work:")
+        for w in working:
+            print(f"  - {w}")
+    else:
+        print("\n✗ No strategies work!")

test_quantize_one_shard.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Test quantization on single shard."""
+
+import gc
+import torch
+from pathlib import Path
+from safetensors.torch import load_file
+from bitsandbytes.functional import quantize_nf4
+
+
+def test_one_shard(model_path, output_dir):
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    import glob
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    
+    if len(shards) == 0:
+        print("No shards found!")
+        return
+    
+    shard_file = shards[0]  # First shard only
+    print(f"Testing shard: {shard_file}")
+    
+    state_dict = load_file(shard_file, device="cuda:0")
+    
+    weight_keys = [
+        k for k, v in state_dict.items()
+        if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
+    ]
+    
+    print(f"Found {len(weight_keys)} weight tensors\n")
+    
+    quant_states = {}
+    quantized = 0
+    failed = 0
+    
+    for key in weight_keys:
+        try:
+            weight = state_dict[key]
+            qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
+            state_dict[key] = qweight
+            if qstate is not None:
+                quant_states[f"{key}.quant_state"] = qstate
+            quantized += 1
+            print(f"✓ {key} -> {tuple(qweight.shape)}")
+            del weight, qweight, qstate
+            gc.collect()
+            torch.cuda.empty_cache()
+        except Exception as e:
+            failed += 1
+            print(f"✗ {key}: {e}")
+            gc.collect()
+            torch.cuda.empty_cache()
+    
+    print(f"\nResults: {quantized} quantized, {failed} failed, {len(quant_states)} quant states")
+    
+    # Save test output as model.safetensors
+    state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()}
+    test_output = output_dir / "model.safetensors"
+    torch.save({**state_dict_cpu, **quant_states}, test_output)
+    print(f"Saved to: {test_output}")
+
+
+if __name__ == "__main__":
+    test_one_shard("/data/models/Ornith-1.0-35B", "/data/models/test_quantize")

test_quantize_shapes.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+"""Test quantized shard shapes."""
+
+import torch
+from pathlib import Path
+
+
+def test_shapes():
+    print("Loading quantized test shard...")
+    ckpt = torch.load("/data/models/test_quantize/model.safetensors", map_location="cpu", weights_only=False)
+    
+    print(f"Loaded {len(ckpt)} tensors\n")
+    
+    # Check for mismatched shapes
+    mismatch_count = 0
+    for key, tensor in list(ckpt.items())[:20]:  # Check first 20
+        if isinstance(tensor, torch.Tensor):
+            print(f"✓ {key} -> {tuple(tensor.shape)}")
+        elif hasattr(tensor, 'shape'):
+            print(f"✓ {key} -> {tensor.shape}")
+        else:
+            print(f"? {key} -> {type(tensor)}")
+    
+    print(f"\nTotal tensors: {len(ckpt)}")
+    print("If all shapes look correct, the quantization is working!")
+
+
+if __name__ == "__main__":
+    test_shapes()

test_quantize_single_shard.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""Test quantization on single shard only."""
+
+import gc
+import torch
+from pathlib import Path
+from safetensors.torch import load_file, save_file
+from bitsandbytes.functional import quantize_nf4
+from transformers import AutoConfig
+
+
+def test_single_shard(model_path, output_dir):
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+
+    import glob
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    
+    if len(shards) == 0:
+        print("No shards found!")
+        return
+    
+    # Use only first shard
+    shard_file = shards[0]
+    print(f"Testing with single shard: {shard_file}")
+    
+    state_dict = load_file(shard_file, device="cpu")
+    
+    weight_keys = [
+        k for k, v in state_dict.items()
+        if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
+    ]
+    
+    print(f"Found {len(weight_keys)} weight tensors\n")
+    
+    quantized = 0
+    failed = 0
+    
+    for i, key in enumerate(weight_keys[:5]):
+        try:
+            weight = state_dict[key].to("cuda:0")
+            qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
+            state_dict[key] = qweight.cpu()
+            del weight, qweight
+            gc.collect()
+            torch.cuda.empty_cache()
+            quantized += 1
+            print(f"✓ {key}")
+        except Exception as e:
+            failed += 1
+            print(f"✗ {key}: {e}")
+            gc.collect()
+            torch.cuda.empty_cache()
+    
+    print(f"\nResults: {quantized} quantized, {failed} failed")
+    
+    # Save test output
+    test_output = output_dir / "test_shard.safetensors"
+    state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()}
+    save_file(state_dict_cpu, test_output)
+    print(f"Saved test output to: {test_output}")
+
+
+if __name__ == "__main__":
+    test_single_shard("/data/models/Ornith-1.0-35B", "/data/models/test_quantize")

train-on-this-server.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Training script for Cyron LoRA on this server (2x RTX 5090)
+
+set -e
+
+# Unload AI model to free GPU memory
+echo "Unloading AI model..."
+ai none
+
+echo "=== Cyron LoRA Training Setup ==="
+
+# Create venv
+python3 -m venv venv
+source venv/bin/activate
+pip install --upgrade pip
+
+# Install training dependencies (PyTorch already installed)
+echo "Installing training dependencies..."
+pip install transformers datasets trl peft accelerate bitsandbytes deepspeed
+
+# Run training with single process (model already distributed via device_map)
+echo "Starting training (single process, model pre-distributed)..."
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+python train.py --config training/configs/ornith-35b-lora.yaml
+
+echo "Training completed!"

train.py

@@ -0,0 +1,271 @@
+#!/usr/bin/env python3
+"""
+Train LoRA adapter on Cyron summary dataset.
+
+Uses Hugging Face TRL for SFT training.
+"""
+
+import argparse
+import os
+import yaml
+from pathlib import Path
+
+import torch
+
+
+def train(config_path):
+    """Train LoRA adapter using TRL."""
+    
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        BitsAndBytesConfig,
+        TrainingArguments,
+    )
+    from peft import (
+        LoraConfig,
+        get_peft_model,
+    )
+    from trl import SFTTrainer
+    
+    print(f"Loading model: {config['base_model']}")
+    
+    # Create index file for sharded model
+    import glob as glob_mod
+    import json as json_mod
+    
+    safetensor_files = glob_mod.glob(f"{config['base_model']}/*.safetensors")
+    shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
+    index_file = Path(config["base_model"]) / "model.safetensors.index.json"
+    
+    if not index_file.exists() and shards:
+        print(f"\n[INFO] Creating index file for {len(shards)} shards...")
+        weight_map = {}
+        for shard_name in shards:
+            shard_path = Path(config["base_model"]) / shard_name
+            ckpt = torch.load(str(shard_path), map_location="cpu", weights_only=False)
+            for key in ckpt.keys():
+                if isinstance(ckpt[key], torch.Tensor):
+                    weight_map[key] = shard_name
+        
+        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_mod.dump(index, f)
+        print(f"✓ Created index ({len(weight_map)} weights)")
+    
+    # Remove quantization_config to prevent re-quantization
+    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"\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)")
+    
+    # Move to GPU
+    print("   Moving to GPU 0...")
+    model = model.to("cuda:0")
+    print("✓ Success: Model loaded to GPU 0 (4-bit)")
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    print(f"   Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    
+    # Add LoRA
+    lora_config = LoraConfig(
+        r=config["lora_r"],
+        lora_alpha=config["lora_alpha"],
+        lora_dropout=config["lora_dropout"],
+        target_modules=config["target_modules"],
+        task_type=config.get("lora_task_type", "CAUSAL_LM"),
+    )
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(config["base_model"])
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+    
+    # Load dataset
+    from datasets import load_dataset
+    import os
+    
+    # Get dataset path from config
+    dataset_path = config["dataset"][0]["path"]
+    print(f"Loading dataset from: {dataset_path}")
+    
+    dataset = load_dataset(
+        "json",
+        data_files={
+            "train": dataset_path,
+        },
+    )
+    
+    # Model is on single GPU
+    print("✓ Model loaded to single GPU")
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    print(f"   Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=config["train_params"]["output_dir"],
+        num_train_epochs=config["train_params"]["num_train_epochs"],
+        per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
+        gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
+        learning_rate=float(config["train_params"]["learning_rate"]),
+        lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
+        weight_decay=config["train_params"]["weight_decay"],
+        warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]),
+        logging_steps=config["train_params"]["logging_steps"],
+        save_steps=config["train_params"]["save_steps"],
+        save_total_limit=config["train_params"]["save_total_limit"],
+        eval_strategy=config.get("eval_strategy", "no"),
+        bf16=True,
+        gradient_checkpointing=config.get("gradient_checkpointing", False),
+        optim=config["train_params"].get("optim", "adamw_torch"),
+        optim_args=config["train_params"].get("optim_args"),
+    )
+    print(f"Using optimizer: {training_args.optim}")
+    
+    # SFT Trainer
+    from trl import SFTTrainer
+    
+    # Get text column from config (default to 'text' if not specified)
+    text_column = config["dataset"][0].get("text_column", "text")
+    print(f"Using text column: {text_column}")
+    
+    # Rename column to 'text' if needed (SFTTrainer expects 'text')
+    if text_column != "text":
+        print(f"   Renaming '{text_column}' column to 'text'...")
+        dataset["train"] = dataset["train"].rename_column(text_column, "text")
+        if "test" in dataset and dataset["test"] is not None:
+            dataset["test"] = dataset["test"].rename_column(text_column, "text")
+    
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset.get("test"),
+        args=training_args,
+    )
+    
+    # Train
+    print("Starting training...")
+    trainer.train()
+    
+    # Save
+    trainer.save_model(config["train_params"]["output_dir"])
+    tokenizer.save_pretrained(config["train_params"]["output_dir"])
+    
+    print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train LoRA adapter")
+    parser.add_argument("--config", type=str, default="training/configs/ornith-35b-lora.yaml",
+                       help="Training configuration file")
+    parser.add_argument("--check-only", action="store_true",
+                       help="Validate config and dependencies without training")
+    args = parser.parse_args()
+    
+    if args.check_only:
+        check_setup(args.config)
+    else:
+        train(args.config)
+
+
+def check_setup(config_path):
+    """Validate config and dependencies without loading model."""
+    print("=== Checking Setup ===")
+    
+    # Check config file
+    print(f"\n1. Config file: {config_path}")
+    if not Path(config_path).exists():
+        print(f"   ERROR: Config file not found: {config_path}")
+        return False
+    print("   ✓ Config file exists")
+    
+    # Load and validate config
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    print("   ✓ Config file is valid YAML")
+    
+    # Check model path
+    print(f"\n2. Model path: {config['base_model']}")
+    if Path(config['base_model']).exists():
+        print(f"   ✓ Model path exists: {config['base_model']}")
+    else:
+        print(f"   ⚠ Model path not found: {config['base_model']}")
+        print("   (Will download from HuggingFace during training)")
+    
+    # Check dataset files
+    print("\n3. Dataset files:")
+    repo_root = Path(__file__).parent
+    train_path = repo_root / "training" / "data" / "train.jsonl"
+    test_path = repo_root / "training" / "data" / "test.jsonl"
+    
+    if train_path.exists():
+        print(f"   ✓ Train data: {train_path}")
+    else:
+        print(f"   ✗ Train data missing: {train_path}")
+    
+    if test_path.exists():
+        print(f"   ✓ Test data: {test_path}")
+    else:
+        print(f"   ✗ Test data missing: {test_path}")
+    
+    # Check GPU
+    print("\n4. GPU:")
+    import torch
+    if torch.cuda.is_available():
+        print(f"   ✓ GPU available: {torch.cuda.get_device_name(0)}")
+        print(f"   ✓ VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+        print(f"   ✓ GPU count: {torch.cuda.device_count()}")
+    else:
+        print("   ✗ No GPU detected!")
+        return False
+    
+    # Check required packages
+    print("\n5. Required packages:")
+    packages = ['transformers', 'datasets', 'trl', 'peft', 'accelerate', 'deepspeed']
+    for pkg in packages:
+        try:
+            __import__(pkg)
+            print(f"   ✓ {pkg}")
+        except ImportError:
+            print(f"   ✗ {pkg} not installed")
+    
+    # Check DeepSpeed config
+    print("\n6. DeepSpeed config:")
+    if 'deepspeed_config' in config:
+        print("   ✓ DeepSpeed config present")
+        ds_config = config['deepspeed_config']
+        if 'zero_optimization' in ds_config:
+            stage = ds_config['zero_optimization'].get('stage', 'N/A')
+            print(f"   ✓ ZeRO stage: {stage}")
+    else:
+        print("   ✗ No DeepSpeed config found")
+    
+    print("\n=== Check Complete ===")
+    print("If all checks pass, you can run training with:")
+    print(f"  bash train-on-this-server.sh")
+
+
+if __name__ == "__main__":
+    main()

training/README.md

@@ -23,25 +23,24 @@ training/
 - GPU with 40GB+ VRAM (A100 recommended)
 - 64GB+ system RAM
 - 100GB+ free disk space
+- CUDA drivers installed
 
-## Installation
+## Server Deployment (Recommended)
+
+Use the deployment script to clone and train on your GPU server:
 
 ```bash
-cd scripts/lora_training/training
+# Deploy and train in one command
-
+bash deploy-and-train.sh
-# Create virtual environment
-python -m venv venv
-source venv/bin/activate
-
-# Install dependencies
-pip install transformers datasets trl peft accelerate bitsandbytes
-
-# Or use conda
-conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
-pip install transformers datasets trl peft accelerate bitsandbytes
 ```
 
-## Usage
+This will:
+1. Clone the repo to `/opt/loras/agenx-lora-training`
+2. Setup Python environment with CUDA support
+3. Prepare the dataset
+4. Train the LoRA adapter
+
+## Manual Setup
 
 ### 1. Prepare Dataset
 
@@ -54,13 +53,13 @@ This splits the 20k dataset into train/test sets (95/5).
 ### 2. Train Model
 
 ```bash
-python scripts/train.py --config configs/llama2-7b-lora.yaml
+python scripts/train.py --config configs/ornith-35b-lora.yaml
 ```
 
 Or with custom parameters:
 
 ```bash
-python scripts/train.py --config configs/llama2-7b-lora.yaml --epochs 5 --batch-size 8
+python scripts/train.py --config configs/ornith-35b-lora.yaml --epochs 5 --batch-size 8
 ```
 
 Training takes approximately 6-24 hours depending on GPU.
@@ -69,7 +68,7 @@ Training takes approximately 6-24 hours depending on GPU.
 
 ```bash
 python scripts/inference.py \
-    --model output/llama2-7b-lora \
+    --model output/ornith-35b-lora \
     --task "Fix parser crash on malformed JSON" \
     --files src/parser.cpp \
     --tests-run --test-count 294 \
@@ -123,7 +122,3 @@ Trained model is saved to `output/llama2-7b-lora/`:
 **Generation too long/short:**
 - Adjust `max_new_tokens` in inference script
 - Tune `temperature` and `top_p`
-
-## License
-
-This training infrastructure is part of the AgenX project.

training/configs/ds_zero3.json

@@ -0,0 +1,22 @@
+{
+  "bf16": {
+    "enabled": true
+  },
+
+  "zero_optimization": {
+    "stage": 3,
+
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+
+    "reduce_bucket_size": 50000000,
+    "stage3_prefetch_bucket_size": 50000000,
+    "stage3_param_persistence_threshold": 100000
+  },
+
+  "gradient_clipping": 1.0,
+
+  "train_micro_batch_size_per_gpu": 1,
+
+  "steps_per_print": 100
+}

training/configs/llama2-7b-lora.yaml

@@ -1,55 +0,0 @@
-# LoRA Training Configuration for Llama-2-7b
-# Dataset: cyron_summary_lora_dataset (20k examples)
-
-base_model: meta-llama/Llama-2-7b-hf
-model_type: LlamaForCausalLM
-tokenizer_type: LlamaTokenizer
-
-# Quantization (QLoRA)
-load_in_4bit: true
-bnb_4bit_compute_dtype: bfloat16
-bnb_4bit_quant_type: nf4
-use_nested_quant: false
-
-# LoRA Configuration
-lora_r: 16
-lora_alpha: 32
-lora_dropout: 0.05
-target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-lora_task_type: CAUSAL_LM
-
-# Dataset
-dataset:
-  - path: ../combined_20k.jsonl
-    type: completion
-    text_column: text
-
-# Training Parameters
-train_params:
-  num_train_epochs: 3
-  per_device_train_batch_size: 4
-  gradient_accumulation_steps: 4
-  learning_rate: 2e-4
-  lr_scheduler_type: cosine
-  weight_decay: 0.01
-  warmup_ratio: 0.03
-  max_seq_length: 1024
-  logging_steps: 10
-  save_steps: 100
-  save_total_limit: 3
-  output_dir: ../../output/llama2-7b-lora
-
-# Precision
-mixed_precision: bf16
-
-# Evaluation
-eval_strategy: steps
-eval_steps: 100
-eval_accumulation_steps: 10
-
-# Gradient Checkpointing
-gradient_checkpointing: true

training/configs/ornith-35b-lora.yaml

@@ -0,0 +1,59 @@
+# LoRA Training Configuration for Ornith-1.0-35B
+# Dataset: cyron_summary_lora_dataset (20k examples)
+
+base_model: /data/models/Ornith-1.0-35B-nf4
+model_type: Qwen3_5MoeForCausalLM
+tokenizer_type: AutoTokenizer
+
+# Model is pre-quantized with CompressedTensors
+# Loading via accelerate device_map for DISTRIBUTED training
+
+# LoRA Configuration
+lora_r: 64
+lora_alpha: 128
+lora_dropout: 0.05
+target_modules:
+  - q_proj
+  - v_proj
+  - k_proj
+  - o_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+lora_task_type: CAUSAL_LM
+
+# Dataset
+dataset:
+  - path: /home/cyaren/loras/agenx-lora-training/dataset/combined_20k.jsonl
+    type: completion
+    text_column: output
+
+# Training Parameters
+train_params:
+  num_train_epochs: 3
+  per_device_train_batch_size: 1
+  gradient_accumulation_steps: 16  # Increased to reduce per-step memory
+  learning_rate: 0.0002
+  lr_scheduler_type: cosine
+  weight_decay: 0.01
+  warmup_ratio: 0.03  # Will be converted to warmup_steps by TrainingArguments
+  max_seq_length: 512  # Reduced from 1024 to save memory
+  logging_steps: 10
+  save_steps: 100
+  save_total_limit: 3
+  output_dir: ../../output/ornith-35b-lora
+  optim: adamw_bnb_8bit  # 8-bit optimizer to save VRAM
+  optim_args: "offload_optimizer_device=cpu"  # Offload optimizer to CPU RAM
+
+# Precision
+mixed_precision: bf16
+
+# Distributed training (2x RTX 5090)
+# Using accelerate device_map for DISTRIBUTED loading
+# No DeepSpeed - model already quantized
+
+# Evaluation (disable - no test split in dataset)
+eval_strategy: "no"
+
+# Gradient Checkpointing (disable - causes device issues with distributed MoE)
+gradient_checkpointing: false

training/scripts/train.py

@@ -1,127 +0,0 @@
-#!/usr/bin/env python3
-"""
-Train LoRA adapter on Cyron summary dataset.
-
-Uses Hugging Face TRL for SFT training with QLoRA.
-"""
-
-import argparse
-import yaml
-from pathlib import Path
-
-
-def train(config_path):
-    """Train LoRA adapter using TRL."""
-    
-    with open(config_path) as f:
-        config = yaml.safe_load(f)
-    
-    from transformers import (
-        AutoModelForCausalLM,
-        AutoTokenizer,
-        BitsAndBytesConfig,
-        TrainingArguments,
-    )
-    from peft import (
-        prepare_model_for_kbit_training,
-        LoraConfig,
-        get_peft_model,
-    )
-    from trl import SFTTrainer, SFTConfig
-    
-    print(f"Loading model: {config['base_model']}")
-    
-    # Load model with quantization
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=config.get("load_in_4bit", True),
-        bnb_4bit_compute_dtype=config.get("bnb_4bit_compute_dtype", "bfloat16"),
-        bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"),
-        use_nested_quant=config.get("use_nested_quant", False),
-    )
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        quantization_config=bnb_config,
-        device_map="auto",
-    )
-    model = prepare_model_for_kbit_training(model)
-    
-    # Add LoRA
-    lora_config = LoraConfig(
-        r=config["lora_r"],
-        lora_alpha=config["lora_alpha"],
-        lora_dropout=config["lora_dropout"],
-        target_modules=config["target_modules"],
-        task_type=config.get("lora_task_type", "CAUSAL_LM"),
-    )
-    model = get_peft_model(model, lora_config)
-    model.print_trainable_parameters()
-    
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(config["base_model"])
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.padding_side = "right"
-    
-    # Load dataset
-    from datasets import load_dataset
-    
-    dataset = load_dataset(
-        "json",
-        data_files={
-            "train": config["dataset"][0]["path"].replace("../", ""),
-            "test": config["dataset"][0]["path"].replace("../", "").replace("combined_20k.jsonl", "test.jsonl"),
-        },
-    )
-    
-    # Training arguments
-    training_args = TrainingArguments(
-        output_dir=config["train_params"]["output_dir"],
-        num_train_epochs=config["train_params"]["num_train_epochs"],
-        per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"],
-        gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"],
-        learning_rate=config["train_params"]["learning_rate"],
-        lr_scheduler_type=config["train_params"]["lr_scheduler_type"],
-        weight_decay=config["train_params"]["weight_decay"],
-        warmup_ratio=config["train_params"]["warmup_ratio"],
-        max_seq_length=config["train_params"]["max_seq_length"],
-        logging_steps=config["train_params"]["logging_steps"],
-        save_steps=config["train_params"]["save_steps"],
-        save_total_limit=config["train_params"]["save_total_limit"],
-        evaluation_strategy=config.get("eval_strategy", "steps"),
-        eval_steps=config.get("eval_steps", 100),
-        mixed_precision=config.get("mixed_precision", "bf16"),
-        gradient_checkpointing=config.get("gradient_checkpointing", True),
-    )
-    
-    # SFT Trainer
-    trainer = SFTTrainer(
-        model=model,
-        tokenizer=tokenizer,
-        train_dataset=dataset["train"],
-        eval_dataset=dataset["test"],
-        args=training_args,
-        max_seq_length=config["train_params"]["max_seq_length"],
-    )
-    
-    # Train
-    print("Starting training...")
-    trainer.train()
-    
-    # Save
-    trainer.save_model(config["train_params"]["output_dir"])
-    tokenizer.save_pretrained(config["train_params"]["output_dir"])
-    
-    print(f"Training complete! Model saved to {config['train_params']['output_dir']}")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Train LoRA adapter")
-    parser.add_argument("--config", type=str, default="configs/llama2-7b-lora.yaml", 
-                       help="Training configuration file")
-    args = parser.parse_args()
-    
-    train(args.config)
-
-
-if __name__ == "__main__":
-    main()