From 6826807ff33e7ab8ce3dcc99c6baa21a2ad0400d Mon Sep 17 00:00:00 2001 From: Aswin Alexander Sam Date: Sun, 2 Aug 2026 18:20:54 +0530 Subject: [PATCH 1/5] export a checkpoint as gguf M5. `bloomery export` writes a GGUF and an Ollama Modelfile, at f16, q8_0 or q4_0. Written here rather than handed to llama.cpp's convert_hf_to_gguf.py. That converter identifies a tokenizer by hashing its output against a table of known models, and a tokenizer this project trained is a new hash by construction -- measured 74a7f913c36b5d2f on a real checkpoint. So the stock path refuses precisely the from-scratch models bloomery leads with. Refusing is right for a general tool, since the hash picks the pre-tokenizer regex and the wrong one yields a model that loads and talks nonsense. But we hold the answer it is trying to recover: these are byte-level BPE with the standard regex, so the field is set outright. The README claimed the opposite -- "GGUF conversion, vLLM and Ollama all work without a bespoke converter" -- since before any of this existed. It now says what is true. What the round trip proved, and what it did not Tensors, vocabulary, metadata and the tied output head all verified by reading the file back. Every one of those passed while the artefact was unloadable: BPE is its merge rules, not just its vocabulary, and llama.cpp refuses a file without them -- "cannot find tokenizer merges in model file". That was found by running the export through Ollama, not by any test here, and there is now a test for it. Verified end to end: `ollama create` accepts the file and the model generates, 16 tokens, done_reason length. Details worth recording Embeddings are tied, so the safetensors hold no lm_head. Loading materialises one, and it is derived from the embedding when it does not. Normalisation weights stay f32 and the embedding f16 even under q4_0, which is what llama.cpp's own quantizer does: a small share of the bytes and a large share of the behaviour. The K-quants are absent because the gguf package raises NotImplementedError when asked to write one. It can read them, which is a different thing. The command names llama-quantize rather than pretending. LoRA checkpoints are merged first. GGUF has no notion of an adapter, so without the merge the export is silently of the untouched base. context_length is the sequence length the model was trained at, not a capability, and a runtime holds it as a hard limit. It is reported. Also fixes a pre-existing flake The memory pre-flight derives its budget from *available* RAM, and a depth-1 model estimates ~0.8 GiB once fixed overhead is counted. On a constrained machine the budget dips under that partway through a suite run, and `train` or `adapt` refuses -- so a test about tokenizers fails for reasons of its host. It showed up as this file passing alone and failing inside the full suite, a different test each time. 32 of 34 invocations were exposed. The budget is now pinned for CLI tests, with TestMemoryGuard opting out because it is what tests the real one. --- .github/workflows/ci.yml | 7 +- README.md | 21 +- pyproject.toml | 6 + src/bloomery/cli.py | 114 +++++++++ src/bloomery/export.py | 345 ++++++++++++++++++++++++++ src/bloomery/jobs/runner.py | 10 + src/bloomery/jobs/types.py | 3 + src/bloomery/paths.py | 5 + src/bloomery/server/static/app.js | 8 + src/bloomery/server/static/index.html | 1 + src/bloomery/train/loop.py | 18 ++ tests/test_cli_train.py | 179 +++++++++++++ tests/test_export.py | 225 +++++++++++++++++ 13 files changed, 935 insertions(+), 7 deletions(-) create mode 100644 src/bloomery/export.py create mode 100644 tests/test_export.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 243d37f..36dd9fd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,7 +64,7 @@ jobs: - name: install shell: bash - run: uv pip install --torch-backend=auto -e ".[dev,train,serve,adapt]" + run: uv pip install --torch-backend=auto -e ".[dev,train,serve,adapt,export]" # Assert the interpreter and the training stack directly. This replaces an # earlier trick of failing the build whenever pytest reported any skip, @@ -82,6 +82,9 @@ jobs: # Same reasoning for the adapter tests: they importorskip on peft, # so a broken adapt install would skip every LoRA test silently. python -c "import peft; print('adapt extra ok', peft.__version__)" + # And the export tests, for the same reason: a broken gguf install + # would skip every round-trip test rather than failing. + python -c "import gguf; print('export extra ok')" # No retry. One was added here on the theory that the illegal-instruction # crashes on Windows were an intermittent fault in the PyTorch wheel that @@ -147,7 +150,7 @@ jobs: # The train extra is needed here too: mypy resolves numpy's bundled stubs, # and dropping it would turn a type error into an unanalysable import. - - run: uv pip install --torch-backend=auto -e ".[dev,train,serve,adapt]" + - run: uv pip install --torch-backend=auto -e ".[dev,train,serve,adapt,export]" - name: ruff check run: ruff check . diff --git a/README.md b/README.md index a1fefb3..f40bdc6 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,10 @@ a milestone are not built yet — see [Roadmap](#roadmap). - **CPU thread caps** — `--cores` works identically on every platform. - **Replay mixtures** — weighted, versioned dataset blends with per-component forgetting detection. +- **Export for llama.cpp and Ollama** — `export` writes a GGUF and a Modelfile, + at f16, q8_0 or q4_0. LoRA adapters are folded in first, since GGUF has no + notion of an adapter. The K-quants need llama.cpp's `llama-quantize`, and the + command says so rather than pretending otherwise. - **Fine-tune on conversations** — `prepare --chat` packs a corpus of chat or prompt/completion records and masks the prompts, so the model is scored only on what it was meant to produce. There is no separate command: the dataset @@ -99,8 +103,7 @@ a milestone are not built yet — see [Roadmap](#roadmap). them. Per-job CPU, memory and GPU limits. It binds to localhost by default. ### Not built yet -- **Preference optimization** — DPO and friends (M5+). -- **Export to GGUF, Ollama, MLX** (M5). +- **Preference optimization** — DPO and friends. Not scheduled. ## Hardware @@ -128,7 +131,7 @@ hardware, `bloomery doctor --json` in an issue is genuinely useful. | **M3** | Web UI, job queue, cancel/resume, resource limits | done | | **M4** | Continued pretraining on existing models, full or LoRA | done | | **M4.1** | Supervised fine-tuning on conversations | done | -| **M5** | Export — GGUF, quantization, Ollama | | +| **M5** | Export — GGUF, quantization, Ollama | done | | **M6** | AMD hardening on real hardware, Windows, evaluation | | M2 came before the web UI because weighted replay is what makes "keep adding @@ -179,8 +182,16 @@ instead. Checkpoints are ordinary Hugging Face directories. `AutoModelForCausalLM.from_pretrained` loads them, which is the whole reason bloomery emits a Llama-architecture model -rather than inventing one — GGUF conversion, vLLM and Ollama all work without a -bespoke converter. +rather than inventing one — vLLM and anything else in the ecosystem read them +directly. + +GGUF is the exception, and `export` writes it here rather than shelling out. +llama.cpp's converter identifies a tokenizer by hashing its output against a +table of known models, and a tokenizer bloomery trained is a new hash by +construction — so the stock path refuses precisely the models this project +exists to produce. Refusing is right for a general tool, since the hash picks +the pre-tokenizer and the wrong one yields a model that loads and talks +nonsense. Bloomery already knows the answer, so it fills the field in itself. `train` estimates memory before it starts and refuses a configuration that cannot fit, because an out-of-memory error arrives whenever the allocator diff --git a/pyproject.toml b/pyproject.toml index 2d37ee4..d543d38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,12 @@ serve = [ "pydantic>=2.9", "websockets>=13.0", ] +# Writing GGUF. The official library from the llama.cpp repository: pure Python, +# and it ships its own type information. Kept out of core, which stays at three. +export = [ + "gguf>=0.10", + "numpy>=1.26", +] # Quantized training. CUDA and ROCm only — no macOS wheels exist. quant = [ "bitsandbytes>=0.44; sys_platform != 'darwin'", diff --git a/src/bloomery/cli.py b/src/bloomery/cli.py index f3516f2..4ac1958 100644 --- a/src/bloomery/cli.py +++ b/src/bloomery/cli.py @@ -7,6 +7,7 @@ import json import os import sys +from dataclasses import replace from pathlib import Path from typing import Any, NoReturn @@ -15,6 +16,7 @@ from bloomery import __version__, paths from bloomery.capability import LADDER_BY_KEY, Method, assess, format_params +from bloomery.export import QUANTIZATIONS from bloomery.probe import probe_host_report from bloomery.probe.types import GIB from bloomery.render import render_report @@ -1119,6 +1121,118 @@ def chat( console.print() +# --------------------------------------------------------------------------- # +# export +# --------------------------------------------------------------------------- # + + +@app.command() +def export( + run: str | None = typer.Option(None, "--run", "-r", help="Run name to export."), + checkpoint: Path | None = typer.Option( + None, "--checkpoint", "-c", help="Path to a checkpoint directory." + ), + name: str | None = typer.Option( + None, "--name", "-n", help="Name for the export. Defaults to the run's." + ), + quantize: str = typer.Option( + "f16", "--quantize", "-q", help=f"One of: {', '.join(QUANTIZATIONS)}." + ), + as_json: bool = typer.Option(False, "--json"), +) -> None: + """Write a checkpoint as GGUF, for llama.cpp and Ollama. + + Written here rather than by llama.cpp's converter, which cannot read the + models this project exists to produce: it identifies a tokenizer by hashing + its output against a table of known models, and a tokenizer bloomery trained + is a new hash by construction. + + LoRA adapters are folded into the weights first, since GGUF has no notion of + an adapter and a runtime would otherwise read the untouched base. + """ + _quiet_transformers() + import shutil + + from bloomery.export import GGUF_NAME as gguf_name + from bloomery.export import ExportError, to_gguf, write_modelfile + from bloomery.train import checkpoint as ckpt + from bloomery.train.loop import ModelLoadError, load_model, merge_adapters + + if bool(run) == bool(checkpoint): + _die("give exactly one of --run or --checkpoint") + if quantize not in QUANTIZATIONS: + _die(f"unknown quantization {quantize!r}; choose one of: {', '.join(QUANTIZATIONS)}") + + target = checkpoint if checkpoint else ckpt.checkpoint_dir(paths.run_dir(run or "")) + if not target.is_dir(): + _die(f"{target} does not exist") + + out_name = name or run or target.parent.name + destination = paths.export_dir(out_name) + + with console.status(f"loading {target}"): + try: + from bloomery.data import load_tokenizer + + model = merge_adapters(load_model(target)) + tokenizer = load_tokenizer(target) + except ModelLoadError as exc: + _die(str(exc)) + except Exception as exc: # noqa: BLE001 - tokenizer loading raises many types + _die(f"could not read {target}: {exc}") + + # Staged then renamed, the same way a checkpoint is written. A half-written + # GGUF looks loadable and is not, and export can run while the checkpoint it + # is reading is being replaced by a training step. + staging = destination.with_name(destination.name + ".tmp") + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + + try: + with console.status(f"writing {quantize} gguf"): + result = to_gguf(model, tokenizer, staging / gguf_name, quantization=quantize) + write_modelfile(staging, tokenizer) + except ExportError as exc: + shutil.rmtree(staging, ignore_errors=True) + _die(str(exc)) + + if destination.exists(): + shutil.rmtree(destination) + staging.rename(destination) + result = replace(result, path=destination / gguf_name) + + if as_json: + console.print_json(json.dumps(result.to_dict())) + return + + console.print( + f"model [bold]{result.architecture}[/bold] {format_params(result.parameters)} params" + ) + console.print(f"format [bold]{result.quantization}[/bold] {result.tensors} tensors") + console.print(f"size [bold]{result.bytes_written / 1e6:,.1f} MB[/bold] → {result.path}") + console.print(f"vocab {result.vocab_size:,} tokens") + # The number most likely to surprise: it is the sequence length the model was + # trained at, and a runtime will treat it as a hard limit. + console.print( + f"context {result.context_length:,} tokens [dim]the length it was trained at[/dim]" + ) + if result.unquantized: + console.print( + f"[yellow] {len(result.unquantized)} tensor(s) kept at f16; their shape " + "cannot be blocked for this format[/yellow]" + ) + console.print() + console.print( + f"next: [bold]ollama create {paths.slug(out_name)} -f {destination / 'Modelfile'}[/bold]" + ) + if quantize != "q4_0": + console.print( + "[dim] smaller still with llama.cpp's llama-quantize, which has the " + "K-quants this cannot write[/dim]" + ) + + # --------------------------------------------------------------------------- # # bench # --------------------------------------------------------------------------- # diff --git a/src/bloomery/export.py b/src/bloomery/export.py new file mode 100644 index 0000000..220d1f7 --- /dev/null +++ b/src/bloomery/export.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: 2026 Aswin Alexander Sam +# SPDX-License-Identifier: AGPL-3.0-or-later +"""Writing a checkpoint out as GGUF, for llama.cpp and Ollama. + +Written here rather than handed to llama.cpp's ``convert_hf_to_gguf.py``, which +cannot read the models this project exists to produce. That converter identifies +a tokenizer by hashing its output on a probe string and matching a table of known +models; a tokenizer bloomery trained is a new hash by construction, and the +converter refuses it rather than guessing. + +Refusing is the right call for a general tool — the hash decides which +pre-tokenizer regex to apply, and applying the wrong one produces a model that +loads and generates rubbish. But bloomery already knows the answer the hash is +trying to recover: its tokenizers are byte-level BPE with the standard regex. So +the field is set directly, and the file is written here. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: # pragma: no cover - typing only + import numpy as np + +# Architectures whose tensor layout this writer knows. `adapt --method full` can +# produce a checkpoint of whatever base it was given, and those are not all +# Llama — mapping one wrongly yields a file that loads and generates nonsense, +# which is far worse than refusing to write it. +# +# Mistral is included because its tensor layout is Llama's; the name differs and +# nothing else does. +SUPPORTED_ARCHITECTURES = frozenset({"llama", "mistral"}) + +# What `--quantize` accepts, mapped to the ggml type. Everything here is +# implemented in the `gguf` package's own quantizer, which is the reference one. +# The K-quants (q4_K, q5_K, q6_K) are deliberately absent: that package can read +# them but raises NotImplementedError when asked to write one, so they need +# llama.cpp's `llama-quantize` and the command says so. +QUANTIZATIONS: tuple[str, ...] = ("f16", "q8_0", "q4_0") + +# Quantized formats work in blocks along the last dimension. A tensor whose last +# dimension is not a whole number of blocks cannot be quantized and is written at +# f16 instead — correct, slightly larger, and reported rather than silent. +_BLOCK = 32 + + +class ExportError(RuntimeError): + """A checkpoint could not be written as GGUF.""" + + +@dataclass(frozen=True, slots=True) +class ExportResult: + path: Path + architecture: str + quantization: str + tensors: int + parameters: int + bytes_written: int + context_length: int + vocab_size: int + # Tensors that could not take the requested quantization and fell back to + # f16, by name. Empty for every shape this project produces. + unquantized: tuple[str, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "path": str(self.path), + "architecture": self.architecture, + "quantization": self.quantization, + "tensors": self.tensors, + "parameters": self.parameters, + "bytes": self.bytes_written, + "context_length": self.context_length, + "vocab_size": self.vocab_size, + "unquantized": list(self.unquantized), + } + + +def architecture_of(config: Any) -> str: + """The model family this config describes, refusing what cannot be mapped.""" + model_type = getattr(config, "model_type", None) + if model_type not in SUPPORTED_ARCHITECTURES: + known = ", ".join(sorted(SUPPORTED_ARCHITECTURES)) + raise ExportError( + f"this checkpoint is a {model_type!r} model, and GGUF export understands " + f"{known}.\n" + "Writing it as one of those would produce a file that loads and generates " + "nonsense, so it is refused instead. A checkpoint from `train` is always " + "supported; one from `adapt` is whatever model it continued." + ) + return str(model_type) + + +def _tensor_type(name: str, array: np.ndarray, quantization: str) -> Any: + """The ggml type to store one tensor as. + + Normalisation weights stay at f32 and the token embedding at f16, which is + what llama.cpp's own quantizer does: they are a small share of a model's + bytes and a large share of its behaviour, so quantizing them costs quality + for almost no size. + """ + from gguf import GGMLQuantizationType as QT + + if array.ndim == 1: + return QT.F32 + if quantization == "f16" or name == "token_embd.weight": + return QT.F16 + if array.shape[-1] % _BLOCK: + return QT.F16 + return {"q8_0": QT.Q8_0, "q4_0": QT.Q4_0}[quantization] + + +def to_gguf( + model: Any, + tokenizer: Any, + out_path: Path, + *, + quantization: str = "f16", +) -> ExportResult: + """Write a loaded model and its tokenizer as a single GGUF file.""" + import numpy as np + import torch + from gguf import GGMLQuantizationType as QT + from gguf import GGUFWriter, get_tensor_name_map, quants + from gguf.constants import MODEL_ARCH + + if quantization not in QUANTIZATIONS: + raise ExportError( + f"unknown quantization {quantization!r}; choose one of: {', '.join(QUANTIZATIONS)}" + ) + + config = model.config + architecture = architecture_of(config) + + state = model.state_dict() + # Tying means the file on disk holds no lm_head, but loading materialises it, + # so it is usually here. Deriving it from the embedding when it is not keeps + # both paths working: GGUF always wants an output projection. + if "lm_head.weight" not in state: + state = dict(state) + state["lm_head.weight"] = state["model.embed_tokens.weight"] + + # The official mapping rather than a hand-written table, so a rename upstream + # is not something this has to notice. + names = get_tensor_name_map(MODEL_ARCH.LLAMA, config.num_hidden_layers) + + out_path.parent.mkdir(parents=True, exist_ok=True) + writer = GGUFWriter(str(out_path), architecture) + try: + writer.add_block_count(config.num_hidden_layers) + # The sequence length the model was trained at, not a capability of the + # architecture. For a bloomery run this is often 512 or 1024, and a + # runtime will hold it as a hard limit. + writer.add_context_length(config.max_position_embeddings) + writer.add_embedding_length(config.hidden_size) + writer.add_feed_forward_length(config.intermediate_size) + writer.add_head_count(config.num_attention_heads) + writer.add_head_count_kv(getattr(config, "num_key_value_heads", config.num_attention_heads)) + writer.add_layer_norm_rms_eps(config.rms_norm_eps) + writer.add_file_type(_file_type(quantization)) + _add_rope(writer, config) + _add_tokenizer(writer, tokenizer, config) + + parameters = 0 + written = 0 + unquantized: list[str] = [] + for source, tensor in state.items(): + mapped = names.get_name(source.removesuffix(".weight")) + if mapped is None: + # Buffers such as rotary inverse frequencies are recomputed by + # the runtime and have no GGUF name; skipping them is correct. + continue + name = f"{mapped}.weight" + array = tensor.to(torch.float32).numpy() + parameters += int(array.size) + + wanted = _tensor_type(name, array, quantization) + # Only a shape that could not be blocked counts as a fallback. The + # embedding is held at f16 by policy, not by failure. + if ( + quantization != "f16" + and array.ndim > 1 + and wanted is QT.F16 + and name != "token_embd.weight" + ): + unquantized.append(name) + data = array.astype(np.float32) if wanted is QT.F32 else quants.quantize(array, wanted) + writer.add_tensor(name, data, raw_dtype=wanted) + written += 1 + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + finally: + writer.close() + + return ExportResult( + path=out_path, + architecture=architecture, + quantization=quantization, + tensors=written, + parameters=parameters, + bytes_written=out_path.stat().st_size, + context_length=config.max_position_embeddings, + vocab_size=config.vocab_size, + unquantized=tuple(unquantized), + ) + + +MODELFILE = "Modelfile" +GGUF_NAME = "model.gguf" + + +def write_modelfile(directory: Path, tokenizer: Any, *, gguf_name: str = GGUF_NAME) -> Path: + """Describe the exported model to Ollama. + + Written rather than run: `ollama create` needs Ollama installed, and an + export that fails because a tool the user has not installed is missing would + be a poor trade for a file that is six lines of text. + + The template is carried across when the checkpoint has one. Without it a + fine-tuned model gets prompted as raw text, which is the same mismatch + `bloomery chat` avoids by applying the template itself. + """ + lines = [ + f"FROM ./{gguf_name}", + "", + ] + template = getattr(tokenizer, "chat_template", None) + if template: + # Ollama's own template syntax is Go's, not Jinja, so the tokenizer's + # template cannot be handed over directly. The GGUF carries it in its + # metadata, where Ollama reads it; this note says where it went so the + # absence of a TEMPLATE line does not read as an omission. + lines += [ + "# This model's chat template travels inside the GGUF metadata,", + "# which Ollama reads directly. Nothing to declare here.", + "", + ] + + eos = getattr(tokenizer, "eos_token", None) + if eos: + lines.append(f'PARAMETER stop "{eos}"') + lines.append("") + + path = directory / MODELFILE + path.write_text("\n".join(lines), encoding="utf-8") + return path + + +def _file_type(quantization: str) -> Any: + from gguf import LlamaFileType + + return { + "f16": LlamaFileType.MOSTLY_F16, + "q8_0": LlamaFileType.MOSTLY_Q8_0, + "q4_0": LlamaFileType.MOSTLY_Q4_0, + }[quantization] + + +def _add_rope(writer: Any, config: Any) -> None: + """Record the RoPE base, reading it rather than assuming the default. + + `to_llama_config` deliberately does not pass this, so it takes whatever the + transformers default is — and transformers 5 renamed ``rope_theta`` to + ``rope_parameters``. Hardcoding 10000.0 here would be right until it was not. + """ + theta = getattr(config, "rope_theta", None) + if theta is None: + parameters = getattr(config, "rope_parameters", None) + if isinstance(parameters, dict): + theta = parameters.get("rope_theta") + if theta is not None: + writer.add_rope_freq_base(float(theta)) + + +def _merges(tokenizer: Any) -> list[str]: + """The BPE merge rules, as GGUF wants them: one space-joined pair per entry. + + Read from the backend's own serialisation rather than from a file beside the + checkpoint, because a tokenizer saved by ``save_pretrained`` may write + ``tokenizer.json`` and nothing else — there is no ``merges.txt`` to find. + + The tokenizers library has emitted these as both ``"a b"`` strings and + ``["a", "b"]`` pairs across versions, so both are accepted. + """ + import json + + backend = getattr(tokenizer, "backend_tokenizer", None) + if backend is None: + raise ExportError( + "this tokenizer has no fast backend, so its merge rules cannot be read. " + "GGUF needs them: a BPE vocabulary without merges cannot tokenize." + ) + + model = json.loads(backend.to_str()).get("model", {}) + raw = model.get("merges") + if not raw: + raise ExportError( + "this tokenizer reports no BPE merges. A byte-level BPE tokenizer always " + "has them, so the checkpoint's tokenizer.json is probably incomplete." + ) + + merges: list[str] = [] + for entry in raw: + merges.append(" ".join(entry) if isinstance(entry, list) else str(entry)) + return merges + + +def _add_tokenizer(writer: Any, tokenizer: Any, config: Any) -> None: + """Write the vocabulary, and say which pre-tokenizer produced it. + + ``gpt2`` and ``default`` are the byte-level BPE settings, which is what every + tokenizer this project trains is — see the module docstring for why naming it + outright is safe here and is not for a general converter. + """ + vocab = tokenizer.get_vocab() + tokens = [token for token, _ in sorted(vocab.items(), key=lambda item: item[1])] + added = {token for token in getattr(tokenizer, "all_special_tokens", []) or []} + + writer.add_tokenizer_model("gpt2") + writer.add_tokenizer_pre("default") + writer.add_token_list(tokens) + # BPE is the merge rules, not just the vocabulary. Without them llama.cpp + # refuses the file outright — "cannot find tokenizer merges in model file" — + # and a writer that emits only the token list produces something every test + # about tensors and vocabulary passes and no runtime will load. + writer.add_token_merges(_merges(tokenizer)) + # 1 is a normal token, 3 is a control token. Marking the specials keeps a + # runtime from printing them back to the user as text. + writer.add_token_types([3 if token in added else 1 for token in tokens]) + + for setter, value in ( + (writer.add_bos_token_id, getattr(config, "bos_token_id", None)), + (writer.add_eos_token_id, getattr(config, "eos_token_id", None)), + (writer.add_pad_token_id, getattr(config, "pad_token_id", None)), + ): + if value is not None: + setter(int(value)) + + template = getattr(tokenizer, "chat_template", None) + if template: + writer.add_chat_template(template) diff --git a/src/bloomery/jobs/runner.py b/src/bloomery/jobs/runner.py index 571cb56..aa992aa 100644 --- a/src/bloomery/jobs/runner.py +++ b/src/bloomery/jobs/runner.py @@ -158,6 +158,12 @@ def created_at(self) -> float | None: "device": "--device", "seed": "--seed", }, + JobKind.EXPORT: { + "run": "--run", + "checkpoint": "--checkpoint", + "name": "--name", + "quantize": "--quantize", + }, JobKind.BENCH: { "size": "--size", "depth": "--depth", @@ -183,6 +189,10 @@ def created_at(self) -> float | None: "resume": "--resume", "force": "--force", }, + # Required even though it is empty: _SWITCHES is subscripted directly + # below, so a missing kind raises KeyError inside the supervisor rather + # than a clean launch error. + JobKind.EXPORT: {}, JobKind.BENCH: {"grad_checkpoint": "--grad-checkpoint"}, } diff --git a/src/bloomery/jobs/types.py b/src/bloomery/jobs/types.py index 09019e2..d8d588d 100644 --- a/src/bloomery/jobs/types.py +++ b/src/bloomery/jobs/types.py @@ -27,6 +27,7 @@ class JobKind(StrEnum): TRAIN = "train" ADAPT = "adapt" BENCH = "bench" + EXPORT = "export" class JobStatus(StrEnum): @@ -53,6 +54,8 @@ def terminal(self) -> bool: # other because VRAM cannot be partitioned on consumer hardware: two training # processes sharing one card do not each get half, they contend for all of it # and the second one dies. +# EXPORT is deliberately absent: it reads a checkpoint and writes a file, so +# it can run beside a training job rather than queueing behind one. EXCLUSIVE_KINDS = frozenset({JobKind.TRAIN, JobKind.ADAPT, JobKind.BENCH}) diff --git a/src/bloomery/paths.py b/src/bloomery/paths.py index c57e2c7..2830b3b 100644 --- a/src/bloomery/paths.py +++ b/src/bloomery/paths.py @@ -66,6 +66,11 @@ def run_dir(name: str) -> Path: return runs_dir() / slug(name) +def export_dir(name: str) -> Path: + """Everything produced by `export` for one model: the GGUF and its Modelfile.""" + return exports_dir() / slug(name) + + def slug(name: str) -> str: """Make a user-supplied name safe to use as a directory. diff --git a/src/bloomery/server/static/app.js b/src/bloomery/server/static/app.js index d95ea9c..28f987e 100644 --- a/src/bloomery/server/static/app.js +++ b/src/bloomery/server/static/app.js @@ -65,6 +65,12 @@ const FIELDS = { { key: "device", label: "Device", choices: ["cuda", "mps", "cpu"] }, { key: "seed", label: "Seed", type: "number" }, ], + export: [ + { key: "run", label: "Run", hint: "a run name from train or adapt" }, + { key: "checkpoint", label: "Checkpoint", hint: "a path, instead of a run" }, + { key: "name", label: "Export name", hint: "defaults to the run's" }, + { key: "quantize", label: "Format", choices: ["f16", "q8_0", "q4_0"] }, + ], bench: [ { key: "size", label: "Size", choices: SIZES, hint: "named preset" }, { key: "depth", label: "Depth", type: "number", min: 1 }, @@ -90,6 +96,8 @@ const SWITCHES = { { key: "resume", label: "Resume from latest checkpoint" }, { key: "force", label: "Force" }, ], + // Empty, but required: SWITCHES[kind] is subscripted directly below. + export: [], bench: [{ key: "grad_checkpoint", label: "Gradient checkpointing" }], }; diff --git a/src/bloomery/server/static/index.html b/src/bloomery/server/static/index.html index 725912e..7664981 100644 --- a/src/bloomery/server/static/index.html +++ b/src/bloomery/server/static/index.html @@ -34,6 +34,7 @@

