Laptop fine-tuning (Mac and Linux), backbone policy and order

README.md

@@ -128,18 +128,20 @@ decisionformer serve vllm-sr/d3-edge --backend vllm
 ## Supported backbones
 
 Every family below is fine-tuned, exported, reloaded and checked against the Decision Index's response rules in
-`tests/test_zoo.py`. Newest first:
+`tests/test_zoo.py`, newest first within each group:
 
 | | families |
 |---|---|
-| **Decoders, hybrid linear attention** | Qwen3.5 (saved whole, with its vision tower), Qwen3-Next |
-| **Decoders, mixture of experts** | GPT-OSS, Llama 4, DeepSeek-V3, Qwen3-MoE, Granite-MoE, OLMoE, Phi-MoE, Mixtral |
 | **Decoders, dense** | Gemma 4, Gemma 3, Qwen3, SmolLM3, Seed-OSS, EXAONE 4, GLM-4, Ministral, Cohere 2, Hunyuan, OLMo 2, Granite, Phi-3, Qwen2, Gemma 2, Gemma, Mistral, Llama, Phi, StableLM, Falcon, GPT-NeoX, GPT-2 |
+| **Decoders, mixture of experts** | GPT-OSS, Llama 4, DeepSeek-V3, Qwen3-MoE, Granite-MoE, OLMoE, Phi-MoE, Mixtral |
+| **Decoders, hybrid linear attention** | Qwen3.5 (saved whole, with its vision tower), Qwen3-Next |
 | **Encoders** | ModernBERT, EuroBERT, DeBERTa-v3 (and v1/v2), XLM-RoBERTa, RoBERTa, ELECTRA, ALBERT, DistilBERT, BERT, T5 and mT5 (the encoder) |
 
