"""Hugging Face Space training core: SFT / DPO / KTO, one step at a time, resumable. Design goals ------------ * **One bucket + one model per deployment (env, not per request).** The ``BUCKET_ID`` and ``MODEL_ID`` environment variables configure the whole Space — they are *not* re-entered on every request. Loading the base model once and keeping it resident avoids re-downloading / re-loading it on every call. * **Model kept resident in RAM/VRAM.** The base model is loaded once and kept in a module-level cache across calls (the expensive part — download + load — happens a single time). The LoRA adapter weights are tiny and are reset from the resumed checkpoint on each call, so reuse is safe: the frozen base weights never change, only the adapter + optimizer state do. * **GPU-when-available, CPU otherwise.** The device is chosen from ``torch.cuda.is_available()``. On a ZeroGPU Space ``@spaces.GPU`` attaches an A10G for the call (so we train on GPU); off ZeroGPU the decorator is a no-op and we train on CPU. Because ZeroGPU's GPU is *transient* (it detaches between calls), the resident cache is only used when the device is persistent (a dedicated-GPU or CPU Space); on ZeroGPU the model is loaded per call from the local HF cache instead. * **No data persistence.** The uploaded "trace" is parsed into an in-memory ``datasets`` object for the duration of the call and discarded afterwards. Only weights + optimizer/scheduler state are ever written to the bucket. * **Dynamic gradient accumulation.** The number of optimizer steps a call runs is derived from the number of uploaded examples and the deployment's gradient-accumulation size ``G`` (``GRAD_ACCUM`` env): ``steps = ceil(num_examples / G)`` (SFT/DPO). A single trace -> 1 step; 10 traces with ``G=4`` -> 3 steps (4 + 4 + 2). The caller no longer chooses a fixed "steps per call" — it just uploads the trace. * **Incremental, resumable training.** Each call advances training by the dynamic number of steps above. The HF ``Trainer`` checkpoint (model/adapter weights + ``optimizer.pt`` + ``scheduler.pt`` + ``rng`` + ``trainer_state.json``) is persisted to an HF Storage Bucket, so the next call resumes exactly where it left off — the optimizer does **not** restart from scratch. Shut the Space down, come back later, keep training. Method data schemas ------------------- SFT (``trl.SFTTrainer``) — prompt/completion or plain text: {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]} or {"text": "a single string"} DPO (``trl.DPOTrainer``) — chosen/rejected pairs: {"prompt": "...", "chosen": "...", "rejected": "..."} (messages-format ``prompt``/``chosen``/``rejected`` also accepted) KTO (``trl.KTOTrainer``) — binary desirable/undesirable feedback: {"prompt": "...", "completion": "...", "desirable": true|false} """ from __future__ import annotations import json import os import tempfile from dataclasses import dataclass, field from pathlib import Path from typing import Any import torch from datasets import Dataset from huggingface_hub import snapshot_download from peft import LoraConfig, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer # `spaces` lets us detect whether we are running on a ZeroGPU Space. On # ZeroGPU the GPU is *transient* (attached only for the duration of a # ``@spaces.GPU`` call), so a model cannot be kept resident across calls — it # would be freed and any cached reference would dangle. We therefore only keep # the model resident when the device is *persistent* (a dedicated-GPU or CPU # Space). The import is optional; the platform ships its own `spaces` and, on # any non-ZeroGPU Space, `spaces.Config.zero_gpu` is simply False. try: from spaces.config import Config as _SpacesConfig _ZERO_GPU = bool(getattr(_SpacesConfig, "zero_gpu", False)) except Exception: # pragma: no cover _ZERO_GPU = False from store import ( checkpoint_paths, pull_checkpoint, push_checkpoint, save_state_index, ) HERE = Path(__file__).resolve().parent LOCAL_CACHE = HERE / ".cache" # Small, CPU-fine base model used for defaults / smoke tests. DEFAULT_MODEL_ID = "HuggingFaceTB/SmolLM2-135M" # Injected when a base model's tokenizer has no chat template, so SFT with # chat-format ("messages") traces works on *any* base model, not just -Instruct # ones. A simple, valid template transformers can apply. DEFAULT_CHAT_TEMPLATE = ( "{%- for message in messages -%}\n" "{% if message['role'] == 'user' %}User: {{ message['content'] }}\n" "{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n" "{% elif message['role'] == 'system' %}System: {{ message['content'] }}\n" "{% endif %}{%- endfor -%}" ) # --------------------------------------------------------------------------- # # Configuration # --------------------------------------------------------------------------- # @dataclass class TrainConfig: # One model + one bucket per deployment. These come from the Space's # environment variables (MODEL_ID, BUCKET_ID) and are NOT re-entered per # request. `resolve_*` fills them from the env when left empty, so the # whole deployment is configured once. model_id: str = "" # e.g. "HuggingFaceTB/SmolLM2-135M" — from MODEL_ID env method: str = "sft" # sft | dpo | kto # LoRA — keeps CPU training tractable on a ZeroGPU container. lora_r: int = 8 lora_alpha: int = 16 lora_dropout: float = 0.0 target_modules: tuple = ("q_proj", "k_proj", "v_proj", "o_proj") # Optimizer / training. learning_rate: float = 1e-4 # Constant LR keeps "one step at a time" stable across resumable calls; # a decaying schedule would require a fixed horizon and is not the goal here. lr_scheduler_type: str = "constant" optim: str = "adamw_torch" weight_decay: float = 0.0 max_length: int = 256 # Bucket + run. bucket_id: str = "" # e.g. "username/my-bucket" — from BUCKET_ID env resume: bool = True # resume from latest checkpoint if one exists # Gradient accumulation: how many micro-batches to accumulate per optimizer # step. The number of optimizer steps this call runs is derived from the # number of uploaded examples: steps_this_call = ceil(num_examples / # grad_accum). 10 examples with GRAD_ACCUM=4 -> 3 steps (4+4+2). # 1 example -> 1 step (single trace, single step). grad_accum: int = 4 save_total_limit: int = 3 # keep only the N most-recent checkpoints # --------------------------------------------------------------------------- # # Result # --------------------------------------------------------------------------- # @dataclass class TrainResult: method: str model_id: str global_step: int loss: float | None = None checkpoint_id: str = "" resumed_from: str | None = None num_examples: int = 0 grad_accum: int = 1 steps_run: int = 0 # optimizer steps this call advanced notes: list = field(default_factory=list) elapsed: float = 0.0 def as_dict(self) -> dict: return { "method": self.method, "model_id": self.model_id, "global_step": self.global_step, "loss": round(self.loss, 6) if self.loss is not None else None, "checkpoint_id": self.checkpoint_id, "resumed_from": self.resumed_from, "num_examples": self.num_examples, "grad_accum": self.grad_accum, "steps_run": self.steps_run, "notes": self.notes, "elapsed_sec": round(self.elapsed, 2), } # --------------------------------------------------------------------------- # # Trace parsing (the ONLY place the uploaded data is touched) # --------------------------------------------------------------------------- # def parse_trace(rows: list[dict], method: str) -> Dataset: """Turn raw uploaded rows into a ``datasets.Dataset`` for the method. Rows come in as a list of dicts (one JSON object per line, or a small JSON array). Anything that doesn't parse is dropped (logged, not fatal). """ m = method.lower() if m == "sft": return _parse_sft(rows) if m == "dpo": return _parse_dpo(rows) if m == "kto": return _parse_kto(rows) raise ValueError(f"Unknown method: {method!r}") def _parse_sft(rows: list[dict]) -> Dataset: # Prefer chat format; fall back to a single "text" field. conv = [] for r in rows: if "messages" in r and isinstance(r["messages"], list): conv.append({"messages": r["messages"]}) if conv: return Dataset.from_list(conv) texts = [] for r in rows: if "text" in r: texts.append({"text": r["text"]}) elif "completion" in r and "prompt" in r: # prompt/completion -> single text for SFT texts.append({"text": f"{r['prompt']}{r['completion']}"}) if texts: return Dataset.from_list(texts) raise ValueError("SFT trace needs rows with 'messages' (chat) or 'text'.") def _parse_dpo(rows: list[dict]) -> Dataset: out = [] for r in rows: if "chosen" in r and "rejected" in r: out.append( { "prompt": r.get("prompt", ""), "chosen": r["chosen"], "rejected": r["rejected"], } ) if not out: raise ValueError("DPO trace needs rows with 'chosen' and 'rejected'.") return Dataset.from_list(out) def _parse_kto(rows: list[dict]) -> Dataset: # trl's KTOTrainer expects a boolean ``label`` column (True = desirable); # the trace uses the friendlier ``desirable`` name, so map it through. out = [] for r in rows: completion = r.get("completion") desirable = r.get("desirable", r.get("label")) if completion is not None and desirable is not None: out.append( { "prompt": r.get("prompt", ""), "completion": completion, "label": bool(desirable), } ) if not out: raise ValueError("KTO trace needs rows with 'completion' and 'desirable'.") return Dataset.from_list(out) # --------------------------------------------------------------------------- # # Env + resident model cache # --------------------------------------------------------------------------- # def _resolve_token() -> str | None: # On a Space the token is set as a secret (HF_TOKEN / HUGGING_FACE_HUB_TOKEN). return ( os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_HUB_TOKEN") ) def resolve_model_id(cfg: TrainConfig) -> str: """The model for this deployment, from MODEL_ID env (or the passed config). One model per deployment — not re-chosen per request.""" return (cfg.model_id or os.environ.get("MODEL_ID") or DEFAULT_MODEL_ID).strip() def resolve_bucket_id(cfg: TrainConfig) -> str: """The bucket for this deployment, from BUCKET_ID env (or the passed config). One bucket per deployment — not re-chosen per request.""" return (cfg.bucket_id or os.environ.get("BUCKET_ID") or "").strip() def resolve_lora_r(cfg: TrainConfig) -> int: """LoRA rank for this deployment, from LORA_R env (or the config default).""" raw = os.environ.get("LORA_R") return int(raw) if raw else cfg.lora_r def resolve_lora_alpha(cfg: TrainConfig) -> int: """LoRA alpha for this deployment, from LORA_ALPHA env (or the config default).""" raw = os.environ.get("LORA_ALPHA") return int(raw) if raw else cfg.lora_alpha def resolve_lora_dropout(cfg: TrainConfig) -> float: """LoRA dropout for this deployment, from LORA_DROPOUT env (or the config default).""" raw = os.environ.get("LORA_DROPOUT") return float(raw) if raw else cfg.lora_dropout def resolve_target_modules(cfg: TrainConfig): """Target modules for this deployment, from LORA_TARGET_MODULES env (comma-separated, e.g. ``q_proj,k_proj,v_proj,o_proj``) or the config default tuple.""" raw = os.environ.get("LORA_TARGET_MODULES") if raw: return tuple(p.strip() for p in raw.split(",") if p.strip()) return tuple(cfg.target_modules) def resolve_learning_rate(cfg: TrainConfig) -> float: """Learning rate for this deployment, from LEARNING_RATE env (or the config default).""" raw = os.environ.get("LEARNING_RATE") return float(raw) if raw else cfg.learning_rate def resolve_max_length(cfg: TrainConfig) -> int: """Max sequence length for this deployment, from MAX_LENGTH env (or the config default).""" raw = os.environ.get("MAX_LENGTH") return int(raw) if raw else cfg.max_length def resolve_grad_accum(cfg: TrainConfig) -> int: """Gradient-accumulation size for this deployment, from GRAD_ACCUM env (or the config default). The number of optimizer steps a call runs is derived from the number of uploaded examples: ceil(num_examples / grad_accum). 10 examples with GRAD_ACCUM=4 -> 3 steps (4+4+2); 1 example -> 1 step.""" raw = os.environ.get("GRAD_ACCUM") val = int(raw) if raw else cfg.grad_accum return max(1, val) def _device_name() -> str: return "cuda" if torch.cuda.is_available() else "cpu" def _device_map(): """Where to load the model: auto-shard across every available GPU, else a single GPU, else CPU. ``device_map="auto"`` lets accelerate spread a large model across all visible GPUs (e.g. a 27B base model across 4 GPUs) so the *same code path* trains a 27B model on 4 GPUs as a 1B model on 1 GPU — no manual setup. With 0/1 GPUs accelerate returns a single device (no sharding, so the Trainer's single-device path is used). The model only spans multiple devices when >1 GPU is present. """ if torch.cuda.is_available(): return "auto" # accelerate distributes across all visible GPUs return "cpu" def _dtype_for_shard(model_id: str, is_sharded: bool): """Load a multi-GPU-sharded model in float16 so the per-GPU footprint roughly halves (a 27B model = ~54GB in fp32 vs ~27GB in fp16, splitting across 4 GPUs becomes trivially affordable). Single-device / CPU models stay float32 for accuracy. LoRA training is robust to a fp16 base.""" if is_sharded and torch.cuda.is_available(): return torch.float16 return torch.float32 # A module-level cache of the loaded base model. Populated once on the first # call and reused on every subsequent call, so the expensive download + load # happens a single time per deployment. The base (frozen) weights never # change between calls; only the LoRA adapter + optimizer state do, and those # are reset from the resumed checkpoint on each call — so reuse is safe. _MODEL_CACHE: dict = {} def _model_cache_key(model_id: str) -> str: return model_id def _load_model_fresh(model_id: str, cfg: TrainConfig): """Download + load the base model and wrap it in a fresh LoRA adapter. The device is chosen automatically: ``device_map="auto"`` shards the model across every visible GPU (so a 27B model on 4 GPUs loads the same way a 1B model loads on 1 GPU), otherwise it uses a single GPU or CPU. A multi-GPU-sharded model is loaded in fp16 to halve the per-GPU footprint; single-device / CPU models stay fp32 for accuracy. """ num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0 is_sharded = num_gpus > 1 device_map = _device_map() torch_dtype = _dtype_for_shard(model_id, is_sharded) tok = AutoTokenizer.from_pretrained(model_id, token=_resolve_token()) model = AutoModelForCausalLM.from_pretrained( model_id, token=_resolve_token(), torch_dtype=torch_dtype, device_map=device_map, ) if tok.pad_token is None: tok.pad_token = tok.eos_token # A base (non-instruct) model may have no chat template; SFT with chat # traces needs one. Inject a simple default so any base model works. if getattr(tok, "chat_template", None) is None: tok.chat_template = DEFAULT_CHAT_TEMPLATE lora = LoraConfig( r=cfg.lora_r, lora_alpha=cfg.lora_alpha, lora_dropout=cfg.lora_dropout, target_modules=list(cfg.target_modules), bias="none", task_type="CAUSAL_LM", ) model = get_peft_model(model, lora) model.print_trainable_parameters() if is_sharded: # Log the sharding so the user can see how the model was spread. devices = sorted({str(d) for d in set(model.hf_device_map.values())}) print(f"[train] model {model_id} sharded across {len(devices)} device(s): " f"{', '.join(devices)} ({num_gpus} GPU, {torch_dtype})") return model, tok def _model_is_cached(model_id: str) -> bool: return model_id in _MODEL_CACHE and _MODEL_CACHE[model_id] is not None def get_model_and_tokenizer(model_id: str, cfg: TrainConfig): """Return a model (PEFT-wrapped) + tokenizer for the deployment. When the device is *persistent* (not ZeroGPU) the loaded base model is cached in a module-level dict and reused across calls — the whole point: the base model stays resident in RAM/VRAM instead of being re-downloaded / re-loaded on every request. The frozen base weights never change between calls; only the LoRA adapter + optimizer state do, and those are reset from the resumed checkpoint on each call, so reuse is safe. On ZeroGPU the GPU is *transient* (detached between calls), so we skip the cache and load per-call from the local HF cache (which still avoids a network re-download — only the load is repeated). Note: LoRA hyper-params (rank / alpha) are fixed at first-load time. To change them mid-deployment, restart the Space (which re-runs the cache). """ # ZeroGPU's GPU is transient: never keep a cached device ref across calls. if _ZERO_GPU: return _load_model_fresh(model_id, cfg) cached = _MODEL_CACHE.get(model_id) if cached is not None: return cached model, tok = _load_model_fresh(model_id, cfg) _MODEL_CACHE[model_id] = (model, tok) return model, tok def _build_config(cfg: TrainConfig, output_dir: Path, num_examples: int, current_step: int): """Build a TRL training config for a single call. **Dynamic gradient accumulation.** The number of optimizer steps this call runs is derived from how many examples the uploaded trace contains and the deployment's gradient-accumulation size ``G`` (``GRAD_ACCUM`` env), *not* from a fixed "steps per call": steps_this_call = ceil(num_examples / (per_device_batch * accum)) with ``accum = ceil(G / per_device_batch)``. With ``G=4`` this is exactly the requested behaviour: a single trace -> 1 step; 10 traces -> 3 steps (4 + 4 + 2). SFT/DPO use a per-device batch of 1 (so one optimizer step == ``G`` examples, matching the example precisely); KTO uses a per-device batch of 2 (its KL term is undefined for a single example), so its per-step example count is ``2 * ceil(G/2)``. ``max_steps`` is the *absolute* target (``current_step + steps_this_call``); with ``resume_from_checkpoint`` the Trainer runs the remaining steps, so the call continues exactly where the last call left off. """ import math use_gpu = torch.cuda.is_available() use_cpu = not use_gpu m = cfg.method.lower() # KTO needs an actual batch > 1 (its KL term is undefined for a single # example); SFT/DPO use a per-device batch of 1. per_device_batch = 2 if m == "kto" else 1 accum = max(1, math.ceil(cfg.grad_accum / per_device_batch)) examples_per_step = per_device_batch * accum steps_this_call = max(1, math.ceil(num_examples / examples_per_step)) if num_examples else 1 target = current_step + steps_this_call common = dict( output_dir=str(output_dir), per_device_train_batch_size=per_device_batch, gradient_accumulation_steps=accum, max_steps=target, # absolute target; resume continues from here learning_rate=cfg.learning_rate, lr_scheduler_type=cfg.lr_scheduler_type, optim=cfg.optim, weight_decay=cfg.weight_decay, max_grad_norm=1.0, logging_strategy="steps", logging_steps=1, save_strategy="steps", save_steps=1, save_total_limit=cfg.save_total_limit, save_only_model=False, # MUST keep optimizer.pt / scheduler.pt / rng use_cpu=use_cpu, # GPU when available, else CPU report_to="none", dataloader_num_workers=0, remove_unused_columns=False, ignore_data_skip=True, # robust resume for tiny/arbitrary datasets ) if m == "sft": from trl import SFTConfig args = SFTConfig(**common, max_length=cfg.max_length) elif m == "dpo": from trl import DPOConfig args = DPOConfig(**common, max_length=cfg.max_length) elif m == "kto": from trl import KTOConfig args = KTOConfig(**common, max_length=cfg.max_length) else: raise ValueError(f"Unknown method {cfg.method!r}") return args, steps_this_call def _build_trainer(method: str, model, tok, dataset, args): if method == "sft": from trl import SFTTrainer return SFTTrainer(model=model, args=args, train_dataset=dataset, processing_class=tok) if method == "dpo": from trl import DPOTrainer # ref_model=None + a PEFT model => trl uses the adapter-disabled model # as the reference (standard LoRA-DPO pattern). return DPOTrainer(model=model, args=args, train_dataset=dataset, processing_class=tok) if method == "kto": from trl import KTOTrainer return KTOTrainer(model=model, args=args, train_dataset=dataset, processing_class=tok) raise ValueError(f"Unknown method {method!r}") # --------------------------------------------------------------------------- # # The main entry point # --------------------------------------------------------------------------- # def _steps_for_call(cfg: TrainConfig, num_examples: int) -> int: """Dynamic optimizer-step count for a single call. ``steps = ceil(num_examples / (per_device_batch * accum))`` with ``accum = ceil(cfg.grad_accum / per_device_batch)``. SFT/DPO use a per-device batch of 1, so ``steps = ceil(num_examples / grad_accum)`` — exactly the requested behaviour (10 examples, grad_accum=4 -> 3 steps; 1 example -> 1 step). KTO uses a per-device batch of 2. A call with no examples still runs 1 step (keeps the optimizer warm / makes the call idempotent).""" import math m = cfg.method.lower() per_device_batch = 2 if m == "kto" else 1 accum = max(1, math.ceil(cfg.grad_accum / per_device_batch)) examples_per_step = per_device_batch * accum if not num_examples: return 1 return max(1, math.ceil(num_examples / examples_per_step)) def run_one_call( rows: list[dict], cfg: TrainConfig, *, log=print, ) -> TrainResult: """Parse the uploaded trace, run ``steps_per_call`` optimizer steps for the chosen method (resuming from the latest checkpoint if one exists), then persist the resulting weights + optimizer state to the bucket. The uploaded ``rows`` (the training data) are never written anywhere — only the trained weights and optimizer/scheduler state are persisted. """ import time start = time.time() # One bucket + one model per deployment: resolve from the env when the # caller did not pin them (the UI no longer asks for them per request). cfg.model_id = resolve_model_id(cfg) cfg.bucket_id = resolve_bucket_id(cfg) # LoRA hyper-params are also fixed per deployment via env (one model + one # bucket + one LoRA config per deployment). They are resolved here so the # resident model cache (keyed by model_id) wraps the *deployment's* LoRA # config at first-load, and every later call reuses the same cached model. cfg.lora_r = resolve_lora_r(cfg) cfg.lora_alpha = resolve_lora_alpha(cfg) cfg.lora_dropout = resolve_lora_dropout(cfg) cfg.target_modules = resolve_target_modules(cfg) cfg.learning_rate = resolve_learning_rate(cfg) cfg.max_length = resolve_max_length(cfg) cfg.grad_accum = resolve_grad_accum(cfg) m = cfg.method.lower() dataset = parse_trace(rows, m) num_examples = len(dataset) log(f"[train] model={cfg.model_id} bucket={cfg.bucket_id or '(none)'} " f"parsed {num_examples} {m} examples from trace " f"(grad_accum={cfg.grad_accum})") notes: list[str] = [] # --- resolve resume checkpoint -------------------------------------- # local_cache = LOCAL_CACHE local_cache.mkdir(parents=True, exist_ok=True) resume_dir = None resumed_from = None if cfg.resume: resume_dir, resumed_from = _resolve_resume( cfg, local_cache, notes ) if resume_dir: log(f"[train] resuming from {resumed_from}") current_step = _read_global_step(resume_dir) if resume_dir else 0 # --- dynamic gradient accumulation: how many steps this call runs -- # steps_this_call = _steps_for_call(cfg, num_examples) target = current_step + steps_this_call output_dir = local_cache / f"run-{target}" log( f"[train] {num_examples} example(s) / grad_accum={cfg.grad_accum} " f"-> {steps_this_call} step(s): {current_step} -> {target}" ) args, steps_this_call = _build_config(cfg, output_dir, num_examples, current_step) # --- build model + trainer (resident cache keeps the base model in # RAM/VRAM across calls; the adapter+optimizer state is reset from the # resumed checkpoint below) ------------------------- # model, tok = get_model_and_tokenizer(cfg.model_id, cfg) trainer = _build_trainer(m, model, tok, dataset, args) # --- train (resume continues from current_step to target) ---------- # train_args = {} if resume_dir: train_args["resume_from_checkpoint"] = resume_dir log(f"[train] training steps {current_step} -> {target}") train_res = trainer.train(**train_args) # --- read back the new global step + loss -------------------------- # new_step = current_step loss = None try: st = trainer.state new_step = int(getattr(st, "global_step", current_step) or current_step) except Exception: pass logs = getattr(train_res, "log", None) or [] if logs: for entry in logs: if "loss" in entry: loss = float(entry["loss"]) if loss is None: try: metrics = getattr(train_res, "metrics", None) or {} train_metrics = metrics.get("train", {}) if "train_loss" in train_metrics: loss = float(train_metrics["train_loss"]) except Exception: pass # --- persist the latest checkpoint to the bucket ------------------- # checkpoint_id = f"checkpoint-{new_step}" saved_dir = _find_checkpoint_dir(output_dir, new_step) or output_dir if cfg.bucket_id: try: push_checkpoint(cfg.bucket_id, saved_dir, checkpoint_id) save_state_index( cfg.bucket_id, { "global_step": new_step, "method": m, "model_id": cfg.model_id, "loss": loss, }, ) log(f"[train] pushed {checkpoint_id} to bucket {cfg.bucket_id}") except Exception as exc: notes.append(f"push failed: {exc}") log(f"[train] PUSH FAILED: {exc}") else: notes.append("no bucket_id set — weights NOT persisted") result = TrainResult( method=m, model_id=cfg.model_id, global_step=new_step, loss=loss, checkpoint_id=checkpoint_id, resumed_from=resumed_from, num_examples=num_examples, grad_accum=cfg.grad_accum, steps_run=new_step - current_step, notes=notes, elapsed=time.time() - start, ) log(f"[train] done: {result.as_dict()}") return result # --------------------------------------------------------------------------- # # helpers # --------------------------------------------------------------------------- # def _resolve_resume(cfg: TrainConfig, local_cache: Path, notes: list) -> tuple[str | None, str | None]: """Decide which checkpoint to resume from. 1. Prefer a fresh local checkpoint from a previous call (same session) — so back-to-back calls resume even without a bucket. 2. Otherwise pull the latest checkpoint from the bucket. """ # 1) fresh local checkpoint from a prior call local = _find_any_checkpoint(local_cache) if local is not None: return str(local), local.name # 2) latest from the bucket if cfg.bucket_id: try: latest = _latest_checkpoint_prefix(cfg.bucket_id) except Exception as exc: # pragma: no cover latest = None notes.append(f"could not list bucket: {exc}") if latest: local_dir = local_cache / f"pull-{latest}" if pull_checkpoint(cfg.bucket_id, latest, local_dir): return str(local_dir), latest notes.append(f"could not download {latest}") return None, None def _find_any_checkpoint(root: Path): """Return the local ``.../checkpoint-N`` dir with the highest N under ``root`` (searches run-*/checkpoint-* and top-level checkpoint-*).""" best = None for cand in list(root.glob("checkpoint-*")) + list(root.glob("run-*/checkpoint-*")): s = _dir_step(cand) if best is None or s > _dir_step(best): best = cand return best def _latest_checkpoint_prefix(bucket_id: str) -> str | None: try: from store import list_latest_checkpoint return list_latest_checkpoint(bucket_id) except Exception: return None def _read_global_step(checkpoint_dir: str | None) -> int: if not checkpoint_dir: return 0 p = Path(checkpoint_dir) / "trainer_state.json" try: data = json.loads(p.read_text()) return int(data.get("global_step", 0)) except Exception: return 0 def _find_checkpoint_dir(output_dir: Path, step: int) -> Path | None: cand = output_dir / f"checkpoint-{step}" if cand.exists(): return cand # the trainer may have saved under a different name; take the max one best = None for d in output_dir.glob("checkpoint-*"): try: s = int(d.name.split("-")[-1]) except Exception: continue if best is None or s > int(_dir_step(best)): best = d return best def _dir_step(d: Path) -> int: try: return int(d.name.split("-")[-1]) except Exception: return -1 def current_progress(bucket_id: str) -> dict: """Lightweight status: latest checkpoint + global step from the index.""" try: from store import load_state_index idx = load_state_index(bucket_id) except Exception: idx = {} latest = None try: latest = _latest_checkpoint_prefix(bucket_id) except Exception: pass return { "global_step": idx.get("global_step", 0), "latest_checkpoint": latest, "checkpoints": checkpoint_paths(bucket_id), }