File size: 13,744 Bytes

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🔐 LFM2.5 Unsloth QLoRA — Root-Fixed Kaggle/Colab Notebook\n",
    "\n",
    "This notebook keeps **Unsloth** for low VRAM, but fixes the repeated Kaggle error from the root:\n",
    "\n",
    "```text\n",
    "AttributeError: 'int' object has no attribute 'mean'\n",
    "```\n",
    "\n",
    "## Root fix\n",
    "\n",
    "The bug is caused by mismatched `unsloth` / `unsloth_zoo` / `transformers` / `trl` versions and Kaggle reusing `/kaggle/working/unsloth_compiled_cache`.\n",
    "\n",
    "This notebook now:\n",
    "\n",
    "1. Deletes Unsloth compiled cache.\n",
    "2. Installs/updates `unsloth` and `unsloth_zoo` together.\n",
    "3. Forces a kernel restart once after install.\n",
    "4. Uses current TRL `SFTConfig` API.\n",
    "5. Sets `trainer.model_accepts_loss_kwargs = False` to prevent `num_items_in_batch` leakage.\n",
    "\n",
    "Default model: `unsloth/LFM2.5-1.2B-Instruct` — best fit for free T4.\n",
    "\n",
    "> Use only for ethical, authorized cybersecurity education and defensive research.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 1. Clean install Unsloth stack — run first\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, shutil, subprocess, pathlib\n",
    "\n",
    "work_dir = pathlib.Path('/kaggle/working') if pathlib.Path('/kaggle/working').exists() else pathlib.Path('/content')\n",
    "marker = work_dir / '.bex_unsloth_env_ready_v3'\n",
    "\n",
    "# Always remove stale compiled trainer cache. This is where the old bug kept coming from.\n",
    "shutil.rmtree(str(work_dir / 'unsloth_compiled_cache'), ignore_errors=True)\n",
    "print('✅ Removed stale unsloth_compiled_cache')\n",
    "\n",
    "if not marker.exists():\n",
    "    print('Installing/updating Unsloth stack. Kernel will restart after this cell. Run all cells again after restart.')\n",
    "    subprocess.check_call([sys.executable, '-m', 'pip', 'uninstall', '-y', 'unsloth', 'unsloth_zoo'])\n",
    "    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-U', '--no-cache-dir', 'unsloth', 'unsloth_zoo'])\n",
    "    marker.write_text('ready')\n",
    "    print('✅ Installed Unsloth. Restarting kernel now...')\n",
    "    os.kill(os.getpid(), 9)\n",
    "else:\n",
    "    print('✅ Unsloth stack already prepared for this session')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 2. Optional Hugging Face login\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "# login(token='hf_YOUR_WRITE_TOKEN')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 3. Imports and version check\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel, is_bfloat16_supported\n",
    "from trl import SFTTrainer, SFTConfig\n",
    "from datasets import load_dataset, concatenate_datasets\n",
    "import torch, random, os, time\n",
    "import transformers, trl, peft, accelerate\n",
    "\n",
    "try:\n",
    "    import unsloth\n",
    "    print('unsloth', getattr(unsloth, '__version__', 'unknown'))\n",
    "except Exception as e:\n",
    "    print('unsloth version unavailable', e)\n",
    "print('transformers', transformers.__version__)\n",
    "print('trl', trl.__version__)\n",
    "print('peft', peft.__version__)\n",
    "print('accelerate', accelerate.__version__)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print('GPU:', torch.cuda.get_device_name(0))\n",
    "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB')\n",
    "else:\n",
    "    print('⚠️ No GPU found. Enable GPU runtime.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 4. Configuration\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_ID = 'unsloth/LFM2.5-1.2B-Instruct'\n",
    "RUN_NAME = 'lfm25-cyber-unsloth-rootfixed'\n",
    "\n",
    "DATASET_CHOICE = 'cybersecurity'  # cybersecurity | ultrachat | openhermes | code_corpus\n",
    "SAMPLE_SIZE = 50000\n",
    "MAX_SEQ_LENGTH = 4096\n",
    "\n",
    "LORA_R = 64\n",
    "LORA_ALPHA = 128\n",
    "BATCH_SIZE = 4\n",
    "GRAD_ACCUM = 2\n",
    "MAX_STEPS = 2000\n",
    "LEARNING_RATE = 2e-4\n",
    "WARMUP_STEPS = 100\n",
    "SAVE_STEPS = 500\n",
    "LOGGING_STEPS = 10\n",
    "PACKING = False  # keep false until your stack is stable\n",
    "SEED = 3407\n",
    "\n",
    "OUTPUT_DIR = './outputs_lfm25_rootfixed'\n",
    "HUB_MODEL_ID = 'your-username/lfm25-cyber-unsloth-rootfixed'\n",
    "PUSH_TO_HUB = False\n",
    "\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 5. Load model with Unsloth 4-bit QLoRA\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=MODEL_ID,\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    dtype=None,\n",
    "    load_in_4bit=True,\n",
    "    device_map={'': torch.cuda.current_device()} if torch.cuda.is_available() else None,\n",
    ")\n",
    "\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = 'right'\n",
    "\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r=LORA_R,\n",
    "    target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n",
    "    lora_alpha=LORA_ALPHA,\n",
    "    lora_dropout=0,\n",
    "    bias='none',\n",
    "    use_gradient_checkpointing='unsloth',\n",
    "    random_state=SEED,\n",
    "    use_rslora=True,\n",
    "    loftq_config=None,\n",
    ")\n",
    "\n",
    "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "total = sum(p.numel() for p in model.parameters())\n",
    "print(f'Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)')\n",
    "if torch.cuda.is_available():\n",
    "    print(f'VRAM after model load: {torch.cuda.memory_allocated()/1e9:.2f} GB')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 6. Load and format dataset\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_SAFE = 'You are a cybersecurity education assistant. Provide defensive, ethical, and authorized security guidance only. Refuse harmful or unauthorized requests.'\n",
    "\n",
    "def convert_sua(example):\n",
    "    return {'messages': [\n",
    "        {'role': 'system', 'content': example.get('system') or SYSTEM_SAFE},\n",
    "        {'role': 'user', 'content': example['user']},\n",
    "        {'role': 'assistant', 'content': example['assistant']},\n",
    "    ]}\n",
    "\n",
    "def convert_ultrachat(example):\n",
    "    return {'messages': example['messages']}\n",
    "\n",
    "def convert_conversations(example):\n",
    "    msgs = []\n",
    "    sys_prompt = example.get('system_prompt', '') or example.get('system', '')\n",
    "    if sys_prompt: msgs.append({'role': 'system', 'content': sys_prompt})\n",
    "    for turn in example['conversations']:\n",
    "        role = 'user' if turn.get('from') in ('human', 'user') else 'assistant'\n",
    "        msgs.append({'role': role, 'content': turn.get('value', '')})\n",
    "    return {'messages': msgs}\n",
    "\n",
    "def convert_code(example):\n",
    "    return {'messages': [\n",
    "        {'role': 'user', 'content': 'Explain and improve this code with attention to safety and correctness.'},\n",
    "        {'role': 'assistant', 'content': example['text']},\n",
    "    ]}\n",
    "\n",
    "if DATASET_CHOICE == 'cybersecurity':\n",
    "    ds1 = load_dataset('AlicanKiraz0/Cybersecurity-Dataset-Fenrir-v2.1', split='train')\n",
    "    ds1 = ds1.map(convert_sua, remove_columns=ds1.column_names)\n",
    "    ds2 = load_dataset('Trendyol/Trendyol-Cybersecurity-Instruction-Tuning-Dataset', split='train')\n",
    "    ds2 = ds2.map(convert_sua, remove_columns=ds2.column_names)\n",
    "    dataset = concatenate_datasets([ds1, ds2])\n",
    "elif DATASET_CHOICE == 'ultrachat':\n",
    "    dataset = load_dataset('HuggingFaceH4/ultrachat_200k', split='train_sft').map(convert_ultrachat, remove_columns=['prompt','prompt_id','messages'])\n",
    "elif DATASET_CHOICE == 'openhermes':\n",
    "    dataset = load_dataset('teknium/OpenHermes-2.5', split='train')\n",
    "    dataset = dataset.map(convert_conversations, remove_columns=dataset.column_names)\n",
    "elif DATASET_CHOICE == 'code_corpus':\n",
    "    dataset = load_dataset('krystv/code-corpus-llm-training', split='train')\n",
    "    dataset = dataset.map(convert_code, remove_columns=dataset.column_names)\n",
    "else:\n",
    "    raise ValueError(DATASET_CHOICE)\n",
    "\n",
    "print('Rows:', len(dataset))\n",
    "if SAMPLE_SIZE and len(dataset) > SAMPLE_SIZE:\n",
    "    dataset = dataset.shuffle(seed=SEED).select(range(SAMPLE_SIZE))\n",
    "    print('Subsampled:', len(dataset))\n",
    "\n",
    "print(dataset[0]['messages'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 7. Convert messages to text\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_text(example):\n",
    "    try:\n",
    "        text = tokenizer.apply_chat_template(example['messages'], tokenize=False, add_generation_prompt=False)\n",
    "    except Exception:\n",
    "        text = '\\n'.join([f\"<{m['role']}>\\n{m['content']}\\n</{m['role']}>\" for m in example['messages']]) + tokenizer.eos_token\n",
    "    return {'text': text}\n",
    "\n",
    "dataset = dataset.map(to_text, remove_columns=['messages'])\n",
    "print(dataset[0]['text'][:500])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 8. Train with current SFTConfig API\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = SFTConfig(\n",
    "    output_dir=OUTPUT_DIR,\n",
    "    dataset_text_field='text',\n",
    "    max_length=MAX_SEQ_LENGTH,\n",
    "    packing=PACKING,\n",
    "    per_device_train_batch_size=BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
    "    warmup_steps=WARMUP_STEPS,\n",
    "    max_steps=MAX_STEPS,\n",
    "    learning_rate=LEARNING_RATE,\n",
    "    fp16=not is_bfloat16_supported(),\n",
    "    bf16=is_bfloat16_supported(),\n",
    "    logging_steps=LOGGING_STEPS,\n",
    "    optim='adamw_8bit',\n",
    "    weight_decay=0.01,\n",
    "    lr_scheduler_type='linear',\n",
    "    seed=SEED,\n",
    "    save_strategy='steps',\n",
    "    save_steps=SAVE_STEPS,\n",
    "    save_total_limit=2,\n",
    "    report_to='none',\n",
    "    disable_tqdm=True,\n",
    "    logging_first_step=True,\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    processing_class=tokenizer,\n",
    "    train_dataset=dataset,\n",
    "    args=args,\n",
    ")\n",
    "\n",
    "# Root compatibility guard for Transformers num_items_in_batch API.\n",
    "# This is recommended for mixed stacks and prevents the old int.mean crash.\n",
    "trainer.model_accepts_loss_kwargs = False\n",
    "if hasattr(trainer, 'model'):\n",
    "    trainer.model.config.use_cache = False\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(f'VRAM before train: {torch.cuda.memory_allocated()/1e9:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB')\n",
    "\n",
    "trainer_stats = trainer.train()\n",
    "print(trainer_stats)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": ["## 9. Save adapter and test\n"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained('./lfm25-lora-adapter')\n",
    "tokenizer.save_pretrained('./lfm25-lora-adapter')\n",
    "print('Saved adapter to ./lfm25-lora-adapter')\n",
    "\n",
    "if PUSH_TO_HUB:\n",
    "    model.push_to_hub(HUB_MODEL_ID)\n",
    "    tokenizer.push_to_hub(HUB_MODEL_ID)\n",
    "    print('Pushed to', HUB_MODEL_ID)\n",
    "\n",
    "FastLanguageModel.for_inference(model)\n",
    "messages = [{'role':'system','content':SYSTEM_SAFE}, {'role':'user','content':'Explain parameterized queries for preventing SQL injection with safe Python.'}]\n",
    "inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors='pt').to(model.device)\n",
    "with torch.no_grad():\n",
    "    outputs = model.generate(input_ids=inputs, max_new_tokens=384, temperature=0.7, top_p=0.9, do_sample=True, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)\n",
    "print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True))\n"
   ]
  }
 ],
 "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10"}},
 "nbformat": 4,
 "nbformat_minor": 5
}