5080deploy: merge upstream #614 and #593, read compressed-tensors NVFP4 experts for qwen4_exp

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)

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<layer>\d+)\.mlp\.experts\.(?P<expert>\d+)\."
-    r"(?P<proj>gate_proj|up_proj|down_proj)\.(?P<kind>weight|weight_scale|weight_scale_2)$"
+    r"(?P<proj>gate_proj|up_proj|down_proj)\.(?P<kind>{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__ = [

python/freetoken/server/launch.py

@@ -70,6 +70,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
 

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:

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

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)

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)

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)

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

tests/server/test_supervisor.py

@@ -213,6 +213,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."""