refactor: replace unsloth with plain transformers+peft for GRPO training

Drop unsloth/vllm dependency chain to eliminate huggingface-hub version
conflicts. Use AutoModelForCausalLM + BitsAndBytesConfig + LoraConfig
directly. Flash attention falls back to SDPA if flash-attn is missing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

training/notebooks/fusion_design_lab_training.ipynb

@@ -35,23 +35,13 @@
    "id": "9a63283cbaf04dbcab1f6479b197f3a8",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "%%capture\n",
-    "# Build deps for constellaration (booz-xform compiles from source)\n",
-    "!apt-get update -qq && apt-get install -y -qq cmake ninja-build g++ gfortran libnetcdf-dev libnetcdff-dev > /dev/null\n",
-    "\n",
-    "!pip install unsloth vllm\n",
-    "!pip install --no-deps trl\n",
-    "!pip install matplotlib requests"
-   ]
+   "source": "%%capture\n# Build deps for constellaration (booz-xform compiles from source)\n!apt-get update -qq && apt-get install -y -qq cmake ninja-build g++ gfortran libnetcdf-dev libnetcdff-dev > /dev/null\n\n!pip install trl peft bitsandbytes datasets matplotlib accelerate\n!pip install \"transformers>=4.51\" \"huggingface-hub<1.0\""
   },
   {
    "cell_type": "markdown",
    "id": "8dd0d8092fe74a7c96281538738b07e2",
    "metadata": {},
-   "source": [
-    "## 2. Load Model with Unsloth"
-   ]
+   "source": "## 2. Load Model with LoRA"
   },
   {
    "cell_type": "code",
@@ -59,7 +49,7 @@
    "id": "72eea5119410473aa328ad9291626812",
    "metadata": {},
    "outputs": [],
-   "source": "from unsloth import FastLanguageModel\n\nMODEL_NAME = \"unsloth/Qwen3.5-4B\"\nMAX_SEQ_LENGTH = 2048\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    model_name=MODEL_NAME,\n    max_seq_length=MAX_SEQ_LENGTH,\n    load_in_4bit=True,\n    fast_inference=True,\n)\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=32,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n    lora_alpha=32,\n    use_gradient_checkpointing=\"unsloth\",\n)\n\nprint(f\"Model loaded: {MODEL_NAME}\")"
+   "source": "import importlib\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\nfrom peft import LoraConfig, get_peft_model\n\nMODEL_NAME = \"Qwen/Qwen3.5-4B\"\nMAX_SEQ_LENGTH = 2048\n\nbnb_config = BitsAndBytesConfig(\n    load_in_4bit=True,\n    bnb_4bit_quant_type=\"nf4\",\n    bnb_4bit_use_double_quant=True,\n    bnb_4bit_compute_dtype=torch.bfloat16,\n)\n\nattn_impl = \"flash_attention_2\" if importlib.util.find_spec(\"flash_attn\") else \"sdpa\"\n\nmodel = AutoModelForCausalLM.from_pretrained(\n    MODEL_NAME,\n    quantization_config=bnb_config,\n    torch_dtype=torch.bfloat16,\n    device_map=\"auto\",\n    attn_implementation=attn_impl,\n)\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\nlora_config = LoraConfig(\n    r=32,\n    lora_alpha=32,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n    lora_dropout=0.0,\n    task_type=\"CAUSAL_LM\",\n)\nmodel = get_peft_model(model, lora_config)\nmodel.gradient_checkpointing_enable()\nmodel.print_trainable_parameters()\nprint(f\"Model loaded: {MODEL_NAME} (attn: {attn_impl})\")"
   },
   {
    "cell_type": "markdown",
@@ -368,99 +358,7 @@
    "id": "3ed186c9a28b402fb0bc4494df01f08d",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import random\n",
-    "\n",
-    "FastLanguageModel.for_inference(model)\n",
-    "\n",
-    "\n",
-    "def reward_term_summary(step_or_obs: object) -> str:\n",
-    "    breakdown_obj = getattr(step_or_obs, \"reward_breakdown\")\n",
-    "    breakdown = (\n",
-    "        breakdown_obj.model_dump() if hasattr(breakdown_obj, \"model_dump\") else breakdown_obj\n",
-    "    )\n",
-    "    terms = []\n",
-    "    for key, value in breakdown.items():\n",
-    "        if key in {\n",
-    "            \"intent\",\n",
-    "            \"total\",\n",
-    "            \"evaluation_failed\",\n",
-    "            \"recovered_from_failure\",\n",
-    "            \"reference_constraints_satisfied\",\n",
-    "            \"reference_score\",\n",
-    "            \"reference_feasibility\",\n",
-    "            \"reference_max_elongation\",\n",
-    "            \"initial_reference_score\",\n",
-    "            \"terminal_score_ratio\",\n",
-    "        }:\n",
-    "            continue\n",
-    "        if isinstance(value, (int, float)) and float(value) != 0.0:\n",
-    "            terms.append(f\"{key}={float(value):+.3f}\")\n",
-    "    return \", \".join(terms) if terms else \"none\"\n",
-    "\n",
-    "\n",
-    "def run_episode_with_model(seed_idx: int) -> tuple[float, list[str]]:\n",
-    "    \"\"\"Run one episode using the trained model.\"\"\"\n",
-    "    env = StellaratorEnvironment()\n",
-    "    obs = env.reset(seed=seed_idx)\n",
-    "    prompt = build_prompt(obs)\n",
-    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
-    "    outputs = model.generate(\n",
-    "        **inputs,\n",
-    "        max_new_tokens=MAX_COMPLETION_LENGTH,\n",
-    "        temperature=0.7,\n",
-    "        do_sample=True,\n",
-    "    )\n",
-    "    completion = tokenizer.decode(\n",
-    "        outputs[0][inputs[\"input_ids\"].shape[1] :], skip_special_tokens=True\n",
-    "    )\n",
-    "    actions = parse_action_plan(completion)\n",
-    "    episode = run_episode_with_actions(actions, seed_idx=seed_idx)\n",
-    "    trace = [\n",
-    "        (\n",
-    "            f\"{step.action_label} → reward={step.reward:.3f} \"\n",
-    "            f\"score={step.p1_score:.4f} feasible={step.constraints_satisfied} \"\n",
-    "            f\"terms={reward_term_summary(step)}\"\n",
-    "        )\n",
-    "        for step in episode.steps\n",
-    "    ]\n",
-    "    return episode.total_reward, trace\n",
-    "\n",
-    "\n",
-    "def run_random_episode(seed_idx: int) -> float:\n",
-    "    \"\"\"Run one episode with random actions for comparison.\"\"\"\n",
-    "    actions = [StellaratorAction(**random.choice(RUN_ACTION_SPECS)) for _ in range(BUDGET)]\n",
-    "    return run_episode_with_actions(actions, seed_idx=seed_idx).total_reward\n",
-    "\n",
-    "\n",
-    "# Evaluate\n",
-    "print(\"=\" * 60)\n",
-    "print(\"TRAINED MODEL EPISODES\")\n",
-    "print(\"=\" * 60)\n",
-    "trained_rewards = []\n",
-    "for seed in range(len(RESET_SEEDS)):\n",
-    "    reward, trace = run_episode_with_model(seed)\n",
-    "    trained_rewards.append(reward)\n",
-    "    print(f\"\\nSeed {seed} — Total reward: {reward:.3f}\")\n",
-    "    for line in trace:\n",
-    "        print(f\"  {line}\")\n",
-    "\n",
-    "print(f\"\\nMean trained reward: {sum(trained_rewards) / len(trained_rewards):.3f}\")\n",
-    "\n",
-    "print(\"\\n\" + \"=\" * 60)\n",
-    "print(\"RANDOM BASELINE (10 episodes per seed)\")\n",
-    "print(\"=\" * 60)\n",
-    "random_rewards = []\n",
-    "for seed in range(len(RESET_SEEDS)):\n",
-    "    seed_rewards = [run_random_episode(seed) for _ in range(10)]\n",
-    "    random_rewards.extend(seed_rewards)\n",
-    "    print(\n",
-    "        f\"Seed {seed} — Mean: {sum(seed_rewards) / len(seed_rewards):.3f}, Best: {max(seed_rewards):.3f}\"\n",
-    "    )\n",
-    "\n",
-    "print(f\"\\nMean random reward: {sum(random_rewards) / len(random_rewards):.3f}\")\n",
-    "print(f\"Mean trained reward: {sum(trained_rewards) / len(trained_rewards):.3f}\")"
-   ]
+   "source": "import random\n\nmodel.eval()\n\n\ndef reward_term_summary(step_or_obs: object) -> str:\n    breakdown_obj = getattr(step_or_obs, \"reward_breakdown\")\n    breakdown = (\n        breakdown_obj.model_dump() if hasattr(breakdown_obj, \"model_dump\") else breakdown_obj\n    )\n    terms = []\n    for key, value in breakdown.items():\n        if key in {\n            \"intent\",\n            \"total\",\n            \"evaluation_failed\",\n            \"recovered_from_failure\",\n            \"reference_constraints_satisfied\",\n            \"reference_score\",\n            \"reference_feasibility\",\n            \"reference_max_elongation\",\n            \"initial_reference_score\",\n            \"terminal_score_ratio\",\n        }:\n            continue\n        if isinstance(value, (int, float)) and float(value) != 0.0:\n            terms.append(f\"{key}={float(value):+.3f}\")\n    return \", \".join(terms) if terms else \"none\"\n\n\ndef run_episode_with_model(seed_idx: int) -> tuple[float, list[str]]:\n    \"\"\"Run one episode using the trained model.\"\"\"\n    env = StellaratorEnvironment()\n    obs = env.reset(seed=seed_idx)\n    prompt = build_prompt(obs)\n    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n    with torch.no_grad():\n        outputs = model.generate(\n            **inputs,\n            max_new_tokens=MAX_COMPLETION_LENGTH,\n            temperature=0.7,\n            do_sample=True,\n        )\n    completion = tokenizer.decode(\n        outputs[0][inputs[\"input_ids\"].shape[1] :], skip_special_tokens=True\n    )\n    actions = parse_action_plan(completion)\n    episode = run_episode_with_actions(actions, seed_idx=seed_idx)\n    trace = [\n        (\n            f\"{step.action_label} → reward={step.reward:.3f} \"\n            f\"score={step.p1_score:.4f} feasible={step.constraints_satisfied} \"\n            f\"terms={reward_term_summary(step)}\"\n        )\n        for step in episode.steps\n    ]\n    return episode.total_reward, trace\n\n\ndef run_random_episode(seed_idx: int) -> float:\n    \"\"\"Run one episode with random actions for comparison.\"\"\"\n    actions = [StellaratorAction(**random.choice(RUN_ACTION_SPECS)) for _ in range(BUDGET)]\n    return run_episode_with_actions(actions, seed_idx=seed_idx).total_reward\n\n\n# Evaluate\nprint(\"=\" * 60)\nprint(\"TRAINED MODEL EPISODES\")\nprint(\"=\" * 60)\ntrained_rewards = []\nfor seed in range(len(RESET_SEEDS)):\n    reward, trace = run_episode_with_model(seed)\n    trained_rewards.append(reward)\n    print(f\"\\nSeed {seed} — Total reward: {reward:.3f}\")\n    for line in trace:\n        print(f\"  {line}\")\n\nprint(f\"\\nMean trained reward: {sum(trained_rewards) / len(trained_rewards):.3f}\")\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"RANDOM BASELINE (10 episodes per seed)\")\nprint(\"=\" * 60)\nrandom_rewards = []\nfor seed in range(len(RESET_SEEDS)):\n    seed_rewards = [run_random_episode(seed) for _ in range(10)]\n    random_rewards.extend(seed_rewards)\n    print(\n        f\"Seed {seed} — Mean: {sum(seed_rewards) / len(seed_rewards):.3f}, Best: {max(seed_rewards):.3f}\"\n    )\n\nprint(f\"\\nMean random reward: {sum(random_rewards) / len(random_rewards):.3f}\")\nprint(f\"Mean trained reward: {sum(trained_rewards) / len(trained_rewards):.3f}\")"
   },
   {
    "cell_type": "markdown",
@@ -474,59 +372,7 @@
    "id": "379cbbc1e968416e875cc15c1202d7eb",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import requests\n",
-    "\n",
-    "from fusion_lab.client import FusionLabClient\n",
-    "\n",
-    "HF_SPACE_URL = \"https://creativeengineer-fusion-design-lab.hf.space\"\n",
-    "\n",
-    "# Check health\n",
-    "health = requests.get(f\"{HF_SPACE_URL}/health\").json()\n",
-    "print(f\"HF Space status: {health['status']}\")\n",
-    "\n",
-    "# Get task description\n",
-    "task = requests.get(f\"{HF_SPACE_URL}/task\").json()\n",
-    "print(f\"\\nTask: {task['description']}\")\n",
-    "print(f\"Constraints: {task['constraints']}\")\n",
-    "print(f\"Budget: {task['budget']}\")\n",
-    "\n",
-    "with FusionLabClient(base_url=HF_SPACE_URL) as env:\n",
-    "    reset_result = env.reset(seed=42)\n",
-    "    remote_obs = reset_result.observation\n",
-    "    print(f\"\\nRemote reset — max_elongation: {remote_obs.max_elongation:.4f}\")\n",
-    "    print(f\"  aspect_ratio: {remote_obs.aspect_ratio:.4f}\")\n",
-    "    print(f\"  constraints_satisfied: {remote_obs.constraints_satisfied}\")\n",
-    "    print(f\"  budget_remaining: {remote_obs.budget_remaining}\")\n",
-    "\n",
-    "    # Generate an action plan from the trained model\n",
-    "    prompt = build_prompt(remote_obs)\n",
-    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
-    "    outputs = model.generate(\n",
-    "        **inputs, max_new_tokens=MAX_COMPLETION_LENGTH, temperature=0.7, do_sample=True\n",
-    "    )\n",
-    "    completion = tokenizer.decode(\n",
-    "        outputs[0][inputs[\"input_ids\"].shape[1] :], skip_special_tokens=True\n",
-    "    )\n",
-    "    actions = parse_action_plan(completion)\n",
-    "\n",
-    "    print(f\"\\nTrained model generated {len(actions)} actions for remote env:\")\n",
-    "    for i, action in enumerate(actions[:BUDGET], start=1):\n",
-    "        if action.intent == \"submit\":\n",
-    "            continue\n",
-    "        result = env.step(action)\n",
-    "        step_obs = result.observation\n",
-    "        reward = float(result.reward) if result.reward is not None else 0.0\n",
-    "        print(\n",
-    "            f\"  Step {i}: {action.intent} {action.parameter or ''} \"\n",
-    "            f\"{action.direction or ''} {action.magnitude or ''} \"\n",
-    "            f\"→ reward={reward:.3f}, score={step_obs.p1_score:.4f}, terms={reward_term_summary(step_obs)}\"\n",
-    "        )\n",
-    "        if result.done:\n",
-    "            print(f\"  Episode done. Final score: {step_obs.p1_score:.4f}\")\n",
-    "            break\n",
-    "print(\"\\nEnvironment is live and accessible for training and evaluation.\")"
-   ]
+   "source": "import requests\n\nfrom fusion_lab.client import FusionLabClient\n\nHF_SPACE_URL = \"https://creativeengineer-fusion-design-lab.hf.space\"\n\n# Check health\nhealth = requests.get(f\"{HF_SPACE_URL}/health\").json()\nprint(f\"HF Space status: {health['status']}\")\n\n# Get task description\ntask = requests.get(f\"{HF_SPACE_URL}/task\").json()\nprint(f\"\\nTask: {task['description']}\")\nprint(f\"Constraints: {task['constraints']}\")\nprint(f\"Budget: {task['budget']}\")\n\nwith FusionLabClient(base_url=HF_SPACE_URL) as env:\n    reset_result = env.reset(seed=42)\n    remote_obs = reset_result.observation\n    print(f\"\\nRemote reset — max_elongation: {remote_obs.max_elongation:.4f}\")\n    print(f\"  aspect_ratio: {remote_obs.aspect_ratio:.4f}\")\n    print(f\"  constraints_satisfied: {remote_obs.constraints_satisfied}\")\n    print(f\"  budget_remaining: {remote_obs.budget_remaining}\")\n\n    # Generate an action plan from the trained model\n    prompt = build_prompt(remote_obs)\n    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n    with torch.no_grad():\n        outputs = model.generate(\n            **inputs, max_new_tokens=MAX_COMPLETION_LENGTH, temperature=0.7, do_sample=True\n        )\n    completion = tokenizer.decode(\n        outputs[0][inputs[\"input_ids\"].shape[1] :], skip_special_tokens=True\n    )\n    actions = parse_action_plan(completion)\n\n    print(f\"\\nTrained model generated {len(actions)} actions for remote env:\")\n    for i, action in enumerate(actions[:BUDGET], start=1):\n        if action.intent == \"submit\":\n            continue\n        result = env.step(action)\n        step_obs = result.observation\n        reward = float(result.reward) if result.reward is not None else 0.0\n        print(\n            f\"  Step {i}: {action.intent} {action.parameter or ''} \"\n            f\"{action.direction or ''} {action.magnitude or ''} \"\n            f\"→ reward={reward:.3f}, score={step_obs.p1_score:.4f}, terms={reward_term_summary(step_obs)}\"\n        )\n        if result.done:\n            print(f\"  Episode done. Final score: {step_obs.p1_score:.4f}\")\n            break\nprint(\"\\nEnvironment is live and accessible for training and evaluation.\")"
   }
  ],
  "metadata": {