Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,119 @@
+# /// script
+# dependencies = ["unsloth[colab-new]", "trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "xformers"]
+# ///
+"""
+Fine-tune FunctionGemma for llama-agent on HuggingFace Jobs.
+
+Submit with:
+    hf_jobs("uv", {
+        "script": "<this script content>",
+        "flavor": "a10g-large",
+        "timeout": "2h",
+        "secrets": {"HF_TOKEN": "$HF_TOKEN"}
+    })
+"""
+
+import os
+
+# Config - override via environment variables
+MODEL_NAME = os.environ.get("MODEL_NAME", "unsloth/functiongemma-270m-it")
+DATASET_NAME = os.environ.get("DATASET_NAME", "victor/functiongemma-agent-sft")
+OUTPUT_REPO = os.environ.get("OUTPUT_REPO", "victor/functiongemma-agent-finetuned")
+MAX_SEQ_LENGTH = int(os.environ.get("MAX_SEQ_LENGTH", "4096"))
+LORA_R = int(os.environ.get("LORA_R", "128"))
+LORA_ALPHA = int(os.environ.get("LORA_ALPHA", "256"))
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", "3"))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "4"))
+GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "2"))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "2e-4"))
+
+# Imports
+from unsloth import FastLanguageModel
+from unsloth.chat_templates import train_on_responses_only
+from datasets import load_dataset
+from trl import SFTTrainer, SFTConfig
+import trackio
+
+print(f"Loading model: {MODEL_NAME}")
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=MODEL_NAME,
+    max_seq_length=MAX_SEQ_LENGTH,
+    load_in_4bit=False,
+    load_in_8bit=False,
+    load_in_16bit=True,
+    full_finetuning=False,
+)
+
+print(f"Adding LoRA adapters (r={LORA_R}, alpha={LORA_ALPHA})")
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=LORA_R,
+    lora_alpha=LORA_ALPHA,
+    lora_dropout=0,
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+    ],
+    bias="none",
+    use_gradient_checkpointing="unsloth",
+    random_state=3407,
+    use_rslora=False,
+    loftq_config=None,
+)
+
+print(f"Loading dataset: {DATASET_NAME}")
+dataset = load_dataset(DATASET_NAME, split="train")
+print(f"Dataset size: {len(dataset)} examples")
+
+# SFTConfig with Trackio monitoring
+sft_config = SFTConfig(
+    dataset_text_field="text",
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    warmup_steps=5,
+    num_train_epochs=NUM_EPOCHS,
+    learning_rate=LEARNING_RATE,
+    logging_steps=10,
+    optim="adamw_8bit",
+    weight_decay=0.001,
+    lr_scheduler_type="linear",
+    seed=3407,
+    output_dir="./outputs",
+    save_steps=500,
+    save_total_limit=3,
+    max_seq_length=MAX_SEQ_LENGTH,
+    # Trackio monitoring
+    report_to="trackio",
+    run_name="functiongemma-agent-sft",
+    # Hub push (CRITICAL - environment is ephemeral!)
+    push_to_hub=True,
+    hub_model_id=OUTPUT_REPO,
+    hub_strategy="every_save",
+)
+
+# Create trainer
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=dataset,
+    eval_dataset=None,
+    args=sft_config,
+)
+
+# CRITICAL: Only train on model responses, not instructions
+print("Applying train_on_responses_only (masking instruction tokens)...")
+trainer = train_on_responses_only(
+    trainer,
+    instruction_part="<start_of_turn>user\n",
+    response_part="<start_of_turn>model\n",
+)
+
+print("Starting training...")
+trainer_stats = trainer.train()
+
+# Final push to hub
+print(f"Pushing final model to {OUTPUT_REPO}...")
+trainer.push_to_hub()
+
+print("Training complete!")
+print(f"Model saved to: https://huggingface.co/{OUTPUT_REPO}")