train_lora

  1"""QLoRA fine-tune of one preset base model on marola's dataset — Tier 2 in README.md.
  2
  3STATUS: the `tiny` preset has been run end to end (see finetune/README.md); every larger preset is
  4written against the documented peft/transformers/trl APIs and NOT RUN here. Treat a first run on a
  5new preset as a debugging session, not a build step.
  6
  7Usage:
  8    pip install -r requirements.txt
  9    huggingface-cli login                      # only the gated `base` preset needs this
 10    python build_dataset.py                    # or: just finetune-dataset
 11    python train_lora.py --preset qwen-7b --epochs 3
 12    # then convert out/<preset>/adapter with llama.cpp's convert_lora_to_gguf.py and
 13    #   just finetune-adapter-model preset=qwen-7b
 14
 15Presets (see PRESETS): `tiny` (SmolLM2-360M, ungated, Apache-2.0, trains on CPU in minutes — **the
 16default**), `small`/`base` (Llama 3.2 1B/3B; the 3B is gated), `qwen-4b`/`qwen-7b`/`qwen-14b`/
 17`qwen-27b` (Apache-2.0 throughout). tiny is the default rather than small so an unqualified run
 18never silently produces a *Llama derivative*: the Llama 3.2 Community Licence requires such a
 19model's name to begin with "Llama" and to ship the agreement plus a "Built with Llama" notice (see
 20finetune/README.md). Both Llama presets remain available and are fine to use — they just have to
 21be chosen, and their obligations met, on purpose.
 22
 23Switching base is safe by construction: each preset trains into its own `out/<preset>/` directory,
 24and a run refuses to start if the directory it is about to write already holds another base
 25model's work. On CPU pass --no-4bit (bitsandbytes needs CUDA).
 26"""
 27
 28from __future__ import annotations
 29
 30import argparse
 31import json
 32import re
 33from pathlib import Path
 34
 35# Iterate fast on a tiny model, then re-run the same script on the real one: the dataset, the LoRA
 36# config and the GGUF/Ollama steps are identical, only `--preset` changes.
 37#
 38# Every field has exactly one consumer, so a new preset cannot half-exist:
 39#   hf/ollama      the HF id to train from, and the Ollama model an adapter trained on it attaches
 40#                  to — `ollama create` fails loudly on a tag that does not exist, but attaching to
 41#                  the WRONG one of a family loads fine and then answers nonsense, so both the tag
 42#                  and its variant (instruct vs thinking) are part of the identity.
 43#   params_b       preflight's VRAM/RAM/disk estimator, and the batch/accum split below.
 44#   gated          whether a Hugging Face login is needed before the weights download.
 45#   licence        an HF licence *id*, passed through to publish_hf.py --base-license.
 46#   name_prefix    Meta's Community Licence requires a Llama derivative's name to start with
 47#                  "Llama-" (MIP-0025 §5.1(3)); merge_export.py reads this instead of guessing
 48#                  from the id.
 49#   thinking       the base's chat template emits reasoning blocks, which changes what an SFT run
 50#                  teaches and what the Ollama model answers with.
 51#   target_modules which projections LoRA attaches to. peft matches by suffix and silently skips
 52#                  what it does not find, so a model whose attention is not q/k/v/o_proj needs its
 53#                  own list rather than the standard one.
 54ATTENTION_AND_MLP = (
 55    "q_proj",
 56    "k_proj",
 57    "v_proj",
 58    "o_proj",
 59    "gate_proj",
 60    "up_proj",
 61    "down_proj",
 62)
 63
 64_DEFAULTS = {
 65    "gated": False,
 66    "licence": "apache-2.0",
 67    "name_prefix": "",
 68    "thinking": False,
 69    "target_modules": ATTENTION_AND_MLP,
 70}
 71
 72PRESETS = {
 73    "tiny": {
 74        **_DEFAULTS,
 75        "hf": "HuggingFaceTB/SmolLM2-360M-Instruct",
 76        "ollama": "smollm2:360m",
 77        "params_b": 0.36,
 78    },
 79    "small": {
 80        **_DEFAULTS,
 81        "hf": "unsloth/Llama-3.2-1B-Instruct",
 82        "ollama": "llama3.2:1b",
 83        "params_b": 1.24,
 84        "licence": "llama3.2",
 85        "name_prefix": "Llama-",
 86    },
 87    "base": {
 88        **_DEFAULTS,
 89        "hf": "meta-llama/Llama-3.2-3B-Instruct",
 90        "ollama": "llama3.2",
 91        "gated": True,
 92        "params_b": 3.2,
 93        "licence": "llama3.2",
 94        "name_prefix": "Llama-",
 95    },
 96    # Qwen: Apache-2.0 throughout, so no naming obligation of any kind — the reason these are the
 97    # recommended step up from `tiny` rather than the Llama presets above. Sizes and licences
 98    # verified against huggingface.co/api/models/<id>, Ollama tags against ollama.com/library,
 99    # both on 2026-09-12. Qwen2.5-3B-Instruct is deliberately absent: its card says `other`, not
100    # apache-2.0, unlike every other size in the family.
101    "qwen-4b": {
102        **_DEFAULTS,
103        "hf": "Qwen/Qwen3-4B-Instruct-2507",
104        # NOT `qwen3:4b` — that tag is 4b-thinking-2507 (digest 359d7dd4bcda), a different
105        # checkpoint from the Instruct-2507 weights trained here (0edcdef34593).
106        "ollama": "qwen3:4b-instruct",
107        "params_b": 4.02,
108    },
109    "qwen-7b": {
110        **_DEFAULTS,
111        "hf": "Qwen/Qwen2.5-7B-Instruct",
112        "ollama": "qwen2.5:7b",
113        "params_b": 7.62,
114    },
115    "qwen-14b": {
116        **_DEFAULTS,
117        "hf": "Qwen/Qwen2.5-14B-Instruct",
118        "ollama": "qwen2.5:14b",
119        "params_b": 14.77,
120    },
121    # The odd one out, and the reason `target_modules`/`thinking` are preset fields at all.
122    # Qwen3.8-27B is post-trained (not a base checkpoint), multimodal
123    # (`Qwen3_5ForConditionalGeneration`, pipeline image-text-to-text — the merge drops the vision
124    # tower) and hybrid-attention: only 16 of its 64 layers use q/k/v/o_proj, the other 48 are
125    # linear-attention layers with different names, so the standard target list would quietly
126    # train a quarter of the attention stack. Its template also wraps every assistant turn in an
127    # empty <think></think> block. NOT RUN here — treat a first run as an experiment, not a build.
128    "qwen-27b": {
129        **_DEFAULTS,
130        "hf": "Qwen/Qwen3.8-27B",
131        "ollama": "qwen3.8:27b",
132        "params_b": 27.78,
133        "thinking": True,
134        "target_modules": "all-linear",
135    },
136}
137
138PRESET_FIELDS = (
139    "hf",
140    "ollama",
141    "gated",
142    "params_b",
143    "licence",
144    "name_prefix",
145    "thinking",
146    "target_modules",
147)
148
149BASE_MARKER = ".marola-base"
150
151
152def preset_for(hf_id: str) -> str | None:
153    """The preset name that owns this HF id, or None for a `--base` outside the table."""
154    for name, p in PRESETS.items():
155        if p["hf"] == hf_id:
156            return name
157    return None
158
159
160def ollama_base_for(hf_id: str) -> str:
161    """The `FROM` line Modelfile.adapter needs — the adapter only fits the family it was trained on."""
162    name = preset_for(hf_id)
163    return PRESETS[name]["ollama"] if name else "<the Ollama model matching " + hf_id + ">"
164
165
166def run_slug(hf_id: str) -> str:
167    """The per-run directory name: the preset when there is one, else a slug of the model id.
168
169    Artifacts are namespaced by this rather than shared, because everything downstream of training
170    is model-shaped: an adapter, a merged checkpoint and a GGUF built from SmolLM2 mean nothing to
171    a Qwen run, and `save_total_limit` deletes by step number without looking at whose checkpoint
172    it is — the two bases' step counts differ, so a shared directory loses whichever run wrote
173    fewer steps.
174    """
175    name = preset_for(hf_id)
176    if name:
177        return name
178    return re.sub(r"[^a-z0-9]+", "-", hf_id.rsplit("/", 1)[-1].lower()).strip("-") or "custom"
179
180
181def run_dir(hf_id: str) -> Path:
182    """`finetune/out/<preset>` — every artifact of one base's run lives under here."""
183    return Path(__file__).parent / "out" / run_slug(hf_id)
184
185
186def default_out(hf_id: str, kind: str = "adapter") -> Path:
187    """`finetune/out/<preset>/adapter` — see run_slug for why this is not `finetune/out/adapter`."""
188    return run_dir(hf_id) / kind
189
190
191def has_artifacts(out: Path) -> bool:
192    """Does this directory already hold a training run's output?"""
193    return bool(list(out.glob("checkpoint-*"))) or (out / "adapter_config.json").exists()
194
195
196def base_conflict(marker_text: str | None, base: str) -> bool:
197    """True when a directory's recorded base model is not the one about to be trained."""
198    return bool(marker_text) and marker_text.strip() != base
199
200
201def read_marker(out: Path) -> str | None:
202    marker = out / BASE_MARKER
203    return marker.read_text() if marker.exists() else None
204
205
206def check_base(out: Path, base: str) -> None:
207    """Refuse to write a second base model's run into another's directory.
208
209    Optimizer, scheduler and adapter shapes are all model-specific: resuming a SmolLM2 checkpoint
210    into a Qwen run either explodes with a shape error or, worse, trains something meaningless.
211    This runs whether or not `--resume` was passed, because an unguarded run does not resume the
212    other base's work — it overwrites it, and `save_total_limit` can delete it mid-run.
213    """
214    if not out.exists():
215        return
216    previous = read_marker(out)
217    if base_conflict(previous, base) and has_artifacts(out):
218        raise SystemExit(
219            f"train_lora: {out} holds a run trained from {previous.strip()!r}, but this run uses "
220            f"{base!r}. Refusing to mix base models in one directory — delete {out}, or pass "
221            "--out for this base (the default is already per-preset)."
222        )
223
224
225def adapter_base(adapter: Path) -> str | None:
226    """Which base model an existing LoRA adapter was trained on, per its own adapter_config.json."""
227    cfg = Path(adapter) / "adapter_config.json"
228    if not cfg.exists():
229        return None
230    try:
231        return json.loads(cfg.read_text()).get("base_model_name_or_path")
232    except json.JSONDecodeError:
233        return None
234
235
236def check_adapter_base(adapter: Path, base: str) -> None:
237    """Refuse a mismatched adapter *before* the base model loads.
238
239    peft raises on its own eventually, but only after transformers has pulled a full checkpoint
240    into RAM — up to tens of GB for the larger presets. This is one JSON read.
241    """
242    found = adapter_base(adapter)
243    if found and found != base:
244        raise SystemExit(
245            f"the adapter at {adapter} was trained on {found!r}, not {base!r}. A LoRA adapter only "
246            "fits the family it was trained on — point --sft-adapter/--adapter at this base's run "
247            "(finetune/out/<preset>/), or retrain."
248        )
249
250
251def latest_checkpoint(out: Path) -> str | None:
252    """The newest `checkpoint-N` under --out, or None to train from scratch.
253
254    Decided from what is on disk rather than from the flag alone: `--resume` on a clean machine
255    must start from scratch, not fail, or the first run of any CI job breaks.
256    """
257    ckpts = sorted(
258        (d for d in out.glob("checkpoint-*") if d.is_dir()),
259        key=lambda d: int(d.name.rsplit("-", 1)[1]),
260    )
261    return str(ckpts[-1]) if ckpts else None
262
263
264def mark_base(out: Path, base: str) -> None:
265    out.mkdir(parents=True, exist_ok=True)
266    (out / BASE_MARKER).write_text(base + "\n")
267
268
269def self_test() -> int:
270    """Everything above is pure or filesystem-only — no torch, no network, seconds to run."""
271    import tempfile
272
273    fails = 0
274
275    def ok(got, want, label):
276        nonlocal fails
277        if got == want:
278            print(f"  ok   {label}")
279        else:
280            fails += 1
281            print(f"  FAIL {label} — got {got!r}, want {want!r}")
282
283    for name, p in PRESETS.items():
284        ok(sorted(p), sorted(PRESET_FIELDS), f"{name} declares every preset field")
285        ok(
286            p["licence"] in ("apache-2.0", "llama3.2"),
287            True,
288            f"{name}'s licence is an HF id, not prose (publish_hf --base-license takes it)",
289        )
290        ok(
291            p["name_prefix"] == ("Llama-" if "llama" in p["hf"].lower() else ""),
292            True,
293            f"{name}'s name_prefix matches Meta's obligation for its base",
294        )
295
296    ok(preset_for("Qwen/Qwen2.5-7B-Instruct"), "qwen-7b", "a base id resolves back to its preset")
297    ok(preset_for("mistralai/Whatever"), None, "an unknown base id resolves to no preset")
298    ok(ollama_base_for(PRESETS["qwen-27b"]["hf"]), "qwen3.8:27b", "27B's Ollama tag is qwen3.8")
299    ok(run_slug("Qwen/Qwen2.5-7B-Instruct"), "qwen-7b", "the run directory is the preset name")
300    ok(run_slug("acme/Some_Model-v2"), "some-model-v2", "an off-table base still gets a slug")
301    ok(
302        default_out("Qwen/Qwen2.5-7B-Instruct").parts[-2:],
303        ("qwen-7b", "adapter"),
304        "the default --out is namespaced by preset",
305    )
306    ok(
307        default_out(PRESETS["tiny"]["hf"]) != default_out(PRESETS["qwen-7b"]["hf"]),
308        True,
309        "tiny and qwen-7b cannot land in the same directory",
310    )
311
312    ok(base_conflict("a\n", "b"), True, "a different recorded base is a conflict")
313    ok(base_conflict("a\n", "a"), False, "the same base is not a conflict")
314
315    with tempfile.TemporaryDirectory() as tmp:
316        ad = Path(tmp) / "adapter"
317        ad.mkdir()
318        smol, qwen = PRESETS["tiny"]["hf"], PRESETS["qwen-7b"]["hf"]
319        check_adapter_base(ad, qwen)  # no adapter_config.json yet — nothing to contradict
320        (ad / "adapter_config.json").write_text(json.dumps({"base_model_name_or_path": smol}))
321        ok(adapter_base(ad), smol, "an adapter reports the base it was trained on")
322        check_adapter_base(ad, smol)
323        try:
324            check_adapter_base(ad, qwen)
325            ok("no refusal", "SystemExit", "a foreign adapter must be refused before loading")
326        except SystemExit as exc:
327            ok(smol in str(exc), True, "the refusal names the adapter's real base")
328    ok(base_conflict(None, "a"), False, "an unmarked directory is not a conflict")
329
330    with tempfile.TemporaryDirectory() as tmp:
331        out = Path(tmp) / "adapter"
332        smol, qwen = PRESETS["tiny"]["hf"], PRESETS["qwen-7b"]["hf"]
333        check_base(out, smol)  # nothing on disk yet
334        mark_base(out, smol)
335        (out / "checkpoint-40").mkdir()
336        ok(latest_checkpoint(out), str(out / "checkpoint-40"), "the newest checkpoint is found")
337        try:
338            check_base(out, qwen)
339            ok("no refusal", "SystemExit", "switching base in a used directory must be refused")
340        except SystemExit as exc:
341            ok(smol in str(exc) and qwen in str(exc), True, "the refusal names both base models")
342        check_base(out, smol)  # the same base is still fine
343        (out / "checkpoint-40").rmdir()
344        ok(latest_checkpoint(out), None, "an empty directory resumes from scratch")
345
346    print("train_lora self-test:", "ok" if not fails else f"{fails} FAILED")
347    return 1 if fails else 0
348
349
350def main() -> None:
351    ap = argparse.ArgumentParser()
352    ap.add_argument(
353        "--preset",
354        choices=sorted(PRESETS),
355        default=None,
356        help="tiny = SmolLM2-360M (CPU, minutes, ungated — for iterating on the dataset); "
357        "small/base = Llama 3.2 1B/3B (the 3B is gated, and both carry Meta's naming obligation); "
358        "qwen-4b/7b/14b/27b = Apache-2.0 all the way down, the recommended step up from tiny. "
359        "Each preset trains into its own out/<preset>/ directory",
360    )
361    ap.add_argument(
362        "--base",
363        default=None,
364        help="explicit HF model id; overrides --preset. An id in the preset table still gets that "
365        "preset's directory, Ollama tag, licence and LoRA targets",
366    )
367    ap.add_argument(
368        "--self-test",
369        action="store_true",
370        help="check the preset table and the base-switching guards; no torch, no network",
371    )
372    ap.add_argument(
373        "--resume",
374        action="store_true",
375        help="continue from the newest checkpoint in --out if one exists. Trainer checkpoints "
376        "carry optimizer, scheduler, RNG and step state, so this resumes mid-epoch rather than "
377        "restarting the epoch — what makes a multi-day or interrupted run survivable",
378    )
379    ap.add_argument(
380        "--save-steps",
381        type=int,
382        default=200,
383        help="checkpoint every N steps; with --resume this bounds what a crash costs",
384    )
385    ap.add_argument("--data", default=str(Path(__file__).parent / "data"))
386    ap.add_argument(
387        "--out",
388        default=None,
389        help="where the adapter and checkpoints go (default: finetune/out/<preset>/adapter)",
390    )
391    ap.add_argument("--epochs", type=int, default=3)
392    ap.add_argument("--lr", type=float, default=2e-4)
393    ap.add_argument("--rank", type=int, default=16)
394    ap.add_argument("--max-len", type=int, default=2048)
395    ap.add_argument(
396        "--no-4bit", action="store_true", help="skip bitsandbytes 4-bit (CPU / no CUDA)"
397    )
398    ap.add_argument(
399        "--device",
400        choices=("auto", "cuda", "cpu", "hybrid"),
401        default="auto",
402        help="auto: cuda when available. hybrid: fill the GPU and spill the rest to CPU RAM "
403        "(accelerate device_map='auto' + max_memory), for a model too large for VRAM alone",
404    )
405    ap.add_argument(
406        "--gpu-mem-gb",
407        type=float,
408        default=None,
409        help="VRAM ceiling for --device hybrid; default is 90%% of the card, leaving headroom for "
410        "activations",
411    )
412    ap.add_argument(
413        "--batch", type=int, default=None, help="per-device batch size (default: picked by size)"
414    )
415    ap.add_argument(
416        "--grad-accum",
417        type=int,
418        default=None,
419        help="gradient accumulation steps (default: by size)",
420    )
421    ap.add_argument(
422        "--no-packing",
423        action="store_true",
424        help="disable example packing; packing is the single biggest throughput win on a corpus of "
425        "short examples like marola's (~2.8k rows, most far below max-len)",
426    )
427    ap.add_argument(
428        "--no-grad-checkpointing",
429        action="store_true",
430        help="disable gradient checkpointing; it trades ~20%% speed for a large drop in activation "
431        "memory, which is what makes the bigger presets fit at all",
432    )
433    args = ap.parse_args()
434    if args.self_test:
435        raise SystemExit(self_test())
436
437    if not args.base:
438        args.base = PRESETS[args.preset or "tiny"]["hf"]
439    # An explicit --base that happens to be a preset's model still gets that preset's settings:
440    # the table is keyed by what the model IS, not by how it was named on the command line.
441    preset_name = args.preset or preset_for(args.base)
442    preset = PRESETS.get(preset_name or "", {})
443    if not args.out:
444        args.out = str(default_out(args.base))
445    out = Path(args.out)
446    check_base(out, args.base)
447    print(
448        f"base model: {args.base} ({preset_name or 'off-table'}) — Ollama FROM for "
449        f"Modelfile.adapter: {ollama_base_for(args.base)}; adapter -> {out}"
450    )
451    if preset.get("thinking"):
452        print(
453            "NOTE: this base's chat template wraps every assistant turn in a reasoning block, so "
454            "the fine-tune teaches that shape too — check one sample before publishing."
455        )
456
457    import torch
458    from datasets import load_dataset
459    from peft import LoraConfig
460    from transformers import AutoModelForCausalLM, AutoTokenizer
461    from trl import SFTConfig, SFTTrainer
462
463    cuda = torch.cuda.is_available()
464    device = args.device
465    if device == "auto":
466        device = "cuda" if cuda else "cpu"
467    if device in ("cuda", "hybrid") and not cuda:
468        raise SystemExit(
469            f"train_lora: --device {device} needs a working CUDA device; torch.cuda.is_available() "
470            "is False. Check `nvidia-smi` — a present card with a broken driver looks exactly like "
471            "no card at all here. Use --device cpu to run anyway."
472        )
473
474    # TF32 costs nothing on Ampere and later and speeds up every matmul that is not already bf16.
475    if cuda:
476        torch.backends.cuda.matmul.allow_tf32 = True
477        torch.backends.cudnn.allow_tf32 = True
478
479    tok = AutoTokenizer.from_pretrained(args.base)
480    tok.pad_token = tok.pad_token or tok.eos_token
481
482    model_kwargs: dict = {"torch_dtype": torch.bfloat16 if cuda else torch.float32}
483    use_4bit = not args.no_4bit and device in ("cuda", "hybrid")
484    if use_4bit:
485        from transformers import BitsAndBytesConfig
486
487        model_kwargs["quantization_config"] = BitsAndBytesConfig(
488            load_in_4bit=True,
489            bnb_4bit_quant_type="nf4",
490            bnb_4bit_compute_dtype=torch.bfloat16,
491            # Quantizing the quantization constants too — ~0.4 bits/param less memory for no
492            # measurable quality cost, which is free headroom on a 24 GB card at 27B.
493            bnb_4bit_use_double_quant=True,
494        )
495    if device == "cuda":
496        model_kwargs["device_map"] = "auto"
497    elif device == "hybrid":
498        # Fill the GPU to a ceiling, spill the remainder to CPU RAM. Slower per step than pure GPU
499        # — every offloaded layer crosses PCIe twice — but it is the difference between running and
500        # an OOM when the model does not fit in VRAM alone. This machine has 188 GB of RAM, so the
501        # CPU side is effectively unbounded.
502        ceiling = args.gpu_mem_gb or (torch.cuda.get_device_properties(0).total_memory / 1e9 * 0.90)
503        model_kwargs["device_map"] = "auto"
504        model_kwargs["max_memory"] = {0: f"{ceiling:.0f}GiB", "cpu": "160GiB"}
505        print(f"hybrid: up to {ceiling:.0f} GiB on the GPU, the rest offloaded to CPU RAM")
506
507    # SDPA is the fastest attention available without a flash-attn build, and unlike flash-attn it
508    # needs no extra wheel — worth asking for explicitly rather than taking the eager default.
509    model_kwargs["attn_implementation"] = "sdpa"
510    model = AutoModelForCausalLM.from_pretrained(args.base, **model_kwargs)
511    if use_4bit:
512        from peft import prepare_model_for_kbit_training
513
514        model = prepare_model_for_kbit_training(
515            model, use_gradient_checkpointing=not args.no_grad_checkpointing
516        )
517
518    data = load_dataset(
519        "json", data_files={"train": f"{args.data}/train.jsonl", "eval": f"{args.data}/eval.jsonl"}
520    )
521
522    # peft matches target modules by suffix and skips what it cannot find, so a hybrid attention
523    # stack needs "all-linear" rather than the q/k/v/o list every other preset uses.
524    targets = preset.get("target_modules", ATTENTION_AND_MLP)
525    lora = LoraConfig(
526        r=args.rank,
527        lora_alpha=2 * args.rank,
528        lora_dropout=0.05,
529        bias="none",
530        task_type="CAUSAL_LM",
531        target_modules="all-linear" if targets == "all-linear" else list(targets),
532    )
533    # Effective batch stays ~8 regardless of size; only how it is split changes. Bigger models get
534    # a smaller per-device batch and more accumulation, because activation memory scales with both
535    # batch and model width.
536    params_b = preset.get("params_b", 1.0)
537    if args.batch is not None:
538        batch = args.batch
539    elif device == "cpu":
540        batch = 1
541    else:
542        batch = 4 if params_b <= 2 else 2 if params_b <= 9 else 1
543    accum = args.grad_accum if args.grad_accum is not None else max(1, 8 // batch)
544
545    cfg = SFTConfig(
546        output_dir=args.out,
547        num_train_epochs=args.epochs,
548        learning_rate=args.lr,
549        per_device_train_batch_size=batch,
550        gradient_accumulation_steps=accum,
551        logging_steps=5,
552        eval_strategy="epoch",
553        save_strategy="epoch",
554        save_total_limit=1,  # a 27B checkpoint per epoch is ~100 GB of churn nobody reads
555        max_length=args.max_len,
556        bf16=cuda,
557        # marola's corpus is ~2.8k mostly-short rows against a 2048-token window, so without
558        # packing most of every batch is padding. Packing concatenates examples up to max_length
559        # and is the single biggest throughput win available here.
560        packing=not args.no_packing,
561        gradient_checkpointing=not args.no_grad_checkpointing,
562        gradient_checkpointing_kwargs={"use_reentrant": False},
563        # Fused optimizer when CUDA is present: fewer kernel launches per step, no accuracy cost.
564        optim="adamw_torch_fused" if cuda else "adamw_torch",
565        # Only useful when packing is off: it batches similar-length examples so a batch is
566        # mostly content rather than padding. With packing on it is redundant.
567        group_by_length=args.no_packing,
568        dataloader_num_workers=4,
569        report_to=[],
570    )
571    print(
572        f"device={device} batch={batch} accum={accum} (effective {batch * accum}) "
573        f"packing={not args.no_packing} grad_checkpointing={not args.no_grad_checkpointing} "
574        f"4bit={use_4bit}"
575    )
576    resume = latest_checkpoint(out) if args.resume else None
577    print(f"resuming from {resume}" if resume else "training from scratch (no checkpoint found)")
578    mark_base(out, args.base)
579    trainer = SFTTrainer(
580        model=model,
581        processing_class=tok,
582        peft_config=lora,
583        args=cfg,
584        train_dataset=data["train"],
585        eval_dataset=data["eval"],
586    )
587    trainer.train(resume_from_checkpoint=resume)
588    trainer.save_model(args.out)
589    print(
590        f"adapter saved to {args.out} — convert with llama.cpp convert_lora_to_gguf.py, then see Modelfile.adapter"
591    )
592
593
594if __name__ == "__main__":
595    main()