Fix: use Colab default torch, auto-detect TRL API, no version pins

KernelX_Training.ipynb

@@ -45,7 +45,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "!pip install -q \"sympy>=1.13,<1.14\" \"transformers==4.46.3\" \"trl==0.15.2\" \"peft>=0.13,<0.15\" datasets \"accelerate>=0.34,<0.36\" huggingface_hub\n\n# Restart runtime to pick up new sympy version (torch stays as Colab's default)\nimport os\nos.kill(os.getpid(), 9)"
+   "source": "!pip install -q -U sympy transformers trl peft datasets accelerate huggingface_hub\n\n# Restart runtime to pick up updated packages cleanly\nimport os\nos.kill(os.getpid(), 9)"
   },
   {
    "cell_type": "markdown",
@@ -245,45 +245,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# Train World Model\n",
-    "lora_config = LoraConfig(\n",
-    "    r=16, lora_alpha=32,\n",
-    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
-    "                     'gate_proj', 'up_proj', 'down_proj'],\n",
-    "    lora_dropout=0.05, bias='none', task_type='CAUSAL_LM',\n",
-    ")\n",
-    "\n",
-    "training_args = SFTConfig(\n",
-    "    output_dir='./world_model_checkpoints',\n",
-    "    num_train_epochs=2,\n",
-    "    per_device_train_batch_size=16,\n",
-    "    gradient_accumulation_steps=2,\n",
-    "    learning_rate=2e-4,\n",
-    "    lr_scheduler_type='cosine',\n",
-    "    warmup_ratio=0.1,\n",
-    "    logging_steps=10,\n",
-    "    eval_strategy='steps',\n",
-    "    eval_steps=100,\n",
-    "    save_total_limit=1,\n",
-    "    fp16=True,\n",
-    "    max_seq_length=512,\n",
-    "    report_to='none',\n",
-    ")\n",
-    "\n",
-    "trainer = SFTTrainer(\n",
-    "    model=model, args=training_args,\n",
-    "    train_dataset=train_ds, eval_dataset=val_ds,\n",
-    "    peft_config=lora_config,\n",
-    ")\n",
-    "\n",
-    "print('Training World Model...')\n",
-    "trainer.train()\n",
-    "\n",
-    "trainer.save_model('./world_model_final')\n",
-    "tokenizer.save_pretrained('./world_model_final')\n",
-    "print('World Model saved.')"
-   ]
+   "source": "# Train World Model\nimport inspect\n\nlora_config = LoraConfig(\n    r=16, lora_alpha=32,\n    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n                     'gate_proj', 'up_proj', 'down_proj'],\n    lora_dropout=0.05, bias='none', task_type='CAUSAL_LM',\n)\n\n# Auto-detect SFTConfig API (max_seq_length in older TRL, max_length in newer)\nsft_sig = inspect.signature(SFTConfig.__init__)\nseq_len_param = 'max_seq_length' if 'max_seq_length' in sft_sig.parameters else 'max_length'\nprint(f'TRL uses: {seq_len_param}')\n\ntraining_args = SFTConfig(\n    output_dir='./world_model_checkpoints',\n    num_train_epochs=2,\n    per_device_train_batch_size=16,\n    gradient_accumulation_steps=2,\n    learning_rate=2e-4,\n    lr_scheduler_type='cosine',\n    warmup_ratio=0.1,\n    logging_steps=10,\n    eval_strategy='steps',\n    eval_steps=100,\n    save_total_limit=1,\n    fp16=True,\n    report_to='none',\n    **{seq_len_param: 512},\n)\n\ntrainer = SFTTrainer(\n    model=model, args=training_args,\n    train_dataset=train_ds, eval_dataset=val_ds,\n    peft_config=lora_config,\n)\n\nprint('Training World Model...')\ntrainer.train()\n\ntrainer.save_model('./world_model_final')\ntokenizer.save_pretrained('./world_model_final')\nprint('World Model saved.')"
   },
   {
    "cell_type": "markdown",
@@ -359,45 +321,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# Reload fresh base model for Strategist\n",
-    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
-    "model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map='auto')\n",
-    "if tokenizer.pad_token is None:\n",
-    "    tokenizer.pad_token = tokenizer.eos_token\n",
-    "\n",
-    "lora_config = LoraConfig(\n",
-    "    r=16, lora_alpha=32,\n",
-    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
-    "                     'gate_proj', 'up_proj', 'down_proj'],\n",
-    "    lora_dropout=0.05, bias='none', task_type='CAUSAL_LM',\n",
-    ")\n",
-    "\n",
-    "ws_args = SFTConfig(\n",
-    "    output_dir='./strategist_warmstart',\n",
-    "    num_train_epochs=2,\n",
-    "    per_device_train_batch_size=16,\n",
-    "    gradient_accumulation_steps=2,\n",
-    "    learning_rate=2e-4,\n",
-    "    fp16=True,\n",
-    "    max_seq_length=512,\n",
-    "    logging_steps=5,\n",
-    "    save_total_limit=1,\n",
-    "    report_to='none',\n",
-    ")\n",
-    "\n",
-    "trainer = SFTTrainer(\n",
-    "    model=model, args=ws_args,\n",
-    "    train_dataset=ws_dataset, peft_config=lora_config,\n",
-    ")\n",
-    "\n",
-    "print('Training Strategist (warm-start)...')\n",
-    "trainer.train()\n",
-    "\n",
-    "trainer.save_model('./strategist_final')\n",
-    "tokenizer.save_pretrained('./strategist_final')\n",
-    "print('Strategist saved.')"
-   ]
+   "source": "# Reload fresh base model for Strategist\ntokenizer = AutoModelForCausalLM  # just to trigger import if needed\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map='auto')\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\nlora_config = LoraConfig(\n    r=16, lora_alpha=32,\n    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n                     'gate_proj', 'up_proj', 'down_proj'],\n    lora_dropout=0.05, bias='none', task_type='CAUSAL_LM',\n)\n\nws_args = SFTConfig(\n    output_dir='./strategist_warmstart',\n    num_train_epochs=2,\n    per_device_train_batch_size=16,\n    gradient_accumulation_steps=2,\n    learning_rate=2e-4,\n    fp16=True,\n    logging_steps=5,\n    save_total_limit=1,\n    report_to='none',\n    **{seq_len_param: 512},\n)\n\ntrainer = SFTTrainer(\n    model=model, args=ws_args,\n    train_dataset=ws_dataset, peft_config=lora_config,\n)\n\nprint('Training Strategist (warm-start)...')\ntrainer.train()\n\ntrainer.save_model('./strategist_final')\ntokenizer.save_pretrained('./strategist_final')\nprint('Strategist saved.')"
   },
   {
    "cell_type": "markdown",