fix(ftw): preserve auxiliary expert banks

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 <phase> <done> <total>` 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
 

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,
     )
 
 

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

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]
+    )