Add proposer training job script

train_proposer_job.py

@@ -0,0 +1,58 @@
+"""
+HF Jobs script: Train cheap proposer (Qwen3-1.7B) for next-action prediction.
+Uses SFTTrainer with LoRA on speculative-actions-proposer-sft dataset.
+"""
+import torch
+from datasets import load_dataset
+from trl import SFTTrainer, SFTConfig
+from peft import LoraConfig
+
+MODEL = "Qwen/Qwen3-1.7B"
+DATASET = "narcolepticchicken/speculative-actions-proposer-sft"
+OUTPUT = "narcolepticchicken/speculative-proposer-qwen3-1.7b"
+
+def main():
+    ds = load_dataset(DATASET)
+    train_ds = ds["train"]
+    eval_ds = ds["test"]
+
+    peft_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        modules_to_save=["embed_tokens", "lm_head"],
+    )
+
+    config = SFTConfig(
+        output_dir="/tmp/proposer-out",
+        hub_model_id=OUTPUT,
+        push_to_hub=True,
+        learning_rate=2e-4,
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        num_train_epochs=3,
+        max_seq_length=4096,
+        bf16=True,
+        gradient_checkpointing=True,
+        logging_strategy="steps",
+        logging_steps=10,
+        logging_first_step=True,
+        disable_tqdm=True,
+        report_to="trackio",
+        run_name="proposer-sft-qwen3-1.7b",
+    )
+
+    trainer = SFTTrainer(
+        model=MODEL,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        args=config,
+        peft_config=peft_config,
+    )
+
+    trainer.train()
+    trainer.push_to_hub()
+    print("Proposer training complete.")
+
+if __name__ == "__main__":
+    main()