Stage 1 on HF Jobs: Unsloth job script, PEFT to MLX converter, base valid loss 0.849, Qwen base model docs

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

train/hf_train.py

@@ -0,0 +1,60 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["unsloth", "datasets", "trl", "huggingface_hub"]
+# ///
+"""Stage 1 LoRA training with Unsloth on Hugging Face Jobs (1x A100 80GB).
+
+Run (pipeline test, 10 steps):
+  hf jobs uv run --flavor a100-large --timeout 45m --secrets HF_TOKEN train/hf_train.py -- --max-steps 10
+Full run: no --max-steps (2 epochs). Settings follow train/config.yaml; alpha is a proposal, see STATE.md.
+"""
+import argparse, json, os, time
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--model", default="unsloth/Qwen3.8-27B-unsloth-bnb-4bit")
+ap.add_argument("--data", default="erhankeseli/abap-stage1-data")
+ap.add_argument("--out", default="erhankeseli/abap-stage1-adapter-test")
+ap.add_argument("--max-steps", type=int, default=-1)
+ap.add_argument("--epochs", type=float, default=2.0)
+ap.add_argument("--rank", type=int, default=16)
+ap.add_argument("--alpha", type=int, default=16)
+ap.add_argument("--lr", type=float, default=5e-5)
+ap.add_argument("--max-seq-length", type=int, default=16384)
+a = ap.parse_args()
+
+from unsloth import FastLanguageModel  # noqa: E402 (import first)
+from datasets import load_dataset  # noqa: E402
+from trl import SFTConfig, SFTTrainer  # noqa: E402
+
+tok_hf = os.environ["HF_TOKEN"]
+model, tok = FastLanguageModel.from_pretrained(
+    a.model, max_seq_length=a.max_seq_length, load_in_4bit=True, token=tok_hf)
+model = FastLanguageModel.get_peft_model(
+    model, r=a.rank, lora_alpha=a.alpha, lora_dropout=0.0, bias="none",
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",
+                    "in_proj_qkv", "in_proj_z", "out_proj"],
+    use_gradient_checkpointing="unsloth", random_state=20261003)
+
+ds = load_dataset(a.data, data_files={"train": "train.jsonl", "valid": "valid.jsonl"}, token=tok_hf)
+cfg = SFTConfig(
+    output_dir="out", per_device_train_batch_size=1, per_device_eval_batch_size=1,
+    gradient_accumulation_steps=1, num_train_epochs=a.epochs, max_steps=a.max_steps,
+    learning_rate=a.lr, lr_scheduler_type="cosine", warmup_steps=min(30, max(1, a.max_steps // 3)) if a.max_steps > 0 else 30,
+    optim="adamw_8bit", weight_decay=0.0, logging_steps=1, eval_strategy="no", save_strategy="no",
+    max_length=a.max_seq_length, dataset_text_field="text", packing=False, seed=20261003, report_to="none")
+tr = SFTTrainer(model=model, processing_class=tok, train_dataset=ds["train"], eval_dataset=ds["valid"], args=cfg)
+
+t0 = time.time()
+tr.train()
+train_s = time.time() - t0
+steps = tr.state.global_step
+ev = tr.evaluate()  # valid loss of the adapter, for the conversion check
+info = {"steps": steps, "train_seconds": round(train_s, 1), "sec_per_step": round(train_s / max(steps, 1), 2),
+        "eval_loss": ev.get("eval_loss"), "rank": a.rank, "alpha": a.alpha, "lr": a.lr,
+        "peak_gpu_gb": round(__import__("torch").cuda.max_memory_allocated() / 2**30, 1)}
+print("RESULT", json.dumps(info))
+model.save_pretrained("adapter")
+json.dump(info, open("adapter/job_result.json", "w"))
+from huggingface_hub import HfApi  # noqa: E402
+HfApi(token=tok_hf).upload_folder(folder_path="adapter", repo_id=a.out, repo_type="model")
+print("PUSHED", a.out)