From c7be5d068b3c273220876dfbd48c6a9d6d6ae138 Mon Sep 17 00:00:00 2001 From: Erik Bogado Date: Mon, 31 Aug 2026 20:59:48 -0300 Subject: [PATCH] fix(ftw): preserve auxiliary expert banks --- python/freetoken/checkpoint/convert.py | 28 ++++++++++ python/freetoken/checkpoint/ftw.py | 22 +++++++- python/freetoken/moe/expert_banks.py | 5 ++ tests/test_ftw_auxiliary.py | 76 ++++++++++++++++++++++++++ 4 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 tests/test_ftw_auxiliary.py diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 420faf2c2..c0fc54d40 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -25,6 +25,8 @@ from .ftw import DEFAULT_SHARD_LIMIT, FTWWriter, layer_bank_entry_name +_AUXILIARY_DIR = "auxiliary" + # Machine-readable convert progress for a supervising process (e.g. a GUI frontend parses # these `FTCONVERT ` stdout lines to drive its convert bar). Gated by # FREETOKEN_CONVERT_PROGRESS=1 so plain CLI use isn't spammed; the human tqdm bars stay on @@ -160,6 +162,29 @@ def num_layers(self) -> int: return len(self._seen) +def _write_auxiliary_ftw(out_dir: str, banks, *, shard_limit: int) -> dict: + sources = banks.auxiliary_sources + if sources is None: + return {} + layer_ids = list(banks.auxiliary_layer_ids) + if not banks.auxiliary_quant_format or not layer_ids: + raise ValueError("auxiliary expert banks require a quant format and layer mapping") + writer = FTWWriter(os.path.join(out_dir, _AUXILIARY_DIR), shard_limit=shard_limit) + count = 0 + for name, per_layer in sources.items(): + if len(per_layer) != len(layer_ids): + raise ValueError(f"auxiliary bank {name!r} has {len(per_layer)} layers, expected {len(layer_ids)}") + for layer_id, tensor in enumerate(per_layer): + writer.add_tensor(layer_bank_entry_name(name, layer_id), tensor, kind="experts_bank") + count += 1 + writer.finalize({ + "quant_format": banks.auxiliary_quant_format, + "expert_bank_num_layers": len(layer_ids), + "counts": {"weight": 0, "experts_bank": count}, + }) + return {"auxiliary_checkpoint": _AUXILIARY_DIR, "auxiliary_layer_ids": layer_ids} + + def convert_checkpoint( model_path: str, out_dir: str, @@ -204,6 +229,7 @@ def convert_checkpoint( writer = FTWWriter(out_dir, shard_limit=shard_limit) n_weight = n_bank = n_alpha = 0 + auxiliary_meta = {} # 1) dense weights (host tensors; load straight to CPU to avoid GPU pressure) _progress("dense", 0, 0) # phase start; per-tensor cumulative bytes follow (total unknown) @@ -273,6 +299,7 @@ def convert_checkpoint( n_bank += name not in ("gate_up_alpha", "down_alpha") n_alpha += name in ("gate_up_alpha", "down_alpha") bar.close() + auxiliary_meta = _write_auxiliary_ftw(out_dir, banks, shard_limit=shard_limit) _progress("finalize") # writing shard index + copying config/tokenizer copied = _copy_metadata(model_path, out_dir) @@ -297,6 +324,7 @@ def convert_checkpoint( "expert_bank_num_layers": num_layers, "counts": {"weight": n_weight, "experts_bank": n_bank + n_alpha}, "copied_metadata": copied, + **auxiliary_meta, }) return index diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index a365e3d5b..fe8a0538a 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -462,6 +462,11 @@ def _backing(layer_id: int) -> str: return "mmap" reader = FTWReader(path) + auxiliary_checkpoint = reader.meta("auxiliary_checkpoint") + auxiliary_layer_ids = tuple(reader.meta("auxiliary_layer_ids", ())) + if bool(auxiliary_checkpoint) != bool(auxiliary_layer_ids): + reader.close() + raise RuntimeError("FTW auxiliary checkpoint and layer mapping must be recorded together") bank_entries = reader.entries("experts_bank") if not bank_entries: reader.close() @@ -627,9 +632,24 @@ def _read_layer(job): # alphas are the small per-expert scale vectors, distinguished by their reserved names # (not a separate kind); everything else under experts_bank is a weight source. alpha_kw = {n: alpha_hb[n].tensor for n in alpha_hb} + auxiliary_kw = {} + if auxiliary_checkpoint: + auxiliary = load_ftw_banks( + os.path.join(path, auxiliary_checkpoint), + num_layers=len(auxiliary_layer_ids), + workers=workers, + chunk=chunk, + ) + if auxiliary is None: + raise RuntimeError("FTW auxiliary checkpoint contains no expert banks") + auxiliary_kw = { + "auxiliary_quant_format": auxiliary.quant_format, + "auxiliary_sources": auxiliary.sources, + "auxiliary_layer_ids": auxiliary_layer_ids, + } return ExpertBanks( reader.meta("quant_format"), sources, **alpha_kw, - layer_residency=applied, + layer_residency=applied, **auxiliary_kw, ) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba8..ef9032c20 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -49,6 +49,11 @@ class ExpertBanks: # streamed straight to its sink instead of staying materialized here) -- set by # convert.py's per-format streaming gate; ``sources`` may hold released tensors. streamed: bool = False + # Optional banks with a different row geometry. They are serialized as a nested FTW + # checkpoint and mapped back to their model layers by ``auxiliary_layer_ids``. + auxiliary_quant_format: str | None = None + auxiliary_sources: dict[str, list[torch.Tensor]] | None = None + auxiliary_layer_ids: tuple[int, ...] = () _PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader diff --git a/tests/test_ftw_auxiliary.py b/tests/test_ftw_auxiliary.py new file mode 100644 index 000000000..38d20b1b9 --- /dev/null +++ b/tests/test_ftw_auxiliary.py @@ -0,0 +1,76 @@ +import json +from types import SimpleNamespace + +import pytest +import torch + + +def test_ftw_writes_auxiliary_banks_as_per_layer_checkpoint(tmp_path): + from freetoken.checkpoint.convert import _write_auxiliary_ftw + + banks = SimpleNamespace( + auxiliary_quant_format="q6_k_down", + auxiliary_sources={ + "down": [ + torch.arange(6, dtype=torch.uint8).reshape(2, 3), + torch.ones((2, 3), dtype=torch.uint8), + ] + }, + auxiliary_layer_ids=(7, 9), + ) + + metadata = _write_auxiliary_ftw(str(tmp_path), banks, shard_limit=4096) + + assert metadata == { + "auxiliary_checkpoint": "auxiliary", + "auxiliary_layer_ids": [7, 9], + } + index = json.loads((tmp_path / "auxiliary" / "freetoken_weight.json").read_text()) + assert index["quant_format"] == "q6_k_down" + assert index["expert_bank_num_layers"] == 2 + assert [entry["name"] for entry in index["tensors"]] == [ + "down#L00000", + "down#L00001", + ] + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="FTW expert banks require pinned GPU memory" +) +def test_ftw_loads_nested_auxiliary_banks(tmp_path): + from freetoken.checkpoint.convert import _write_auxiliary_ftw + from freetoken.checkpoint.ftw import ( + FTWWriter, + layer_bank_entry_name, + load_ftw_banks, + ) + + auxiliary = SimpleNamespace( + auxiliary_quant_format="q6_k_down", + auxiliary_sources={"down": [torch.arange(6, dtype=torch.uint8).reshape(2, 3)]}, + auxiliary_layer_ids=(9,), + ) + metadata = _write_auxiliary_ftw(str(tmp_path), auxiliary, shard_limit=4096) + writer = FTWWriter(str(tmp_path), shard_limit=4096) + for layer_id in range(2): + writer.add_tensor( + layer_bank_entry_name("gate_up", layer_id), + torch.full((2, 3), layer_id, dtype=torch.uint8), + kind="experts_bank", + ) + writer.finalize( + { + "quant_format": "q4_k_q5_k", + "expert_bank_num_layers": 2, + **metadata, + } + ) + + banks = load_ftw_banks(str(tmp_path), num_layers=2) + + assert banks.auxiliary_quant_format == "q6_k_down" + assert banks.auxiliary_layer_ids == (9,) + assert len(banks.auxiliary_sources["down"]) == 1 + torch.testing.assert_close( + banks.auxiliary_sources["down"][0], auxiliary.auxiliary_sources["down"][0] + )