-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.
+Backbones are the architectures built into Hugging Face transformers, loaded with its `Auto` classes and saved with
+`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.
 
 ## Fine-tune
 
@@ -174,6 +176,25 @@ A run writes `runs/<name>/best` and `last`: finished checkpoints in the standard
 `from_pretrained`, `serve` and `push`. It also writes `metrics.jsonl`, per-example predictions at every
 evaluation, resumable state and `summary.json`.
 
+### On a laptop (Mac or Linux)
+
+A small model trains on a laptop. This recipe fully fine-tunes SmolLM2-135M-Instruct on deciformer-15k for 300
+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 serve runs/local-smollm2-135m/best        # try it on /v1/systemone
+```
+
+The device is picked automatically: a CUDA GPU, then Apple silicon (MPS), then the CPU. Set
+`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.
+
 ### Training data
 
 Training data is the request with the gold answers, one JSON object per line. One record, spread out here for
@@ -272,7 +293,9 @@ its own pooler (`decision_token`, `last_token`, `cls`, `mean`, `attention`, `opt
 
 `data.prompt.style` picks the prompt: `d3` (the default: Decision 3.0's byte-exact prompt, servable by vLLM) or
 `autojev` (AutoJev's). Decoders use the tokenizer's own chat template, and `chat_template_kwargs` defaults to
-`{enable_thinking: false}`. Register your own format with `decisionformer.register_format`.
+`{enable_thinking: false}`. Register your own format with `decisionformer.register_format`. Training examples
+longer than `data.max_length` have the end of their state cut (`data.overflow`); loaded and served models never
+truncate and refuse what does not fit instead.
 
 ### Decoder attention
 

configs/recipes/local_smollm2.yaml

@@ -0,0 +1,34 @@
+# 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
+# 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
+  head: {type: code_readout}
+
+finetune:
+  mode: full
+  gradient_checkpointing: true
+
+data:
+  train: data/deciformer-15k/train.jsonl
+  calibration: data/deciformer-15k/calibration.jsonl
+  eval: {dev: data/deciformer-15k/dev.jsonl}
+  max_length: 1024
+
+train:
+  output_dir: runs/local-smollm2-135m
+  max_steps: 300
+  global_batch_size: 16
+  micro_batch_size: 4
+  token_budget: 4096
+  lr: 5.0e-5
+  head_lr: 1.0e-4
+  eval_every: 100
+  save_every: 100
+
+eval:
+  micro_batch_size: 8

src/decisionformer/api.py

@@ -21,6 +21,7 @@
 from decisionformer.backbone import Backbone, context_limit, load_backbone, resolve_kind
 from decisionformer.config import HeadConfig, PoolerConfig
 from decisionformer.data import make_batch, split_micro_batches
+from decisionformer.distributed import pick_device
 from decisionformer.formats import get_format
 from decisionformer.heads import CodeReadoutHead, build_head
 from decisionformer.model import DECISION_CONFIG, HEADS, READOUT, DecisionModel, model_card
@@ -90,7 +91,7 @@ def from_pretrained(cls, name_or_path: str | Path, *, revision: str | None = Non
         config = json.loads(config_file.read_text())
         ext = config.get("decisionformer") or {}
         style = get_format(prompt or _prompt_name(config))
-        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        device = torch.device(device) if device else pick_device()
         dtype = dtype or _default_dtype(device)
         tokenizer_dir = root if (root / "tokenizer_config.json").is_file() else root / "processor"
         tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=trust_remote_code)

src/decisionformer/config.py

@@ -76,7 +76,9 @@ class DataConfig(_Strict):
     calibration: str | list[str] | None = None
     prompt: PromptConfig = Field(default_factory=PromptConfig)
     max_length: int = 2048
-    overflow: Literal["error", "truncate_state_end", "truncate_state_start"] = "error"
+    # 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"
     max_options: int = Field(255, ge=2, le=702)
     shuffle_options: bool = True
 

src/decisionformer/distributed.py

@@ -27,17 +27,25 @@ def enabled(self) -> bool:
         return self.world_size > 1
 
 
+def pick_device(local_rank: int = 0, world_size: int = 1) -> torch.device:
+    """`DECISIONFORMER_DEVICE` (cuda, mps or cpu) when set; otherwise CUDA, then Apple's MPS for single-process
+    runs, then CPU."""
+    wanted = os.environ.get("DECISIONFORMER_DEVICE", "")
+    if wanted in ("", "cuda") and torch.cuda.is_available():
+        index = local_rank % torch.cuda.device_count()
+        torch.cuda.set_device(index)
+        return torch.device("cuda", index)
+    if wanted in ("", "mps") and world_size == 1 and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
 def setup(backend: str = "auto", timeout_minutes: int = 30) -> DistributedState:
     world_size = int(os.environ.get("WORLD_SIZE", "1"))
     rank = int(os.environ.get("RANK", "0"))
     local_rank = int(os.environ.get("LOCAL_RANK", "0"))
-    use_cuda = torch.cuda.is_available() and os.environ.get("DECISIONFORMER_DEVICE", "") != "cpu"
-    if use_cuda:
-        index = local_rank % torch.cuda.device_count()
-        torch.cuda.set_device(index)
-        device = torch.device("cuda", index)
-    else:
-        device = torch.device("cpu")
+    device = pick_device(local_rank, world_size)
+    use_cuda = device.type == "cuda"
     if world_size > 1 and not dist.is_initialized():
         chosen = backend if backend != "auto" else ("nccl" if use_cuda else "gloo")
         kwargs: dict[str, Any] = {"backend": chosen, "timeout": datetime.timedelta(minutes=timeout_minutes)}

src/decisionformer/trainer.py

@@ -138,8 +138,8 @@ 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 GPU(s)/process(es)", self.total_steps,
-                    self.steps_per_epoch, state.world_size)
+        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)
         if config.eval.at_start and self.step == 0:
             self.evaluate()
         window_loss, window_count, window_start = 0.0, 0.0, time.monotonic()

tests/conftest.py

@@ -3,13 +3,16 @@
 from __future__ import annotations
 
 import json
+import os
 import random
 import string
 from pathlib import Path
 
 import pytest
 import torch
 
+os.environ.setdefault("DECISIONFORMER_DEVICE", "cpu")  # deterministic CPU runs, also on machines with a GPU or MPS
+
 COLORS = ["red", "blue", "green", "yellow"]
 SHAPES = ["circle", "square", "triangle"]
 SIZES = ["tiny", "small", "medium", "large", "huge"]

tests/test_units.py

@@ -392,3 +392,15 @@ def test_training_precision_defaults_to_bf16(monkeypatch):
     assert resolve_dtype(FinetuneConfig(mode="qlora"), "fp16") == torch.float16
     assert resolve_dtype(FinetuneConfig(mode="qlora", quantization={"compute_dtype": "float32"}), "bf16") == torch.float32
     assert resolve_dtype(FinetuneConfig(mode="full", model_dtype="bfloat16"), "bf16") == torch.bfloat16
+
+
+def test_device_choice(monkeypatch):
+    from decisionformer.distributed import pick_device
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+    monkeypatch.delenv("DECISIONFORMER_DEVICE", raising=False)
+    assert pick_device().type == "mps", "Apple GPUs train single-process runs"
+    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"