Drop unsloth: use standard bitsandbytes 4-bit + peft LoRA + TRL GRPOTrainer

Dockerfile

@@ -21,7 +21,6 @@ ENV CXX=/usr/bin/g++
 RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
-ENV LD_LIBRARY_PATH="/home/user/.local/lib/python3.10/site-packages/nvidia/nvjitlink/lib:/home/user/.local/lib/python3.10/site-packages/nvidia/cuda_runtime/lib:$LD_LIBRARY_PATH"
 
 WORKDIR /app
 
@@ -29,19 +28,13 @@ WORKDIR /app
 COPY --chown=user requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# 2. PyTorch 2.6.0 + CUDA 12.1
-RUN pip install --no-cache-dir torch==2.6.0 --extra-index-url https://download.pytorch.org/whl/cu121
+# 2. PyTorch with CUDA
+RUN pip install --no-cache-dir torch --extra-index-url https://download.pytorch.org/whl/cu121
 
-# 3. Training deps (no unsloth here)
+# 3. Training deps (standard stack, no unsloth)
 COPY --chown=user requirements-training.txt .
 RUN pip install --no-cache-dir -r requirements-training.txt
 
-# 4. Unsloth --no-deps (avoids torchao>=0.13 conflict)
-RUN pip install --no-cache-dir --no-deps unsloth==2025.11.1 unsloth_zoo==2025.11.1
-
-# 5. Remove torchao if pulled in (incompatible with torch 2.6, crashes transformers)
-RUN pip uninstall -y torchao 2>/dev/null; true
-
 # --- App code ---
 COPY --chown=user src/ /app/src/
 COPY --chown=user training/ /app/training/

requirements-training.txt

@@ -1,4 +1,4 @@
-# Training deps (torch, unsloth installed separately in Dockerfile)
+# Training deps — standard stack, no unsloth
 transformers>=4.46.0,<4.52
 trl>=0.12.0,<0.16
 peft>=0.13.0

run_training.py

@@ -58,26 +58,30 @@ def run_grpo_training():
     )
 
     # ── 1. Load model ──
-    print("\n[1/6] Loading model with Unsloth...")
-    try:
-        from unsloth import FastLanguageModel
-        MODEL_NAME = "unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit"
-        model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=MODEL_NAME, max_seq_length=2048, load_in_4bit=True,
-        )
-        model = FastLanguageModel.get_peft_model(
-            model, r=16, lora_alpha=16, lora_dropout=0,
-            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
-                            "gate_proj", "up_proj", "down_proj"],
-        )
-        print(f"  Model: {MODEL_NAME}")
-        print(f"  Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
-    except ImportError:
-        print("WARNING: Unsloth not available, using standard loading")
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    print("\n[1/6] Loading model with bitsandbytes 4-bit...")
+    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+    from peft import LoraConfig, get_peft_model
+
+    MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+        bnb_4bit_use_double_quant=True,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, quantization_config=bnb_config, device_map="auto",
+    )
+    lora_config = LoraConfig(
+        r=16, lora_alpha=16, lora_dropout=0,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                        "gate_proj", "up_proj", "down_proj"],
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora_config)
+    print(f"  Model: {MODEL_NAME}")
+    print(f"  Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
 
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -239,10 +243,7 @@ def run_grpo_training():
 
     # ── 5. Post-training evaluation ──
     print("\n[5/6] Evaluating trained model...")
-    try:
-        FastLanguageModel.for_inference(model)
-    except Exception:
-        pass
+    model.eval()
 
     def trained_generate(prompt):
         messages = [