Start a job

+ diff --git a/src/bloomery/train/loop.py b/src/bloomery/train/loop.py index 740d227..b808115 100644 --- a/src/bloomery/train/loop.py +++ b/src/bloomery/train/loop.py @@ -204,6 +204,24 @@ def _load_adapted(path: Path) -> Any: ) from exc +def merge_adapters(model: Any) -> Any: + """Fold any LoRA adapters into the weights they modify. + + Export needs one set of weights, not a base plus a diff: GGUF has no notion + of an adapter, and a runtime reading the file would get the untouched base. + + A model with no adapters is returned unchanged, so a caller does not have to + know which shape it was handed. + """ + merge = getattr(model, "merge_and_unload", None) + if merge is None: + return model + # eval() first: dropout is active on a freshly loaded adapter, and merging + # while it is would fold a randomly masked version of the update. + model.eval() + return merge() + + def attach_adapter(model: Any, settings: LoraSettings) -> Any: """Freeze the model and train low-rank adapters against it instead. diff --git a/tests/test_cli_train.py b/tests/test_cli_train.py index 10669e7..9ea9d0b 100644 --- a/tests/test_cli_train.py +++ b/tests/test_cli_train.py @@ -9,7 +9,9 @@ from __future__ import annotations import json +from dataclasses import replace from pathlib import Path +from typing import Any from unittest import mock import pytest @@ -28,6 +30,34 @@ def isolated_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: return home +@pytest.fixture(autouse=True) +def _stable_memory_budget(monkeypatch: pytest.MonkeyPatch) -> None: + """Stop ambient RAM from deciding whether these tests pass. + + The pre-flight budget is derived from *available* system memory, and even a + depth-1 model estimates around 0.8 GiB once the fixed context overhead is + counted. On a machine with little free memory — or simply partway through a + long suite run — the budget dips under that and `train` and `adapt` refuse, + so a test about tokenizers or exports fails for reasons of its own host. + + Observed as exactly that: this file passes alone and fails inside the full + suite, with a different test each time. + + The budget is pinned generously rather than the check disabled, so a run that + genuinely cannot fit is still refused — TestMemoryGuard asks for shapes that + need terabytes and is unaffected. + """ + from bloomery import capability + + real = capability.derive_budget + + def generous(report: Any, *, device_type: str | None = None) -> Any: + budget = real(report, device_type=device_type) + return replace(budget, total=max(budget.total, 64 * 1024**3)) + + monkeypatch.setattr(capability, "derive_budget", generous) + + def _invoke(*args: str): # noqa: ANN202 - typer's Result type is internal result = runner.invoke(app, list(args)) if result.exception is not None and not isinstance(result.exception, SystemExit): @@ -652,6 +682,14 @@ def test_pinning_a_mixture_version(self) -> None: class TestMemoryGuard: """The pre-flight check that stops a run before it OOMs partway through.""" + @pytest.fixture(autouse=True) + def _stable_memory_budget(self) -> None: + """Opt out of the module-wide pin: this class is what tests the budget. + + Everywhere else a pinned budget stops ambient RAM deciding unrelated + outcomes. Here the real one is the subject. + """ + @pytest.fixture(autouse=True) def prepared(self, isolated_home: Path) -> None: assert ( @@ -1291,3 +1329,144 @@ def record(loaded, prompt, config=None, *, raw=False): # noqa: ANN001, ANN202 ) assert seen == [True, False], seen + + +class TestExport: + """Writing a checkpoint out for llama.cpp and Ollama.""" + + @pytest.fixture + def trained(self, isolated_home: Path) -> Path: + assert ( + _invoke("prepare", "--name", "c", "--synthetic", "400", "--vocab", "300").exit_code == 0 + ) + assert ( + _invoke( + "train", + "--data", + "c", + "--name", + "r", + "--depth", + "1", + "--steps", + "2", + "--batch", + "4", + "--seq", + "32", + "--device", + "cpu", + ).exit_code + == 0 + ) + return isolated_home / "runs/r/latest" + + def test_writes_a_gguf_and_a_modelfile(self, trained: Path, isolated_home: Path) -> None: + result = _invoke("export", "--run", "r") + assert result.exit_code == 0, plain(result.stdout) + out = isolated_home / "exports/r" + assert (out / "model.gguf").is_file() + assert (out / "Modelfile").is_file() + + @pytest.mark.parametrize("quantize", ["f16", "q8_0", "q4_0"]) + def test_each_format_writes(self, trained: Path, isolated_home: Path, quantize: str) -> None: + result = _invoke("export", "--run", "r", "--quantize", quantize, "--name", quantize) + assert result.exit_code == 0, plain(result.stdout) + assert (isolated_home / f"exports/{quantize}/model.gguf").is_file() + + def test_it_reports_the_context_it_was_trained_at(self, trained: Path) -> None: + """The number most likely to surprise: a runtime treats it as a hard limit.""" + result = _invoke("export", "--run", "r", "--json") + assert result.exit_code == 0, plain(result.stdout) + payload = json.loads(plain(result.stdout)) + assert payload["context_length"] == 32 + assert payload["architecture"] == "llama" + + def test_an_unknown_format_is_refused(self, trained: Path) -> None: + result = _invoke("export", "--run", "r", "--quantize", "q4_k_m") + assert result.exit_code == 1 + assert "unknown quantization" in plain(result.stdout) + + def test_a_missing_run_is_reported(self) -> None: + result = _invoke("export", "--run", "nope") + assert result.exit_code == 1 + + def test_exactly_one_target(self, trained: Path) -> None: + assert _invoke("export").exit_code == 1 + assert _invoke("export", "--run", "r", "--checkpoint", str(trained)).exit_code == 1 + + def test_a_lora_checkpoint_is_merged_before_export(self, isolated_home: Path) -> None: + """GGUF has no notion of an adapter. + + Without the merge the export would silently be of the untouched base: + it writes, it loads, and none of the fine-tuning is in it. + """ + pytest.importorskip("peft", reason="adapters need the adapt extra") + from transformers import AutoTokenizer + + assert ( + _invoke("prepare", "--name", "c", "--synthetic", "400", "--vocab", "300").exit_code == 0 + ) + assert ( + _invoke( + "train", + "--data", + "c", + "--name", + "b", + "--depth", + "1", + "--steps", + "2", + "--batch", + "4", + "--seq", + "32", + "--device", + "cpu", + ).exit_code + == 0 + ) + base = isolated_home / "runs/b/latest" + AutoTokenizer.from_pretrained(str(base)).save_pretrained(str(base)) + assert ( + _invoke( + "adapt", + "--from", + str(base), + "--data", + "c", + "--name", + "a", + "--steps", + "4", + "--batch", + "2", + "--seq", + "32", + "--device", + "cpu", + ).exit_code + == 0 + ) + + assert _invoke("export", "--run", "a", "--name", "adapted").exit_code == 0 + assert _invoke("export", "--checkpoint", str(base), "--name", "plain").exit_code == 0 + + from gguf import GGUFReader + + def tensor(name: str, export: str): # noqa: ANN202 + reader = GGUFReader(str(isolated_home / f"exports/{export}/model.gguf")) + return next(t.data for t in reader.tensors if t.name == name) + + merged = tensor("blk.0.attn_q.weight", "adapted") + plain_base = tensor("blk.0.attn_q.weight", "plain") + assert not (merged == plain_base).all(), "the adapters were not folded in" + + def test_a_half_written_export_is_not_left_behind( + self, trained: Path, isolated_home: Path + ) -> None: + """A staging directory must not survive a failure looking like an export.""" + result = _invoke("export", "--run", "r", "--name", "ok") + assert result.exit_code == 0, plain(result.stdout) + assert not (isolated_home / "exports/ok.tmp").exists() diff --git a/tests/test_export.py b/tests/test_export.py new file mode 100644 index 0000000..8070141 --- /dev/null +++ b/tests/test_export.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: 2026 Aswin Alexander Sam +# SPDX-License-Identifier: AGPL-3.0-or-later +"""Tests for GGUF export. + +The round trip is what matters. A GGUF that writes without error and holds the +wrong numbers is the failure this has to exclude: it loads, it generates, and +what comes out is noise with no indication that the file is at fault. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +pytest.importorskip("gguf", reason="export tests need the [export] extra") + +import numpy as np # noqa: E402 +import torch # noqa: E402 +from gguf import GGMLQuantizationType as QT # noqa: E402 +from gguf import GGUFReader, quants # noqa: E402 + +from bloomery.arch import spec_from_depth # noqa: E402 +from bloomery.export import ( # noqa: E402 + QUANTIZATIONS, + ExportError, + architecture_of, + to_gguf, + write_modelfile, +) +from bloomery.train.loop import build_model # noqa: E402 + + +@pytest.fixture(scope="module") +def model(tokenizer: Any) -> Any: + """A real model, small enough to export in a moment.""" + spec = spec_from_depth(2, vocab=len(tokenizer), seq=64) + return build_model(spec, eos_token_id=0, seed=0) + + +def read(path: Path) -> dict[str, Any]: + reader = GGUFReader(str(path)) + return {tensor.name: tensor for tensor in reader.tensors} + + +def as_float(tensor: Any) -> np.ndarray: + data = np.asarray(tensor.data) + if tensor.tensor_type not in (QT.F32, QT.F16): + data = quants.dequantize(data, tensor.tensor_type) + return data.astype(np.float32) + + +class TestRoundTrip: + """What was written must be what the model held.""" + + @pytest.mark.parametrize( + ("quantization", "tolerance"), + [("f16", 1e-2), ("q8_0", 5e-2), ("q4_0", 1.0)], + ) + def test_weights_survive( + self, model: Any, tokenizer: Any, tmp_path: Path, quantization: str, tolerance: float + ) -> None: + out = tmp_path / f"{quantization}.gguf" + to_gguf(model, tokenizer, out, quantization=quantization) + written = read(out) + state = model.state_dict() + + for source, name in ( + ("model.embed_tokens.weight", "token_embd.weight"), + ("model.layers.0.self_attn.q_proj.weight", "blk.0.attn_q.weight"), + ("model.layers.0.mlp.down_proj.weight", "blk.0.ffn_down.weight"), + ("model.layers.1.mlp.gate_proj.weight", "blk.1.ffn_gate.weight"), + ("model.norm.weight", "output_norm.weight"), + ): + expected = state[source].to(torch.float32).numpy() + got = as_float(written[name]).reshape(expected.shape) + error = float(np.abs(got - expected).max()) + assert error <= tolerance, f"{name} drifted by {error} at {quantization}" + + def test_the_output_head_is_written_though_it_is_tied( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """The safetensors on disk hold no lm_head, and GGUF needs one anyway. + + Tying means the file has a single embedding matrix serving both ends. + A GGUF without an output projection loads and produces nothing useful. + """ + assert model.config.tie_word_embeddings + out = tmp_path / "tied.gguf" + to_gguf(model, tokenizer, out) + written = read(out) + + assert "output.weight" in written + embedding = as_float(written["token_embd.weight"]) + assert np.array_equal(as_float(written["output.weight"]), embedding) + + def test_normalisation_weights_are_not_quantized( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """A small share of the bytes and a large share of the behaviour.""" + out = tmp_path / "q4.gguf" + to_gguf(model, tokenizer, out, quantization="q4_0") + written = read(out) + assert written["output_norm.weight"].tensor_type == QT.F32 + assert written["blk.0.attn_norm.weight"].tensor_type == QT.F32 + # And the thing that should be quantized, was. + assert written["blk.0.ffn_down.weight"].tensor_type == QT.Q4_0 + + def test_quantizing_actually_shrinks_the_file( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + sizes = {} + for quantization in QUANTIZATIONS: + out = tmp_path / f"size-{quantization}.gguf" + sizes[quantization] = to_gguf( + model, tokenizer, out, quantization=quantization + ).bytes_written + assert sizes["f16"] > sizes["q8_0"] > sizes["q4_0"], sizes + + +class TestTokenizerTravels: + """A GGUF carries its own vocabulary; a wrong one is silent nonsense.""" + + def test_every_token_is_written_in_order( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + out = tmp_path / "vocab.gguf" + to_gguf(model, tokenizer, out) + field = GGUFReader(str(out)).fields["tokenizer.ggml.tokens"] + + expected = [ + token for token, _ in sorted(tokenizer.get_vocab().items(), key=lambda kv: kv[1]) + ] + assert len(field.data) == len(expected) + first = bytes(field.parts[field.data[0]]).decode("utf-8", "replace") + assert first == expected[0] + + def test_the_merge_rules_are_written(self, model: Any, tokenizer: Any, tmp_path: Path) -> None: + """BPE is the merges, not just the vocabulary. + + Found by loading an export into Ollama, not by any test here: llama.cpp + refuses a file without them outright — "cannot find tokenizer merges in + model file" — and every assertion about tensors, vocabulary and metadata + passed while the artefact would not load at all. + """ + out = tmp_path / "merges.gguf" + to_gguf(model, tokenizer, out) + fields = GGUFReader(str(out)).fields + + assert "tokenizer.ggml.merges" in fields + assert len(fields["tokenizer.ggml.merges"].data) > 0 + + def test_a_merge_is_a_space_joined_pair( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """The shape llama.cpp expects, and the one the library has changed. + + tokenizers has emitted merges as both "a b" strings and ["a", "b"] pairs + across versions; either is read, and only the joined form is written. + """ + out = tmp_path / "shape.gguf" + to_gguf(model, tokenizer, out) + field = GGUFReader(str(out)).fields["tokenizer.ggml.merges"] + first = bytes(field.parts[field.data[0]]).decode("utf-8", "replace") + assert " " in first, first + assert len(first.split(" ")) == 2, first + + def test_the_pre_tokenizer_is_declared( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """The field llama.cpp's converter refuses to guess for our tokenizers. + + It picks this by hashing a tokenizer's output against a table of known + models; ours is a new hash by construction. We know the answer, so we + state it — and a GGUF without it cannot be tokenized correctly at all. + """ + out = tmp_path / "pre.gguf" + to_gguf(model, tokenizer, out) + fields = GGUFReader(str(out)).fields + assert "tokenizer.ggml.pre" in fields + assert "tokenizer.ggml.model" in fields + + def test_a_chat_template_is_carried_across( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """Otherwise a fine-tuned model is prompted in a shape it never saw.""" + tokenizer.chat_template = ( + "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}" + ) + try: + out = tmp_path / "chat.gguf" + to_gguf(model, tokenizer, out) + assert "tokenizer.chat_template" in GGUFReader(str(out)).fields + finally: + tokenizer.chat_template = None + + +class TestRefusals: + def test_an_architecture_we_cannot_map_is_refused_by_name(self) -> None: + """`adapt --method full` can produce one, and a wrong mapping is silent.""" + + class NotLlama: + model_type = "mamba" + + with pytest.raises(ExportError, match="mamba"): + architecture_of(NotLlama()) + + def test_an_unknown_quantization_is_refused( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + with pytest.raises(ExportError, match="unknown quantization"): + to_gguf(model, tokenizer, tmp_path / "x.gguf", quantization="q3_k_xxl") + + +class TestModelfile: + def test_it_points_at_the_gguf_beside_it(self, tokenizer: Any, tmp_path: Path) -> None: + path = write_modelfile(tmp_path, tokenizer) + text = path.read_text(encoding="utf-8") + assert text.startswith("FROM ./model.gguf") + + def test_it_declares_a_stop_token(self, tokenizer: Any, tmp_path: Path) -> None: + """Without one a runtime generates past the end of the reply.""" + text = write_modelfile(tmp_path, tokenizer).read_text(encoding="utf-8") + assert "PARAMETER stop" in text From 366cee0639697b2b379826e0ef37e541bba965e3 Mon Sep 17 00:00:00 2001 From: Aswin Alexander Sam Date: Sun, 2 Aug 2026 18:35:40 +0530 Subject: [PATCH 2/5] address the review on gguf export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ten findings. Three were the same class as the missing merges: metadata that writes without complaint and produces a file no runtime will load. A mistral checkpoint was written with general.architecture = "mistral". llama.cpp registers this family once, under llama, and there is no such entry — the file would have failed with "unknown model architecture". Its tensor layout is Llama's, which is why the checkpoint is accepted at all; the name is not. Reported as mistral, written as llama. The token list was built by sorting the vocabulary, which reproduces the ids only when they are contiguous. Ours are; a published checkpoint's need not be, because they reserve blocks of ids and are routinely shorter than vocab_size. A gap shifted every later token onto the wrong id, and a short vocabulary left fewer entries than the embedding has rows. It is built by index now, sized to vocab_size, with unnamed rows filled. A checkpoint with neither an output projection nor an embedding raised KeyError past every handler. It is refused like everything else here. Never a moment with no checkpoint checkpoint.save deleted the old directory and then renamed the new one in. Between those calls there is nothing on disk, and a kill or a failed rename — across filesystems, say — takes the good checkpoint and puts nothing back. The old one is moved aside and discarded only once the new one is in place, with the move reversed if the rename fails. That also settles the export race properly. The plan noted that export runs concurrently with training and reads runs//latest while it is being replaced, and then did nothing about it. The window is what made that unsafe, and the window is gone, so the two can run together because of a property rather than a warning. The same fix is applied to export's own directory. Export cleanup ran only for ExportError, so a full disk or an interrupt mid-quantization left a partial GGUF in exports/.tmp. It is a finally now, guarded by whether the result was placed. Tests The q4_0 round trip allowed an absolute error of 1.0 against weights whose standard deviation is about 0.02 — it would have passed a writer that emitted zeros, which is exactly what it exists to catch. Tolerances scale to each tensor's own spread now, with a correlation check that the values are that tensor and not merely something small. Verified by making the writer emit zeros and watching it fail. And a test mutated the shared tokenizer fixture, restoring the chat template to None rather than to what was there — the same fault flagged in M4.1, reintroduced. monkeypatch now. Also scoped a README claim to the runtimes this repo actually tests, and made the web form say that run and checkpoint are alternatives, which the CLI enforces and the form did not mention. Re-verified in Ollama after the metadata changes: still loads, still generates. --- README.md | 5 +- src/bloomery/cli.py | 31 +++++++- src/bloomery/export.py | 54 ++++++++++++-- src/bloomery/jobs/types.py | 6 ++ src/bloomery/server/static/app.js | 4 +- src/bloomery/train/checkpoint.py | 23 +++++- tests/test_export.py | 120 +++++++++++++++++++++++++++++- tests/test_train.py | 66 ++++++++++++++++ 8 files changed, 289 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index f40bdc6..cf9f1b7 100644 --- a/README.md +++ b/README.md @@ -182,8 +182,9 @@ instead. Checkpoints are ordinary Hugging Face directories. `AutoModelForCausalLM.from_pretrained` loads them, which is the whole reason bloomery emits a Llama-architecture model -rather than inventing one — vLLM and anything else in the ecosystem read them -directly. +rather than inventing one — vLLM and other runtimes that read Llama-architecture +Hugging Face checkpoints load them directly. (A LoRA run writes adapters rather +than a whole model; `export` folds those in.) GGUF is the exception, and `export` writes it here rather than shelling out. llama.cpp's converter identifies a tokenizer by hashing its output against a diff --git a/src/bloomery/cli.py b/src/bloomery/cli.py index 4ac1958..e2e0e7e 100644 --- a/src/bloomery/cli.py +++ b/src/bloomery/cli.py @@ -1189,17 +1189,40 @@ def export( shutil.rmtree(staging) staging.mkdir(parents=True) + # Cleared on every exit, not only on ExportError. A full disk raises OSError + # and a long quantization pass can take a KeyboardInterrupt, and either would + # otherwise leave a partial GGUF sitting in exports/.tmp. + placed = False try: with console.status(f"writing {quantize} gguf"): result = to_gguf(model, tokenizer, staging / gguf_name, quantization=quantize) write_modelfile(staging, tokenizer) + + # The old export is moved aside rather than deleted, so there is never a + # moment with no export at all — a kill between the two, or a rename that + # fails across filesystems, would otherwise take the previous good GGUF + # and put nothing in its place. + previous = destination.with_name(destination.name + ".previous") + if previous.exists(): + shutil.rmtree(previous) + if destination.exists(): + destination.rename(previous) + try: + staging.rename(destination) + except OSError: + if previous.exists() and not destination.exists(): + previous.rename(destination) + raise + placed = True + shutil.rmtree(previous, ignore_errors=True) except ExportError as exc: - shutil.rmtree(staging, ignore_errors=True) _die(str(exc)) + except OSError as exc: + _die(f"could not write the export: {exc}") + finally: + if not placed: + shutil.rmtree(staging, ignore_errors=True) - if destination.exists(): - shutil.rmtree(destination) - staging.rename(destination) result = replace(result, path=destination / gguf_name) if as_json: diff --git a/src/bloomery/export.py b/src/bloomery/export.py index 220d1f7..750f78d 100644 --- a/src/bloomery/export.py +++ b/src/bloomery/export.py @@ -33,6 +33,12 @@ # nothing else does. SUPPORTED_ARCHITECTURES = frozenset({"llama", "mistral"}) +# What goes in `general.architecture`, which is not the same question as which +# checkpoints can be read. llama.cpp has one registry entry for this family and +# it is called llama — there is no `mistral` in it, so writing the model_type +# through produces a file that fails to load with "unknown model architecture". +GGUF_ARCHITECTURE = "llama" + # What `--quantize` accepts, mapped to the ggml type. Everything here is # implemented in the `gguf` package's own quantizer, which is the reference one. # The K-quants (q4_K, q5_K, q6_K) are deliberately absent: that package can read @@ -139,15 +145,26 @@ def to_gguf( # so it is usually here. Deriving it from the embedding when it is not keeps # both paths working: GGUF always wants an output projection. if "lm_head.weight" not in state: + embedding = state.get("model.embed_tokens.weight") + if embedding is None: + # Refused rather than left to raise KeyError, which would reach the + # user as a traceback past every handler this module's callers have. + raise ExportError( + "this checkpoint has neither an output projection nor a token " + "embedding, so there is nothing to write as GGUF's output layer. " + "It does not look like a causal language model." + ) state = dict(state) - state["lm_head.weight"] = state["model.embed_tokens.weight"] + state["lm_head.weight"] = embedding # The official mapping rather than a hand-written table, so a rename upstream # is not something this has to notice. names = get_tensor_name_map(MODEL_ARCH.LLAMA, config.num_hidden_layers) out_path.parent.mkdir(parents=True, exist_ok=True) - writer = GGUFWriter(str(out_path), architecture) + # GGUF_ARCHITECTURE, not the model_type: llama.cpp registers this family + # under one name, and `architecture` is kept for what we report back. + writer = GGUFWriter(str(out_path), GGUF_ARCHITECTURE) try: writer.add_block_count(config.num_hidden_layers) # The sequence length the model was trained at, not a capability of the @@ -276,6 +293,35 @@ def _add_rope(writer: Any, config: Any) -> None: writer.add_rope_freq_base(float(theta)) +def _vocabulary(tokenizer: Any, config: Any) -> tuple[list[str], set[str]]: + """The token list, indexed by id, sized to the embedding table. + + Built by position rather than by sorting the vocabulary, because sorting only + reproduces the ids when they happen to be contiguous. Tokenizers this project + trains are; ones that arrive with a published checkpoint need not be. They + reserve blocks of ids, and they are routinely shorter than ``vocab_size`` — + the embedding has rows nothing is named for. + + Sorting a sparse vocabulary silently shifts every token after the first gap + by one, so the model tokenizes to ids that mean something else. Nothing + errors; the output is simply wrong. + """ + vocab = tokenizer.get_vocab() + size = int(getattr(config, "vocab_size", 0)) or (max(vocab.values(), default=-1) + 1) + + by_id: dict[int, str] = {} + for token, index in vocab.items(): + position = int(index) + if 0 <= position < size: + by_id[position] = token + + # An unnamed row still needs an entry, or the token list is shorter than the + # embedding and llama.cpp reads the two as disagreeing about the vocabulary. + tokens = [by_id.get(index, f"[UNUSED_{index}]") for index in range(size)] + added = {token for token in getattr(tokenizer, "all_special_tokens", []) or [] if token} + return tokens, added + + def _merges(tokenizer: Any) -> list[str]: """The BPE merge rules, as GGUF wants them: one space-joined pair per entry. @@ -316,9 +362,7 @@ def _add_tokenizer(writer: Any, tokenizer: Any, config: Any) -> None: tokenizer this project trains is — see the module docstring for why naming it outright is safe here and is not for a general converter. """ - vocab = tokenizer.get_vocab() - tokens = [token for token, _ in sorted(vocab.items(), key=lambda item: item[1])] - added = {token for token in getattr(tokenizer, "all_special_tokens", []) or []} + tokens, added = _vocabulary(tokenizer, config) writer.add_tokenizer_model("gpt2") writer.add_tokenizer_pre("default") diff --git a/src/bloomery/jobs/types.py b/src/bloomery/jobs/types.py index d8d588d..6c1088e 100644 --- a/src/bloomery/jobs/types.py +++ b/src/bloomery/jobs/types.py @@ -56,6 +56,12 @@ def terminal(self) -> bool: # and the second one dies. # EXPORT is deliberately absent: it reads a checkpoint and writes a file, so # it can run beside a training job rather than queueing behind one. +# +# That concurrency used to be unsafe. checkpoint.save deleted `latest` before +# renaming the new one into place, so an export starting in that window found +# nothing there. The save now moves the old checkpoint aside instead, leaving +# no moment when the directory is absent — which is what makes running the two +# together sound, rather than a note saying to be careful. EXCLUSIVE_KINDS = frozenset({JobKind.TRAIN, JobKind.ADAPT, JobKind.BENCH}) diff --git a/src/bloomery/server/static/app.js b/src/bloomery/server/static/app.js index 28f987e..25cfebe 100644 --- a/src/bloomery/server/static/app.js +++ b/src/bloomery/server/static/app.js @@ -66,8 +66,8 @@ const FIELDS = { { key: "seed", label: "Seed", type: "number" }, ], export: [ - { key: "run", label: "Run", hint: "a run name from train or adapt" }, - { key: "checkpoint", label: "Checkpoint", hint: "a path, instead of a run" }, + { key: "run", label: "Run", hint: "a run name — or a checkpoint below, not both" }, + { key: "checkpoint", label: "Checkpoint", hint: "a path — or a run above, not both" }, { key: "name", label: "Export name", hint: "defaults to the run's" }, { key: "quantize", label: "Format", choices: ["f16", "q8_0", "q4_0"] }, ], diff --git a/src/bloomery/train/checkpoint.py b/src/bloomery/train/checkpoint.py index 03ce657..5c9e7cb 100644 --- a/src/bloomery/train/checkpoint.py +++ b/src/bloomery/train/checkpoint.py @@ -101,10 +101,27 @@ def save( + "\n" ) - # Replace only once everything is on disk. + # Replace only once everything is on disk — and never leave a moment with no + # checkpoint at all. Deleting the old one first opens a window where a kill, + # or a rename that fails, loses the good checkpoint and puts nothing in its + # place. The old one is moved aside instead, and only discarded once the new + # one is where it belongs. + # + # Not academic: `export` reads runs//latest while a run may be saving, + # and that window is exactly when it finds nothing there. + previous = directory.with_name(directory.name + ".previous") + if previous.exists(): + shutil.rmtree(previous) if directory.exists(): - shutil.rmtree(directory) - staging.rename(directory) + directory.rename(previous) + try: + staging.rename(directory) + except OSError: + # Put the old one back rather than leaving the caller with neither. + if previous.exists() and not directory.exists(): + previous.rename(directory) + raise + shutil.rmtree(previous, ignore_errors=True) return directory diff --git a/tests/test_export.py b/tests/test_export.py index 8070141..3f79a7d 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -55,11 +55,15 @@ class TestRoundTrip: """What was written must be what the model held.""" @pytest.mark.parametrize( - ("quantization", "tolerance"), - [("f16", 1e-2), ("q8_0", 5e-2), ("q4_0", 1.0)], + ("quantization", "share"), + # As a fraction of each tensor's own spread. An absolute bound is + # meaningless here: a freshly built model initialises to a standard + # deviation around 0.02, so a tolerance of 1.0 would admit all zeros, + # or the wrong tensor entirely. + [("f16", 0.01), ("q8_0", 0.05), ("q4_0", 0.6)], ) def test_weights_survive( - self, model: Any, tokenizer: Any, tmp_path: Path, quantization: str, tolerance: float + self, model: Any, tokenizer: Any, tmp_path: Path, quantization: str, share: float ) -> None: out = tmp_path / f"{quantization}.gguf" to_gguf(model, tokenizer, out, quantization=quantization) @@ -75,8 +79,23 @@ def test_weights_survive( ): expected = state[source].to(torch.float32).numpy() got = as_float(written[name]).reshape(expected.shape) + + spread = float(np.abs(expected).max()) error = float(np.abs(got - expected).max()) - assert error <= tolerance, f"{name} drifted by {error} at {quantization}" + assert error <= share * spread, ( + f"{name} drifted by {error} at {quantization}, " + f"more than {share:.0%} of its {spread:.4f} range" + ) + # And it must actually be this tensor, not merely something small. + # Correlation is undefined for a constant tensor — the norms + # initialise to all ones — so those are checked for equality + # instead, which is the stronger claim anyway. + if float(expected.std()) > 0: + assert np.corrcoef(got.ravel(), expected.ravel())[0, 1] > 0.9, ( + f"{name} does not track the weights it came from at {quantization}" + ) + else: + assert np.allclose(got, expected), f"{name} changed at {quantization}" def test_the_output_head_is_written_though_it_is_tied( self, model: Any, tokenizer: Any, tmp_path: Path @@ -196,6 +215,87 @@ def test_a_chat_template_is_carried_across( tokenizer.chat_template = None +class TestMetadataLlamaCppWillAccept: + """Fields whose wrong value produces a file that writes and will not load. + + The same class of failure as the missing merges: nothing here errors at + write time, and llama.cpp refuses the result. + """ + + def test_the_architecture_is_the_one_llama_cpp_registers( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """A mistral checkpoint must still say llama. + + Its tensor layout is Llama's, and llama.cpp has one registry entry for + the family. Writing `general.architecture = "mistral"` produces + "unknown model architecture" on load — there is no such entry. + """ + out = tmp_path / "arch.gguf" + model.config.model_type = "mistral" + try: + result = to_gguf(model, tokenizer, out) + finally: + model.config.model_type = "llama" + + field = GGUFReader(str(out)).fields["general.architecture"] + written = bytes(field.parts[field.data[0]]).decode() + assert written == "llama", written + # Reported as what it actually is, which is a different question. + assert result.architecture == "mistral" + + def test_the_token_list_matches_the_embedding_rows( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """Fewer entries than rows and llama.cpp reads the two as disagreeing.""" + out = tmp_path / "rows.gguf" + to_gguf(model, tokenizer, out) + reader = GGUFReader(str(out)) + tokens = reader.fields["tokenizer.ggml.tokens"] + embedding = next(t for t in reader.tensors if t.name == "token_embd.weight") + + assert len(tokens.data) == model.config.vocab_size + assert len(tokens.data) == int(embedding.shape[1]) + + def test_a_sparse_vocabulary_keeps_every_token_at_its_own_id( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """Sorting only reproduces ids when they happen to be contiguous. + + Published checkpoints reserve blocks of ids. Sorting a vocabulary with a + gap shifts every token after it by one, so the model tokenizes to ids + that mean something else — with nothing reporting it. + """ + + class Sparse: + """Ids 0, 1 and 9, as a reserved block would leave them.""" + + chat_template = None + all_special_tokens: list[str] = [] + is_fast = True + backend_tokenizer = tokenizer.backend_tokenizer + + def get_vocab(self) -> dict[str, int]: + return {"a": 0, "b": 1, "z": 9} + + model.config.vocab_size = 10 + original = model.get_input_embeddings().weight.shape[0] + try: + out = tmp_path / "sparse.gguf" + to_gguf(model, tokenizer, out) # a sanity export, real tokenizer + from bloomery.export import _vocabulary + + tokens, _ = _vocabulary(Sparse(), Sparse) + finally: + model.config.vocab_size = original + + assert len(tokens) == 10 + assert tokens[0] == "a" + assert tokens[1] == "b" + assert tokens[9] == "z", "the gap shifted a token onto the wrong id" + assert tokens[5].startswith("["), "an unnamed row still needs an entry" + + class TestRefusals: def test_an_architecture_we_cannot_map_is_refused_by_name(self) -> None: """`adapt --method full` can produce one, and a wrong mapping is silent.""" @@ -206,6 +306,18 @@ class NotLlama: with pytest.raises(ExportError, match="mamba"): architecture_of(NotLlama()) + def test_a_model_with_no_output_layer_is_refused(self, tokenizer: Any, tmp_path: Path) -> None: + """Refused, rather than reaching the user as a KeyError traceback.""" + + class Hollow: + config = type("C", (), {"model_type": "llama", "num_hidden_layers": 1})() + + def state_dict(self) -> dict[str, Any]: + return {} + + with pytest.raises(ExportError, match="nothing to write"): + to_gguf(Hollow(), tokenizer, tmp_path / "hollow.gguf") + def test_an_unknown_quantization_is_refused( self, model: Any, tokenizer: Any, tmp_path: Path ) -> None: diff --git a/tests/test_train.py b/tests/test_train.py index 4060109..e622419 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -939,3 +939,69 @@ def test_a_step_that_scored_nothing_does_not_report_a_perfect_loss(self) -> None reported = step_loss if contributed else float("nan") assert contributed == 0 assert math.isnan(reported), "an unscored step must not look like a perfect one" + + +class TestReplacingACheckpointLeavesNoGap: + """There must never be a moment with no checkpoint at all. + + Deleting the old one before renaming the new one into place opens a window + where a kill, or a rename that fails, takes the good checkpoint and puts + nothing back. `export` reads runs//latest while a run may be saving, + and that window is exactly when it finds nothing there. + """ + + def _save(self, directory: Path, tokenizer: Any, *, step: int) -> Path: + import torch + + from bloomery.train import checkpoint as ckpt + + model = build_model(spec_from_depth(1, vocab=64, seq=16), eos_token_id=0, seed=0) + return ckpt.save( + directory, + model=model, + tokenizer=tokenizer, + optimizer=torch.optim.AdamW(model.parameters(), lr=1e-3), + step=step, + tokens_seen=step * 10, + best_val_loss=None, + ) + + def test_a_second_save_replaces_the_first(self, tmp_path: Path, tokenizer: Any) -> None: + from bloomery.train import checkpoint as ckpt + + directory = tmp_path / "latest" + self._save(directory, tokenizer, step=1) + self._save(directory, tokenizer, step=2) + + assert ckpt.load_resume_state(directory).step == 2 + # And nothing is left lying about that a reader could mistake for one. + assert not directory.with_name("latest.previous").exists() + assert not directory.with_name("latest.tmp").exists() + + def test_the_old_checkpoint_survives_a_failed_rename( + self, tmp_path: Path, tokenizer: Any, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The case the window existed for: something goes wrong mid-swap. + + A cross-filesystem rename raises, and the caller must still have the + checkpoint it had before rather than neither. + """ + from bloomery.train import checkpoint as ckpt + + directory = tmp_path / "latest" + self._save(directory, tokenizer, step=1) + + real = Path.rename + + def refuse(self: Path, target: Any) -> None: + if self.name.endswith(".tmp"): + raise OSError("cross-device link") + real(self, target) + + monkeypatch.setattr(Path, "rename", refuse) + with pytest.raises(OSError): + self._save(directory, tokenizer, step=2) + monkeypatch.undo() + + assert directory.is_dir(), "the previous checkpoint was lost" + assert ckpt.load_resume_state(directory).step == 1 From 14f6ca5b39d3b0581f5d3d3f9c253a46df45a4a1 Mon Sep 17 00:00:00 2001 From: Aswin Alexander Sam Date: Sun, 2 Aug 2026 18:51:11 +0530 Subject: [PATCH 3/5] recover a checkpoint left aside by an interrupted save MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fix in the previous commit narrowed the window where no checkpoint exists, and opened a worse path in doing so. A kill between the two renames leaves the good checkpoint at .previous with nothing at — and the next save began by deleting .previous to clear its way, destroying the only copy there was. The promotion is now preceded by a recovery: if the canonical path is missing and .previous holds a checkpoint, it is put back. Readers do the same, so `--resume` does not report nothing to resume while a complete checkpoint sits beside the path it looked at, and `export` finds one when it arrives mid-save — which it is built to do, since it runs concurrently with training by design. The window itself is not closed. An atomic directory exchange would do it and the syscall is Linux-only, so instead the state a kill leaves is made unambiguous and recoverable rather than impossible. That is the honest claim and it is what the docstring says. The test for this passed against the broken version, which is worth recording. A save that merely succeeds after an interruption hides the loss, because its own result replaces what was destroyed; only a save that then fails leaves nothing at all. The test makes the second rename raise and asserts the interrupted checkpoint is still there. It fails against the version without recovery. --- src/bloomery/cli.py | 4 ++ src/bloomery/train/checkpoint.py | 34 +++++++++++++++++ tests/test_train.py | 63 ++++++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+) diff --git a/src/bloomery/cli.py b/src/bloomery/cli.py index e2e0e7e..03b049d 100644 --- a/src/bloomery/cli.py +++ b/src/bloomery/cli.py @@ -1164,6 +1164,10 @@ def export( _die(f"unknown quantization {quantize!r}; choose one of: {', '.join(QUANTIZATIONS)}") target = checkpoint if checkpoint else ckpt.checkpoint_dir(paths.run_dir(run or "")) + # A run saving right now may have been killed mid-promotion, leaving the + # checkpoint beside this path rather than at it. Export runs concurrently + # with training by design, so it is the command most likely to arrive then. + ckpt.restore_interrupted(target) if not target.is_dir(): _die(f"{target} does not exist") diff --git a/src/bloomery/train/checkpoint.py b/src/bloomery/train/checkpoint.py index 5c9e7cb..c1e9a82 100644 --- a/src/bloomery/train/checkpoint.py +++ b/src/bloomery/train/checkpoint.py @@ -110,8 +110,15 @@ def save( # Not academic: `export` reads runs//latest while a run may be saving, # and that window is exactly when it finds nothing there. previous = directory.with_name(directory.name + ".previous") + + # Recover before clearing. A save killed between the two renames below + # leaves the checkpoint under `previous` and nothing at `directory`; going + # straight to rmtree here would then delete the only copy that exists, which + # is a worse outcome than the window this whole dance is closing. + restore_interrupted(directory) if previous.exists(): shutil.rmtree(previous) + if directory.exists(): directory.rename(previous) try: @@ -125,6 +132,28 @@ def save( return directory +def restore_interrupted(directory: Path) -> bool: + """Put back a checkpoint left aside by a save that did not finish. + + The promotion above is two renames, and a process killed between them leaves + the good checkpoint at ``.previous`` with nothing at ````. An + atomic directory exchange would remove even that gap, but the syscall for it + is Linux-only, and this project runs on three platforms. + + So the gap is made recoverable instead of impossible: the state it leaves is + unambiguous — a complete checkpoint under a known name — and this puts it + back. Called before any save clears the way, and available to a reader that + finds nothing where it expected a checkpoint. + + Returns whether anything was restored. + """ + previous = directory.with_name(directory.name + ".previous") + if directory.exists() or not previous.is_dir(): + return False + previous.rename(directory) + return True + + def load_resume_state( directory: Path, optimizer: torch.optim.Optimizer | None = None ) -> ResumeState: @@ -157,11 +186,16 @@ def load_resume_state( def is_resumable(directory: Path) -> bool: """Whether a run can be picked up from this directory. + Recovers first: a save killed mid-promotion leaves the checkpoint beside + this path rather than at it, and reporting "nothing to resume" then would + discard a complete checkpoint that is sitting right there. + Needs the optimizer state, plus weights in one of the two shapes a run writes: a whole model, or the adapters a LoRA run produced. Checking only for ``config.json`` would report every adapter checkpoint as unresumable, which is the shape a long adaptation run leaves behind. """ + restore_interrupted(directory) if not (directory / TRAINER_STATE).is_file(): return False return (directory / "config.json").is_file() or (directory / "adapter_config.json").is_file() diff --git a/tests/test_train.py b/tests/test_train.py index e622419..e959ebd 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1005,3 +1005,66 @@ def refuse(self: Path, target: Any) -> None: assert directory.is_dir(), "the previous checkpoint was lost" assert ckpt.load_resume_state(directory).step == 1 + + def test_a_save_killed_mid_promotion_is_recovered(self, tmp_path: Path, tokenizer: Any) -> None: + """The window that cannot be closed on three platforms, made harmless. + + An atomic directory exchange would remove it, but that syscall is + Linux-only. So the state a kill leaves is unambiguous — a complete + checkpoint under a known name — and it gets put back. + """ + from bloomery.train import checkpoint as ckpt + + directory = tmp_path / "latest" + self._save(directory, tokenizer, step=1) + # Exactly what a kill between the two renames leaves behind. + directory.rename(directory.with_name("latest.previous")) + assert not directory.exists() + + assert ckpt.restore_interrupted(directory) is True + assert ckpt.load_resume_state(directory).step == 1 + assert not directory.with_name("latest.previous").exists() + + def test_the_next_save_does_not_destroy_an_interrupted_one( + self, tmp_path: Path, tokenizer: Any, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The hole the first version of this fix opened. + + After an interrupted promotion the only checkpoint is under `.previous`. + Clearing the way for the next save by deleting it destroys that copy — + and if the new save then fails, nothing is left at all. A save that + merely succeeds hides this, because its own result replaces what was + lost; the failure is what exposes it. + """ + from bloomery.train import checkpoint as ckpt + + directory = tmp_path / "latest" + self._save(directory, tokenizer, step=1) + # Exactly what a kill between the two renames leaves. + directory.rename(directory.with_name("latest.previous")) + + real = Path.rename + + def refuse(self: Path, target: Any) -> None: + if self.name.endswith(".tmp"): + raise OSError("cross-device link") + real(self, target) + + monkeypatch.setattr(Path, "rename", refuse) + with pytest.raises(OSError): + self._save(directory, tokenizer, step=2) + monkeypatch.undo() + + assert directory.is_dir(), "the interrupted checkpoint was destroyed" + assert ckpt.load_resume_state(directory).step == 1 + + def test_a_reader_finds_an_interrupted_checkpoint(self, tmp_path: Path, tokenizer: Any) -> None: + """Otherwise `--resume` reports nothing to resume, beside a complete one.""" + from bloomery.train import checkpoint as ckpt + + directory = tmp_path / "latest" + self._save(directory, tokenizer, step=3) + directory.rename(directory.with_name("latest.previous")) + + assert ckpt.is_resumable(directory) + assert ckpt.load_resume_state(directory).step == 3 From 4c01f0da4db16f737b465d0d05216eb6f48ae7f0 Mon Sep 17 00:00:00 2001 From: Aswin Alexander Sam Date: Sun, 2 Aug 2026 19:08:32 +0530 Subject: [PATCH 4/5] let readers find an aside checkpoint without moving it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adding recovery to readers in the previous commit created the race it was meant to guard against. A save mid-promotion is indistinguishable from an interrupted one — the canonical directory is absent and .previous holds a checkpoint — so a concurrent export doing the restore takes .previous out from under the save. The save's own rename then lands on a path that is suddenly occupied, and its recovery finds nothing to put back. A microsecond window became a failed training run, and export runs alongside training by design, so it is the command most likely to be there. Reading is safe; renaming is what races. `resolve` returns wherever the checkpoint can be read from and moves nothing, and is_resumable, load_resume_state and export use it. `restore_interrupted` still exists and is documented as the writer's alone. The exists-then-rmtree in save is now rmtree(ignore_errors=True): two operations against one path, and anything clearing it between them would raise FileNotFoundError out of an otherwise successful save. The test for this did not exist when I first claimed it did. The string replacement that was supposed to add it silently matched nothing, and five tests passing read as six. It asserts the directory listing is unchanged across a read, and fails against a reader that restores. --- src/bloomery/cli.py | 6 ++-- src/bloomery/train/checkpoint.py | 55 +++++++++++++++++++++++++------- tests/test_train.py | 35 +++++++++++++++++++- 3 files changed, 81 insertions(+), 15 deletions(-) diff --git a/src/bloomery/cli.py b/src/bloomery/cli.py index 03b049d..58522fa 100644 --- a/src/bloomery/cli.py +++ b/src/bloomery/cli.py @@ -1166,8 +1166,10 @@ def export( target = checkpoint if checkpoint else ckpt.checkpoint_dir(paths.run_dir(run or "")) # A run saving right now may have been killed mid-promotion, leaving the # checkpoint beside this path rather than at it. Export runs concurrently - # with training by design, so it is the command most likely to arrive then. - ckpt.restore_interrupted(target) + # with training by design, so it is the command most likely to arrive then — + # and for the same reason it must only read the aside copy, never move it: + # a save in mid-promotion looks exactly like an interrupted one. + target = ckpt.resolve(target) if not target.is_dir(): _die(f"{target} does not exist") diff --git a/src/bloomery/train/checkpoint.py b/src/bloomery/train/checkpoint.py index c1e9a82..f79c3a7 100644 --- a/src/bloomery/train/checkpoint.py +++ b/src/bloomery/train/checkpoint.py @@ -116,8 +116,10 @@ def save( # straight to rmtree here would then delete the only copy that exists, which # is a worse outcome than the window this whole dance is closing. restore_interrupted(directory) - if previous.exists(): - shutil.rmtree(previous) + # ignore_errors rather than exists-then-remove: the check and the call + # are two operations, and anything that clears the path between them + # would otherwise raise FileNotFoundError out of a successful save. + shutil.rmtree(previous, ignore_errors=True) if directory.exists(): directory.rename(previous) @@ -138,29 +140,56 @@ def restore_interrupted(directory: Path) -> bool: The promotion above is two renames, and a process killed between them leaves the good checkpoint at ``.previous`` with nothing at ````. An atomic directory exchange would remove even that gap, but the syscall for it - is Linux-only, and this project runs on three platforms. + is Linux-only and this project runs on three platforms, so the gap is made + recoverable rather than impossible. + + **Only :func:`save` may call this.** It moves a directory, and a save in + mid-promotion is in exactly the state this looks for — so a concurrent + caller doing the restore would take ``.previous`` out from under the save, + whose own rename then fails onto a path that is suddenly occupied and whose + recovery finds nothing to put back. That turns a harmless window into a + failed training run. - So the gap is made recoverable instead of impossible: the state it leaves is - unambiguous — a complete checkpoint under a known name — and this puts it - back. Called before any save clears the way, and available to a reader that - finds nothing where it expected a checkpoint. + A reader wanting the same robustness uses :func:`resolve`, which reads the + aside copy where it lies and moves nothing. Returns whether anything was restored. """ previous = directory.with_name(directory.name + ".previous") if directory.exists() or not previous.is_dir(): return False - previous.rename(directory) + try: + previous.rename(directory) + except OSError: + # Something else got there first. Nothing to do and nothing broken. + return False return True +def resolve(directory: Path) -> Path: + """Where this checkpoint can actually be read from, without moving anything. + + Normally the path given. After a save that was killed mid-promotion, the + complete checkpoint is at ``.previous`` instead, and a reader that + only looked at the canonical path would report nothing while a usable + checkpoint sat beside it. + + Deliberately read-only. Putting it back is the writer's job — see + :func:`restore_interrupted` for what happens when a reader tries. + """ + if directory.exists(): + return directory + previous = directory.with_name(directory.name + ".previous") + return previous if (previous / TRAINER_STATE).is_file() else directory + + def load_resume_state( directory: Path, optimizer: torch.optim.Optimizer | None = None ) -> ResumeState: """Restore step counters, and optimizer state if an optimizer is given.""" import torch - state_path = directory / TRAINER_STATE + state_path = resolve(directory) / TRAINER_STATE if not state_path.is_file(): raise FileNotFoundError(f"no {TRAINER_STATE} in {directory}") @@ -186,16 +215,18 @@ def load_resume_state( def is_resumable(directory: Path) -> bool: """Whether a run can be picked up from this directory. - Recovers first: a save killed mid-promotion leaves the checkpoint beside + Resolved first: a save killed mid-promotion leaves the checkpoint beside this path rather than at it, and reporting "nothing to resume" then would - discard a complete checkpoint that is sitting right there. + overlook a complete checkpoint sitting right there. Read where it lies — + moving it is the writer's job, and a reader that moved it would break a + save that is running. Needs the optimizer state, plus weights in one of the two shapes a run writes: a whole model, or the adapters a LoRA run produced. Checking only for ``config.json`` would report every adapter checkpoint as unresumable, which is the shape a long adaptation run leaves behind. """ - restore_interrupted(directory) + directory = resolve(directory) if not (directory / TRAINER_STATE).is_file(): return False return (directory / "config.json").is_file() or (directory / "adapter_config.json").is_file() diff --git a/tests/test_train.py b/tests/test_train.py index e959ebd..9e15977 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1021,8 +1021,14 @@ def test_a_save_killed_mid_promotion_is_recovered(self, tmp_path: Path, tokenize directory.rename(directory.with_name("latest.previous")) assert not directory.exists() - assert ckpt.restore_interrupted(directory) is True + # A reader finds it where it lies, without moving anything. + assert ckpt.resolve(directory).name == "latest.previous" assert ckpt.load_resume_state(directory).step == 1 + assert directory.with_name("latest.previous").is_dir() + + # The writer is what puts it back. + assert ckpt.restore_interrupted(directory) is True + assert directory.is_dir() assert not directory.with_name("latest.previous").exists() def test_the_next_save_does_not_destroy_an_interrupted_one( @@ -1068,3 +1074,30 @@ def test_a_reader_finds_an_interrupted_checkpoint(self, tmp_path: Path, tokenize assert ckpt.is_resumable(directory) assert ckpt.load_resume_state(directory).step == 3 + + def test_a_reader_does_not_move_a_checkpoint_out_from_under_a_save( + self, tmp_path: Path, tokenizer: Any + ) -> None: + """Why readers resolve rather than restore. + + A save mid-promotion looks exactly like an interrupted one: the canonical + directory is absent and `.previous` holds a checkpoint. A reader that + restored would take it, the save's own rename would then land on an + occupied path, and its recovery would find nothing to put back — turning + a microsecond window into a failed run. + """ + from bloomery.train import checkpoint as ckpt + + directory = tmp_path / "latest" + self._save(directory, tokenizer, step=1) + # The exact state a save is in between its two renames. + directory.rename(directory.with_name("latest.previous")) + + before = sorted(path.name for path in tmp_path.iterdir()) + assert ckpt.resolve(directory).is_dir() + assert ckpt.is_resumable(directory) + assert ckpt.load_resume_state(directory).step == 1 + + assert sorted(path.name for path in tmp_path.iterdir()) == before, ( + "a read moved something on disk" + ) From 80ab2430463c7f2155e73cf9f2301adc0a78d059 Mon Sep 17 00:00:00 2001 From: Aswin Alexander Sam Date: Sun, 2 Aug 2026 19:22:48 +0530 Subject: [PATCH 5/5] size the token list from the embedding, not the config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two findings from the review that carried no comment id, so no thread ever tracked them and both went unimplemented. The vocabulary was sized from `config.vocab_size` and never checked against the embedding it has to describe. llama.cpp builds `token_embd.weight` as {n_embd, n_vocab} from the token list it read and rejects a file whose tensor disagrees, so a stale config writes something whose own two halves contradict each other. The two numbers match in a checkpoint nobody resized; a padded embedding leaves the config low, and a resize on an older transformers leaves it behind entirely. Both arrive through `adapt` on someone else's checkpoint. Sized from the rows now. A token whose id is past the last row is refused rather than dropped: dropping exports cleanly and crashes the first time the tokenizer produces one. The reported vocab is what was written rather than what was claimed. The Modelfile interpolated the eos token into a quoted value with no escaping, so a token containing a quote ended the value early and left the rest as stray arguments. It is left out rather than escaped — whether that grammar reads a backslash as an escape is not something this verifies, and a stop parameter that is quietly the wrong string is worse than an absent one. The eos id travels in the GGUF metadata, which is what a runtime stops on regardless. Each of the three tests was run against the old behaviour first. The malformed line they produce is `PARAMETER stop "<|end"of"text|>"`. --- src/bloomery/export.py | 67 ++++++++++++++++++++++++++++++++----- tests/test_export.py | 76 +++++++++++++++++++++++++++++++++++++----- 2 files changed, 125 insertions(+), 18 deletions(-) diff --git a/src/bloomery/export.py b/src/bloomery/export.py index 750f78d..f39552e 100644 --- a/src/bloomery/export.py +++ b/src/bloomery/export.py @@ -157,6 +157,11 @@ def to_gguf( state = dict(state) state["lm_head.weight"] = embedding + # The token list is sized from this rather than from config.vocab_size. See + # _vocabulary: llama.cpp checks the two against each other and refuses a file + # where they disagree. + rows = _embedding_rows(state) + # The official mapping rather than a hand-written table, so a rename upstream # is not something this has to notice. names = get_tensor_name_map(MODEL_ARCH.LLAMA, config.num_hidden_layers) @@ -178,7 +183,7 @@ def to_gguf( writer.add_layer_norm_rms_eps(config.rms_norm_eps) writer.add_file_type(_file_type(quantization)) _add_rope(writer, config) - _add_tokenizer(writer, tokenizer, config) + vocabulary = _add_tokenizer(writer, tokenizer, config, rows) parameters = 0 written = 0 @@ -221,7 +226,8 @@ def to_gguf( parameters=parameters, bytes_written=out_path.stat().st_size, context_length=config.max_position_embeddings, - vocab_size=config.vocab_size, + # What was written, which is not always what the config claims. + vocab_size=vocabulary, unquantized=tuple(unquantized), ) @@ -259,7 +265,17 @@ def write_modelfile(directory: Path, tokenizer: Any, *, gguf_name: str = GGUF_NA eos = getattr(tokenizer, "eos_token", None) if eos: - lines.append(f'PARAMETER stop "{eos}"') + # Quoted, so a token carrying a quote or a newline would end the value + # early and leave the rest as stray arguments. No escape sequence is + # written instead: the Modelfile grammar's handling of one is not + # something this can verify, and a stop parameter that is silently the + # wrong string is worse than none. It is belt-and-braces anyway — the eos + # id travels in the GGUF metadata, which is what Ollama stops on. + if any(character in str(eos) for character in '"\r\n'): + lines.append(f"# The eos token is not quotable here, so no stop parameter: {eos!r}") + lines.append("# Ollama stops on the eos id in the GGUF metadata regardless.") + else: + lines.append(f'PARAMETER stop "{eos}"') lines.append("") path = directory / MODELFILE @@ -293,21 +309,50 @@ def _add_rope(writer: Any, config: Any) -> None: writer.add_rope_freq_base(float(theta)) -def _vocabulary(tokenizer: Any, config: Any) -> tuple[list[str], set[str]]: +def _embedding_rows(state: Any) -> int: + """How many tokens the model can actually embed, or 0 if that is unknowable.""" + for name in ("model.embed_tokens.weight", "lm_head.weight"): + tensor = state.get(name) + if tensor is not None and getattr(tensor, "ndim", 0) == 2: + return int(tensor.shape[0]) + return 0 + + +def _vocabulary(tokenizer: Any, config: Any, rows: int = 0) -> tuple[list[str], set[str]]: """The token list, indexed by id, sized to the embedding table. Built by position rather than by sorting the vocabulary, because sorting only reproduces the ids when they happen to be contiguous. Tokenizers this project trains are; ones that arrive with a published checkpoint need not be. They - reserve blocks of ids, and they are routinely shorter than ``vocab_size`` — - the embedding has rows nothing is named for. + reserve blocks of ids, and they are routinely shorter than the embedding — + it has rows nothing is named for. Sorting a sparse vocabulary silently shifts every token after the first gap by one, so the model tokenizes to ids that mean something else. Nothing errors; the output is simply wrong. + + Sized from the embedding's own row count, not ``config.vocab_size``, because + the row count is the number llama.cpp checks: it builds ``token_embd.weight`` + as ``{n_embd, n_vocab}`` from the token list it read and rejects a file whose + tensor disagrees. The two are the same number in a checkpoint nobody resized, + and ``config.vocab_size`` is stale in one that was — a padded embedding leaves + it low, and `resize_token_embeddings` on an older transformers leaves it + behind entirely. Trusting the config there writes a file no runtime loads. """ vocab = tokenizer.get_vocab() - size = int(getattr(config, "vocab_size", 0)) or (max(vocab.values(), default=-1) + 1) + size = rows or int(getattr(config, "vocab_size", 0)) or (max(vocab.values(), default=-1) + 1) + + # Tokens the model has no row for. Dropping them quietly would leave a + # tokenizer that emits ids the embedding cannot index — a crash at generation + # time, from a file that exported cleanly. + overflow = sorted(int(index) for index in vocab.values() if int(index) >= size) + if overflow: + raise ExportError( + f"this tokenizer has {len(overflow)} token(s) the model cannot embed: ids " + f"{overflow[:5]}{'…' if len(overflow) > 5 else ''} against {size} embedding rows.\n" + "Tokens were probably added without resizing the embedding. Export is " + "refused because the resulting model would fail the moment one was produced." + ) by_id: dict[int, str] = {} for token, index in vocab.items(): @@ -355,14 +400,16 @@ def _merges(tokenizer: Any) -> list[str]: return merges -def _add_tokenizer(writer: Any, tokenizer: Any, config: Any) -> None: +def _add_tokenizer(writer: Any, tokenizer: Any, config: Any, rows: int = 0) -> int: """Write the vocabulary, and say which pre-tokenizer produced it. ``gpt2`` and ``default`` are the byte-level BPE settings, which is what every tokenizer this project trains is — see the module docstring for why naming it outright is safe here and is not for a general converter. + + Returns how many tokens were written. """ - tokens, added = _vocabulary(tokenizer, config) + tokens, added = _vocabulary(tokenizer, config, rows) writer.add_tokenizer_model("gpt2") writer.add_tokenizer_pre("default") @@ -387,3 +434,5 @@ def _add_tokenizer(writer: Any, tokenizer: Any, config: Any) -> None: template = getattr(tokenizer, "chat_template", None) if template: writer.add_chat_template(template) + + return len(tokens) diff --git a/tests/test_export.py b/tests/test_export.py index 3f79a7d..d26059d 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -257,6 +257,30 @@ def test_the_token_list_matches_the_embedding_rows( assert len(tokens.data) == model.config.vocab_size assert len(tokens.data) == int(embedding.shape[1]) + def test_a_stale_vocab_size_does_not_shrink_the_token_list( + self, model: Any, tokenizer: Any, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The embedding's rows are the truth; the config is what goes stale. + + llama.cpp sizes `token_embd.weight` as {n_embd, n_vocab} from the token + list it read, so a list sized from a config that no longer matches the + embedding writes a file whose own two halves disagree — it does not load. + A padded embedding or a resize on an older transformers leaves exactly + this gap, and both arrive through `adapt` on someone else's checkpoint. + """ + rows = int(model.get_input_embeddings().weight.shape[0]) + monkeypatch.setattr(model.config, "vocab_size", rows - 3) + + out = tmp_path / "stale.gguf" + result = to_gguf(model, tokenizer, out) + + reader = GGUFReader(str(out)) + embedding = next(t for t in reader.tensors if t.name == "token_embd.weight") + assert len(reader.fields["tokenizer.ggml.tokens"].data) == rows + assert int(embedding.shape[1]) == rows + # Reported as what was written, not as what the config claimed. + assert result.vocab_size == rows + def test_a_sparse_vocabulary_keeps_every_token_at_its_own_id( self, model: Any, tokenizer: Any, tmp_path: Path ) -> None: @@ -278,16 +302,9 @@ class Sparse: def get_vocab(self) -> dict[str, int]: return {"a": 0, "b": 1, "z": 9} - model.config.vocab_size = 10 - original = model.get_input_embeddings().weight.shape[0] - try: - out = tmp_path / "sparse.gguf" - to_gguf(model, tokenizer, out) # a sanity export, real tokenizer - from bloomery.export import _vocabulary + from bloomery.export import _vocabulary - tokens, _ = _vocabulary(Sparse(), Sparse) - finally: - model.config.vocab_size = original + tokens, _ = _vocabulary(Sparse(), Sparse, 10) assert len(tokens) == 10 assert tokens[0] == "a" @@ -318,6 +335,29 @@ def state_dict(self) -> dict[str, Any]: with pytest.raises(ExportError, match="nothing to write"): to_gguf(Hollow(), tokenizer, tmp_path / "hollow.gguf") + def test_a_token_the_model_cannot_embed_is_refused( + self, model: Any, tokenizer: Any, tmp_path: Path + ) -> None: + """`add_tokens` without a resize leaves ids the embedding cannot index. + + Dropping them quietly would export cleanly and crash the first time the + tokenizer produced one. + """ + rows = int(model.get_input_embeddings().weight.shape[0]) + + class Overflowing: + """The real tokenizer, plus one token nothing has a row for.""" + + chat_template = None + all_special_tokens: list[str] = [] + backend_tokenizer = tokenizer.backend_tokenizer + + def get_vocab(self) -> dict[str, int]: + return {**tokenizer.get_vocab(), "<|extra|>": rows} + + with pytest.raises(ExportError, match="cannot embed"): + to_gguf(model, Overflowing(), tmp_path / "overflow.gguf") + def test_an_unknown_quantization_is_refused( self, model: Any, tokenizer: Any, tmp_path: Path ) -> None: @@ -335,3 +375,21 @@ def test_it_declares_a_stop_token(self, tokenizer: Any, tmp_path: Path) -> None: """Without one a runtime generates past the end of the reply.""" text = write_modelfile(tmp_path, tokenizer).read_text(encoding="utf-8") assert "PARAMETER stop" in text + + def test_an_eos_token_that_cannot_be_quoted_is_left_out( + self, tokenizer: Any, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A quote inside the value ends it early and strands the rest as arguments. + + Left out rather than escaped: whether this grammar reads a backslash as an + escape is not something this project verifies, and a stop parameter that is + quietly the wrong string is worse than an absent one. The eos id travels in + the GGUF metadata, which is what a runtime actually stops on. + + monkeypatch, not assignment: the tokenizer fixture is shared across tests. + """ + monkeypatch.setattr(tokenizer, "eos_token", '<|end"of"text|>', raising=False) + text = write_modelfile(tmp_path, tokenizer).read_text(encoding="utf-8") + + assert "PARAMETER stop" not in text + assert "not quotable" in text