From 26d11bd7efdbf5137919b3cec98ed37eaf3abd44 Mon Sep 17 00:00:00 2001 From: Abhay Kumar Date: Sat, 24 Oct 2026 01:25:27 +0530 Subject: [PATCH] Developer experience: bundled recipes, safe defaults, Python training API, server batching - Recipes ship in the package (src/decisionformer/recipes): train, config and adapt accept a recipe name; decisionformer recipes lists them and decisionformer init copies one to edit. Configs can inherit a recipe. - configs/base.yaml is gone: its values are the code defaults (resumable state and evaluation every 500 steps, resume auto, max_length 1024); the learning rate follows the fine-tuning mode. Resumable state is removed when a run finishes, and rerunning a finished run warns first. - decisionformer.train() takes a recipe name, a YAML file or a dict, with dotted overrides; load_config accepts the same. - New starting recipe deciformer15k_qwen3 (Qwen3-0.6B + LoRA): 0.824 dev accuracy on deciformer-15k in one epoch. The Qwen3.5 recipe is removed until Qwen3.5 fine-tuning learns reliably (docs/benchmarks.md). - predict reads every accepted record shape, batches requests and reports records it cannot answer as error lines instead of crashing. - serve batches requests that arrive while the model is busy (--max-batch) and documents the request and response schemas at /docs. - Logs: Hub request lines are silenced, precision is named, and training steps show progress and the time left. Training warns about overlap between the training data and the evaluation or calibration data. - Local run exports answer with their run name instead of best/last. - CLI help shows examples and the recipes; the README installs from GitHub until the first PyPI release. --- CONTRIBUTING.md | 7 + README.md | 55 ++++--- configs/base.yaml | 26 --- configs/benchmarks/mnli_8gb/_base.yaml | 1 - configs/recipes/deciformer15k_qwen35.yaml | 31 ---- docs/benchmarks.md | 22 +++ src/decisionformer/__init__.py | 8 +- src/decisionformer/adapt.py | 4 +- src/decisionformer/api.py | 9 +- src/decisionformer/cli.py | 116 ++++++++++++-- src/decisionformer/config.py | 72 +++++++-- .../decisionformer}/recipes/adapt_mntp.yaml | 3 +- .../recipes/deciformer15k_qwen3.yaml | 25 +++ .../recipes/decoder_bidirectional.yaml | 3 +- .../recipes/decoder_lora_code_readout.yaml | 5 +- .../recipes/decoder_qlora_multitask.yaml | 1 - .../decisionformer}/recipes/encoder_full.yaml | 3 +- .../recipes/local_smollm2.yaml | 3 +- src/decisionformer/schema.py | 18 +++ src/decisionformer/serve.py | 148 ++++++++++++++++-- src/decisionformer/trainer.py | 85 ++++++++-- src/decisionformer/vllm_backend.py | 4 +- tests/test_api.py | 58 +++++++ tests/test_training.py | 16 ++ tests/test_units.py | 69 ++++++++ 25 files changed, 641 insertions(+), 151 deletions(-) delete mode 100644 configs/base.yaml delete mode 100644 configs/recipes/deciformer15k_qwen35.yaml rename {configs => src/decisionformer}/recipes/adapt_mntp.yaml (92%) create mode 100644 src/decisionformer/recipes/deciformer15k_qwen3.yaml rename {configs => src/decisionformer}/recipes/decoder_bidirectional.yaml (86%) rename {configs => src/decisionformer}/recipes/decoder_lora_code_readout.yaml (84%) rename {configs => src/decisionformer}/recipes/decoder_qlora_multitask.yaml (98%) rename {configs => src/decisionformer}/recipes/encoder_full.yaml (88%) rename {configs => src/decisionformer}/recipes/local_smollm2.yaml (92%) 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"