From 1d5ce5459b6632adc60feb50b07cd4bbf73af1c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Sat, 3 Oct 2026 09:45:33 -0400 Subject: [PATCH 1/4] fix(checkpoint): refuse converting into a directory that already holds an FTW --- python/freetoken/checkpoint/convert.py | 20 ++++++++ tests/checkpoint/test_convert_refusal.py | 60 ++++++++++++++++++++++++ 2 files changed, 80 insertions(+) create mode 100644 tests/checkpoint/test_convert_refusal.py diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index a875d77f34..2e6ebbad93 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -103,6 +103,25 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]: return copied +def _refuse_occupied_output(out_dir: str) -> None: + """Refuse an output dir that already holds FTW output: the index of a finished conversion, + or shards left by an interrupted one. There is no resume; the writer would open shard + 00000 with "wb" and a stale index could then describe the wrong bytes.""" + from .ftw import INDEX_NAME, is_ftw_checkpoint + + if is_ftw_checkpoint(out_dir): + found = INDEX_NAME + else: + shards = glob.glob(os.path.join(glob.escape(out_dir), "freetoken-*.ftw")) + if not shards: + return + found = os.path.basename(min(shards)) + raise SystemExit( + f"{out_dir} already holds an FTW checkpoint ({found}); " + f"convert into an empty directory or remove the old one first" + ) + + class _ConvertSink: """Layer-completion sink for ``load_expert_banks(layer_sink=...)``: writes each completed layer's banks as their own FTW entries immediately (name @@ -184,6 +203,7 @@ def convert_checkpoint( if is_ftw_checkpoint(model_path): raise SystemExit(f"{model_path} is already an FTW checkpoint") + _refuse_occupied_output(out_dir) tp = try_get_tp_info() if tp is None: set_tp_info(rank=0, size=1) diff --git a/tests/checkpoint/test_convert_refusal.py b/tests/checkpoint/test_convert_refusal.py new file mode 100644 index 0000000000..6c09f79ad8 --- /dev/null +++ b/tests/checkpoint/test_convert_refusal.py @@ -0,0 +1,60 @@ +"""convert_checkpoint refuses an output directory that already holds FTW output, before it writes anything. + +No GPU: the refusal sits ahead of the TP/CUDA setup, so a CPU box reaches it with an empty source dir.""" + +import os + +import pytest +import torch + +from freetoken.checkpoint.convert import _refuse_occupied_output, convert_checkpoint +from freetoken.checkpoint.ftw import INDEX_NAME, FTWWriter + + +def _listing(d): + return sorted((name, os.stat(os.path.join(d, name)).st_size, os.stat(os.path.join(d, name)).st_mtime_ns) + for name in os.listdir(d)) + + +def _refused(src, out): + before = _listing(out) + with pytest.raises(SystemExit, match="already holds an FTW checkpoint") as ei: + convert_checkpoint(str(src), str(out)) + assert _listing(out) == before # nothing written, nothing touched + return str(ei.value) + + +def test_finished_ftw_in_the_output_dir_is_refused_by_name(tmp_path): + out = tmp_path / "out" + writer = FTWWriter(str(out)) + writer.add_tensor("model.a.weight", torch.zeros(4, dtype=torch.bfloat16)) + writer.finalize({}) + src = tmp_path / "src" + src.mkdir() + msg = _refused(src, out) + assert str(out) in msg and INDEX_NAME in msg + + +def test_shards_from_an_interrupted_conversion_are_refused_by_name(tmp_path): + out = tmp_path / "out" + out.mkdir() + (out / "freetoken-00001.ftw").write_bytes(b"\0" * 64) + (out / "freetoken-00000.ftw").write_bytes(b"\0" * 64) + src = tmp_path / "src" + src.mkdir() + msg = _refused(src, out) + assert str(out) in msg and "freetoken-00000.ftw" in msg and INDEX_NAME not in msg + + +def test_empty_absent_or_unrelated_output_dirs_pass(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + assert _refuse_occupied_output(str(empty)) is None + absent = tmp_path / "absent" + assert _refuse_occupied_output(str(absent)) is None + assert not absent.exists() # the check creates nothing + other = tmp_path / "other" + other.mkdir() + (other / "config.json").write_text("{}") # copied metadata alone is not FTW output + (other / "notes.ftw.bak").write_bytes(b"") + assert _refuse_occupied_output(str(other)) is None From a8498c7e96e25fad7ed5e61968188fdbb66b35c6 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 6 Oct 2026 13:27:13 +0300 Subject: [PATCH 2/4] fix(server): leak no semaphores when a signal stops the server --- python/freetoken/server/launch.py | 8 ++++++ python/freetoken/server/supervisor.py | 7 +++++ tests/server/test_launch.py | 37 +++++++++++++++++++++++++++ tests/server/test_supervisor.py | 31 ++++++++++++++++++++++ 4 files changed, 83 insertions(+) create mode 100644 tests/server/test_launch.py diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index bb0c256f46..5179e8e238 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -59,6 +59,14 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: if args.shell_mode: _detach_process_group() + import threading + + from tqdm import tqdm + + # Only rank 0 draws a bar, so no bar lock spans processes; tqdm's default one is a multiprocessing + # semaphore, which this worker leaves to resource_tracker as leaked when a signal ends it. + tqdm.set_lock(threading.RLock()) + # published (not bound) here: the engine binds it after the allocator setup from freetoken.gpu_select import set_assigned_gpu diff --git a/python/freetoken/server/supervisor.py b/python/freetoken/server/supervisor.py index 322419a65f..20deb7d132 100644 --- a/python/freetoken/server/supervisor.py +++ b/python/freetoken/server/supervisor.py @@ -55,6 +55,12 @@ class BackendHandle: processes: List[Any] = field(default_factory=list) expected_acks: int = 0 + def release_acks(self) -> None: + """Drop the ack queue once readiness has read it for the last time. A multiprocessing queue's + semaphores are unlinked only when it is collected, and a main process that a signal ends never + collects it, so resource_tracker reports them as leaked.""" + self.ack_queue = None + class WorkerDied(Exception): def __init__(self, message: str) -> None: @@ -172,6 +178,7 @@ def _shutting_down() -> bool: on_failure(exc.message) return + handle.release_acks() on_ready() while True: diff --git a/tests/server/test_launch.py b/tests/server/test_launch.py new file mode 100644 index 0000000000..defba4f307 --- /dev/null +++ b/tests/server/test_launch.py @@ -0,0 +1,37 @@ +"""The scheduler worker's own process setup, before it builds the engine.""" + +from __future__ import annotations + +import queue +import threading + +import pytest +import torch +from tqdm import tqdm + +import freetoken.gpu_select +import freetoken.scheduler +from freetoken.distributed import DistributedInfo +from freetoken.server import launch +from freetoken.server.args import ServerArgs +from freetoken.utils import progress + + +def test_the_scheduler_worker_draws_its_bars_under_a_lock_of_its_own(monkeypatch): + """Only rank 0 draws a bar, so no bar lock spans processes. tqdm's default one is a multiprocessing + semaphore, which a worker ended by a signal leaves to resource_tracker as leaked.""" + + class EngineBuilt(Exception): + pass + + def scheduler(args): + raise EngineBuilt(tqdm.get_lock()) + + monkeypatch.setattr(freetoken.scheduler, "Scheduler", scheduler) + monkeypatch.setattr(freetoken.gpu_select, "set_assigned_gpu", lambda target: None) + monkeypatch.setattr(progress, "_PROGRESS_SINK", None) + monkeypatch.setattr(tqdm, "_lock", tqdm.get_lock()) + args = ServerArgs(model_path="unused", tp_info=DistributedInfo(0, 1), dtype=torch.bfloat16) + with pytest.raises(EngineBuilt) as built: + launch._run_scheduler(args, queue.Queue()) + assert type(built.value.args[0]) is type(threading.RLock()) diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..1ee1ca17bf 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -211,6 +211,37 @@ def on_ready() -> None: assert "failure" not in seen # graceful stop: the death was not reported +def test_supervisor_releases_the_ack_queue_once_the_backend_is_ready(): + """Readiness is the ack queue's last read. A multiprocessing queue's semaphores are unlinked only when + it is collected, and a main process ended by a signal never collects it.""" + import multiprocessing + import weakref + + from freetoken.server.supervisor import run_backend_supervisor + + class Proc: + name = "freetoken-TP0-scheduler" + alive = True + + def is_alive(self) -> bool: + return self.alive + + proc = Proc() + ack_queue = multiprocessing.get_context("spawn").Queue() + ack_queue.put("scheduler ready") + queue_at_ready = weakref.ref(ack_queue) + handle = BackendHandle(ack_queue=ack_queue, processes=[proc], expected_acks=1) + del ack_queue + seen: dict = {} + + def on_ready() -> None: + seen["queue at ready"] = queue_at_ready() + proc.alive = False + + run_backend_supervisor(handle, LoadProgress(), on_ready=on_ready, poll=0.01, is_shutting_down=lambda: True) + assert seen == {"queue at ready": None} + + def test_supervisor_silent_on_startup_death_during_shutdown(): """A worker dying mid-load while an orderly stop is already in progress must not be reported as a load failure either.""" From a3206a944d9ea06a60b836c3eb017c176e27c5e7 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:45:44 -0400 Subject: [PATCH 3/4] feat(qwen4_exp): read llm-compressor (compressed-tensors) NVFP4 routed experts The expert spec now takes its tensor names and global-scale direction from the checkpoint's quant config, as qwen3_5_moe does, so RedHatAI/Qwen3.8-Flash-Next-NVFP4 loads and converts; the nvidia modelopt checkpoint keeps the same spec. Assisted-by: Claude Opus 5.5 --- python/freetoken/models/qwen4_exp/weight.py | 35 +++++++----- tests/models/qwen4_exp/test_weight.py | 59 +++++++++++++++++++++ tests/models/qwen4_exp/test_weight_ckpt.py | 5 +- 3 files changed, 83 insertions(+), 16 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 91c8c1bd17..26f47af508 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -28,7 +28,7 @@ from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) -from freetoken.layers.quantization import get_quant_config +from freetoken.layers.quantization import QuantKind, get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight, init_logger @@ -37,20 +37,16 @@ logger = init_logger(__name__) -# Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW -# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's -# stacked ``mtp.layers.N.mlp.experts.*`` tensors. -_EXPERT_KEY_RE = re.compile( +# Routed NVFP4 experts: per-expert, un-fused, under the checkpoint's dialect names (nvidia modelopt +# or llm-compressor). Matched against the RAW weight_map key in nvfp4_banks. The +# ``model.language_model.`` anchor excludes the MTP head's stacked ``mtp.layers.N.mlp.experts.*``. +_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" ) +# role -> the expert bank reader's canonical (modelopt) tensor kind +_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") -_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( - key_pattern=_EXPERT_KEY_RE, - proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, - layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc="Qwen3.8-Flash-Next NVFP4 experts", -) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") @@ -543,8 +539,19 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + """The per-expert NVFP4 layout under the checkpoint's dialect names (modelopt or llm-compressor).""" + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1a9a100edb..1167eb1b4b 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -518,3 +518,62 @@ def test_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, quan (tmp_path / "config.json").write_text(json.dumps(_config_json(quantization_config))) with pytest.raises(ValueError, match=match): _load(str(tmp_path)) + + +_CT_NVFP4_EXPERTS = { + "quant_method": "compressed-tensors", "format": "nvfp4-pack-quantized", "quantization_status": "compressed", + "config_groups": {"group_0": { + "format": "nvfp4-pack-quantized", "targets": [r"re:.*mlp\.experts\..*(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "float", "symmetric": True, "group_size": 16, + "strategy": "tensor_group", "dynamic": False, "scale_dtype": "torch.float8_e4m3fn"}, + "input_activations": {"num_bits": 4, "type": "float", "symmetric": True, "group_size": 16, + "strategy": "tensor_group", "dynamic": "local", "scale_dtype": "torch.float8_e4m3fn"}, + }}, + "ignore": ["lm_head"], +} + + +def test_compressed_tensors_experts_load_like_their_modelopt_twin(tmp_path, monkeypatch): + """An llm-compressor export (RedHatAI) names the routed experts weight_packed / weight_global_scale, + stores the quant-side global (the reciprocal of modelopt's weight_scale_2) and adds an unused + input_global_scale: its pieces must equal the modelopt twin's bit for bit.""" + from freetoken.layers.quantization import QuantConfig, set_quant_config + from freetoken.layers.quantization import configs as qconfigs + from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces + from freetoken.models.qwen4_exp.weight import nvfp4_expert_spec + + monkeypatch.setattr(qconfigs, "_QUANT_CONFIG", qconfigs._QUANT_CONFIG) # restored at teardown + layers, experts, hid, inter = 2, 3, 32, 16 + g = torch.Generator().manual_seed(0) + ct, mo = {}, {} + for layer in range(layers): + for e in range(experts): + for proj, (rows, cols) in {"gate_proj": (inter, hid), "up_proj": (inter, hid), "down_proj": (hid, inter)}.items(): + base = f"{LM}.layers.{layer}.mlp.experts.{e}.{proj}" + packed = torch.randint(0, 256, (rows, cols // 2), dtype=torch.uint8, generator=g) + scale = (torch.rand(rows, cols // 16, generator=g) * 4 + 0.1).to(torch.float8_e4m3fn) + glob = torch.rand(1, generator=g) * 1000 + 1 + ct |= {f"{base}.weight_packed": packed, f"{base}.weight_scale": scale, + f"{base}.weight_global_scale": glob, f"{base}.input_global_scale": torch.rand(1, generator=g)} + mo |= {f"{base}.weight": packed, f"{base}.weight_scale": scale, + f"{base}.weight_scale_2": 1.0 / glob, f"{base}.input_scale": torch.rand(1, generator=g)} + config = SimpleNamespace(num_experts=experts, num_moe_layers=layers, moe_intermediate_size=inter) + + def load(name, tensors, quant): + folder = tmp_path / name + folder.mkdir() + save_file(tensors, str(folder / "model.safetensors")) + set_quant_config(QuantConfig.from_hf({"quantization_config": quant})) + spec = nvfp4_expert_spec(str(folder), config) + pieces = iter_nvfp4_expert_pieces(str(folder), config, spec, drop_page_cache=lambda path: None, primary=False) + return spec, {(lay, e): piece for lay, e, _, piece in pieces} + + ct_spec, ct_pieces = load("ct", ct, _CT_NVFP4_EXPERTS) + mo_spec, mo_pieces = load("mo", mo, {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}) + assert ct_spec.global_reciprocal and not mo_spec.global_reciprocal + assert sorted(ct_pieces) == sorted(mo_pieces) == [(lay, e) for lay in range(layers) for e in range(experts)] + for key, piece in ct_pieces.items(): + assert sorted(piece) == sorted(mo_pieces[key]) + for role, tensor in piece.items(): + twin = mo_pieces[key][role] + assert tensor.dtype == twin.dtype and torch.equal(tensor.view(torch.uint8), twin.view(torch.uint8)), (key, role) diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py index 605c40f18b..38d32bd1af 100644 --- a/tests/models/qwen4_exp/test_weight_ckpt.py +++ b/tests/models/qwen4_exp/test_weight_ckpt.py @@ -22,10 +22,10 @@ from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.weight import ( - _NVFP4_SOURCE_SPEC, _ZERO_CENTERED_NORM_SUFFIXES, iter_weights, load_ple_table, + nvfp4_expert_spec, ) from freetoken.moe.expert_banks import build_expert_banks from freetoken.moe.host_banks import HostResidency @@ -271,8 +271,9 @@ def layer0_expert_banks(): """The real NVFP4 source-bank loader, restricted to layer 0 (1.32 GiB instead of 63.5).""" if not torch.cuda.is_available(): pytest.skip("expert bank pinning needs CUDA") + install_quant_config(MODEL_PATH) spec = dataclasses.replace( - _NVFP4_SOURCE_SPEC, layer_to_bank=lambda layer, config: 0 if layer == 0 else None + nvfp4_expert_spec(MODEL_PATH, None), layer_to_bank=lambda layer, config: 0 if layer == 0 else None ) config = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, num_moe_layers=1) From ecb52158346aaeab871b4190950a327d9f263e70 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 18:46:57 -0400 Subject: [PATCH 4/4] test(moe): let the prefetch counter test's copy spies pass #610's blocks_per_bank through Assisted-by: Claude Opus 5.5 --- tests/models/qwen4_exp/test_skeleton.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 2b47748f53..0803cbca28 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -1117,13 +1117,13 @@ def spy_ensure(layer, ids, **kw): calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone())) return real_ensure(layer, ids, **kw) - def spy_slim(*args): + def spy_slim(*args, **kw): streams["slim"].add(torch.cuda.current_stream().cuda_stream) - return real_slim(*args) + return real_slim(*args, **kw) - def spy_multi(*args): + def spy_multi(*args, **kw): streams["demand"].add(torch.cuda.current_stream().cuda_stream) - return real_multi(*args) + return real_multi(*args, **kw) def install(): monkeypatch.setattr(cache, "prefetch_ensure", spy_pf)