File size: 3,434 Bytes
ec3854f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | # /// script
# requires-python = ">=3.10"
# dependencies = [
# "unsloth",
# "trl>=0.12.0",
# "transformers>=4.45.0",
# "datasets>=3.0.0",
# "peft>=0.13.0",
# "accelerate>=0.34.0",
# "trackio",
# "huggingface_hub>=0.25.0",
# ]
# ///
"""H1 cloud SFT — Qwen3-32B-Instruct-2507 + Unsloth + LoRA r=16 on ATC parser dataset.
Submit via:
hf jobs uv run --flavor a100-large --timeout 3h --secrets HF_TOKEN \\
"https://huggingface.co/<USER>/atc-parser-train-script/resolve/main/train_h1_cloud.py"
Goal: break M3-local 65% intent_em ceiling. Target 85%+.
Inputs (env, HF Jobs auto-injects HF_TOKEN):
HF_DATASET_REPO default kinglyai/atc-parser-spike-v0
HF_OUTPUT_REPO default kinglyai/qwen3-32b-atc-parser-v1
HF_TRACKIO_PROJECT default atc-parser
RUN_NAME default qwen3-32b-r16-h1-v1
"""
from __future__ import annotations
import os
from datasets import load_dataset
from peft import LoraConfig
from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig
import trackio
BASE_MODEL = os.environ.get("HF_BASE_MODEL", "Qwen/Qwen3-32B-Instruct-2507")
DATASET_REPO = os.environ.get("HF_DATASET_REPO", "kinglyai/atc-parser-spike-v0")
OUTPUT_REPO = os.environ.get("HF_OUTPUT_REPO", "kinglyai/qwen3-32b-atc-parser-v1")
TRACKIO_PROJECT = os.environ.get("HF_TRACKIO_PROJECT", "atc-parser")
RUN_NAME = os.environ.get("RUN_NAME", "qwen3-32b-r16-h1-v1")
def main() -> None:
trackio.init(project=TRACKIO_PROJECT, name=RUN_NAME)
# Load with Unsloth (60% less VRAM, 2× faster)
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=BASE_MODEL,
max_seq_length=1024,
dtype=None, # auto bf16
load_in_4bit=False, # use 8-bit or full bf16 on a100-80GB; LoRA only adapter
)
# LoRA r=16 — 2× capacity vs M3 spike's r=8
model = FastLanguageModel.get_peft_model(
model,
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"],
use_gradient_checkpointing="unsloth",
random_state=42,
)
# Load dataset (3 JSONL files in HF dataset repo)
ds = load_dataset(DATASET_REPO, data_files={
"train": "train.jsonl",
"valid": "valid.jsonl",
"test": "test.jsonl",
})
cfg = SFTConfig(
output_dir="/tmp/sft_out",
num_train_epochs=2,
per_device_train_batch_size=4,
per_device_eval_batch_size=4,
gradient_accumulation_steps=4,
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.05,
bf16=True,
eval_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=200,
save_total_limit=3,
logging_steps=10,
report_to="trackio",
run_name=RUN_NAME,
push_to_hub=True,
hub_model_id=OUTPUT_REPO,
hub_strategy="every_save",
hub_private_repo=True,
save_safetensors=True,
max_length=1024,
)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=ds["train"],
eval_dataset=ds["valid"],
args=cfg,
)
trainer.train()
trainer.save_model()
trainer.push_to_hub()
trackio.finish()
if __name__ == "__main__":
main()
|