diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f6ff935..75e89fe 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -34,6 +34,13 @@ class MyHead(Head): Unknown or missing constructor arguments are reported when the config loads. +## Adding a recipe + +Recipes ship with the package in `src/decisionformer/recipes/`, so `decisionformer train NAME` works after a plain +`pip install`. The first comment line is the description `decisionformer recipes` shows. Put only what differs from +the defaults in a recipe, and add a recipe only after it has learned in a real run. The unit tests load every +recipe. + ## Development loop ```bash diff --git a/README.md b/README.md index 75228bb..8dc6887 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ generation. Three lines to decide, one command to train, one command to serve. Python 3.10 or newer, with PyTorch 2.6+ and transformers 5.20+. -**With pip** +**With pip** (from the first public release; until then, install from GitHub) ```bash pip install decisionformer @@ -114,8 +114,10 @@ curl localhost:8000/v1/systemone -H 'content-type: application/json' -d '{ ``` The server takes the same `state` and `questions` and returns the same JSON as `decide`, with `id`, `object` and -`created` added. It is the System One API used by Jev, vLLM and the Decision Index. Over-long requests are -refused with HTTP 422, never truncated, and `--api-key-env MY_KEY_VAR` requires `Authorization: Bearer `. +`created` added. It is the System One API used by Jev, vLLM and the Decision Index; the request and response +schemas are at `/docs`. Over-long requests are refused with HTTP 422, never truncated, and +`--api-key-env MY_KEY_VAR` requires `Authorization: Bearer `. A request that arrives while the model is idle +runs at once; requests that arrive while it is busy are answered together, up to `--max-batch` (16) at a time. For code-readout checkpoints on Qwen3.5 (the Decision 3.0 layout), vLLM can run the model behind the same API. Other checkpoints fail with a clear error that points to the default backend: @@ -141,22 +143,26 @@ Backbones are the architectures built into Hugging Face transformers, loaded wit `save_pretrained`; decisionformer has no model code of its own. Hub models that need their own modeling code (`trust_remote_code`) are not supported. Decoders read the answer as a code token (`A`, `B`, ...) at the end of the prompt, which works from step 0; encoders score each option. Encoder-decoder models without a standalone text -encoder in transformers (BART) are refused with a clear error. +encoder in transformers (BART) are refused with a clear error. Fine-tuning Qwen3.5 with the answer-code readout +has not learned reliably in our runs yet (`docs/benchmarks.md`); loading and serving Qwen3.5 checkpoints such as +`vllm-sr/d3-edge` works. ## Fine-tune ```bash decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k -decisionformer train configs/recipes/deciformer15k_qwen35.yaml -torchrun --nproc-per-node 8 -m decisionformer train configs/recipes/deciformer15k_qwen35.yaml # multi-GPU +decisionformer train deciformer15k_qwen3 # Qwen3-0.6B + LoRA, one epoch +torchrun --nproc-per-node 8 -m decisionformer train deciformer15k_qwen3 # multi-GPU ``` -Every setting lives in a YAML config, and any value can be overridden on the command line with dotted keys -(`train.lr=1e-4`). `decisionformer config FILE` prints the resolved configuration. A minimal config: +`deciformer15k_qwen3` is a bundled recipe: `decisionformer recipes` lists them all, and +`decisionformer init deciformer15k_qwen3` copies one to a YAML file you can edit. Every setting lives in that YAML, +and any value can be overridden on the command line with dotted keys (`train.lr=1e-4`). +`decisionformer config FILE` prints the resolved configuration. A minimal config: ```yaml model: - backbone: Qwen/Qwen3.5-0.8B # any Hub id or local path + backbone: Qwen/Qwen3-0.6B # any Hub id or local path attention: causal # causal | bidirectional | prefix_bidirectional head: {type: code_readout} finetune: @@ -168,13 +174,24 @@ data: eval: {dev: data/dev.jsonl} train: output_dir: runs/my-decider - lr: 2.0e-4 push_to_hub: my-org/my-decider # optional: uploads best/ when training ends (private by default) ``` -A run writes `runs//best` and `last`: finished checkpoints in the standard Hugging Face layout, ready for -`from_pretrained`, `serve` and `push`. It also writes `metrics.jsonl`, per-example predictions at every -evaluation, resumable state and `summary.json`. +Everything else has a default. The learning rate follows the fine-tuning mode (2e-5 full, 2e-4 LoRA and QLoRA, +1e-3 head only). Training evaluates and saves resumable state every 500 steps, and rerunning the same command +resumes where it stopped; the resumable state is removed once the run finishes. A run writes `runs//best` and `last`: finished checkpoints in the standard Hugging +Face layout, ready for `from_pretrained`, `serve` and `push`. It also writes `metrics.jsonl`, per-example +predictions at every evaluation and `summary.json`. Training warns when the training data shares ids, families or +identical states with the evaluation or calibration data. + +The same from Python, with a recipe name, a YAML file or a dict: + +```python +import decisionformer + +summary = decisionformer.train("deciformer15k_qwen3", {"train.max_steps": 200}) +model = decisionformer.AutoDecisionModel.from_pretrained("runs/deciformer15k-qwen3-0.6b/best") +``` ### On a laptop (Mac or Linux) @@ -184,7 +201,7 @@ steps: ```bash pip install "decisionformer[data]" decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k -decisionformer train configs/recipes/local_smollm2.yaml +decisionformer train local_smollm2 decisionformer serve runs/local-smollm2-135m/best # try it on /v1/systemone ``` @@ -192,8 +209,8 @@ The device is picked automatically: a CUDA GPU, then Apple silicon (MPS), then t `DECISIONFORMER_DEVICE=cpu` to force the CPU. On MPS and CPU, training runs in float32. The recipe keeps memory low with micro-batches of 4, at most 4,096 tokens per micro-batch and gradient checkpointing: on a 6-core CPU it peaked at 3.7 GB of memory and took about 15 s per step. If memory is still tight, lower them: -`train.micro_batch_size=2 train.token_budget=2048`. A checkpoint is saved every 100 steps, and -rerunning the same command resumes from it. +`train.micro_batch_size=2 train.token_budget=2048`. Resumable state is saved every 100 steps, and rerunning the +same command resumes from it. ### Training data @@ -354,7 +371,9 @@ python -m decision_index pipeline --engine decisionformer.integrations.decision_ ## Measured results -On one 8 GB RTX 2060 SUPER in FP16, on MNLI: Qwen3-0.6B with LoRA and `code_readout` went from 0.36 zero-shot to +On one 8 GB RTX 2060 SUPER in FP16, the starting recipe `deciformer15k_qwen3` (Qwen3-0.6B with LoRA, one epoch of +deciformer-15k, 30 minutes) reaches **0.824** dev accuracy across its 9 tasks, from 0.75 (VitaminC) to 0.90 (Civil +Comments), with calibrated ECE 0.023. On MNLI: Qwen3-0.6B with LoRA and `code_readout` went from 0.36 zero-shot to **0.83** dev accuracy in 250 steps (calibrated ECE 0.042), and Qwen3.5-0.8B with QLoRA and bidirectional attention reached **0.84** in 150 steps. The vLLM backend matched the transformers backend to within 0.001 per probability, with identical choices, at a median of 30 ms per 3-question request. Configs and full tables are in [`docs/benchmarks.md`](docs/benchmarks.md). @@ -369,7 +388,7 @@ ruff check src tests ``` Experimental: `decisionformer adapt` runs masked next-token prediction (as in LLM2Vec) before bidirectional -training (`configs/recipes/adapt_mntp.yaml`). A short stage did not help in our tests. +training (the `adapt_mntp` recipe). A short stage did not help in our tests. ## Roadmap diff --git a/configs/base.yaml b/configs/base.yaml deleted file mode 100644 index 5227d06..0000000 --- a/configs/base.yaml +++ /dev/null @@ -1,26 +0,0 @@ -# Shared defaults. Recipes inherit this file with `base: ../base.yaml` and override what they need. -data: - max_length: 1024 - overflow: truncate_state_end # training only: served models refuse over-long requests - max_options: 255 - shuffle_options: true - -train: - seed: 42 - epochs: 1 - global_batch_size: 32 - micro_batch_size: 8 - weight_decay: 0.01 - max_grad_norm: 1.0 - scheduler: {type: cosine, warmup_ratio: 0.05, min_lr_ratio: 0.1} - log_every: 10 - save_every: 500 - keep_last: 2 - resume: auto - -eval: - ece_bins: 15 - micro_batch_size: 32 - selection_metric: accuracy - selection_mode: max - use_calibrated: true diff --git a/configs/benchmarks/mnli_8gb/_base.yaml b/configs/benchmarks/mnli_8gb/_base.yaml index 7b260fe..b565ad5 100644 --- a/configs/benchmarks/mnli_8gb/_base.yaml +++ b/configs/benchmarks/mnli_8gb/_base.yaml @@ -1,5 +1,4 @@ # 8 GB Turing GPU (RTX 2060 SUPER): fp16 autocast, no bf16 / flash-attention. -base: ../../base.yaml data: train: data/mnli/train.jsonl calibration: data/mnli/calibration_small.jsonl diff --git a/configs/recipes/deciformer15k_qwen35.yaml b/configs/recipes/deciformer15k_qwen35.yaml deleted file mode 100644 index f9fe3b0..0000000 --- a/configs/recipes/deciformer15k_qwen35.yaml +++ /dev/null @@ -1,31 +0,0 @@ -# Demo: Qwen3.5-0.8B on FrontiersMind/deciformer-15k (9 tasks: safety, toxicity, QA, NLI, paraphrase, fact -# checking and intent routing in English and German; choice and noul questions), LoRA and the answer-code readout. -# decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k -# decisionformer train configs/recipes/deciformer15k_qwen35.yaml -# The export is a Decision 3.0 code readout on Qwen3.5, so `decisionformer serve ... --backend vllm` can serve it. -# Trains in bf16; GPUs without bf16 (Turing, V100) fall back to fp16 automatically. -base: ../base.yaml - -model: - backbone: Qwen/Qwen3.5-0.8B - attention: causal - head: {type: code_readout} - -finetune: - mode: lora - gradient_checkpointing: true - lora: {r: 16, lora_alpha: 32, lora_dropout: 0.05, target_modules: all-linear} - -data: - train: data/deciformer-15k/train.jsonl - calibration: data/deciformer-15k/calibration.jsonl - eval: {dev: data/deciformer-15k/dev.jsonl} - max_length: 2048 # 6 of 11,982 training prompts are longer; their states are truncated in training - -train: - output_dir: runs/deciformer15k-qwen3.5-0.8b - lr: 2.0e-4 - head_lr: 2.0e-5 - mixed_precision: bf16 - token_budget: 8192 - eval_every: 125 diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 4b1f821..330a84e 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -13,6 +13,28 @@ records each (`dev_small.jsonl`, `calibration_small.jsonl`, separated by whole f ¹ Before v0.2's final pooler default: code-readout heads now read the last token in every attention mode (the Decision 3.0 layout). +**deciformer-15k, the starting recipe.** `decisionformer train deciformer15k_qwen3` (Qwen3-0.6B, LoRA r16, causal, +one epoch of 11,982 records = 375 updates, fp16 on the 2060 SUPER, 30 minutes with evaluations) on the family-level +split from `decisionformer data split` (1,509 dev records): + +| step | dev accuracy (calibrated) | ECE | Brier | NLL | +|---|---|---|---|---| +| 125 | 0.792 | 0.039 | 0.314 | 0.585 | +| 250 | 0.805 | 0.030 | 0.281 | 0.533 | +| 375 | **0.824** | 0.023 | 0.252 | 0.467 | + +Per task at step 375: Civil Comments 0.904, PAWS 0.875, MASSIVE en-US 0.864, SQuAD 2.0 0.860, MASSIVE de-DE 0.823, +Aegis 2.0 0.790, MultiNLI 0.774, BoolQ 0.772, VitaminC 0.750. + +**Known issue: Qwen3.5 with causal attention.** Qwen3.5-0.8B fine-tuned with causal attention stayed at chance on +MNLI (0.34-0.37) in four runs: LoRA with an fp16 base (150 steps, twice), QLoRA (60 steps), and LoRA on an fp32 base +(60 steps). Each run's loss rose near step 25 and then settled at the class prior. Before training, the code +readout's logits equal the language model's own next-token logits for the answer codes exactly, so the prompt, +answer position and readout rows are correct. Qwen3-0.6B learns with causal attention, and Qwen3.5 learned with +bidirectional attention and a mean pooler (0.84 above). With bidirectional attention and the last-token readout +(the Decision 3.0 layout, the default since v0.2), a 40-step QLoRA check on deciformer-15k showed the same loss spike +near step 20 (1.55 → 2.04 → 1.23) and 0.41 dev accuracy on 128 records, so no Qwen3.5 recipe ships yet. + **Parity with the Decision 3.0 reference runtime.** `AutoDecisionModel.from_pretrained("vllm-sr/d3-edge")` against the checkpoint's own `trust_remote_code` runtime, BF16, four requests (choice, noul and score questions, an empty state, a JSON state, a 32-option choice): one question per forward pass, the probabilities are identical (maximum diff --git a/src/decisionformer/__init__.py b/src/decisionformer/__init__.py index ad484f8..0791f7e 100644 --- a/src/decisionformer/__init__.py +++ b/src/decisionformer/__init__.py @@ -4,17 +4,21 @@ model = AutoDecisionModel.from_pretrained("org/my-decider") model.decide("My card was charged twice.", {"team": {"type": "choice", "criteria": {"billing": None, "shipping": None}}}) + + import decisionformer + decisionformer.train("deciformer15k_qwen3") # fine-tune from a bundled recipe, a YAML file or a dict """ __version__ = "0.2.0.dev0" from decisionformer.api import AutoDecisionModel, Decider, UnsupportedInput # noqa: E402 -from decisionformer.config import RunConfig, load_config # noqa: E402 +from decisionformer.config import RunConfig, list_recipes, load_config # noqa: E402 from decisionformer.formats import PromptFormat, get_format, register_format # noqa: E402 from decisionformer.prompting import ContextLengthError # noqa: E402 from decisionformer.schema import DecisionInputError # noqa: E402 +from decisionformer.trainer import train # noqa: E402 __all__ = [ "AutoDecisionModel", "ContextLengthError", "Decider", "DecisionInputError", "PromptFormat", "RunConfig", - "UnsupportedInput", "__version__", "get_format", "load_config", "register_format", + "UnsupportedInput", "__version__", "get_format", "list_recipes", "load_config", "register_format", "train", ] diff --git a/src/decisionformer/adapt.py b/src/decisionformer/adapt.py index 026b5be..229059f 100644 --- a/src/decisionformer/adapt.py +++ b/src/decisionformer/adapt.py @@ -33,6 +33,7 @@ from decisionformer.inference import autocast from decisionformer.optim import lr_lambda from decisionformer.schema import read_records +from decisionformer.trainer import _setup_logging logger = logging.getLogger("decisionformer.adapt") @@ -68,8 +69,7 @@ def adapt(config: RunConfig) -> Path: if config.finetune.mode not in ("full", "lora"): raise ValueError("adapt supports finetune.mode 'full' or 'lora'") state = dist_utils.setup(config.distributed.backend, config.distributed.timeout_minutes) - logging.basicConfig(level=logging.INFO if state.is_main else logging.WARNING, - format="%(asctime)s %(levelname)s %(name)s: %(message)s", force=True) + _setup_logging(state.is_main) torch.manual_seed(config.train.seed) model_cfg = config.model tokenizer = load_tokenizer(model_cfg.backbone, model_cfg.revision, model_cfg.trust_remote_code) diff --git a/src/decisionformer/api.py b/src/decisionformer/api.py index fc112ca..462b485 100644 --- a/src/decisionformer/api.py +++ b/src/decisionformer/api.py @@ -43,6 +43,13 @@ def _local_dir(name_or_path: str | Path, revision: str | None, token: str | None return Path(snapshot_download(str(name_or_path), revision=revision, token=token)) +def display_name(name_or_path: str | Path, root: Path) -> str: + """How a model names itself in responses: its Hub id, or its run name for a local `runs//best`.""" + if not Path(str(name_or_path)).expanduser().is_dir(): + return str(name_or_path) + return root.resolve().parent.name if root.name in ("best", "last") else root.name + + def _prompt_name(config: dict[str, Any]) -> str: name = config.get("prompt") or config.get("prompt_format") if name: @@ -132,7 +139,7 @@ def from_pretrained(cls, name_or_path: str | Path, *, revision: str | None = Non builder = PromptBuilder(tokenizer, style, kind, max_length=limit, overflow="error", codes=codes, code_token_ids=config.get("token_ids") or None, chat_template_kwargs=ext.get("chat_template_kwargs")) - name = ext.get("name") or (str(name_or_path) if not Path(str(name_or_path)).is_dir() else root.name) + name = ext.get("name") or display_name(name_or_path, root) return cls(model, builder, tokenizer, config, device=device, name=name, max_input_tokens=limit, batch_size=batch_size, path=root) diff --git a/src/decisionformer/cli.py b/src/decisionformer/cli.py index fdc6c82..951dc90 100644 --- a/src/decisionformer/cli.py +++ b/src/decisionformer/cli.py @@ -1,4 +1,10 @@ -"""decisionformer command line: train, eval, predict, serve, push, compare, data check/split, config (adapt: experimental).""" +"""decisionformer: fine-tune, evaluate, serve and publish one-pass decision models on Hugging Face transformers. + +Start with a bundled recipe: + decisionformer recipes list them + decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k + decisionformer train deciformer15k_qwen3 fine-tune (a recipe name or your own YAML file) + decisionformer serve runs/deciformer15k-qwen3-0.6b/best answer requests on /v1/systemone""" from __future__ import annotations @@ -9,26 +15,57 @@ from typing import Any -def _train(args: argparse.Namespace) -> None: +def _run_config(args: argparse.Namespace, **options: Any) -> Any: from decisionformer.config import load_config + + try: + return load_config(args.config, args.overrides, **options) + except FileNotFoundError as error: + raise SystemExit(f"error: {error}") from None + + +def _train(args: argparse.Namespace) -> None: from decisionformer.trainer import train - summary = train(load_config(args.config, args.overrides)) + summary = train(_run_config(args)) if summary: print(json.dumps(summary, indent=2)) def _adapt(args: argparse.Namespace) -> None: from decisionformer.adapt import adapt - from decisionformer.config import load_config - print(adapt(load_config(args.config, args.overrides, validate_components=False))) + print(adapt(_run_config(args, validate_components=False))) def _config(args: argparse.Namespace) -> None: - from decisionformer.config import dump_config, load_config + from decisionformer.config import dump_config + + print(dump_config(_run_config(args))) + + +def _recipes(args: argparse.Namespace) -> None: + from decisionformer.config import list_recipes - print(dump_config(load_config(args.config, args.overrides))) + recipes = list_recipes() + width = max(map(len, recipes)) + for name, description in recipes.items(): + print(f"{name:<{width}} {description}") + print("\nRun one with `decisionformer train NAME`, or copy it to edit: `decisionformer init NAME`.") + + +def _init(args: argparse.Namespace) -> None: + from decisionformer.config import config_path + + try: + source = config_path(args.recipe) + except FileNotFoundError as error: + raise SystemExit(f"error: {error}") from None + target = Path(args.path or f"{source.stem}.yaml") + if target.exists() and not args.force: + raise SystemExit(f"error: {target} exists (use --force to overwrite)") + target.write_text(source.read_text()) + print(f"Wrote {target}. Edit it, then: decisionformer train {target}") def _load(args: argparse.Namespace) -> Any: @@ -76,14 +113,41 @@ def _eval(args: argparse.Namespace) -> None: def _predict(args: argparse.Namespace) -> None: - from decisionformer.schema import read_records + from decisionformer.schema import as_request, read_records decider = _load(args) - lines = [] - for number, record in enumerate(r for path in args.data for r in read_records(path)): - response = decider.decide(record.get("state"), record["questions"], images=record.get("images")) - lines.append(json.dumps({"id": record.get("id", number), **response}, ensure_ascii=False) + "\n") + records = [r for path in args.data for r in read_records(path)] + lines: list[str] = [] + for start in range(0, len(records), 64): # requests in a chunk share forward passes + chunk = list(enumerate(records[start:start + 64], start)) + requests: list[dict[str, Any] | Exception] = [] + for number, record in chunk: + try: + state, questions = as_request(record, number) + requests.append({"state": state, "questions": questions, "images": record.get("images")}) + except ValueError as error: + requests.append(error) + valid = [request for request in requests if not isinstance(request, Exception)] + try: + answers = iter(decider.decide_many(valid)) + results: list[dict[str, Any] | Exception] = [r if isinstance(r, Exception) else next(answers) + for r in requests] + except ValueError: # one request does not fit: answer the others and report it alone + results = [r if isinstance(r, Exception) else _decide_one(decider, r) for r in requests] + for (number, record), result in zip(chunk, results): + body = {"error": {"message": str(result)}} if isinstance(result, Exception) else result + lines.append(json.dumps({"id": record.get("id", number), **body}, ensure_ascii=False) + "\n") _write(lines, args.output) + failed = sum('"error": {"message"' in line for line in lines) + if failed: + print(f"{failed} of {len(lines)} requests could not be answered; see their `error` lines", file=sys.stderr) + + +def _decide_one(decider: Any, request: dict[str, Any]) -> dict[str, Any] | Exception: + try: + return decider.decide_many([request])[0] + except ValueError as error: + return error def _serve(args: argparse.Namespace) -> None: @@ -92,8 +156,8 @@ def _serve(args: argparse.Namespace) -> None: try: serve(args.model, host=args.host, port=args.port, device=args.device, backend=args.backend, - api_key_env=args.api_key_env, batch_size=args.batch_size, max_length=args.max_length, - revision=args.revision) + api_key_env=args.api_key_env, max_batch=args.max_batch, batch_size=args.batch_size, + max_length=args.max_length, revision=args.revision) except UnsupportedBackend as error: raise SystemExit(f"error: {error}") from None @@ -161,17 +225,33 @@ def _model_options(command: argparse.ArgumentParser) -> None: def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(prog="decisionformer", description=__doc__) + raw = argparse.RawDescriptionHelpFormatter + parser = argparse.ArgumentParser(prog="decisionformer", description=__doc__, formatter_class=raw) commands = parser.add_subparsers(dest="command", required=True) + examples = """examples: + decisionformer {name} deciformer15k_qwen3 + decisionformer {name} my_config.yaml train.lr=1e-4 train.max_steps=200 + torchrun --nproc-per-node 8 -m decisionformer {name} my_config.yaml (multi-GPU training) + +`decisionformer recipes` lists the bundled recipes; `decisionformer init NAME` copies one to edit.""" for name, handler, help_text in (("train", _train, "Fine-tune a decision model (torchrun for multi-GPU)"), ("config", _config, "Print the fully resolved configuration"), ("adapt", _adapt, "Experimental: MNTP adaptation of a decoder")): - command = commands.add_parser(name, help=help_text) - command.add_argument("config") + command = commands.add_parser(name, help=help_text, description=help_text, formatter_class=raw, + epilog=examples.format(name=name)) + command.add_argument("config", help="a YAML config file or a bundled recipe name") command.add_argument("overrides", nargs="*", help="dotted overrides, e.g. train.lr=1e-5") command.set_defaults(handler=handler) + recipes = commands.add_parser("recipes", help="List the bundled training recipes") + recipes.set_defaults(handler=_recipes) + init = commands.add_parser("init", help="Copy a bundled recipe to a YAML file you can edit") + init.add_argument("recipe", help="recipe name (see `decisionformer recipes`)") + init.add_argument("path", nargs="?", help="where to write it (default: RECIPE.yaml)") + init.add_argument("--force", action="store_true", help="overwrite an existing file") + init.set_defaults(handler=_init) + evaluate = commands.add_parser("eval", help="Score a decision model on labeled requests (JSONL)") _model_options(evaluate) evaluate.add_argument("data", nargs="+") @@ -193,6 +273,8 @@ def build_parser() -> argparse.ArgumentParser: serve.add_argument("--port", type=int, default=8000) serve.add_argument("--backend", choices=("transformers", "vllm"), default="transformers", help="vllm: vLLM computes the readout (Qwen3.5/3.8 code-readout checkpoints only)") + serve.add_argument("--max-batch", type=int, default=16, + help="requests arriving together are answered in one batch of up to this many") serve.add_argument("--api-key-env", default="DECISIONFORMER_API_KEY", help="require `Authorization: Bearer ` when this variable is set") serve.set_defaults(handler=_serve) diff --git a/src/decisionformer/config.py b/src/decisionformer/config.py index 97c4cbe..95ce80d 100644 --- a/src/decisionformer/config.py +++ b/src/decisionformer/config.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Any, Literal @@ -75,7 +76,7 @@ class DataConfig(_Strict): eval: dict[str, str | list[str]] = Field(default_factory=dict) calibration: str | list[str] | None = None prompt: PromptConfig = Field(default_factory=PromptConfig) - max_length: int = 2048 + max_length: int = 1024 # Training and evaluation prompts longer than max_length: cut the state (default) or fail. Loaded and served # models never truncate; they refuse what does not fit. overflow: Literal["error", "truncate_state_end", "truncate_state_start"] = "truncate_state_end" @@ -129,7 +130,7 @@ class TrainConfig(_Strict): global_batch_size: int = Field(32, ge=1) micro_batch_size: int = Field(8, ge=1) token_budget: int | None = Field(None, ge=1) - lr: float = 2e-5 + lr: float | None = None # None: 2e-5 for full fine-tuning, 2e-4 for LoRA and QLoRA, 1e-3 for head_only head_lr: float | None = None weight_decay: float = 0.01 adam_betas: tuple[float, float] = (0.9, 0.999) @@ -140,10 +141,10 @@ class TrainConfig(_Strict): scheduler: SchedulerConfig = Field(default_factory=SchedulerConfig) loss: LossConfig = Field(default_factory=LossConfig) log_every: int = 10 - eval_every: int | None = None - save_every: int | None = None + eval_every: int | None = 500 + save_every: int | None = 500 # resumable state; rerunning the same command resumes (train.resume: auto) keep_last: int = 2 - resume: str | None = None + resume: str | None = "auto" shard_optimizer_states: bool = False find_unused_parameters: bool | Literal["auto"] = "auto" drop_last: bool = False @@ -160,7 +161,7 @@ def _yaml_no(cls, value: Any) -> Any: class EvalConfig(_Strict): at_start: bool = False # also evaluate before the first update (the step-0 baseline) ece_bins: int = 15 - micro_batch_size: int = 16 + micro_batch_size: int = 32 token_budget: int | None = None selection_set: str | None = None selection_metric: str = "accuracy" @@ -184,6 +185,9 @@ class AdaptConfig(_Strict): max_length: int = 512 +DEFAULT_LR = {"full": 2e-5, "lora": 2e-4, "qlora": 2e-4, "head_only": 1e-3} + + class RunConfig(_Strict): model: ModelConfig finetune: FinetuneConfig = Field(default_factory=FinetuneConfig) @@ -201,26 +205,64 @@ def _consistency(self) -> RunConfig: raise ValueError("train.global_batch_size must be at least train.micro_batch_size") if self.eval.selection_set is not None and self.eval.selection_set not in self.data.eval: raise ValueError(f"eval.selection_set '{self.eval.selection_set}' is not one of data.eval") + if self.train.lr is None: + self.train.lr = DEFAULT_LR[self.finetune.mode] return self -def _load_with_bases(path: Path, seen: tuple[Path, ...] = ()) -> Any: - path = path.resolve() - if path in seen: - raise ValueError(f"Config inheritance cycle through {path}") - node = OmegaConf.load(path) +RECIPES = Path(__file__).parent / "recipes" + + +def list_recipes() -> dict[str, str]: + """The bundled recipes: name -> the first line of the recipe's description.""" + recipes = {} + for path in sorted(RECIPES.glob("*.yaml")): + comments = [line[1:].strip() for line in path.read_text().splitlines() if line.startswith("#")] + recipes[path.stem] = comments[0] if comments else "" + return recipes + + +def config_path(source: str | Path) -> Path: + """A YAML config file, or the name of a bundled recipe (`decisionformer recipes` lists them).""" + path = Path(source).expanduser() + if path.is_file(): + return path + if len(path.parts) == 1 and path.suffix in ("", ".yaml", ".yml") and (RECIPES / f"{path.stem}.yaml").is_file(): + return RECIPES / f"{path.stem}.yaml" + raise FileNotFoundError(f"{source} is neither a config file nor a bundled recipe. Recipes: " + f"{', '.join(list_recipes())}. Copy one to edit it: decisionformer init NAME") + + +def _with_bases(node: Any, parent: Path, seen: tuple[Path, ...]) -> Any: bases = node.pop("base", None) if "base" in node else None if bases is None: return node merged = OmegaConf.create() for base in [bases] if isinstance(bases, str) else list(bases): - merged = OmegaConf.merge(merged, _load_with_bases(path.parent / str(base), (*seen, path))) + local = parent / str(base) + merged = OmegaConf.merge(merged, _load_file(local if local.is_file() else config_path(str(base)), seen)) return OmegaConf.merge(merged, node) -def load_config(path: str | Path, overrides: list[str] | None = None, *, validate_components: bool = True) -> RunConfig: - node = _load_with_bases(Path(path)) - if overrides: +def _load_file(path: Path, seen: tuple[Path, ...] = ()) -> Any: + path = path.resolve() + if path in seen: + raise ValueError(f"Config inheritance cycle through {path}") + return _with_bases(OmegaConf.load(path), path.parent, (*seen, path)) + + +def load_config(source: str | Path | Mapping[str, Any], overrides: Sequence[str] | Mapping[str, Any] | None = None, + *, validate_components: bool = True) -> RunConfig: + """A run config from a YAML file, a bundled recipe name or a dict (each may name a `base` to inherit from), + with overrides as dotted `key=value` strings or a `{dotted.key: value}` dict.""" + if isinstance(source, Mapping): + node = _with_bases(OmegaConf.create(dict(source)), Path.cwd(), ()) + else: + node = _load_file(config_path(source)) + if isinstance(overrides, Mapping): + for key, value in overrides.items(): + OmegaConf.update(node, key, value, merge=True) + elif overrides: node = OmegaConf.merge(node, OmegaConf.from_dotlist(list(overrides))) config = RunConfig.model_validate(OmegaConf.to_container(node, resolve=True)) if validate_components: diff --git a/configs/recipes/adapt_mntp.yaml b/src/decisionformer/recipes/adapt_mntp.yaml similarity index 92% rename from configs/recipes/adapt_mntp.yaml rename to src/decisionformer/recipes/adapt_mntp.yaml index 55e352b..230a24b 100644 --- a/configs/recipes/adapt_mntp.yaml +++ b/src/decisionformer/recipes/adapt_mntp.yaml @@ -1,10 +1,9 @@ # Step 1 of a bidirectional decoder: masked next-token prediction (LLM2Vec-style) on in-domain text. -# decisionformer adapt configs/recipes/adapt_mntp.yaml +# decisionformer adapt adapt_mntp # Then train with model.backbone=runs/smollm2-mntp/adapted and model.attention=bidirectional. # Budget generously: in docs/benchmarks.md, 150 steps restored the MNTP loss but not task information. # Use plenty of in-domain text (more than the decision training set alone), bf16 or fp32, and compare the # result against a causal baseline before adopting it. -base: ../base.yaml model: backbone: HuggingFaceTB/SmolLM2-360M-Instruct diff --git a/src/decisionformer/recipes/deciformer15k_qwen3.yaml b/src/decisionformer/recipes/deciformer15k_qwen3.yaml new file mode 100644 index 0000000..f64e98f --- /dev/null +++ b/src/decisionformer/recipes/deciformer15k_qwen3.yaml @@ -0,0 +1,25 @@ +# Start here: Qwen3-0.6B with LoRA on FrontiersMind/deciformer-15k (9 decision tasks, English and German). +# The tasks: safety, toxicity, QA, NLI, paraphrase, fact checking and intent routing. +# decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k +# decisionformer train deciformer15k_qwen3 +# One epoch (375 updates). Trains in bf16; GPUs without bf16 (Turing, V100) fall back to fp16 automatically. + +model: + backbone: Qwen/Qwen3-0.6B + head: {type: code_readout} + +finetune: + mode: lora + gradient_checkpointing: true + lora: {r: 16, lora_alpha: 32, lora_dropout: 0.05, target_modules: all-linear} + +data: + train: data/deciformer-15k/train.jsonl + calibration: data/deciformer-15k/calibration.jsonl + eval: {dev: data/deciformer-15k/dev.jsonl} + +train: + output_dir: runs/deciformer15k-qwen3-0.6b + head_lr: 2.0e-5 # the readout starts from the LM-head rows; keep its step small + token_budget: 8192 + eval_every: 125 diff --git a/configs/recipes/decoder_bidirectional.yaml b/src/decisionformer/recipes/decoder_bidirectional.yaml similarity index 86% rename from configs/recipes/decoder_bidirectional.yaml rename to src/decisionformer/recipes/decoder_bidirectional.yaml index 04d40f2..10f78df 100644 --- a/configs/recipes/decoder_bidirectional.yaml +++ b/src/decisionformer/recipes/decoder_bidirectional.yaml @@ -1,8 +1,7 @@ # Decoder with bidirectional attention: every prompt token sees every other token, like an encoder. # Each option is pooled from its own tokens and scored, so any number of options works in any order. -# Experimental for decoders: without a long adaptation stage (configs/recipes/adapt_mntp.yaml) this underperforms +# Experimental for decoders: without a long adaptation stage (the adapt_mntp recipe) this underperforms # causal attention, often staying at chance; see docs/benchmarks.md. Prefer bf16/fp32: fp16 diverged here. -base: ../base.yaml model: backbone: HuggingFaceTB/SmolLM2-360M-Instruct diff --git a/configs/recipes/decoder_lora_code_readout.yaml b/src/decisionformer/recipes/decoder_lora_code_readout.yaml similarity index 84% rename from configs/recipes/decoder_lora_code_readout.yaml rename to src/decisionformer/recipes/decoder_lora_code_readout.yaml index 9bbf38e..de9e2dc 100644 --- a/configs/recipes/decoder_lora_code_readout.yaml +++ b/src/decisionformer/recipes/decoder_lora_code_readout.yaml @@ -1,7 +1,6 @@ # Causal decoder (default attention) with LoRA and the answer-code readout initialized from the LM head. -# decisionformer train configs/recipes/decoder_lora_code_readout.yaml -# torchrun --nproc-per-node 8 -m decisionformer train configs/recipes/decoder_lora_code_readout.yaml -base: ../base.yaml +# decisionformer train decoder_lora_code_readout +# torchrun --nproc-per-node 8 -m decisionformer train decoder_lora_code_readout model: backbone: Qwen/Qwen3-0.6B diff --git a/configs/recipes/decoder_qlora_multitask.yaml b/src/decisionformer/recipes/decoder_qlora_multitask.yaml similarity index 98% rename from configs/recipes/decoder_qlora_multitask.yaml rename to src/decisionformer/recipes/decoder_qlora_multitask.yaml index bd69225..8056ccc 100644 --- a/configs/recipes/decoder_qlora_multitask.yaml +++ b/src/decisionformer/recipes/decoder_qlora_multitask.yaml @@ -1,6 +1,5 @@ # QLoRA (4-bit base) with one head per task. Records choose their head with the `task` field. # Requires `pip install decisionformer[qlora]` and a CUDA GPU. -base: ../base.yaml model: backbone: Qwen/Qwen3-1.7B diff --git a/configs/recipes/encoder_full.yaml b/src/decisionformer/recipes/encoder_full.yaml similarity index 88% rename from configs/recipes/encoder_full.yaml rename to src/decisionformer/recipes/encoder_full.yaml index e679ede..0d2fbd3 100644 --- a/configs/recipes/encoder_full.yaml +++ b/src/decisionformer/recipes/encoder_full.yaml @@ -1,7 +1,6 @@ # Encoder baseline: full fine-tuning of ModernBERT on MNLI with a fixed label set. # python examples/prepare_datasets.py mnli --output data/mnli --train-size 50000 -# decisionformer train configs/recipes/encoder_full.yaml -base: ../base.yaml +# decisionformer train encoder_full model: backbone: answerdotai/ModernBERT-base diff --git a/configs/recipes/local_smollm2.yaml b/src/decisionformer/recipes/local_smollm2.yaml similarity index 92% rename from configs/recipes/local_smollm2.yaml rename to src/decisionformer/recipes/local_smollm2.yaml index 032c463..cf4447f 100644 --- a/configs/recipes/local_smollm2.yaml +++ b/src/decisionformer/recipes/local_smollm2.yaml @@ -1,9 +1,8 @@ # Fine-tune on a laptop (Mac or Linux): SmolLM2-135M-Instruct on FrontiersMind/deciformer-15k, full fine-tuning. # decisionformer data split hf:FrontiersMind/deciformer-15k:train --out data/deciformer-15k -# decisionformer train configs/recipes/local_smollm2.yaml +# decisionformer train local_smollm2 # The device is picked automatically: a CUDA GPU, else Apple silicon (MPS), else CPU; set DECISIONFORMER_DEVICE=cpu # to force the CPU. Small micro-batches, a per-batch token budget and gradient checkpointing keep memory low. -base: ../base.yaml model: backbone: HuggingFaceTB/SmolLM2-135M-Instruct diff --git a/src/decisionformer/schema.py b/src/decisionformer/schema.py index 3fad2fc..f814651 100644 --- a/src/decisionformer/schema.py +++ b/src/decisionformer/schema.py @@ -159,6 +159,24 @@ def _distribution(question: Question, value: Any, where: str) -> list[float] | N return [v / total for v in values] +def as_request(record: dict[str, Any], index: int = 0) -> tuple[Any, dict[str, Any]]: + """The `state` and System One `questions` of a record in any accepted shape: a request with `questions`, a single + `question` object, or a flat one-question row.""" + request_id = str(record.get("id", index)) + if "state" not in record: + raise DecisionInputError(f"record {request_id}: needs `state`") + key = str(record.get("question_key", "decision")) + if "questions" in record: + return record["state"], record["questions"] + if isinstance(record.get("question"), str) and "options" in record: + question, _, _ = _flat_question(key, record, request_id) + return record["state"], {key: {"type": question.type, "instructions": question.instructions, + "criteria": question.criteria}} + if isinstance(record.get("question"), dict): + return record["state"], {key: record["question"]} + raise DecisionInputError(f"record {request_id}: needs `questions` (or a single `question`)") + + FLAT_FIELDS = ("question_type", "options", "option_descriptions", "label_index", "label_code", "soft_target") diff --git a/src/decisionformer/serve.py b/src/decisionformer/serve.py index 00899d8..3e0e82b 100644 --- a/src/decisionformer/serve.py +++ b/src/decisionformer/serve.py @@ -6,16 +6,22 @@ Capacity refusals use HTTP 422 with the phrases the Decision Index `http` engine recognizes ("maximum context length", "options per choice"), so over-long requests count as unsupported instead of failing the run. -Requests run one at a time: the latency the Index measures is single-request latency. +A request that arrives while the model is idle runs at once, alone, so single-request latency is unchanged; requests +that arrive while it is busy are answered together in one batch. The request and response schemas are in /docs. """ +import asyncio import hmac import os +import queue import threading import time import uuid -from typing import Any +from concurrent.futures import Future +from typing import Any, Literal + +from pydantic import BaseModel, Field from decisionformer import __version__ from decisionformer.api import Decider, UnsupportedInput @@ -23,15 +29,102 @@ from decisionformer.schema import DecisionInputError -def create_app(decider: Any, *, api_key: str | None = None) -> Any: +# ----------------------------------------------------------------------------------------------- API schema (docs) +class QuestionSpec(BaseModel): + type: Literal["choice", "noul", "score"] = Field(description="choice: pick one option; noul: yes/no; " + "score: pick one of 2-10 ordered levels") + instructions: Any = Field(None, description="the question, in words") + criteria: dict[str, Any] | list[Any] | None = Field( + None, description="choice: option -> description (or null), or a list of options; score: the ordered levels; " + "noul: optional descriptions of 'false' and 'true'") + + +class SystemOneRequest(BaseModel): + state: Any = Field(description="what to decide about: text or any JSON value") + questions: dict[str, QuestionSpec] = Field(description="question id -> question") + + model_config = {"json_schema_extra": {"examples": [{ + "state": "I was charged twice for my subscription this month.", + "questions": {"angry": {"type": "noul", "instructions": "Is the customer angry?"}, + "team": {"type": "choice", "criteria": {"billing": "Payments and refunds", "technical": None}}}, + }]}} + + +class Answer(BaseModel): + type: Literal["choice", "noul", "score"] + choice: str | None = Field(None, description="choice: the most probable option") + noul: float | None = Field(None, description="noul: the probability of yes") + score: float | None = Field(None, description="score: the expected level, 0-based") + probabilities: dict[str, float] | None = Field(None, description="choice and score: per option or level") + confidence: float | None = None + legend: dict[str, Any] | None = Field(None, description="score: level index -> level") + + +class SystemOneResponse(BaseModel): + id: str + object: Literal["structured_decision"] + created: int + model: str + answers: dict[str, Answer] + usage: dict[str, int] + + +class ErrorResponse(BaseModel): + error: dict[str, str] + + +# ----------------------------------------------------------------------------------------------- batching +class Batcher: + """Answers requests on one worker thread. The worker takes the next request and every request already queued + behind it (up to `max_batch`) and answers them in one `decide_many` call; if that fails, it answers them one by + one so each request gets its own result or error.""" + + def __init__(self, decider: Any, max_batch: int = 16) -> None: + self.decider = decider + self.max_batch = max(1, max_batch) + self.queue: queue.Queue[tuple[dict[str, Any], Future]] = queue.Queue() + threading.Thread(target=self._run, name="decisionformer-batcher", daemon=True).start() + + def submit(self, request: dict[str, Any]) -> Future: + future: Future = Future() + self.queue.put((request, future)) + return future + + def _answer(self, requests: list[dict[str, Any]]) -> list[Any]: + try: + return list(self.decider.decide_many(requests)) + except Exception as error: # noqa: BLE001 - delivered to the request(s) it belongs to + if len(requests) == 1: + return [error] + return [self._answer([request])[0] for request in requests] + + def _run(self) -> None: + while True: + batch = [self.queue.get()] + while len(batch) < self.max_batch: + try: + batch.append(self.queue.get_nowait()) + except queue.Empty: + break + for (_, future), result in zip(batch, self._answer([request for request, _ in batch])): + if isinstance(result, BaseException): + future.set_exception(result) + else: + future.set_result(result) + + +# ----------------------------------------------------------------------------------------------- app +def create_app(decider: Any, *, api_key: str | None = None, max_batch: int = 16) -> Any: try: from fastapi import FastAPI, Request + from fastapi.openapi.utils import get_openapi from fastapi.responses import JSONResponse except ImportError as error: raise ImportError("`decisionformer serve` needs `pip install decisionformer[serve]`") from error - app = FastAPI(title="decisionformer", version=__version__) - lock = threading.Lock() + app = FastAPI(title="decisionformer", version=__version__, + description="Answer System One requests: typed questions about a state, one forward pass each.") + batcher = Batcher(decider, max_batch) def error(status: int, message: str, kind: str) -> Any: return JSONResponse(status_code=status, content={"error": {"message": message, "type": kind}}) @@ -44,13 +137,14 @@ def authorized(request: Request) -> bool: @app.get("/health") def health() -> dict[str, Any]: return {"status": "ready", "model": decider.name, "max_input_tokens": decider.max_input_tokens, - "max_options": len(decider.builder.codes), "prompt_format": decider.builder.format.name} + "max_options": len(decider.builder.codes), "prompt_format": decider.builder.format.name, + "max_batch": batcher.max_batch} @app.get("/v1/models") def models() -> dict[str, Any]: return {"object": "list", "data": [{"id": decider.name, "object": "model", "owned_by": "decisionformer"}]} - @app.post("/v1/systemone") + @app.post("/v1/systemone", summary="Answer every question about a state") async def systemone(request: Request) -> Any: if not authorized(request): return error(401, "missing or invalid API key", "authentication_error") @@ -62,16 +156,45 @@ async def systemone(request: Request) -> Any: return error(400, "a request needs `state` and `questions`", "invalid_request") started = time.perf_counter() try: - with lock: - response = decider.decide(body.get("state"), body["questions"], images=body.get("images")) + response = await asyncio.wrap_future(batcher.submit( + {"state": body.get("state"), "questions": body["questions"], "images": body.get("images")})) except DecisionInputError as exc: return error(400, str(exc), "invalid_request") except (ContextLengthError, UnsupportedInput) as exc: return error(422, str(exc), "unsupported_request") - response.update(id=f"decision-{uuid.uuid4().hex}", object="structured_decision", created=int(time.time())) + response = {**response, "id": f"decision-{uuid.uuid4().hex}", "object": "structured_decision", + "created": int(time.time())} return JSONResponse(content=response, headers={ "server-timing": f"total;dur={(time.perf_counter() - started) * 1000:.1f}"}) + def openapi() -> dict[str, Any]: + # The endpoint parses requests itself (its error codes and messages are part of the protocol), so the + # schemas are added to the generated document for /docs instead of being enforced by FastAPI. + if app.openapi_schema: + return app.openapi_schema + schema = get_openapi(title=app.title, version=app.version, description=app.description, routes=app.routes) + components = schema.setdefault("components", {}).setdefault("schemas", {}) + for model in (SystemOneRequest, SystemOneResponse, ErrorResponse): + body = model.model_json_schema(ref_template="#/components/schemas/{model}") + components.update(body.pop("$defs", {})) + components[model.__name__] = body + + def content(name: str) -> dict[str, Any]: + return {"application/json": {"schema": {"$ref": f"#/components/schemas/{name}"}}} + + operation = schema["paths"]["/v1/systemone"]["post"] + operation["requestBody"] = {"required": True, "content": content("SystemOneRequest")} + operation["responses"] = { + "200": {"description": "One answer per question", "content": content("SystemOneResponse")}, + "400": {"description": "Malformed request", "content": content("ErrorResponse")}, + "401": {"description": "Missing or invalid API key", "content": content("ErrorResponse")}, + "422": {"description": "Too long or too many options for this model (never truncated)", + "content": content("ErrorResponse")}, + } + app.openapi_schema = schema + return schema + + app.openapi = openapi return app @@ -89,8 +212,9 @@ def load_decider(model: str, *, backend: str = "transformers", device: str | Non def serve(model: str, *, host: str = "127.0.0.1", port: int = 8000, device: str | None = None, - backend: str = "transformers", api_key_env: str = "DECISIONFORMER_API_KEY", **load_options: Any) -> None: + backend: str = "transformers", api_key_env: str = "DECISIONFORMER_API_KEY", max_batch: int = 16, + **load_options: Any) -> None: decider = load_decider(model, backend=backend, device=device, **load_options) # model errors first import uvicorn - uvicorn.run(create_app(decider, api_key=os.environ.get(api_key_env)), host=host, port=port) + uvicorn.run(create_app(decider, api_key=os.environ.get(api_key_env), max_batch=max_batch), host=host, port=port) diff --git a/src/decisionformer/trainer.py b/src/decisionformer/trainer.py index 14b48a8..5ca58d8 100644 --- a/src/decisionformer/trainer.py +++ b/src/decisionformer/trainer.py @@ -9,6 +9,7 @@ import random import shutil import time +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Any @@ -17,7 +18,7 @@ from torch.nn.parallel import DistributedDataParallel from decisionformer import distributed as dist_utils -from decisionformer.config import RunConfig, dump_config +from decisionformer.config import RunConfig, dump_config, load_config from decisionformer.data import Collator, DecisionDataset, plan_epoch from decisionformer.finetune import count_parameters, effective_precision from decisionformer.inference import ( @@ -41,6 +42,18 @@ def _setup_logging(is_main: bool) -> None: logging.basicConfig(level=logging.INFO if is_main else logging.WARNING, format="%(asctime)s %(levelname)s %(name)s: %(message)s", force=True) + for name in ("httpx", "httpx2", "httpcore", "httpcore2", "huggingface_hub", "urllib3", "filelock"): # Hub requests + logging.getLogger(name).setLevel(logging.WARNING) + + +PRECISION_NAMES = {"no": "float32", "bf16": "bf16 autocast", "fp16": "fp16 autocast with loss scaling"} + + +def _duration(seconds: float) -> str: + seconds = int(seconds) + if seconds >= 3600: + return f"{seconds // 3600}h{seconds % 3600 // 60:02d}m" + return f"{seconds // 60}m{seconds % 60:02d}s" if seconds >= 60 else f"{seconds}s" def _seed_everything(seed: int) -> None: @@ -59,6 +72,27 @@ def _gather_to_main(value: Any) -> list[Any]: return gathered or [] +def _warn_about_leakage(config: RunConfig) -> None: + """Warn when training data shares ids, families or identical states with the evaluation or calibration data + (`decisionformer data check` also looks for near-duplicates).""" + from decisionformer.datacheck import check_splits + + splits = {name: path for name, path in config.data.eval.items() if isinstance(path, str)} + if isinstance(config.data.calibration, str): + splits["calibration"] = config.data.calibration + if not splits or not isinstance(config.data.train, str): + return + try: + report = check_splits({**splits, "train": config.data.train}, near_duplicates=False) + except Exception as error: # noqa: BLE001 - a check must never stop training + logger.debug("Leakage check skipped: %s", error) + return + for problem in report["problems"]: + logger.warning("%s shares %d %s value(s) with an earlier split (e.g. %s); evaluation on it will be " + "optimistic. Run `decisionformer data check` for details.", problem["split"], + problem["count"], problem["overlap"], ", ".join(map(str, problem["examples"][:3]))) + + class Trainer: def __init__(self, config: RunConfig) -> None: self.config = config @@ -77,6 +111,8 @@ def __init__(self, config: RunConfig) -> None: self.train_set = DecisionDataset.from_paths(config.data.train, builder, labeled_only=True) self.eval_sets = {name: DecisionDataset.from_paths(paths, builder) for name, paths in config.data.eval.items()} self.calibration = DecisionDataset.from_paths(config.data.calibration, builder) if config.data.calibration else None + if self.state.is_main: + _warn_about_leakage(config) pad = self.result.tokenizer.pad_token_id self.train_collator = Collator(builder, pad, shuffle_options=config.data.shuffle_options, seed=config.train.seed) self.eval_collator = Collator(builder, pad) @@ -138,11 +174,12 @@ def train(self) -> dict[str, Any]: logger.info("No trainable parameters: evaluation-only run (zero-shot baseline)") self.evaluate() return self._finish() - logger.info("Training %d steps (%d per epoch) on %d process(es), device %s, precision %s", self.total_steps, - self.steps_per_epoch, state.world_size, state.device.type, self.precision) + logger.info("Training %d steps (%d per epoch) on %d process(es), device %s, %s", self.total_steps, + self.steps_per_epoch, state.world_size, state.device.type, PRECISION_NAMES[self.precision]) if config.eval.at_start and self.step == 0: self.evaluate() window_loss, window_count, window_start = 0.0, 0.0, time.monotonic() + started, first_step = time.monotonic(), self.step while self.step < self.total_steps: epoch, _ = self._plan_for(self.step) loss_sum, count = self.accumulate_gradients(self.step) @@ -152,16 +189,20 @@ def train(self) -> dict[str, Any]: self.step += 1 if self.step % config.train.log_every == 0 or self.step in (1, self.total_steps): loss_sum, count = dist_utils.all_reduce_sum([window_loss, window_count], state.device) - self._log({"event": "train", "step": self.step, "epoch": epoch, "loss": loss_sum / max(count, 1), - "lr": self.scheduler.get_last_lr()[0], "grad_norm": grad_norm, - "seconds": time.monotonic() - window_start}) + per_step = (time.monotonic() - started) / (self.step - first_step) + self._log({"event": "train", "step": self.step, "total_steps": self.total_steps, "epoch": epoch, + "loss": loss_sum / max(count, 1), "lr": self.scheduler.get_last_lr()[0], + "grad_norm": grad_norm, "seconds": time.monotonic() - window_start, + "eta_seconds": per_step * (self.total_steps - self.step)}) window_loss, window_count, window_start = 0.0, 0.0, time.monotonic() evaluate_now = config.train.eval_every and self.step % config.train.eval_every == 0 if evaluate_now or self.step == self.total_steps: self.evaluate() - if config.train.save_every and (self.step % config.train.save_every == 0 or self.step == self.total_steps): - self.save_checkpoint() - if config.train.stop_after is not None and self.step >= config.train.stop_after: + pausing = config.train.stop_after is not None and self.step >= config.train.stop_after + if self.step < self.total_steps and (pausing or (config.train.save_every + and self.step % config.train.save_every == 0)): + self.save_checkpoint() # resumable state; a finished run needs none (best/ and last/ are exported) + if pausing: logger.info("Pausing at step %d (train.stop_after); resume with train.resume=auto", self.step) break return self._finish() @@ -302,7 +343,12 @@ def _resume(self) -> None: found = sorted(p for p in (self.output / "checkpoints").glob("step-*") if p.is_dir() and not p.name.endswith(".tmp")) if not found: - logger.info("resume=auto: no checkpoint found, starting fresh") + finished = self.output / "summary.json" + if finished.exists() and json.loads(finished.read_text()).get("complete"): + logger.warning("%s holds a finished run; training it again from step 0 overwrites best/ and " + "last/ (set train.output_dir to keep it)", self.output) + else: + logger.info("resume=auto: no checkpoint found, starting fresh") return path = found[-1] else: @@ -337,7 +383,9 @@ def _log(self, record: dict[str, Any]) -> None: with (self.output / "metrics.jsonl").open("a") as stream: stream.write(line + "\n") if record.get("event") == "train": - logger.info("step %(step)d loss %(loss).4f lr %(lr).3g grad_norm %(grad_norm).3f", record) + logger.info("step %d/%d loss %.4f lr %.3g grad_norm %.3f, %s left (training only)", record["step"], + record["total_steps"], record["loss"], record["lr"], record["grad_norm"], + _duration(record["eta_seconds"])) else: for name, entry in record.get("sets", {}).items(): for flavor, metrics in entry.items(): @@ -357,6 +405,8 @@ def _finish(self) -> dict[str, Any]: shutil.copytree(self.output / "last", self.output / "best", dirs_exist_ok=True) if config.finetune.export == "merged" and config.finetune.mode in ("lora", "qlora"): self._merge_exports() + if self.step == self.total_steps and (self.output / "checkpoints").exists(): + shutil.rmtree(self.output / "checkpoints") # only needed to resume an unfinished run summary = {"steps": self.step, "complete": self.step == self.total_steps, "best_step": self.best_step, "best": self.best, "temperatures": self.temperatures, "output_dir": str(self.output)} if config.train.push_to_hub and (self.output / "best").exists(): @@ -382,7 +432,18 @@ def _merge_exports(self) -> None: merge_checkpoint(self.output / name, device="cpu", dtype=dtype, save_dtype=export_dtype(self.config)) -def train(config: RunConfig) -> dict[str, Any]: +def train(config: RunConfig | str | Path | Mapping[str, Any], + overrides: Sequence[str] | Mapping[str, Any] | None = None) -> dict[str, Any]: + """Fine-tune and export a decision model; returns the run summary. + + decisionformer.train("deciformer15k_qwen3") # a bundled recipe + decisionformer.train("my.yaml", {"train.lr": 1e-4, "train.max_steps": 200}) # a file, with overrides + decisionformer.train({"model": {...}, "data": {...}, "train": {"output_dir": "runs/x"}}) + """ + if not isinstance(config, RunConfig): + config = load_config(config, overrides) + elif overrides: + config = load_config(config.model_dump(mode="json"), overrides) trainer = Trainer(config) try: return trainer.train() diff --git a/src/decisionformer/vllm_backend.py b/src/decisionformer/vllm_backend.py index 121f257..0f643ee 100644 --- a/src/decisionformer/vllm_backend.py +++ b/src/decisionformer/vllm_backend.py @@ -17,7 +17,7 @@ import torch -from decisionformer.api import Decider, _local_dir, _prompt_name +from decisionformer.api import Decider, _local_dir, _prompt_name, display_name from decisionformer.formats import get_format from decisionformer.model import DECISION_CONFIG, HEADS, READOUT from decisionformer.prompting import Encoded, PromptBuilder, prepare_tokenizer @@ -102,7 +102,7 @@ def from_pretrained(cls, name_or_path: str | Path, *, revision: str | None = Non builder = PromptBuilder(tokenizer, get_format(prompt or _prompt_name(config)), "decoder", max_length=limit, overflow="error", codes=config["codes"], code_token_ids=config["token_ids"], chat_template_kwargs=ext.get("chat_template_kwargs")) - name = ext.get("name") or (str(name_or_path) if not Path(str(name_or_path)).is_dir() else root.name) + name = ext.get("name") or display_name(name_or_path, root) return cls(llm, builder, tokenizer, config, name=name, max_input_tokens=limit, path=root) def code_logits(self, decisions: list[Decision], encoded: list[Encoded]) -> list[torch.Tensor]: diff --git a/tests/test_api.py b/tests/test_api.py index a6ecbf6..f6de28f 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -183,6 +183,10 @@ def test_systemone_server(trained): assert too_long.status_code == 422 and "maximum context length" in too_long.text bad = client.post("/v1/systemone", json={"state": "x", "questions": {"q": {"type": "score", "criteria": ["a"]}}}) assert bad.status_code == 400 and "a score takes 2 to 10 levels" in bad.text + schema = client.get("/openapi.json").json() + request_schema = schema["paths"]["/v1/systemone"]["post"]["requestBody"]["content"]["application/json"]["schema"] + assert request_schema["$ref"].endswith("SystemOneRequest") + assert {"SystemOneRequest", "SystemOneResponse", "QuestionSpec", "Answer"} <= set(schema["components"]["schemas"]) locked = TestClient(create_app(model, api_key="secret")) assert locked.post("/v1/systemone", json=body).status_code == 401 assert locked.post("/v1/systemone", json=body, headers={"Authorization": "Bearer secret"}).status_code == 200 @@ -208,6 +212,20 @@ def test_predict_and_eval_commands(trained, data_dir, tmp_path, capsys): main(["predict", str(trained / "best"), str(requests), "--output", str(tmp_path / "out.jsonl"), "--device", "cpu"]) out = json.loads((tmp_path / "out.jsonl").read_text()) assert out["id"] == "r1" and set(out["answers"]) == set(QUESTIONS) + assert out["model"] == trained.name, "a local runs//best answers as " + # every accepted shape (request, single question, flat row), and a record that cannot be answered + mixed = tmp_path / "mixed.jsonl" + mixed.write_text("".join(json.dumps(r) + "\n" for r in [ + {"id": "flat", "state": "the tiny square is green .", "question": "What is the color ?", + "question_type": "choice", "options": ["red", "green"], "option_descriptions": [None, None], "label": "green"}, + {"id": "single", "state": "the tiny square is green .", "question": {"type": "noul"}}, + {"id": "broken", "state": "x", "questions": {"q": {"type": "choice"}}}, + ])) + main(["predict", str(trained / "best"), str(mixed), "--output", str(tmp_path / "mixed.out"), "--device", "cpu"]) + rows = {r["id"]: r for r in map(json.loads, (tmp_path / "mixed.out").read_text().splitlines())} + assert set(rows["flat"]["answers"]["decision"]["probabilities"]) == {"red", "green"} + assert rows["single"]["answers"]["decision"]["type"] == "noul" + assert "criteria" in rows["broken"]["error"]["message"] and "answers" not in rows["broken"] main(["eval", str(trained / "best"), str(data_dir / "multi_dev.jsonl"), "--device", "cpu"]) report = json.loads(capsys.readouterr().out) assert report["overall"]["count"] == 48 and set(report["types"]) == {"choice", "noul", "score"} @@ -237,3 +255,43 @@ def create_commit(self, repo_id, commit_message, operations): assert push_folder(trained / "best", "org/decider") == "https://huggingface.co/org/decider" assert calls["private"] is True and "README.md" in calls["ignored"] assert 'AutoDecisionModel.from_pretrained("org/decider")' in calls["card"] and "" not in calls["card"] + + +def test_server_batches_concurrent_requests_and_isolates_errors(trained): + """Requests that queue up while the model is busy share one decide_many call; a bad one fails alone.""" + import threading + from concurrent.futures import ThreadPoolExecutor + + from fastapi.testclient import TestClient + + from decisionformer.serve import Batcher, create_app + + model = AutoDecisionModel.from_pretrained(trained / "best", device="cpu", dtype=torch.float32, max_length=200) + original, release, calls = model.decide_many, threading.Event(), [] + + def recording(requests): + calls.append(len(requests)) + release.wait(60) # the first call holds the worker until every other request is queued + return original(requests) + + model.decide_many = recording + batcher = Batcher(model, max_batch=8) + good = {"state": "the huge triangle is yellow .", "questions": {"red": QUESTIONS["red"]}} + too_long = {"state": " ".join(["red"] * 300), "questions": {"q": {"type": "noul"}}} + first = batcher.submit(good) + while not calls: + threading.Event().wait(0.01) + queued = [batcher.submit(request) for request in (good, too_long, good)] + release.set() + assert first.result(timeout=60)["answers"]["red"]["type"] == "noul" + assert queued[0].result(timeout=60)["answers"]["red"]["type"] == "noul" + with pytest.raises(ContextLengthError): + queued[1].result(timeout=60) + assert queued[2].result(timeout=60)["answers"]["red"]["type"] == "noul" + assert calls == [1, 3, 1, 1, 1], "one alone, three queued as one batch, then one by one after the failure" + + model.decide_many = original + client = TestClient(create_app(model, max_batch=4)) + with ThreadPoolExecutor(8) as pool: + statuses = list(pool.map(lambda _: client.post("/v1/systemone", json=good).status_code, range(8))) + assert statuses == [200] * 8 and client.get("/health").json()["max_batch"] == 4 diff --git a/tests/test_training.py b/tests/test_training.py index 485756d..b9cc5f5 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -221,6 +221,9 @@ def run(name: str, stops: list[int | None]): "data": {"eval": {}}}, overrides=["data.calibration=null"]) train(config) + paused = stop is not None + assert (Path(config.train.output_dir) / "checkpoints").exists() == paused, \ + "a pause keeps resumable state; a finished run removes it" return load_file(str(Path(config.train.output_dir) / "last" / "readout.safetensors")), \ load_file(str(Path(config.train.output_dir) / "last" / "model.safetensors")) @@ -300,3 +303,16 @@ def test_exports_are_saved_in_the_export_dtype(model_root, data_dir, tmp_path, m AutoDecisionModel.from_pretrained(best, device="cpu", dtype=torch.float32).decide( "the small circle is blue .", {"red": {"type": "noul", "instructions": "Is it red ?"}}) + + +def test_train_from_python_with_a_dict_config(model_root, data_dir, tmp_path): + import decisionformer + + summary = decisionformer.train({ + "model": {"backbone": str(build_decoder(model_root, "llama")), "head": {"type": "code_readout"}}, + "finetune": {"mode": "lora", "lora": {"r": 4, "lora_dropout": 0.0}}, + "data": {"train": str(data_dir / "train.jsonl"), "eval": {"dev": str(data_dir / "dev.jsonl")}}, + "train": {"output_dir": str(tmp_path / "run"), "global_batch_size": 8, "micro_batch_size": 8}, + }, {"train.max_steps": 2}) + assert summary["steps"] == 2 and summary["complete"] + assert (tmp_path / "run" / "best" / "decision_config.json").exists() diff --git a/tests/test_units.py b/tests/test_units.py index 22b49ed..2df2515 100644 --- a/tests/test_units.py +++ b/tests/test_units.py @@ -404,3 +404,72 @@ def test_device_choice(monkeypatch): assert pick_device(0, world_size=2).type == "cpu", "multi-process runs (gloo) stay on CPU" monkeypatch.setenv("DECISIONFORMER_DEVICE", "cpu") assert pick_device().type == "cpu" + + +def test_bundled_recipes_load_by_name_and_resolve_defaults(): + from decisionformer.config import config_path, list_recipes + + recipes = list_recipes() + assert "deciformer15k_qwen3" in recipes and "local_smollm2" in recipes and all(recipes.values()) + for name in recipes: + config = load_config(name) # every bundled recipe is valid as shipped + assert config.train.lr is not None and config.train.resume == "auto" + assert load_config("deciformer15k_qwen3.yaml").model.backbone == "Qwen/Qwen3-0.6B" + with pytest.raises(FileNotFoundError, match="decisionformer init NAME"): + config_path("no_such_recipe") + + +def test_configs_from_dicts_with_overrides_and_mode_learning_rates(tmp_path): + base = {"model": {"backbone": "x", "head": {"type": "option_scoring"}}, "train": {"output_dir": "out"}} + assert load_config(base).train.lr == 2e-5 # full fine-tuning + assert load_config({**base, "finetune": {"mode": "lora"}}).train.lr == 2e-4 + assert load_config({**base, "finetune": {"mode": "head_only"}}).train.lr == 1e-3 + tuned = load_config(base, {"train.lr": 5e-5, "train.max_steps": 10, "data.eval": {"dev": "d.jsonl"}}) + assert tuned.train.lr == 5e-5 and tuned.train.max_steps == 10 and tuned.data.eval == {"dev": "d.jsonl"} + # a user file can inherit from a bundled recipe + child = _write(tmp_path / "mine.yaml", "base: local_smollm2\ntrain: {output_dir: runs/mine, max_steps: 5}\n") + mine = load_config(child) + assert mine.model.backbone == "HuggingFaceTB/SmolLM2-135M-Instruct" and mine.train.max_steps == 5 + defaults = load_config(base) + assert (defaults.train.save_every, defaults.train.eval_every, defaults.data.max_length) == (500, 500, 1024) + + +def test_recipes_and_init_commands(tmp_path, capsys, monkeypatch): + from decisionformer.cli import main + + main(["recipes"]) + assert "deciformer15k_qwen3" in capsys.readouterr().out + monkeypatch.chdir(tmp_path) + main(["init", "local_smollm2"]) + assert (tmp_path / "local_smollm2.yaml").read_text().startswith("# Fine-tune on a laptop") + with pytest.raises(SystemExit, match="exists"): + main(["init", "local_smollm2"]) + main(["init", "local_smollm2", "my.yaml"]) + assert load_config(tmp_path / "my.yaml").model.backbone.endswith("SmolLM2-135M-Instruct") + with pytest.raises(SystemExit, match="neither a config file nor a bundled recipe"): + main(["train", "missing.yaml"]) + + +def test_training_warns_about_overlap_between_splits(tmp_path, caplog): + from conftest import synthetic_records, write_jsonl + + from decisionformer.trainer import _warn_about_leakage + + records = synthetic_records(16, 1) + write_jsonl(tmp_path / "train.jsonl", records) + write_jsonl(tmp_path / "dev.jsonl", records[:4]) + config = load_config({"model": {"backbone": "x", "head": {"type": "option_scoring"}}, + "data": {"train": str(tmp_path / "train.jsonl"), "eval": {"dev": str(tmp_path / "dev.jsonl")}}, + "train": {"output_dir": "out"}}) + with caplog.at_level("WARNING", logger="decisionformer"): + _warn_about_leakage(config) + assert "train shares 4 id value(s)" in caplog.text + + +def test_local_runs_answer_with_their_run_name(tmp_path): + from decisionformer.api import display_name + + best = tmp_path / "runs" / "my-decider" / "best" + best.mkdir(parents=True) + assert display_name(best, best) == "my-decider" + assert display_name("org/model", best) == "org/model"