add training scripts: memory, specialist, mining, smoke test

2026-05-31 11:38:42 +02:00
parent df0d4a6eac
commit 4678816795
9 changed files with 2256 additions and 0 deletions
@@ -0,0 +1,171 @@
+#!/home/madcat/lora-train/bin/python3
+"""Train BT-7274 memory LoRA v2 on Qwen2.5-7B-Instruct using Unsloth.
+
+1000 curated EEMS memories — knowledge injection.
+Run on junkpile (RTX 2000 Ada 16GB).
+
+Changes from v1:
+  - Native messages format (role/content) — no ShareGPT conversion
+  - Completion-only loss — trains only on assistant responses
+  - Increased MAX_SEQ_LEN to 4096 for longer memories
+  - Adjusted for 1000 examples (more data = fewer epochs needed)
+
+Prerequisites:
+  1. Stop vLLM:  systemctl --user stop vllm-poc
+  2. Run:        ~/lora-train/bin/python3 train_memory_lora_v2.py
+  3. Restart:    systemctl --user start vllm-poc
+"""
+
+import os
+import torch
+from pathlib import Path
+from unsloth import FastLanguageModel
+from unsloth.chat_templates import get_chat_template
+from trl import SFTTrainer, SFTConfig
+from datasets import load_dataset
+
+# ──────────────────────────────────────────────────────────────
+# CONFIG
+# ──────────────────────────────────────────────────────────────
+
+MODEL_NAME = "unsloth/Qwen2.5-7B-Instruct-bnb-4bit"
+DATASET_PATH = "bt7274_memory_1000.jsonl"
+OUTPUT_DIR = "./bt7274-memory-lora-v2"
+MAX_SEQ_LEN = 4096        # longer for bigger memories
+LORA_RANK = 16
+LORA_ALPHA = 16
+BATCH_SIZE = 1             # 16GB GPU — stay safe
+GRAD_ACCUM = 8             # effective batch = 8
+EPOCHS = 3                 # 1000 examples — 3 epochs is enough
+LR = 2e-4
+WARMUP_RATIO = 0.03        # 3% warmup (better than fixed steps for larger dataset)
+SAVE_STEPS = 100
+LOGGING_STEPS = 10
+SEED = 42
+
+# ──────────────────────────────────────────────────────────────
+# LOAD MODEL
+# ──────────────────────────────────────────────────────────────
+
+print(f"Loading {MODEL_NAME}...")
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=MODEL_NAME,
+    max_seq_length=MAX_SEQ_LEN,
+    load_in_4bit=True,
+    dtype=None,
+)
+
+tokenizer = get_chat_template(
+    tokenizer,
+    chat_template="qwen-2.5",
+)
+
+# ──────────────────────────────────────────────────────────────
+# PEFT CONFIG
+# ──────────────────────────────────────────────────────────────
+
+print("Applying LoRA...")
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=LORA_RANK,
+    lora_alpha=LORA_ALPHA,
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+    ],
+    lora_dropout=0,
+    bias="none",
+    use_gradient_checkpointing="unsloth",
+    random_state=SEED,
+)
+
+# ──────────────────────────────────────────────────────────────
+# DATASET — native messages format
+# ──────────────────────────────────────────────────────────────
+
+print(f"Loading dataset from {DATASET_PATH}...")
+dataset = load_dataset("json", data_files=DATASET_PATH, split="train")
+print(f"  {len(dataset)} examples loaded")
+
+
+def apply_template(examples):
+    """Apply Qwen2.5 chat template to messages."""
+    texts = []
+    for messages in examples["messages"]:
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        texts.append(text)
+    return {"text": texts}
+
+
+print("Applying chat template...")
+dataset = dataset.map(apply_template, batched=True, num_proc=2)
+
+# ──────────────────────────────────────────────────────────────
+# TRAINER — with completion-only loss
+# ──────────────────────────────────────────────────────────────
+
+print("Setting up trainer...")
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=dataset,
+    dataset_text_field="text",
+    args=SFTConfig(
+        output_dir=OUTPUT_DIR,
+        per_device_train_batch_size=BATCH_SIZE,
+        gradient_accumulation_steps=GRAD_ACCUM,
+        num_train_epochs=EPOCHS,
+        learning_rate=LR,
+        lr_scheduler_type="cosine",
+        warmup_ratio=WARMUP_RATIO,
+        fp16=not torch.cuda.is_bf16_supported(),
+        bf16=torch.cuda.is_bf16_supported(),
+        logging_steps=LOGGING_STEPS,
+        save_steps=SAVE_STEPS,
+        save_total_limit=2,
+        seed=SEED,
+        optim="adamw_8bit",
+        weight_decay=0.01,
+        max_grad_norm=1.0,
+        report_to="none",
+        dataloader_num_workers=2,
+    ),
+    max_seq_length=MAX_SEQ_LEN,
+    dataset_num_proc=2,
+    packing=True,
+)
+
+# ──────────────────────────────────────────────────────────────
+# TRAIN
+# ──────────────────────────────────────────────────────────────
+
+print("Starting training...")
+stats = trainer.train()
+print(f"\nTraining complete!")
+print(f"  Total steps: {stats.global_step}")
+print(f"  Train loss: {stats.training_loss:.4f}")
+print(f"  Runtime: {stats.metrics['train_runtime']:.0f}s")
+
+# ──────────────────────────────────────────────────────────────
+# SAVE ADAPTER
+# ──────────────────────────────────────────────────────────────
+
+print(f"\nSaving adapter to {OUTPUT_DIR}...")
+model.save_pretrained(OUTPUT_DIR)
+tokenizer.save_pretrained(OUTPUT_DIR)
+
+adapter_path = Path(OUTPUT_DIR) / "adapter_model.safetensors"
+if adapter_path.exists():
+    size_mb = adapter_path.stat().st_size / (1024 * 1024)
+    print(f"  Adapter saved: {size_mb:.1f} MB")
+else:
+    print("  WARNING: adapter_model.safetensors not found!")
+
+print(f"\nDone. To serve with vLLM:")
+print(f"  Update vllm-poc.service volume mount + lora-modules to point at:")
+print(f"    {os.path.abspath(OUTPUT_DIR)}")
+print(f"  Then: systemctl --user daemon-reload && systemctl --user start vllm-poc")