From 69b11f967ccbfe5b36b4a3e5574dd13051242394 Mon Sep 17 00:00:00 2001 From: Abhay Kumar Date: Fri, 23 Oct 2026 19:21:49 +0530 Subject: [PATCH] Laptop fine-tuning (Mac MPS, CPU), truncating training prompts, backbone order - Device choice: DECISIONFORMER_DEVICE (cuda, mps, cpu) when set, else CUDA, then Apple silicon (MPS) for single-process runs, then CPU. Loading for inference follows the same order. Tests stay on CPU. - data.overflow defaults to truncate_state_end for training and evaluation data; loaded and served models still refuse what does not fit. - configs/recipes/local_smollm2.yaml: SmolLM2-135M-Instruct on deciformer-15k with low-memory settings (3.7 GB peak on CPU). - README: backbones are architectures built into transformers (no trust_remote_code models); table ordered dense, mixture of experts, hybrid, encoders; new section for fine-tuning on a laptop. - The training log names the device and precision. --- README.md | 37 ++++++++++++++++++++++++------ configs/recipes/local_smollm2.yaml | 34 +++++++++++++++++++++++++++ src/decisionformer/api.py | 3 ++- src/decisionformer/config.py | 4 +++- src/decisionformer/distributed.py | 22 ++++++++++++------ src/decisionformer/trainer.py | 4 ++-- tests/conftest.py | 3 +++ tests/test_units.py | 12 ++++++++++ 8 files changed, 101 insertions(+), 18 deletions(-) create mode 100644 configs/recipes/local_smollm2.yaml diff --git a/README.md b/README.md index 04cd4fa..75228bb 100644 --- a/README.md +++ b/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//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 diff --git a/configs/recipes/local_smollm2.yaml b/configs/recipes/local_smollm2.yaml new file mode 100644 index 0000000..032c463 --- /dev/null +++ b/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 diff --git a/src/decisionformer/api.py b/src/decisionformer/api.py index 84c6f58..fc112ca 100644 --- a/src/decisionformer/api.py +++ b/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) diff --git a/src/decisionformer/config.py b/src/decisionformer/config.py index 1bb1d4c..97c4cbe 100644 --- a/src/decisionformer/config.py +++ b/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 diff --git a/src/decisionformer/distributed.py b/src/decisionformer/distributed.py index d185883..74701db 100644 --- a/src/decisionformer/distributed.py +++ b/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)} diff --git a/src/decisionformer/trainer.py b/src/decisionformer/trainer.py index 5d440d5..14b48a8 100644 --- a/src/decisionformer/trainer.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py index 732f033..0644a99 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os import random import string from pathlib import Path @@ -10,6 +11,8 @@ 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"] diff --git a/tests/test_units.py b/tests/test_units.py index d0a736b..22b49ed 100644 --- a/tests/test_units.py +++ b/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"