From 2f9e09ca083050e6e652b065de4c1a27f750d923 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:56 -0400 Subject: [PATCH 01/17] fix(qwen4): reconcile PLE router and parser behavior --- python/freetoken/models/qwen4_exp/config.py | 6 +- python/freetoken/models/qwen4_exp/model.py | 42 +++++- python/freetoken/server/args.py | 2 + tests/models/test_qwen4_exp.py | 145 +++++++++++++++++++- tests/models/test_qwen4_exp_raw_config.py | 78 +++++++++++ tests/server/test_parser_auto_selection.py | 8 ++ 6 files changed, 270 insertions(+), 11 deletions(-) create mode 100644 tests/models/test_qwen4_exp_raw_config.py diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index b53811a98..f57f461d6 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -159,9 +159,9 @@ def parse_config(hf_config: Any) -> ModelConfig: num_experts_per_tok=int(text.num_experts_per_tok), moe_intermediate_size=int(text.moe_intermediate_size), shared_expert_intermediate_size=int(text.shared_expert_intermediate_size), - # The released Qwen3.8-Flash-Next configs omit this older Qwen MoE - # field. Omission means that the router weights are not renormalized. - norm_topk_prob=bool(getattr(text, "norm_topk_prob", False)), + # The official Qwen4-Exp config defaults this field to True. Released + # checkpoints may omit it, while an explicit False must remain False. + norm_topk_prob=bool(getattr(text, "norm_topk_prob", True)), model_type=str(hf_config.model_type), architectures=list(hf_config.architectures), moe_enabled=True, diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index d0ba393a6..9b884cc68 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -278,17 +278,28 @@ def build_ngram_ids( return torch.cat(blocks, dim=-1) -def _ple_request_tokens(req, forwarded_ids: torch.Tensor | None = None) -> torch.Tensor: - """Return the complete host token history visible to this forward. +def _ple_request_tokens( + req, + forwarded_ids: torch.Tensor | None = None, + *, + start: int = 0, +) -> torch.Tensor: + """Return host-visible request tokens from ``start`` through ``device_len``. The overlap scheduler advances ``device_len`` before it drains the prior sampled token to ``req.input_ids``. During decode, that one current token is already present in ``batch.input_ids``. Join it to the committed host prefix - so PLE hashes the same history as a non-overlapped forward. + so PLE hashes the same history as a non-overlapped forward. The default + preserves the complete-history contract; callers may select a validated + suffix when only the N-gram dependency prefix is needed. """ + if not 0 <= start <= req.device_len: + raise ValueError( + f"Qwen4-Exp PLE history start {start} is outside [0, {req.device_len}]" + ) host_len = req.input_ids.numel() if host_len >= req.device_len: - return req.input_ids[: req.device_len] + return req.input_ids[start : req.device_len] if host_len != req.cached_len: raise RuntimeError( "Qwen4-Exp PLE host history has an unexpected gap: " @@ -300,7 +311,14 @@ def _ple_request_tokens(req, forwarded_ids: torch.Tensor | None = None) -> torch "Qwen4-Exp PLE needs the current forwarded tokens: " f"got {actual}, expected {req.extend_len}" ) - return torch.cat((req.input_ids[: req.cached_len], forwarded_ids.to(device="cpu"))) + host_start = min(start, req.cached_len) + forwarded_start = max(0, start - req.cached_len) + return torch.cat( + ( + req.input_ids[host_start : req.cached_len], + forwarded_ids[forwarded_start:].to(device="cpu"), + ) + ) class _HostNGramEmbedding(BaseOP): @@ -400,7 +418,8 @@ def _current_ngram_ids(self) -> torch.Tensor: forwarded = forwarded_host[ forwarded_offset : forwarded_offset + extend_len ] - tokens = _ple_request_tokens(req, forwarded) + history_start = max(0, req.cached_len - (self.ngram_size - 1)) + tokens = _ple_request_tokens(req, forwarded, start=history_start) all_ids = build_ngram_ids( tokens, ngram_size=self.ngram_size, @@ -410,8 +429,17 @@ def _current_ngram_ids(self) -> torch.Tensor: vocab_sizes=vocab_sizes, offsets=offsets, ) - pieces.append(all_ids[req.cached_len : req.device_len]) + pieces.append( + all_ids[ + req.cached_len - history_start : req.device_len - history_start + ] + ) forwarded_offset += extend_len + if forwarded_offset != batch.input_ids.numel(): + raise RuntimeError( + f"Qwen4-Exp PLE consumed {forwarded_offset} forwarded tokens, " + f"but the batch carries {batch.input_ids.numel()}" + ) result = torch.cat(pieces, dim=0) if result.shape[0] != batch.input_ids.numel(): raise RuntimeError( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index adef3e393..22567e43c 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -174,6 +174,8 @@ def _infer_tool_call_parser(model_path: str) -> str: if ( "qwen3_5" in marker or "qwen3.5" in marker + or "qwen4_exp" in marker + or "qwen4exp" in marker or ("qwen3" in marker and "coder" in marker) ): return "qwen3_coder" diff --git a/tests/models/test_qwen4_exp.py b/tests/models/test_qwen4_exp.py index 7d9912e42..ce9029971 100644 --- a/tests/models/test_qwen4_exp.py +++ b/tests/models/test_qwen4_exp.py @@ -6,8 +6,13 @@ import torch import freetoken.models.qwen4_exp as qwen4_exp +import freetoken.models.qwen4_exp.model as qwen4_model from freetoken.models.qwen4_exp.config import parse_config -from freetoken.models.qwen4_exp.model import _ple_request_tokens, build_ngram_ids +from freetoken.models.qwen4_exp.model import ( + _HostNGramEmbedding, + _ple_request_tokens, + build_ngram_ids, +) from freetoken.models.qwen4_exp.weight import _rename, _try_fuse from freetoken.models.register import get_model_spec @@ -106,6 +111,13 @@ def test_qwen4_config_accepts_missing_norm_topk_prob(): hf_config = _config() del hf_config.text_config.norm_topk_prob config = parse_config(hf_config) + assert config.norm_topk_prob + + +def test_qwen4_config_preserves_explicit_false_norm_topk_prob(): + hf_config = _config() + hf_config.text_config.norm_topk_prob = False + config = parse_config(hf_config) assert not config.norm_topk_prob @@ -211,6 +223,19 @@ def test_ple_request_tokens_uses_complete_prefill_history(): assert _ple_request_tokens(req).tolist() == [11, 12, 13] +def test_ple_request_tokens_uses_bounded_non_overlap_decode_suffix(): + history = torch.tensor([10, 11, 12, 13, 14, 15]) + req = SimpleNamespace( + input_ids=history, + cached_len=5, + device_len=6, + extend_len=1, + ) + + assert _ple_request_tokens(req, start=3).tolist() == [13, 14, 15] + assert torch.equal(req.input_ids, history) + + def test_ple_request_tokens_joins_overlap_decode_token(): req = SimpleNamespace( input_ids=torch.tensor([11, 12]), @@ -221,6 +246,33 @@ def test_ple_request_tokens_joins_overlap_decode_token(): assert _ple_request_tokens(req, torch.tensor([13])).tolist() == [11, 12, 13] +def test_ple_request_tokens_bounds_overlap_multi_token_extension(): + history = torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + req = SimpleNamespace( + input_ids=history, + cached_len=8, + device_len=11, + extend_len=3, + ) + + tokens = _ple_request_tokens(req, torch.tensor([18, 19, 20]), start=6) + + assert tokens.tolist() == [16, 17, 18, 19, 20] + assert tokens.numel() == 2 + req.extend_len + assert torch.equal(req.input_ids, history) + + +def test_ple_request_tokens_validates_overlap_forwarded_token_count(): + req = SimpleNamespace( + input_ids=torch.tensor([11, 12]), + cached_len=2, + device_len=4, + extend_len=2, + ) + with pytest.raises(RuntimeError, match="needs the current forwarded tokens"): + _ple_request_tokens(req, torch.tensor([13])) + + def test_ple_request_tokens_rejects_noncontiguous_host_history(): req = SimpleNamespace( input_ids=torch.tensor([11]), @@ -230,3 +282,94 @@ def test_ple_request_tokens_rejects_noncontiguous_host_history(): ) with pytest.raises(RuntimeError, match="unexpected gap"): _ple_request_tokens(req, torch.tensor([13])) + + +@pytest.mark.parametrize("start", [-1, 4]) +def test_ple_request_tokens_rejects_invalid_suffix_start(start): + req = SimpleNamespace( + input_ids=torch.tensor([11, 12, 13]), + cached_len=2, + device_len=3, + extend_len=1, + ) + with pytest.raises(ValueError, match="history start"): + _ple_request_tokens(req, start=start) + + +def test_incremental_ngram_ids_match_full_history_across_eos_boundary(): + tokens = torch.tensor([10, 99, 4, 5, 6]) + kwargs = { + "ngram_size": 3, + "heads_per_ngram": 1, + "eos_token_id": 99, + "multipliers": torch.tensor([3, 5, 7]), + "vocab_sizes": torch.tensor([101, 103]), + "offsets": torch.tensor([0, 101]), + } + cached_len = 3 + history_start = cached_len - (kwargs["ngram_size"] - 1) + + full = build_ngram_ids(tokens, **kwargs) + incremental = build_ngram_ids(tokens[history_start:], **kwargs) + + assert torch.equal( + incremental[cached_len - history_start :], + full[cached_len:], + ) + + +def test_current_ngram_ids_keeps_forwarded_offsets_request_local(monkeypatch): + req_a_history = torch.tensor([1, 2]) + req_b_history = torch.tensor([10, 11, 12]) + req_a = SimpleNamespace( + input_ids=req_a_history, + cached_len=2, + device_len=4, + extend_len=2, + ) + req_b = SimpleNamespace( + input_ids=req_b_history, + cached_len=3, + device_len=4, + extend_len=1, + ) + batch = SimpleNamespace( + is_decode=True, + padded_reqs=[req_a, req_b], + input_ids=torch.tensor([3, 4, 13]), + ) + monkeypatch.setattr(qwen4_model, "get_global_ctx", lambda: SimpleNamespace(batch=batch)) + + multipliers = torch.tensor([3, 5, 7]) + vocab_sizes = torch.tensor([101, 103]) + offsets = torch.tensor([0, 101]) + embedding = SimpleNamespace( + _host_constants=(multipliers, vocab_sizes, offsets), + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + ) + + actual = _HostNGramEmbedding._current_ngram_ids(embedding) + expected_a = build_ngram_ids( + torch.tensor([1, 2, 3, 4]), + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + multipliers=multipliers, + vocab_sizes=vocab_sizes, + offsets=offsets, + )[2:] + expected_b = build_ngram_ids( + torch.tensor([10, 11, 12, 13]), + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + multipliers=multipliers, + vocab_sizes=vocab_sizes, + offsets=offsets, + )[3:] + + assert torch.equal(actual, torch.cat((expected_a, expected_b))) + assert torch.equal(req_a.input_ids, req_a_history) + assert torch.equal(req_b.input_ids, req_b_history) diff --git a/tests/models/test_qwen4_exp_raw_config.py b/tests/models/test_qwen4_exp_raw_config.py new file mode 100644 index 000000000..81745aa9d --- /dev/null +++ b/tests/models/test_qwen4_exp_raw_config.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.utils.hf import RawConfigShim + + +def _raw_checkpoint_config() -> RawConfigShim: + """Raw config shape used when installed Transformers predates Qwen4-Exp.""" + return RawConfigShim( + { + "architectures": ["Qwen4ExpForConditionalGeneration"], + "model_type": "qwen4_exp", + "image_token_id": 248056, + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + "text_config": { + "model_type": "qwen4_exp_text", + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "qwen_sparse_attention", + ], + "head_dim": 256, + "rope_parameters": { + "partial_rotary_factor": 0.25, + "rope_theta": 10_000_000, + "rope_type": "default", + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + }, + "indexer_budget": 2048, + "indexer_n_heads": 4, + "indexer_kv_heads": 1, + "indexer_head_dim": 128, + "indexer_compress_ratio": 4, + "max_position_embeddings": 262_144, + "num_key_value_heads": 2, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "eos_token_id": 248044, + "hc_count": 4, + "hc_lowrank": 320, + "ple_layer_ids": [2], + "ple_embed_dim": 2560, + "ple_conv_kernel_size": 4, + "ngram_size": 3, + "heads_per_ngram": 8, + "ngram_vocab_size_base": 20_000_000, + "split_ngram_parts": 128, + "output_gate_type": "sigmoid", + "hidden_act": "silu", + "num_hidden_layers": 4, + "num_attention_heads": 24, + "hidden_size": 2560, + "vocab_size": 248320, + "rms_norm_eps": 1e-6, + "num_experts": 512, + "num_experts_per_tok": 10, + "moe_intermediate_size": 640, + "shared_expert_intermediate_size": 640, + "tie_word_embeddings": False, + }, + } + ) + + +def test_qwen4_raw_config_uses_official_topk_normalization_default(): + config = parse_config(_raw_checkpoint_config()) + + assert config.norm_topk_prob is True + assert config.rotary_config.max_position == 262_144 + assert config.attn_type_for_layer(3).value == "qsa" diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e8..78ceaaffe 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -87,6 +87,14 @@ def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch(): assert _inferred("Qwen3MoeForCausalLM")[0] == "qwen25" +@pytest.mark.parametrize( + "architecture", + ["Qwen4ExpForConditionalGeneration", "qwen4_exp"], +) +def test_qwen4_exp_uses_qwen3_coder_tool_parser(architecture): + assert _inferred(architecture)[0] == "qwen3_coder" + + def test_an_explicit_choice_beats_inference(): config = _Config({"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"}) with patch("freetoken.utils.cached_load_hf_config", lambda _path: config): From 8994cdef5840ed4a4a72068dd93071ea24752d15 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:56 -0400 Subject: [PATCH 02/17] test(qwen4): add QSA differential oracle --- tests/kernels/test_qsa_differential.py | 314 +++++++++++++++ .../test_qsa_incremental_differential.py | 369 ++++++++++++++++++ 2 files changed, 683 insertions(+) create mode 100644 tests/kernels/test_qsa_differential.py create mode 100644 tests/kvcache/test_qsa_incremental_differential.py diff --git a/tests/kernels/test_qsa_differential.py b/tests/kernels/test_qsa_differential.py new file mode 100644 index 000000000..20d06caa8 --- /dev/null +++ b/tests/kernels/test_qsa_differential.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import json +import math +import time + +import pytest +import torch + +from freetoken.attention.qsa import select_qsa_logical_rows +from freetoken.kernel.triton.qsa import qsa_sparse_gqa + + +SEED = 38038 +COMPRESS_RATIO = 4 +TOKEN_BUDGET = 2048 +OUTPUT_WIDTH = TOKEN_BUDGET + COMPRESS_RATIO - 1 +QUERY_POSITIONS = ( + 0, + 1, + 2, + 3, + 4, + 2047, + 2048, + 2049, + 2050, + 2051, + 2052, + 4095, + 8191, + 65535, +) + + +def _official_contiguous_oracle( + index_q: torch.Tensor, + compressed_keys: torch.Tensor, + query_positions: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Independent eager transcription of the pinned official indexer semantics. + + This oracle intentionally does not call any FreeToken selection or compaction + helper. Its scope is the serving topology supported by FreeToken: one + contiguous causal history per request. + """ + rows = torch.full( + (query_positions.numel(), OUTPUT_WIDTH), -1, dtype=torch.int32 + ) + counts = torch.empty(query_positions.numel(), dtype=torch.int32) + offsets = torch.arange(COMPRESS_RATIO, dtype=torch.long) + for query_row, position_tensor in enumerate(query_positions.cpu()): + position = int(position_tensor) + visible = position + 1 + complete_groups = visible // COMPRESS_RATIO + width = min(TOKEN_BUDGET // COMPRESS_RATIO, complete_groups) + if width: + query = index_q[query_row].cpu().float() + keys = compressed_keys[:complete_groups, 0].cpu().float() + scores = torch.relu(query @ keys.transpose(0, 1)).sum(dim=0) + scores /= math.sqrt(query.shape[-1]) + groups = torch.topk(scores, width, sorted=True).indices + selected = (groups[:, None] * COMPRESS_RATIO + offsets).flatten() + else: + selected = torch.empty(0, dtype=torch.long) + tail = torch.arange( + complete_groups * COMPRESS_RATIO, visible, dtype=torch.long + ) + selected = torch.cat((selected, tail)) + rows[query_row, : selected.numel()] = selected.to(torch.int32) + counts[query_row] = selected.numel() + return rows, counts + + +def _score_separated_fixture( + positions: tuple[int, ...] = QUERY_POSITIONS, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + query_positions = torch.tensor(positions, dtype=torch.int64) + complete_groups = (max(positions) + 1) // COMPRESS_RATIO + dim = 8 + index_q = torch.zeros(len(positions), 4, dim, dtype=torch.float32) + index_q[:, :, 0] = torch.tensor([1.0, 1.5, 2.0, 2.5]) + compressed_keys = torch.zeros(complete_groups, 1, dim, dtype=torch.float32) + # Strictly increasing positive scores make exact top-k order authoritative. + compressed_keys[:, 0, 0] = torch.arange( + 1, complete_groups + 1, dtype=torch.float32 + ) + return index_q, compressed_keys, query_positions + + +def _live_rows(rows: torch.Tensor, counts: torch.Tensor, index: int) -> torch.Tensor: + return rows[index, : int(counts[index])].long() + + +def _assert_group_and_tail_invariants( + selected: torch.Tensor, position: int, expected_count: int +) -> None: + assert selected.numel() == expected_count + assert torch.all((selected >= 0) & (selected <= position)) + assert torch.unique(selected).numel() == selected.numel() + visible = position + 1 + complete_groups = visible // COMPRESS_RATIO + tail = torch.arange(complete_groups * COMPRESS_RATIO, visible) + if tail.numel(): + assert torch.equal(selected[-tail.numel() :].cpu(), tail) + selected = selected[: -tail.numel()] + assert selected.numel() % COMPRESS_RATIO == 0 + if selected.numel(): + groups = selected.view(-1, COMPRESS_RATIO) + assert torch.equal( + groups, + groups[:, :1] + torch.arange(COMPRESS_RATIO, device=groups.device), + ) + assert torch.all(groups[:, 0] % COMPRESS_RATIO == 0) + + +def _independent_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + rows: torch.Tensor, + counts: torch.Tensor, + scale: float, +) -> torch.Tensor: + output = torch.zeros_like(q, dtype=torch.float32) + gqa = q.shape[1] // k.shape[1] + for query_row in range(q.shape[0]): + selected = rows[query_row, : int(counts[query_row])].long().cpu() + for kv_head in range(k.shape[1]): + heads = slice(kv_head * gqa, (kv_head + 1) * gqa) + scores = torch.einsum( + "hd,td->ht", + q[query_row, heads].cpu().float(), + k[selected, kv_head].cpu().float(), + ) * scale + probabilities = torch.softmax(scores, dim=-1) + output[query_row, heads] = torch.einsum( + "ht,td->hd", probabilities, v[selected, kv_head].cpu().float() + ) + return output + + +def test_qsa_unique_score_selection_matches_pinned_official_oracle_cpu(): + torch.manual_seed(SEED) + index_q, compressed_keys, positions = _score_separated_fixture() + expected_rows, expected_counts = _official_contiguous_oracle( + index_q, compressed_keys, positions + ) + actual_rows, actual_counts = select_qsa_logical_rows( + index_q, + compressed_keys, + positions, + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + assert torch.equal(actual_counts.cpu(), expected_counts) + assert torch.equal(actual_rows.cpu(), expected_rows) + + for row, position in enumerate(QUERY_POSITIONS): + selected = _live_rows(actual_rows, actual_counts, row) + assert selected.numel() <= OUTPUT_WIDTH + assert torch.all(selected <= position) + if position <= 2050: + assert set(selected.tolist()) == set(range(position + 1)) + + boundary = QUERY_POSITIONS.index(2051) + boundary_rows = _live_rows(actual_rows, actual_counts, boundary) + assert actual_counts[boundary].item() == TOKEN_BUDGET + assert set(boundary_rows.tolist()) == set(range(4, 2052)) + assert not any(token in boundary_rows.tolist() for token in range(4)) + + after = QUERY_POSITIONS.index(2052) + assert actual_counts[after].item() == TOKEN_BUDGET + 1 + assert _live_rows(actual_rows, actual_counts, after)[-1].item() == 2052 + + +@pytest.mark.parametrize("key_value", [0.0, -1.0]) +def test_qsa_tied_scores_obey_order_insensitive_official_invariants_cpu(key_value): + positions = torch.tensor([2051, 2052, 4095], dtype=torch.int64) + index_q = torch.ones(positions.numel(), 4, 8) + compressed_keys = torch.full((1024, 1, 8), key_value) + rows, counts = select_qsa_logical_rows( + index_q, + compressed_keys, + positions, + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + for row, position in enumerate(positions.tolist()): + complete = (position + 1) // COMPRESS_RATIO + tail = (position + 1) % COMPRESS_RATIO + expected_count = min(complete, TOKEN_BUDGET // COMPRESS_RATIO) * 4 + tail + _assert_group_and_tail_invariants( + _live_rows(rows, counts, row), position, expected_count + ) + + +def test_qsa_float32_sparse_attention_matches_independent_eager_oracle_cpu(): + torch.manual_seed(SEED) + positions_tuple = (2051, 2052, 4095) + index_q, compressed_keys, positions = _score_separated_fixture(positions_tuple) + selected, counts = select_qsa_logical_rows( + index_q, + compressed_keys, + positions, + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + q = torch.randn(len(positions_tuple), 4, 16, dtype=torch.float32) + k = torch.randn(max(positions_tuple) + 1, 2, 16, dtype=torch.float32) + v = torch.randn_like(k) + scale = 16**-0.5 + actual = qsa_sparse_gqa(q, k, v, selected, counts, scale) + expected = _independent_sparse_attention(q, k, v, selected, counts, scale) + error = (actual.float() - expected).abs() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "device": "cpu", + "dtype": "float32", + "max_abs_error": error.max().item(), + "mean_abs_error": error.mean().item(), + }, + sort_keys=True, + ), + ) + torch.testing.assert_close(actual.float(), expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_cuda_selection_and_compaction_match_independent_oracle(): + torch.manual_seed(SEED) + index_q, compressed_keys, positions = _score_separated_fixture() + expected_rows, expected_counts = _official_contiguous_oracle( + index_q, compressed_keys, positions + ) + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + actual_rows, actual_counts = select_qsa_logical_rows( + index_q.cuda(), + compressed_keys.cuda(), + positions.cuda(), + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + assert torch.equal(actual_counts.cpu(), expected_counts) + assert torch.equal(actual_rows.cpu(), expected_rows) + peak = torch.cuda.max_memory_allocated() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "case": "cuda_selection_compaction", + "device": torch.cuda.get_device_name(), + "driver_runtime": torch.version.cuda, + "elapsed_seconds": elapsed, + "peak_memory_bytes": peak, + }, + sort_keys=True, + ), + ) + assert peak < 1 << 30 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_cuda_bf16_sparse_gqa_matches_independent_eager_oracle(): + torch.manual_seed(SEED) + positions_tuple = (2051, 65535) + index_q, compressed_keys, positions = _score_separated_fixture(positions_tuple) + selected, counts = _official_contiguous_oracle(index_q, compressed_keys, positions) + dim = 64 + q_cpu = torch.randn(len(positions_tuple), 8, dim, dtype=torch.bfloat16) + k_cpu = torch.randn(max(positions_tuple) + 1, 2, dim, dtype=torch.bfloat16) + v_cpu = torch.randn_like(k_cpu) + scale = dim**-0.5 + expected = _independent_sparse_attention( + q_cpu, k_cpu, v_cpu, selected, counts, scale + ) + + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + actual = qsa_sparse_gqa( + q_cpu.cuda(), + k_cpu.cuda(), + v_cpu.cuda(), + selected.cuda(), + counts.cuda(), + scale, + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + actual_cpu = actual.cpu().float() + error = (actual_cpu - expected).abs() + peak = torch.cuda.max_memory_allocated() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "case": "cuda_bf16_sparse_gqa", + "device": torch.cuda.get_device_name(), + "driver_runtime": torch.version.cuda, + "elapsed_seconds": elapsed, + "peak_memory_bytes": peak, + "max_abs_error": error.max().item(), + "mean_abs_error": error.mean().item(), + }, + sort_keys=True, + ), + ) + assert peak < 1 << 30 + torch.testing.assert_close(actual_cpu, expected, rtol=2e-2, atol=2e-2) diff --git a/tests/kvcache/test_qsa_incremental_differential.py b/tests/kvcache/test_qsa_incremental_differential.py new file mode 100644 index 000000000..85f5aca1c --- /dev/null +++ b/tests/kvcache/test_qsa_incremental_differential.py @@ -0,0 +1,369 @@ +from __future__ import annotations + +import json +import math +import time +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.attention.qsa as qsa_module +from freetoken.attention.qsa import QSAAttnBackend, select_qsa_logical_rows +from freetoken.distributed.info import DistributedInfo +from freetoken.kvcache.qsa_pool import QSAKVCache + + +SEED = 38038 +RATIO = 4 +BUDGET = 2048 +LAYER_ID = 1 +INDEX_DIM = 8 + + +class _RecordingSyntheticIndexer: + """Observable stand-in for official K RMSNorm + block-start RoPE. + + The backend owns pooling and position routing, while the model indexer owns + normalization/RoPE. This deterministic transform makes both inputs visible + without importing model weights or calling a production selection helper. + """ + + def __init__(self): + self.calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def normalize_compressed_keys( + self, keys: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + self.calls.append((keys.detach().cpu().clone(), positions.detach().cpu().clone())) + values = keys.float() + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + 1e-6) + # A small block-start marker makes an incorrect current/end position observable. + marker = positions[0].float().view(-1, 1, 1) / 4096.0 + return (normalized + marker).to(keys.dtype) + + +def _oracle_normalize_and_position( + pooled: torch.Tensor, block_start_rope: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + values = pooled.float() + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + 1e-6) + marker = block_start_rope[0].float().view(-1, 1, 1) / 4096.0 + return (normalized + marker).to(dtype) + + +def _pool(monkeypatch, device: torch.device, num_pages: int = 70) -> QSAKVCache: + monkeypatch.setattr( + "freetoken.kvcache.mha_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + return QSAKVCache( + num_kv_heads=2, + num_layers=2, + head_dim=16, + num_pages=num_pages, + page_size=64, + dtype=torch.bfloat16, + device=device, + index_num_kv_heads=1, + index_head_dim=INDEX_DIM, + compress_ratio=RATIO, + layer_ids=(LAYER_ID,), + ) + + +def _backend_and_context(monkeypatch, pool: QSAKVCache, max_tokens: int = 2112): + stride = max_tokens + page_table = torch.stack( + ( + torch.arange(max_tokens, device=pool.device), + torch.arange(stride, stride + max_tokens, device=pool.device), + ) + ).long() + assert int(page_table.max()) < pool.k_cache(LAYER_ID).numel() // (2 * 16) + context = SimpleNamespace(kv_cache=pool, page_table=page_table) + monkeypatch.setattr(qsa_module, "get_global_ctx", lambda: context) + config = SimpleNamespace( + qwen4_args=SimpleNamespace( + indexer_compress_ratio=RATIO, + indexer_budget=BUDGET, + ) + ) + return QSAAttnBackend(config), context + + +def _raw_keys(length: int, request_id: int, device: torch.device) -> torch.Tensor: + base = torch.arange(length * INDEX_DIM, dtype=torch.float32).view(length, 1, INDEX_DIM) + values = base / 97.0 + 1.0 + request_id * 100.0 + return values.to(device=device, dtype=torch.bfloat16) + + +def _rope_positions(length: int, request_id: int, device: torch.device) -> torch.Tensor: + logical = torch.arange(length, device=device, dtype=torch.int64) + return torch.stack( + (logical, logical + 1000 * request_id, logical + 2000 * request_id) + ) + + +def _request(start: int, end: int, table_idx: int): + return SimpleNamespace( + cached_len=start, + device_len=end, + extend_len=end - start, + table_idx=table_idx, + ) + + +def _batch(reqs, rope_chunks): + lengths = [req.extend_len for req in reqs] + return SimpleNamespace( + reqs=reqs, + padded_reqs=reqs, + rope_positions=torch.cat(rope_chunks, dim=1), + input_ids=torch.empty(sum(lengths), dtype=torch.long), + ) + + +def _oracle_completed( + full_keys: torch.Tensor, full_rope: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + groups = full_keys.shape[0] // RATIO + if not groups: + return ( + full_keys.new_empty((0, 1, INDEX_DIM)), + full_rope.new_empty((3, 0)), + ) + members = full_keys[: groups * RATIO].view(groups, RATIO, 1, INDEX_DIM) + pooled = members.float().mean(dim=1).to(full_keys.dtype) + starts = torch.arange(0, groups * RATIO, RATIO, device=full_keys.device) + rope = full_rope.index_select(1, starts) + return _oracle_normalize_and_position(pooled, rope, full_keys.dtype), rope + + +def _oracle_selection( + query: torch.Tensor, keys: torch.Tensor, position: int +) -> tuple[torch.Tensor, int]: + complete = (position + 1) // RATIO + width = min(BUDGET // RATIO, complete) + if width: + score = torch.relu( + query.float() @ keys[:complete, 0].float().transpose(0, 1) + ).sum(dim=0) / math.sqrt(query.shape[-1]) + groups = torch.topk(score, width, sorted=True).indices + rows = ( + groups[:, None] * RATIO + torch.arange(RATIO, device=groups.device) + ).flatten() + else: + rows = torch.empty(0, dtype=torch.long, device=query.device) + tail = torch.arange(complete * RATIO, position + 1, device=query.device) + rows = torch.cat((rows, tail)) + return rows, rows.numel() + + +def _assert_incremental_state( + backend: QSAAttnBackend, + context, + full_keys: torch.Tensor, + full_rope: torch.Tensor, + end: int, + table_idx: int, +) -> None: + expected, _ = _oracle_completed(full_keys[:end], full_rope[:, :end]) + starts = torch.arange(0, expected.shape[0] * RATIO, RATIO, device=full_keys.device) + physical = context.page_table[table_idx].index_select(0, starts) + compressed_rows = torch.div(physical, RATIO, rounding_mode="floor") + actual = backend.kvcache.compressed_k_cache(LAYER_ID).index_select( + 0, compressed_rows.long() + ) + assert torch.equal(actual, expected) + + latest_start = max(0, end - RATIO) + latest = torch.arange(latest_start, end, device=full_keys.device) + assert torch.equal( + backend.kvcache.pending_group(LAYER_ID, table_idx, latest), + full_keys.index_select(0, latest), + ) + assert torch.equal( + backend.kvcache.pending_rope_group(LAYER_ID, table_idx, latest), + full_rope.index_select(1, latest).transpose(0, 1), + ) + + query = torch.ones(1, 4, INDEX_DIM, device=full_keys.device) + selected, counts = select_qsa_logical_rows( + query, + actual, + torch.tensor([end - 1], device=full_keys.device), + compress_ratio=RATIO, + token_budget=BUDGET, + ) + expected_rows, expected_count = _oracle_selection(query[0], actual, end - 1) + assert int(counts[0]) == expected_count + assert torch.equal(selected[0, :expected_count].long(), expected_rows.long()) + + +def _run_pattern(monkeypatch, pattern: list[int], device: torch.device): + total = sum(pattern) + pool = _pool(monkeypatch, device) + backend, context = _backend_and_context(monkeypatch, pool) + indexer = _RecordingSyntheticIndexer() + full_keys = _raw_keys(total, 0, device) + full_rope = _rope_positions(total, 0, device) + start = 0 + for length in pattern: + end = start + length + req = _request(start, end, 0) + batch = _batch([req], [full_rope[:, start:end]]) + backend.prepare_metadata(batch) + backend._compress_current_keys( + indexer, full_keys[start:end], LAYER_ID, batch + ) + _assert_incremental_state( + backend, context, full_keys, full_rope, end, table_idx=0 + ) + start = end + return backend, context, indexer, full_keys, full_rope + + +@pytest.mark.parametrize( + "pattern", + ([1, 3], [3, 1], [5, 2, 1], [1, 1, 1, 1, 1, 1, 1, 1]), +) +def test_qsa_incremental_compression_matches_full_history_cpu(monkeypatch, pattern): + torch.manual_seed(SEED) + backend, context, indexer, full_keys, full_rope = _run_pattern( + monkeypatch, list(pattern), torch.device("cpu") + ) + expected, expected_rope = _oracle_completed(full_keys, full_rope) + assert indexer.calls + # Across all calls, every emitted completed group uses its block-start RoPE. + recorded_rope = torch.cat([positions for _, positions in indexer.calls], dim=1) + assert torch.equal(recorded_rope, expected_rope.cpu()) + starts = torch.arange(0, expected.shape[0] * RATIO, RATIO) + rows = torch.div(context.page_table[0].cpu().index_select(0, starts), RATIO, rounding_mode="floor") + assert torch.equal( + backend.kvcache.compressed_k_cache(LAYER_ID).cpu().index_select(0, rows), + expected.cpu(), + ) + + +def test_qsa_incremental_crosses_first_sparse_boundary_cpu(monkeypatch): + torch.manual_seed(SEED) + pattern = [2047, 1, 1, 1, 1, 1] + backend, context, _, full_keys, full_rope = _run_pattern( + monkeypatch, pattern, torch.device("cpu") + ) + assert full_keys.shape[0] == 2052 + expected, _ = _oracle_completed(full_keys, full_rope) + assert expected.shape[0] == 513 + query = torch.ones(1, 4, INDEX_DIM) + selected, counts = select_qsa_logical_rows( + query, + expected, + torch.tensor([2051]), + compress_ratio=RATIO, + token_budget=BUDGET, + ) + oracle_rows, oracle_count = _oracle_selection(query[0], expected, 2051) + assert oracle_count == BUDGET + assert int(counts[0]) == BUDGET + assert torch.equal(selected[0, :BUDGET].long(), oracle_rows.long()) + # The least-scoring complete group is the only omitted group. + omitted = set(range(2052)) - set(oracle_rows.tolist()) + assert omitted == set(range(4)) + assert backend.kvcache is context.kv_cache + + +def test_qsa_multi_request_cache_and_pending_state_are_isolated_cpu(monkeypatch): + torch.manual_seed(SEED) + pool = _pool(monkeypatch, torch.device("cpu")) + backend, context = _backend_and_context(monkeypatch, pool) + indexer = _RecordingSyntheticIndexer() + lengths = (5, 7) + keys = [_raw_keys(length, req_id, torch.device("cpu")) for req_id, length in enumerate(lengths)] + rope = [_rope_positions(length, req_id, torch.device("cpu")) for req_id, length in enumerate(lengths)] + reqs = [_request(0, length, req_id) for req_id, length in enumerate(lengths)] + batch = _batch(reqs, rope) + backend.prepare_metadata(batch) + backend._compress_current_keys(indexer, torch.cat(keys), LAYER_ID, batch) + + for req_id, length in enumerate(lengths): + _assert_incremental_state( + backend, context, keys[req_id], rope[req_id], length, req_id + ) + first_expected, _ = _oracle_completed(keys[0], rope[0]) + second_expected, _ = _oracle_completed(keys[1], rope[1]) + assert not torch.equal(first_expected[0], second_expected[0]) + + index_q = torch.ones(sum(lengths), 4, INDEX_DIM) + physical, counts = backend._select_physical_rows(index_q, LAYER_ID, batch) + offset = 0 + for req_id, length in enumerate(lengths): + for local_position in range(length): + row = offset + local_position + assert int(counts[row]) == local_position + 1 + actual = set(physical[row, : int(counts[row])].tolist()) + expected = set( + context.page_table[req_id, : local_position + 1].tolist() + ) + assert actual == expected + offset += length + + +def test_qsa_request_slot_reset_invalidates_pending_state_cpu(monkeypatch): + torch.manual_seed(SEED) + backend, context, _, old_keys, old_rope = _run_pattern( + monkeypatch, [5, 2, 1], torch.device("cpu") + ) + new_keys = _raw_keys(3, 1, torch.device("cpu")) + new_rope = _rope_positions(3, 1, torch.device("cpu")) + req = _request(0, 3, 0) + batch = _batch([req], [new_rope]) + backend.prepare_metadata(batch) + backend._compress_current_keys( + _RecordingSyntheticIndexer(), new_keys, LAYER_ID, batch + ) + assert torch.equal( + backend.kvcache.pending_group(LAYER_ID, 0, torch.arange(3)), new_keys + ) + with pytest.raises(RuntimeError, match="pending-key state is missing"): + backend.kvcache.pending_group(LAYER_ID, 0, torch.tensor([4])) + selected, counts = select_qsa_logical_rows( + torch.ones(1, 4, INDEX_DIM), + backend.kvcache.compressed_k_cache(LAYER_ID)[:0], + torch.tensor([2]), + compress_ratio=RATIO, + token_budget=BUDGET, + ) + assert int(counts[0]) == 3 + assert torch.equal(selected[0, :3], torch.arange(3, dtype=torch.int32)) + assert old_keys.shape[0] == old_rope.shape[1] == 8 + assert backend.kvcache is context.kv_cache + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_incremental_cache_public_surface_matches_cpu_oracle_cuda(monkeypatch): + torch.manual_seed(SEED) + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + _run_pattern( + monkeypatch, + [1, 1, 1, 1, 1, 1, 1, 1], + torch.device("cuda"), + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + peak = torch.cuda.max_memory_allocated() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "case": "cuda_incremental_cache", + "device": torch.cuda.get_device_name(), + "driver_runtime": torch.version.cuda, + "elapsed_seconds": elapsed, + "peak_memory_bytes": peak, + }, + sort_keys=True, + ), + ) + assert peak < 1 << 30 From 068b279d980e90dd5f13a352533b8181ed22cf88 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:56 -0400 Subject: [PATCH 03/17] feat(qwen4): add native q3 ple reader --- python/freetoken/checkpoint/__init__.py | 2 + python/freetoken/checkpoint/q3_ple.py | 276 +++++++++++++++++++++ python/freetoken/models/qwen4_exp/model.py | 59 ++++- tests/checkpoint/test_q3_ple.py | 150 +++++++++++ 4 files changed, 484 insertions(+), 3 deletions(-) create mode 100644 python/freetoken/checkpoint/q3_ple.py create mode 100644 tests/checkpoint/test_q3_ple.py diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index f3ce70fcd..d11cb7fcd 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -12,8 +12,10 @@ load_ftw_banks, ) from .convert import convert_checkpoint +from .q3_ple import Q3PLEReader, Q3PLESegment __all__ = [ "FTWReader", "FTWWriter", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", + "Q3PLEReader", "Q3PLESegment", ] diff --git a/python/freetoken/checkpoint/q3_ple.py b/python/freetoken/checkpoint/q3_ple.py new file mode 100644 index 000000000..0e968fb4a --- /dev/null +++ b/python/freetoken/checkpoint/q3_ple.py @@ -0,0 +1,276 @@ +"""Native reader for the Qwen4 ``Q3_PLE_32`` lookup-table sidecar. + +The reader deliberately knows only the PLE GET_ROWS format. It does not expose a +matmul/dequant path and it never maps or allocates the complete table. A small JSON +directory describes the logical rows and the byte ranges containing each segment; +the payload remains an ordinary read-only file on the project's required ``Z:`` +volume. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import struct +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import torch + + +BLOCK_VALUES = 32 +BLOCK_BYTES = 14 +ROW_VALUES = 160 +BLOCKS_PER_ROW = 5 +ROW_BYTES = BLOCKS_PER_ROW * BLOCK_BYTES +FORMAT = "q3_ple_32" +VERSION = 1 +ALIGN = 4096 + + +def _z_path(path: str | os.PathLike[str]) -> Path: + """Resolve *path* and fail closed unless its physical drive is ``Z:``.""" + + resolved = Path(path).expanduser().resolve(strict=True) + drive, _ = os.path.splitdrive(str(resolved)) + # On Windows splitdrive is authoritative. The second clause keeps the + # check useful in a POSIX test harness mounted as ``/z`` while still never + # accepting an unqualified relative path. + if drive.upper() != "Z:" and not str(resolved).lower().startswith("/z/"): + raise ValueError(f"Q3_PLE_32 backing must resolve to Z:, got {resolved}") + return resolved + + +def _unpack_codes(payload: bytes) -> list[int]: + if len(payload) != 12: + raise ValueError(f"Q3_PLE_32 code payload must be 12 bytes, got {len(payload)}") + bits = int.from_bytes(payload, "little") + return [(bits >> (3 * i)) & 0x7 for i in range(BLOCK_VALUES)] + + +def _decode_row(row: bytes) -> torch.Tensor: + if len(row) != ROW_BYTES: + raise ValueError(f"Q3_PLE_32 row must be {ROW_BYTES} bytes, got {len(row)}") + values: list[float] = [] + for block_start in range(0, ROW_BYTES, BLOCK_BYTES): + block = row[block_start : block_start + BLOCK_BYTES] + # The format authority specifies a little-endian BF16 scalar. Decode + # through an integer bit pattern so host endianness cannot leak in. + scale_bits = int.from_bytes(block[:2], "little") + scale = struct.unpack(" int: + return self.end_row - self.first_row + + +class Q3PLEReader: + """Bounded random-row reader for a validated Q3_PLE_32 sidecar. + + ``manifest_path`` points to ``ple-q3.json`` and ``data_path`` may override + its ``data_file``. Opening validates the JSON schema, segment coverage, + file length, and whole-file/segment hashes in bounded chunks. ``gather`` + then reads only the requested 70-byte rows and returns BF16 values. + """ + + def __init__(self, manifest_path: str | os.PathLike[str], *, data_path: str | os.PathLike[str] | None = None): + self.manifest_path = _z_path(manifest_path) + with self.manifest_path.open("r", encoding="utf-8") as handle: + manifest = json.load(handle) + self.manifest = manifest + if manifest.get("format") != FORMAT or int(manifest.get("version", -1)) != VERSION: + raise ValueError("unsupported Q3_PLE_32 format/version") + if manifest.get("endianness", "little") != "little": + raise ValueError("Q3_PLE_32 requires little-endian metadata") + if int(manifest.get("block_values", BLOCK_VALUES)) != BLOCK_VALUES: + raise ValueError("Q3_PLE_32 block_values mismatch") + if int(manifest.get("block_bytes", BLOCK_BYTES)) != BLOCK_BYTES: + raise ValueError("Q3_PLE_32 block_bytes mismatch") + if int(manifest.get("row_values", ROW_VALUES)) != ROW_VALUES: + raise ValueError("Q3_PLE_32 row_values mismatch") + if int(manifest.get("row_bytes", ROW_BYTES)) != ROW_BYTES: + raise ValueError("Q3_PLE_32 row_bytes mismatch") + + candidate = data_path + if candidate is None: + candidate = self.manifest.get("data_file") + if not candidate: + raise ValueError("Q3_PLE_32 manifest has no data_file") + data_candidate = Path(candidate) + if not data_candidate.is_absolute(): + data_candidate = self.manifest_path.parent / data_candidate + self.data_path = _z_path(data_candidate) + self.row_count = int(manifest.get("rows", 0)) + self.total_payload_bytes = self.row_count * ROW_BYTES + if self.row_count <= 0: + raise ValueError("Q3_PLE_32 rows must be positive") + declared_payload = int(manifest.get("payload_bytes", self.total_payload_bytes)) + if declared_payload != self.total_payload_bytes: + raise ValueError("Q3_PLE_32 payload_bytes does not equal rows * row_bytes") + + raw_segments = manifest.get("segments") + if not isinstance(raw_segments, list) or not raw_segments: + raise ValueError("Q3_PLE_32 segment directory is empty") + self.segments: tuple[Q3PLESegment, ...] = tuple(self._parse_segment(item) for item in raw_segments) + self._validate_segments() + stat = self.data_path.stat() + expected_file_bytes = int(manifest.get("file_bytes", stat.st_size)) + if stat.st_size != expected_file_bytes: + raise ValueError(f"Q3_PLE_32 file length mismatch: {stat.st_size} != {expected_file_bytes}") + self._handle = self.data_path.open("rb") + self._io_lock = threading.Lock() + try: + self._verify_hashes() + except Exception: + self._handle.close() + raise + + self.weight_scale = float(manifest.get("weight_scale", 1.0)) + if not math.isfinite(self.weight_scale): + raise ValueError("Q3_PLE_32 weight_scale must be finite") + + def _parse_segment(self, item: object) -> Q3PLESegment: + if not isinstance(item, dict): + raise ValueError("Q3_PLE_32 segment must be an object") + try: + data_offset = item.get("data_offset") + if data_offset is None: + data_offset = item["offset"] + byte_length = item.get("byte_length") + if byte_length is None: + byte_length = item["length"] + segment = Q3PLESegment( + first_row=int(item["first_row"]), + end_row=int(item["end_row"]), + data_offset=int(data_offset), + byte_length=int(byte_length), + sha256=str(item["sha256"]).lower(), + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("malformed Q3_PLE_32 segment directory") from exc + if len(segment.sha256) != 64 or any(c not in "0123456789abcdef" for c in segment.sha256): + raise ValueError("malformed Q3_PLE_32 segment hash") + return segment + + def _validate_segments(self) -> None: + expected_row = 0 + previous_end = 0 + for segment in self.segments: + if segment.first_row != expected_row or segment.end_row <= segment.first_row: + raise ValueError("Q3_PLE_32 segment rows have a gap, overlap, or bad order") + if segment.data_offset < 0 or segment.data_offset % ALIGN: + raise ValueError("Q3_PLE_32 segment data offset is not 4 KiB aligned") + if segment.byte_length != segment.rows * ROW_BYTES: + raise ValueError("Q3_PLE_32 segment byte length does not match rows") + if segment.data_offset < previous_end: + raise ValueError("Q3_PLE_32 segment byte ranges overlap") + expected_row = segment.end_row + previous_end = segment.data_offset + segment.byte_length + if expected_row != self.row_count: + raise ValueError("Q3_PLE_32 segment rows do not cover the table") + if previous_end > int(self.manifest.get("file_bytes", previous_end)): + raise ValueError("Q3_PLE_32 segment exceeds file length") + + def _read_exact(self, offset: int, length: int) -> bytes: + with self._io_lock: + self._handle.seek(offset) + payload = self._handle.read(length) + if len(payload) != length: + raise OSError(f"short Q3_PLE_32 read at {offset}: {len(payload)}/{length}") + return payload + + def _verify_hashes(self) -> None: + # Hashes are checked incrementally; this never allocates the table. + whole = hashlib.sha256() + with self.data_path.open("rb") as handle: + while True: + chunk = handle.read(8 << 20) + if not chunk: + break + whole.update(chunk) + expected_whole = str(self.manifest.get("sha256", "")).lower() + if len(expected_whole) != 64 or whole.hexdigest() != expected_whole: + raise ValueError("Q3_PLE_32 whole-file hash mismatch") + for segment in self.segments: + digest_ctx = hashlib.sha256() + remaining = segment.byte_length + offset = segment.data_offset + while remaining: + take = min(8 << 20, remaining) + digest_ctx.update(self._read_exact(offset, take)) + offset += take + remaining -= take + digest = digest_ctx.hexdigest() + if digest != segment.sha256: + raise ValueError(f"Q3_PLE_32 segment hash mismatch at row {segment.first_row}") + + def _segment_for_row(self, row: int) -> Q3PLESegment: + if not 0 <= row < self.row_count: + raise IndexError(f"Q3_PLE_32 row {row} outside 0..{self.row_count - 1}") + # Segment count is small (128 in the production sidecar); linear search + # keeps the directory representation transparent and deterministic. + for segment in self.segments: + if segment.first_row <= row < segment.end_row: + return segment + raise AssertionError("validated Q3_PLE_32 directory did not locate row") + + def read_row(self, row: int) -> torch.Tensor: + segment = self._segment_for_row(int(row)) + offset = segment.data_offset + (int(row) - segment.first_row) * ROW_BYTES + return _decode_row(self._read_exact(offset, ROW_BYTES)) + + def gather(self, row_indices: Sequence[int], *, apply_weight_scale: bool = False) -> torch.Tensor: + """Gather rows in the caller's order without deduplication or reordering.""" + + rows = [self.read_row(int(index)) for index in row_indices] + if not rows: + output = torch.empty((0, ROW_VALUES), dtype=torch.bfloat16) + else: + output = torch.stack(rows, dim=0) + if apply_weight_scale: + output = output * self.weight_scale + return output + + def gather16(self, row_indices: Sequence[int], *, apply_weight_scale: bool = False) -> torch.Tensor: + if len(row_indices) != 16: + raise ValueError(f"Qwen4 PLE requires exactly 16 logical rows, got {len(row_indices)}") + return self.gather(row_indices, apply_weight_scale=apply_weight_scale) + + def close(self) -> None: + handle, self._handle = getattr(self, "_handle", None), None + if handle is not None: + handle.close() + + def __enter__(self) -> "Q3PLEReader": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self.close() + + +__all__ = [ + "ALIGN", + "BLOCK_BYTES", + "BLOCK_VALUES", + "FORMAT", + "Q3PLEReader", + "Q3PLESegment", + "ROW_BYTES", + "ROW_VALUES", + "VERSION", +] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 9b884cc68..9f5ec6dc2 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -25,13 +25,13 @@ get_rope, ) from freetoken.models.blocks import BaseLLMModel +from freetoken.models.config import ModelConfig from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet +from freetoken.checkpoint.q3_ple import Q3PLEReader from freetoken.utils import download_hf_weight, nvtx_annotate if TYPE_CHECKING: - from freetoken.models.config import ModelConfig - from .args import Qwen4ExpArgs @@ -341,8 +341,20 @@ def __init__(self, config: ModelConfig, layer_id: int): self._scale = torch.tensor(1.0, dtype=torch.bfloat16) self._host_constants: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None self._dummy = False + self._q3_reader: Q3PLEReader | None = None - def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + def load_host_weights( + self, + model_path: str, + *, + dummy: bool = False, + ple_format: str = "fp8_safetensors", + ) -> None: + if ple_format == "q3_ple_32": + self.load_q3_ple_weights(model_path) + return + if ple_format != "fp8_safetensors": + raise ValueError(f"unsupported Qwen4 PLE format: {ple_format}") if dummy: self._dummy = True return @@ -400,6 +412,28 @@ def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: f"PLE table has {int(self._shard_ends[-1])} rows, needs {expected_rows}" ) + def load_q3_ple_weights(self, manifest_path: str) -> None: + """Opt into the native Q3_PLE_32 sidecar; FP8 Safetensors stays default.""" + + reader = Q3PLEReader(manifest_path) + if self._host_constants is None: + self._host_constants = ( + self.layer_multipliers.cpu(), + self.ngram_heads_vocab_sizes.cpu(), + self.ngram_heads_offsets.cpu(), + ) + expected_rows = int(self._host_constants[1][-1] + self._host_constants[2][-1]) + if reader.row_count < expected_rows: + reader.close() + raise RuntimeError( + f"Q3_PLE_32 table has {reader.row_count} rows, needs {expected_rows}" + ) + self._q3_reader = reader + self._shards = [] + self._handles = [] + self._shard_ends = torch.empty(0, dtype=torch.long) + self._scale = torch.tensor(reader.weight_scale, dtype=torch.bfloat16) + def _current_ngram_ids(self) -> torch.Tensor: if self._host_constants is None: raise RuntimeError("Qwen4-Exp PLE host weights are not loaded") @@ -452,6 +486,10 @@ def forward(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: token_count = get_global_ctx().batch.input_ids.numel() return torch.zeros(token_count, self.embedding_dim, device=device, dtype=dtype) ngram_ids = self._current_ngram_ids().reshape(-1) + if self._q3_reader is not None: + rows = self._q3_reader.gather(ngram_ids.tolist()) + embedded = rows.to(device=device, dtype=dtype) * self._scale.to(device=device, dtype=dtype) + return embedded.view(-1, self.embedding_dim) shard_ids = torch.bucketize(ngram_ids, self._shard_ends, right=True) output = torch.empty( ngram_ids.numel(), @@ -496,6 +534,10 @@ def __init__(self, config: ModelConfig, layer_id: int): def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: self.ple_embedding.load_host_weights(model_path, dummy=dummy) + def load_q3_ple_weights(self, manifest_path: str) -> None: + """Load an explicit Q3_PLE_32 sidecar for this layer.""" + self.ple_embedding.load_q3_ple_weights(manifest_path) + def _short_conv(self, hidden: torch.Tensor) -> torch.Tensor: batch = get_global_ctx().batch reqs = batch.padded_reqs if batch.is_decode else batch.reqs @@ -691,6 +733,14 @@ def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: if layer.ple is not None: layer.ple.load_host_weights(model_path, dummy=dummy) + def load_q3_ple_weights(self, manifest_paths: str | dict[int, str]) -> None: + """Opt into Q3_PLE_32 using one manifest or a layer-id manifest map.""" + for layer_id, layer in enumerate(self.layers.op_list): + if layer.ple is None: + continue + manifest = manifest_paths[layer_id] if isinstance(manifest_paths, dict) else manifest_paths + layer.ple.load_q3_ple_weights(manifest) + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: hidden = self.embed_tokens.forward(input_ids) mm_embeds = getattr(get_global_ctx().batch, "mm_embeds", None) @@ -735,6 +785,9 @@ def encode_images( def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: self.model.load_host_weights(model_path, dummy=dummy) + def load_q3_ple_weights(self, manifest_paths: str | dict[int, str]) -> None: + self.model.load_q3_ple_weights(manifest_paths) + def forward(self) -> torch.Tensor: hidden = self.model.forward(get_global_ctx().batch.input_ids) return self.lm_head.forward(hidden) diff --git a/tests/checkpoint/test_q3_ple.py b/tests/checkpoint/test_q3_ple.py new file mode 100644 index 000000000..76e6d88a9 --- /dev/null +++ b/tests/checkpoint/test_q3_ple.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import hashlib +import json +import os +import sys +from pathlib import Path + +import pytest +import torch + +from freetoken.checkpoint.q3_ple import ( + ALIGN, + BLOCK_BYTES, + BLOCK_VALUES, + Q3PLEReader, + ROW_BYTES, + ROW_VALUES, +) +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT.parent.parent / "scripts")) +from q3_ple_32_reference import dequantize_row, encode_table + + +FIXTURE_ROOT = ROOT / "artifacts" / "stage6_q3ple_fixture" +MANIFEST_PATH = FIXTURE_ROOT / "ple-q3.json" +DATA_PATH = FIXTURE_ROOT / "ple-q3-000.bin" + + +def _write_fixture() -> tuple[Path, Path, int]: + FIXTURE_ROOT.mkdir(parents=True, exist_ok=True) + rows = [] + for row in range(12): + rows.append([((row + 1) * 0.125) * ((i % 17) - 8) for i in range(ROW_VALUES)]) + encoded = encode_table(rows, refinement_passes=2, scale_dtype="bf16") + split = 5 * ROW_BYTES + second_offset = ALIGN + payload = encoded[:split] + bytes(second_offset - split) + encoded[split:] + DATA_PATH.write_bytes(payload) + segments = [] + for first, end, offset in ((0, 5, 0), (5, 12, second_offset)): + segment_bytes = encoded[first * ROW_BYTES : end * ROW_BYTES] + segments.append( + { + "first_row": first, + "end_row": end, + "data_offset": offset, + "byte_length": len(segment_bytes), + "sha256": hashlib.sha256(segment_bytes).hexdigest(), + } + ) + manifest = { + "format": "q3_ple_32", + "version": 1, + "endianness": "little", + "block_values": BLOCK_VALUES, + "block_bytes": BLOCK_BYTES, + "row_values": ROW_VALUES, + "row_bytes": ROW_BYTES, + "rows": len(rows), + "payload_bytes": len(encoded), + "file_bytes": len(payload), + "data_file": DATA_PATH.name, + "weight_scale": 1.25, + "sha256": hashlib.sha256(payload).hexdigest(), + "segments": segments, + } + MANIFEST_PATH.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return MANIFEST_PATH, DATA_PATH, len(rows) + + +@pytest.fixture(scope="module") +def fixture_paths(): + return _write_fixture() + + +def test_q3_reader_constants_and_ordered_gather(fixture_paths): + manifest, _, row_count = fixture_paths + with Q3PLEReader(manifest) as reader: + assert reader.row_count == row_count + assert reader.total_payload_bytes == row_count * 70 + assert reader.gather([11, 0, 11]).shape == (3, ROW_VALUES) + scaled = reader.gather16(list(range(16)) if row_count >= 16 else [0] * 16, apply_weight_scale=True) + assert scaled.dtype == torch.bfloat16 + assert scaled.shape == (16, ROW_VALUES) + assert torch.equal(scaled[0], reader.gather([0], apply_weight_scale=True)[0]) + + +def test_q3_reader_matches_authoritative_codec(fixture_paths): + manifest, data, _ = fixture_paths + # The fixture uses two segments and includes alignment padding between them. + encoded = data.read_bytes() + with Q3PLEReader(manifest) as reader: + for row in (0, 1, 4, 5, 6, 11): + raw_row = (encoded[row * ROW_BYTES : (row + 1) * ROW_BYTES] + if row < 5 else encoded[ALIGN + (row - 5) * ROW_BYTES : ALIGN + (row - 4) * ROW_BYTES]) + expected = torch.tensor( + dequantize_row(raw_row), + dtype=torch.bfloat16, + ) + assert torch.equal(reader.gather([row])[0], expected) + + +@pytest.mark.parametrize( + "field,value", + [("version", 2), ("endianness", "big"), ("row_bytes", 69), ("sha256", "0" * 64)], +) +def test_q3_reader_rejects_bad_manifest(fixture_paths, field, value): + manifest, _, _ = fixture_paths + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw[field] = value + bad = manifest.with_name(f"bad-{field}.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + + +def test_q3_reader_rejects_gap_overlap_and_truncation(fixture_paths): + manifest, data, _ = fixture_paths + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw["segments"][1]["first_row"] = 6 + bad = manifest.with_name("bad-segments.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + + truncated = data.with_name("truncated.bin") + truncated.write_bytes(data.read_bytes()[:-1]) + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw["data_file"] = truncated.name + raw["file_bytes"] -= 1 + bad = manifest.with_name("bad-truncated.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + truncated.unlink(missing_ok=True) + + +def test_q3_reader_requires_z_backing(): + with pytest.raises((ValueError, FileNotFoundError)): + Q3PLEReader("C:\\q3-ple\\ple-q3.json") From 4fdd38a24162c0162b4cf8a590170ce319a038de Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 04/17] feat(qwen4): add file-backed expert source --- python/freetoken/moe/__init__.py | 6 + python/freetoken/moe/expert_source.py | 314 ++++++++++++++++++++++++++ python/freetoken/moe/offload_cache.py | 150 +++++++++++- tests/moe/test_file_expert_source.py | 147 ++++++++++++ 4 files changed, 611 insertions(+), 6 deletions(-) create mode 100644 python/freetoken/moe/expert_source.py create mode 100644 tests/moe/test_file_expert_source.py diff --git a/python/freetoken/moe/__init__.py b/python/freetoken/moe/__init__.py index 5fc41911e..9e0dc1706 100644 --- a/python/freetoken/moe/__init__.py +++ b/python/freetoken/moe/__init__.py @@ -71,4 +71,10 @@ def create_moe_backend(backend: str) -> BaseMoeBackend: "SUPPORTED_MOE_BACKENDS", "OFFLOAD_MOE_BACKENDS", "is_offload_moe_backend", + "FileExpertSource", + "ExpertSourceError", ] + +# Kept at module bottom to avoid importing torch/file-I/O helpers while the +# backend registry is initialized by lightweight callers. +from .expert_source import ExpertSourceError, FileExpertSource # noqa: E402 diff --git a/python/freetoken/moe/expert_source.py b/python/freetoken/moe/expert_source.py new file mode 100644 index 000000000..2ddd311e4 --- /dev/null +++ b/python/freetoken/moe/expert_source.py @@ -0,0 +1,314 @@ +"""Bounded, file-backed routed-expert sources. + +The normal :class:`~freetoken.moe.offload_cache.OffloadMoeCache` source is a +resident HostBank. ``FileExpertSource`` is the explicit alternative used by +the Qwen4 text-only tier: one aligned record per expert and no materialised +full-layer tensor. The format is deliberately small and boring so corruption +is detected before the first request is served. +""" + +from __future__ import annotations + +import hashlib +import os +import struct +import threading +from pathlib import Path +from typing import Iterable + +import torch + + +MAGIC = b"FTEXNV4\0" +VERSION = 1 +HEADER_BYTES = 4096 +RECORD_BYTES = 2_772_992 +RAW_RECORD_BYTES = 2_772_480 +NUM_EXPERTS = 512 +ALIGNMENT = 4096 + +# Native ModelOpt NVFP4 planes for Qwen4's local expert geometry (H=2560, +# I=640). The six planes are kept in this order in every record. +PLANE_LAYOUT: tuple[tuple[str, int, str], ...] = ( + ("gate_up_packed", 1_638_400, "uint8"), + ("gate_up_scale", 204_800, "float8_e4m3fn"), + ("gate_up_global", 2_560, "float16"), + ("down_packed", 819_200, "uint8"), + ("down_scale", 102_400, "float8_e4m3fn"), + ("down_global", 5_120, "float16"), +) +_PLANE_OFFSETS = {} +_cursor = 0 +for _name, _size, _dtype in PLANE_LAYOUT: + _PLANE_OFFSETS[_name] = _cursor + _cursor += _size +assert _cursor == RAW_RECORD_BYTES + +_HEADER_STRUCT = struct.Struct("<8sIIIIQQ16s32s32s") + + +class ExpertSourceError(RuntimeError): + """Raised when a tier cannot be trusted or read exactly.""" + + +def _z_path(path: str | os.PathLike[str]) -> Path: + resolved = Path(path).resolve() + drive, _ = os.path.splitdrive(str(resolved)) + # ``Path.drive`` is reliable on Windows; splitdrive also keeps tests clear + # on environments where pathlib's Windows flavour is not selected. + drive = (resolved.drive or drive).upper() + if drive != "Z:": + raise ExpertSourceError(f"file-backed experts must reside on Z:, got {resolved}") + return resolved + + +def _dtype(name: str) -> torch.dtype: + return { + "uint8": torch.uint8, + "float16": torch.float16, + "float8_e4m3fn": torch.float8_e4m3fn, + }[name] + + +class FileExpertSource: + """Read fixed NVFP4 expert records from one layer sidecar. + + ``read_record`` returns independent CPU tensors, while ``read_into`` copies + directly into the destination slot planes. Calls are synchronous and each + call owns at most one bounded record buffer. ``read_records`` exposes an + explicit queue-depth guard for future asynchronous readers; this reference + implementation is intentionally serial (depth one) and therefore graph + capture incompatible. + """ + + bank_schema = tuple(name for name, _, _ in PLANE_LAYOUT) + record_bytes = RECORD_BYTES + raw_record_bytes = RAW_RECORD_BYTES + num_experts = NUM_EXPERTS + max_queue_depth = 16 + + def __init__( + self, + path: str | os.PathLike[str], + *, + expected_sha256: str | None = None, + expected_source_fingerprint: str | bytes | None = None, + num_experts: int = NUM_EXPERTS, + max_queue_depth: int = 1, + verify_hash: bool = True, + ) -> None: + self.path = _z_path(path) + if not self.path.is_file(): + raise ExpertSourceError(f"missing expert tier: {self.path}") + if not 1 <= int(max_queue_depth) <= 16: + raise ValueError("max_queue_depth must be in [1, 16]") + if not 1 <= int(num_experts) <= NUM_EXPERTS: + raise ValueError(f"num_experts must be in [1, {NUM_EXPERTS}]") + self.num_experts = int(num_experts) + self.requested_queue_depth = int(max_queue_depth) + self.staging_record_bytes = RECORD_BYTES + self.max_staging_records = self.requested_queue_depth + self._lock = threading.Lock() + self._closed = False + self._fd = os.open(str(self.path), os.O_RDONLY | getattr(os, "O_BINARY", 0)) + try: + self._validate_file(expected_source_fingerprint) + self.sha256 = self._hash_file() if verify_hash else None + if expected_sha256 is not None: + expected_sha256 = expected_sha256.lower() + if self.sha256 is None: + self.sha256 = self._hash_file() + if self.sha256 != expected_sha256: + raise ExpertSourceError( + f"expert tier hash mismatch: expected {expected_sha256}, got {self.sha256}" + ) + except Exception: + os.close(self._fd) + raise + self.read_count = 0 + self.bytes_read = 0 + self.max_inflight = 0 + self._inflight = 0 + + @classmethod + def create_synthetic( + cls, + path: str | os.PathLike[str], + *, + num_experts: int = NUM_EXPERTS, + records: Iterable[bytes] | None = None, + source_fingerprint: bytes | None = None, + ) -> str: + """Create a tiny deterministic sidecar for tests (never model data).""" + path = _z_path(path) + if num_experts < 1: + raise ValueError("num_experts must be positive") + source_fingerprint = source_fingerprint or hashlib.sha256(b"synthetic").digest() + source_fingerprint = bytes(source_fingerprint[:32]).ljust(32, b"\0") + rows = iter(records) if records is not None else None + payload_hash = hashlib.sha256() + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("wb") as fh: + header = bytearray(HEADER_BYTES) + _HEADER_STRUCT.pack_into( + header, + 0, + MAGIC, + VERSION, + HEADER_BYTES, + num_experts, + len(PLANE_LAYOUT), + RAW_RECORD_BYTES, + RECORD_BYTES, + b"nvfp4-qwen4-v1\0\0", # 16-byte layout tag + source_fingerprint, + b"\0" * 32, + ) + fh.write(header) + for expert_id in range(num_experts): + raw = next(rows) if rows is not None else bytes([expert_id & 0xFF]) * RAW_RECORD_BYTES + if len(raw) != RAW_RECORD_BYTES: + raise ValueError("synthetic record must contain exactly RAW_RECORD_BYTES") + record = raw + bytes(RECORD_BYTES - RAW_RECORD_BYTES) + payload_hash.update(record) + fh.write(record) + digest = payload_hash.digest() + with path.open("r+b") as fh: + fh.seek(0) + header = bytearray(fh.read(HEADER_BYTES)) + header[struct.calcsize("<8sIIIIQQ16s32s") : _HEADER_STRUCT.size] = digest + fh.seek(0) + fh.write(header) + return hashlib.sha256(path.read_bytes()).hexdigest() + + def _validate_file(self, expected_source_fingerprint: str | bytes | None) -> None: + size = self.path.stat().st_size + expected_size = HEADER_BYTES + self.num_experts * self.record_bytes + header = self._read_exact(HEADER_BYTES, 0) + if len(header) != HEADER_BYTES: + raise ExpertSourceError("truncated expert tier header") + try: + magic, version, hbytes, experts, planes, raw_bytes, rec_bytes, tag, fingerprint, payload_hash = _HEADER_STRUCT.unpack_from(header) + except struct.error as exc: + raise ExpertSourceError("malformed expert tier header") from exc + if magic != MAGIC or version != VERSION or hbytes != HEADER_BYTES: + raise ExpertSourceError("unsupported expert tier magic/version/header") + if experts != self.num_experts or planes != len(PLANE_LAYOUT): + raise ExpertSourceError("expert tier geometry mismatch") + if raw_bytes != RAW_RECORD_BYTES or rec_bytes != RECORD_BYTES or tag.rstrip(b"\0") != b"nvfp4-qwen4-v1": + raise ExpertSourceError("expert tier layout mismatch") + if size != expected_size: + raise ExpertSourceError(f"expert tier length mismatch: {size} != {expected_size}") + if expected_source_fingerprint is not None: + expected = bytes.fromhex(expected_source_fingerprint) if isinstance(expected_source_fingerprint, str) else bytes(expected_source_fingerprint) + if fingerprint != expected[:32].ljust(32, b"\0"): + raise ExpertSourceError("expert tier source fingerprint mismatch") + self.payload_sha256 = payload_hash.hex() + if self.payload_sha256 != "00" * 32: + h = hashlib.sha256() + offset = HEADER_BYTES + with self.path.open("rb") as fh: + fh.seek(offset) + while True: + block = fh.read(8 << 20) + if not block: + break + h.update(block) + if h.digest() != payload_hash: + raise ExpertSourceError("expert tier payload hash mismatch") + + def _read_exact(self, size: int, offset: int) -> bytes: + if self._closed: + raise ExpertSourceError("expert tier is closed") + with self._lock: + if hasattr(os, "pread"): + data = os.pread(self._fd, size, offset) + else: # pragma: no cover - Windows Python fallback + os.lseek(self._fd, offset, os.SEEK_SET) + data = os.read(self._fd, size) + if len(data) != size: + raise ExpertSourceError(f"short expert tier read at offset {offset}: {len(data)} != {size}") + return data + + def _hash_file(self) -> str: + h = hashlib.sha256() + with self.path.open("rb") as fh: + while True: + block = fh.read(8 << 20) + if not block: + break + h.update(block) + return h.hexdigest() + + def _record_bytes(self, expert_id: int) -> bytes: + if self._closed: + raise ExpertSourceError("expert tier is closed") + expert_id = int(expert_id) + if not 0 <= expert_id < self.num_experts: + raise IndexError(f"expert_id {expert_id} outside [0, {self.num_experts})") + self._inflight += 1 + self.max_inflight = max(self.max_inflight, self._inflight) + try: + offset = HEADER_BYTES + expert_id * self.record_bytes + if offset % ALIGNMENT != 0 or self.record_bytes % ALIGNMENT != 0: + raise ExpertSourceError("expert tier record is not aligned to 4096 bytes") + data = self._read_exact(self.record_bytes, offset) + self.read_count += 1 + self.bytes_read += len(data) + return data + finally: + self._inflight -= 1 + + def read_record(self, expert_id: int) -> dict[str, torch.Tensor]: + raw = self._record_bytes(expert_id) + out: dict[str, torch.Tensor] = {} + for name, size, dtype_name in PLANE_LAYOUT: + offset = _PLANE_OFFSETS[name] + dtype = _dtype(dtype_name) + if name == "gate_up_packed": + shape = (1280, 1280) + elif name == "down_packed": + shape = (2560, 320) + elif name == "gate_up_scale": + shape = (1280, 160) + elif name == "gate_up_global": + shape = (1280,) + elif name == "down_scale": + shape = (2560, 40) + else: + shape = (2560,) + out[name] = torch.frombuffer(memoryview(raw)[offset : offset + size], dtype=dtype).clone().reshape(shape) + return out + + def read_records(self, expert_ids: Iterable[int], *, max_concurrency: int = 1) -> list[dict[str, torch.Tensor]]: + if not 1 <= int(max_concurrency) <= self.requested_queue_depth: + raise ValueError(f"max_concurrency must be in [1, {self.requested_queue_depth}]") + # Serial is deliberate: it keeps staging bounded and is graph-safe only + # outside CUDA capture. A future async implementation may use up to 16. + return [self.read_record(eid) for eid in expert_ids] + + def read_into(self, expert_id: int, destinations: dict[str, torch.Tensor], slot: int) -> int: + """Read one record and copy its six planes into a GPU/CPU cache slot.""" + rows = self.read_record(expert_id) + for name in self.bank_schema: + dst = destinations[name] + if dst.ndim < 1 or not 0 <= slot < dst.shape[0]: + raise ValueError(f"destination slot {slot} invalid for {name}") + if tuple(dst.shape[1:]) != tuple(rows[name].shape): + raise ValueError(f"destination shape mismatch for {name}") + dst[slot].copy_(rows[name], non_blocking=False) + return self.record_bytes + + def close(self) -> None: + if not self._closed: + os.close(self._fd) + self._closed = True + + def __enter__(self) -> "FileExpertSource": + return self + + def __exit__(self, *_exc) -> None: + self.close() + + +__all__ = ["FileExpertSource", "ExpertSourceError", "PLANE_LAYOUT", "HEADER_BYTES", "RECORD_BYTES", "RAW_RECORD_BYTES"] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6ee764061..aa588c8b6 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -195,6 +195,14 @@ def __post_init__(self) -> None: self.bank_schema = _BANK_SCHEMAS[self.quant_format] self.bank_sources: dict[str, list[torch.Tensor]] = {} self.bank_caches: dict[str, torch.Tensor] = {} + # Optional true file-backed tier. Unlike ``bank_sources`` these entries + # never point at a full-layer HostBank; one aligned record is fetched per + # miss into the existing slot cache. The source object is deliberately + # kept separate so CPU/hybrid callers cannot accidentally treat a tier as + # pageable resident memory. + self.file_sources = {} + self._pending_file_fetches: list[tuple[int, int]] = [] + self._pending_file_materialize = False # per-layer host residency: the GPU movement paths require "pinned"; LOCKED/PAGEABLE layers decode on the CPU executor and prefill via copy_missing's pageable branch # _unpinned_layers is the derived id set the hot paths test against self.layer_residency: list[str] = [] @@ -277,7 +285,7 @@ def __post_init__(self) -> None: def set_bank_sources( self, - sources: dict[str, list[torch.Tensor]], + sources: dict[str, list[torch.Tensor | None]], layer_residency: list[str] | None = None, ) -> None: """Attach the host (CPU pinned) expert source banks and allocate a GPU slot @@ -321,8 +329,16 @@ def set_bank_sources( for name in self.bank_schema: per_layer = sources[name] assert len(per_layer) == self.num_layers, (name, len(per_layer)) - head = per_layer[0] + heads = [source for source in per_layer if source is not None] + if not heads: + raise ValueError(f"bank {name!r} has no resident shape; attach a file source") + head = heads[0] for layer_id, source in enumerate(per_layer): + if source is None: + # A None row is an explicit file-tier placeholder. The + # corresponding layer must be registered with set_file_sources + # before any request reaches it. + continue assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous" assert source.size(0) == self.num_experts, (name, layer_id, source.shape) assert source.shape == head.shape and source.dtype == head.dtype, ( @@ -339,6 +355,34 @@ def set_bank_sources( if self.prefill_overlap: self._init_prefill_overlap_buffers() + def set_file_sources(self, sources: dict[int, object]) -> None: + """Register GPU-only, fixed-record expert tiers by MoE layer. + + ``set_bank_sources`` must contain ``None`` placeholders for these layers, + which makes the absence of a HostBank explicit. File tiers are never + valid for CPU/hybrid decode or prefill overlap because both paths require + resident host tensors and CUDA graph capture cannot contain synchronous + file I/O. + """ + if self.decode_target != "gpu": + raise ValueError("file-backed expert tiers are GPU-only; use resident HostBanks for CPU/hybrid") + if self.prefill_overlap: + raise ValueError("file-backed expert tiers require prefill_overlap=False") + if not self.bank_sources: + raise ValueError("set_bank_sources must allocate cache planes before file tiers") + for layer_id, source in sources.items(): + layer_id = int(layer_id) + if not 0 <= layer_id < self.num_layers: + raise ValueError(f"file tier layer {layer_id} outside cache geometry") + if not hasattr(source, "bank_schema") or tuple(source.bank_schema) != tuple(self.bank_schema): + raise ValueError("file tier bank schema does not match cache quant_format") + if int(source.num_experts) != self.num_experts: + raise ValueError("file tier expert count does not match cache geometry") + for name in self.bank_schema: + if self.bank_sources[name][layer_id] is not None: + raise ValueError(f"layer {layer_id} already has a resident HostBank for {name}") + self.file_sources[layer_id] = source + def _build_copy_plan(self) -> None: """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes). @@ -364,7 +408,10 @@ def _build_copy_plan(self) -> None: dst_ptrs, feats = [], [] layer_src_ptrs = [[] for _ in range(self.num_layers)] for per_layer, cache in self.banks: - feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size() + head = next((source for source in per_layer if source is not None), None) + if head is None: + return + feat = math.prod(head.shape[1:]) * head.element_size() if feat % 16 != 0 or cache.data_ptr() % 16 != 0: return # leave fused disabled; copy_missing uses the per-bank path for layer_id, source in enumerate(per_layer): @@ -376,6 +423,12 @@ def _build_copy_plan(self) -> None: # The kernel dereferences these on the GPU, so store each host bank's # device alias (== data_ptr() under UVA identity; differs on # Windows/WDDM). + if source is None: + # A file-tier layer has no host pointer and is never passed + # through the CUDA copy kernel; copy_missing dispatches its + # synchronous record reader below. + layer_src_ptrs[layer_id].append(0) + continue src_dev = device_ptr(source) if src_dev % 16 != 0: return @@ -448,7 +501,11 @@ def rebuild(self, cache_size: int) -> None: torch.cuda.empty_cache() # 3. Reallocate the slot cache from the retained host sources. for name in self.bank_schema: - head = self.bank_sources[name][0] + head = next( + (source for source in self.bank_sources[name] if source is not None), None + ) + if head is None: + raise ValueError(f"bank {name!r} has no resident shape for rebuild") self.bank_caches[name] = torch.empty( (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device ) @@ -798,6 +855,9 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_released[buffer_id] = True def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: + if layer_id in self.file_sources: + self.ensure_file_experts(layer_id, expert_ids) + return from freetoken.moe.offload_kernels import ensure_experts if self.collect_decode_freq: @@ -809,6 +869,52 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: self._pending_whole_layer = False ensure_experts(self, layer_id, expert_ids) + def ensure_file_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: + """Host-bookkeep file-tier misses and rewrite IDs to GPU cache slots. + + This path intentionally performs a bounded host sync and is not CUDA-graph + capturable. Qwen4 disables graph capture when a file source is attached; + callers that cannot make that guarantee must reject the configuration. + """ + if layer_id not in self.file_sources: + raise ValueError(f"layer {layer_id} has no file source") + if self.decode_target != "gpu": + raise RuntimeError("file-backed experts cannot serve CPU/hybrid decode") + original_shape = tuple(expert_ids.shape) + raw_ids = [int(value) for value in expert_ids.detach().cpu().reshape(-1).tolist()] + self.step += 1 + step = int(self.step.item()) + mapped: list[int] = [] + pending: list[tuple[int, int]] = [] + for expert_id in raw_ids: + if not 0 <= expert_id < self.num_experts: + raise IndexError(f"expert_id {expert_id} outside cache geometry") + slot = int(self.slot_for_id[layer_id, expert_id].item()) + if slot < 0: + free = torch.nonzero(self.id_of_slot < 0, as_tuple=False).reshape(-1) + if free.numel(): + slot = int(free[0].item()) + else: + slot = int(torch.argmin(self.usage).item()) + old = int(self.id_of_slot[slot].item()) + if old >= 0: + self.slot_for_id.view(-1)[old] = -1 + flat_id = layer_id * self.num_experts + expert_id + self.slot_for_id[layer_id, expert_id] = slot + self.id_of_slot[slot] = flat_id + pending.append((slot, expert_id)) + self.usage[slot] = step + mapped.append(slot) + self._pending_src_layer = layer_id + self._pending_whole_layer = False + self._pending_file_materialize = False + # ``copy_missing`` has no use for the LRU scratch arrays on this path, but + # keep num_indices truthful for diagnostics and callers that inspect it. + self.num_indices[0] = len(pending) + self._pending_file_fetches = pending + replacement = torch.tensor(mapped, dtype=expert_ids.dtype, device=expert_ids.device).reshape(original_shape) + expert_ids.copy_(replacement) + def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. @@ -831,6 +937,12 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None ) def materialize_layer(self, layer_id: int) -> None: + if layer_id in self.file_sources: + self._pending_src_layer = layer_id + self._pending_whole_layer = True + self._pending_file_materialize = True + self._pending_file_fetches = [] + return from freetoken.moe.offload_kernels import materialize_layer self._pending_src_layer = layer_id @@ -838,12 +950,24 @@ def materialize_layer(self, layer_id: int) -> None: materialize_layer(self, layer_id) def reset(self) -> None: - from freetoken.moe.offload_kernels import reset_cache + if self.device.type == "cuda": + from freetoken.moe.offload_kernels import reset_cache - reset_cache(self) + reset_cache(self) + else: + # The production path is CUDA-only, but a CPU synthetic FileExpertSource + # fixture still needs deterministic reset/reuse semantics without + # attempting to launch a Triton kernel on a CPU tensor. + self.slot_for_id.fill_(-1) + self.id_of_slot.fill_(-1) + self.usage.zero_() + self.step.zero_() + self.num_indices.zero_() # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe # it here so a new sequence starts with cold hybrid fetch priorities. self.expert_recency.fill_(-1) + self._pending_file_fetches = [] + self._pending_file_materialize = False def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -969,6 +1093,20 @@ def copy_missing(self) -> None: assert self.banks, "set_bank_sources must register the banks first" layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" + if layer_id in self.file_sources: + source = self.file_sources[layer_id] + if self._pending_file_materialize: + pairs = [(expert_id, expert_id) for expert_id in range(self.num_experts)] + else: + pairs = list(self._pending_file_fetches) + destinations = { + name: cache for name, (_, cache) in zip(self.bank_schema, self.banks) + } + for slot, expert_id in pairs: + source.read_into(expert_id, destinations, slot) + self._pending_file_fetches = [] + self._pending_file_materialize = False + return if layer_id in self._unpinned_layers: if not self._pending_whole_layer: raise RuntimeError( diff --git a/tests/moe/test_file_expert_source.py b/tests/moe/test_file_expert_source.py new file mode 100644 index 000000000..39b9e9130 --- /dev/null +++ b/tests/moe/test_file_expert_source.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch + +from freetoken.moe.expert_source import ( + ExpertSourceError, + FileExpertSource, + PLANE_LAYOUT, + RAW_RECORD_BYTES, +) + + +@pytest.fixture +def z_fixture_dir(): + root = Path.cwd() / ".stage6-test-fixtures" / uuid4().hex + root.mkdir(parents=True, exist_ok=False) + try: + assert (root.drive or "").upper() == "Z:", root + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +def _record(byte: int) -> bytes: + return bytes([byte]) * RAW_RECORD_BYTES + + +def _resident_banks(num_layers: int, experts: int): + shapes = { + "gate_up_packed": ((experts, 1280, 1280), torch.uint8), + "gate_up_scale": ((experts, 1280, 160), torch.float8_e4m3fn), + "gate_up_global": ((experts, 1280), torch.float16), + "down_packed": ((experts, 2560, 320), torch.uint8), + "down_scale": ((experts, 2560, 40), torch.float8_e4m3fn), + "down_global": ((experts, 2560), torch.float16), + } + return { + name: [torch.zeros(shape, dtype=dtype) if layer == 0 else None for layer in range(num_layers)] + for name, (shape, dtype) in shapes.items() + } + + +def test_file_expert_source_reads_exact_planes_and_rejects_tamper(z_fixture_dir): + root = z_fixture_dir + path = root / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i) for i in range(3)]) + with FileExpertSource(path, num_experts=3, expected_sha256=digest) as src: + rows = src.read_record(2) + assert set(rows) == {name for name, _, _ in PLANE_LAYOUT} + assert rows["gate_up_packed"].shape == (1280, 1280) + assert rows["gate_up_scale"].shape == (1280, 160) + assert int(rows["gate_up_packed"].flatten()[0]) == 2 + assert src.read_count == 1 + assert src.max_inflight <= 1 <= 16 + with pytest.raises(ExpertSourceError, match="closed"): + src.read_record(0) + + +def test_file_expert_source_cache_miss_fills_slots_without_host_layer(z_fixture_dir): + path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 11) for i in range(3)]) + src = FileExpertSource(path, num_experts=3, expected_sha256=digest) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=3, + cache_size=3, + device=torch.device("cpu"), + quant_format="nvfp4", + decode_target="gpu", + ) + cache.set_bank_sources(_resident_banks(2, 3)) + cache.set_file_sources({1: src}) + assert all(cache.bank_sources[name][1] is None for name in cache.bank_schema) + ids = torch.tensor([2, 0], dtype=torch.int32) + cache.ensure_experts(1, ids) + assert ids.tolist() == [0, 1] + cache.copy_missing() + assert cache.bank_caches["gate_up_packed"][0, 0, 0].item() == 13 + assert cache.bank_caches["gate_up_packed"][1, 0, 0].item() == 11 + assert src.bytes_read == 2 * src.record_bytes + assert cache._pending_file_fetches == [] + cache.reset() + assert cache.id_of_slot.tolist() == [-1, -1, -1] + assert cache.slot_for_id.tolist() == [[-1, -1, -1], [-1, -1, -1]] + finally: + src.close() + + +def test_file_tier_materialize_streams_complete_layer(z_fixture_dir): + path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 21) for i in range(3)]) + src = FileExpertSource(path, num_experts=3, expected_sha256=digest) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(2, 3, 3, torch.device("cpu"), quant_format="nvfp4") + cache.set_bank_sources(_resident_banks(2, 3)) + cache.set_file_sources({1: src}) + cache.materialize_layer(1) + cache.copy_missing() + values = cache.bank_caches["gate_up_packed"][:, 0, 0].tolist() + assert values == [21, 22, 23] + assert src.read_count == 3 + finally: + src.close() + + +def test_file_tier_payload_hash_fails_closed(z_fixture_dir): + path = z_fixture_dir / "experts-L00.nvfp4" + FileExpertSource.create_synthetic(path, num_experts=1, records=[_record(9)]) + with path.open("r+b") as fh: + fh.seek(4096 + 17) + fh.write(b"x") + with pytest.raises(ExpertSourceError, match="payload hash mismatch"): + FileExpertSource(path, num_experts=1) + + +def test_file_tier_rejects_cpu_hybrid_and_overlap(z_fixture_dir): + path = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=1, records=[_record(7)]) + src = FileExpertSource(path, num_experts=1, expected_sha256=digest) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + for target in ("cpu", "hybrid"): + cache = OffloadMoeCache(1, 1, 1, torch.device("cpu"), decode_target=target, quant_format="nvfp4") + # Source registration itself fails before any source can be used; + # a direct call is sufficient to prove the policy and avoids giant + # synthetic resident allocations in this negative test. + with pytest.raises(ValueError, match="GPU-only"): + cache.set_file_sources({0: src}) + finally: + src.close() + + +def test_file_expert_source_rejects_wrong_volume(monkeypatch): + with pytest.raises(ExpertSourceError, match="Z:"): + FileExpertSource(r"C:\\not-a-tier.nvfp4") From a5f8ea52b9e7800aa6be0a04b953047270d5b4d5 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 05/17] feat(qwen4): add synthetic NVFP4 active-weight path --- python/freetoken/checkpoint/nvfp4.py | 173 ++++++++++++++++++ python/freetoken/models/qwen3_5_moe/gdn.py | 29 ++- python/freetoken/models/qwen4_exp/model.py | 32 +++- python/freetoken/models/qwen4_exp/weight.py | 9 +- tests/kernels/test_qwen4_nvfp4_active.py | 37 ++++ .../models/test_qwen4_exp_nvfp4_components.py | 114 ++++++++++++ 6 files changed, 375 insertions(+), 19 deletions(-) create mode 100644 python/freetoken/checkpoint/nvfp4.py create mode 100644 tests/kernels/test_qwen4_nvfp4_active.py create mode 100644 tests/models/test_qwen4_exp_nvfp4_components.py diff --git a/python/freetoken/checkpoint/nvfp4.py b/python/freetoken/checkpoint/nvfp4.py new file mode 100644 index 000000000..906a912c6 --- /dev/null +++ b/python/freetoken/checkpoint/nvfp4.py @@ -0,0 +1,173 @@ +"""Deterministic host-side NVFP4 W4A16 encoding helpers. + +The native FreeToken dense NVFP4 operators consume three row-major tensors: + +* packed E2M1 codes (two low-bit-first nibbles per byte), +* one positive E4M3 scale for every 16 input values, and +* one FP16 positive global scale per output row. + +This module is intentionally CPU-safe and does not retain a BF16 copy. It is +used by metadata/conversion code and by synthetic component tests; runtime +operators continue to live in :mod:`freetoken.kernel.triton.nvfp4_linear`. +""" + +from __future__ import annotations + +import torch + + +# Keep this table in lock-step with the Triton/native dequant implementations. +# The unsigned codes are magnitudes; bit 3 is the sign bit. +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E2M1_MAGNITUDES = torch.tensor(E2M1_VALUES, dtype=torch.float32) +E2M1_SIGNED = torch.tensor( + E2M1_VALUES + tuple(-v for v in E2M1_VALUES), dtype=torch.float32 +) +_FP16_MAX = float(torch.finfo(torch.float16).max) +_FP16_MIN_SUBNORMAL = 2.0 ** -24 +_E4M3_MAX = 448.0 +_E4M3_MIN_SUBNORMAL = 2.0 ** -9 + + +def _round_e2m1_rne(magnitude: torch.Tensor) -> torch.Tensor: + """Round non-negative values to E2M1 using the shared tie-to-even rule. + + Ties are resolved by the parity of the integer E2M1 code (for example, + 0.5/1.0 resolves to code 2, while 1.0/1.5 resolves to code 2). The + comparison is carried out in float64 so all BF16 inputs have deterministic + behavior at exact midpoints. + """ + + if torch.any(~torch.isfinite(magnitude)) or torch.any(magnitude < 0): + raise ValueError("E2M1 rounding expects finite non-negative values") + grid = E2M1_MAGNITUDES.to(device=magnitude.device, dtype=torch.float64) + x = magnitude.to(torch.float64).unsqueeze(-1) + distance = (x - grid).abs() + minimum = distance.min(dim=-1, keepdim=True).values + candidates = distance == minimum + # Prefer the even code among exact ties. Since candidates are at most two + # adjacent codes, selecting the last even candidate gives the desired rule. + codes = torch.arange(8, device=magnitude.device).expand_as(distance) + even = candidates & ((codes & 1) == 0) + picked = torch.where(even, codes, torch.full_like(codes, -1)).amax(dim=-1) + # Non-ties have no even candidate only for an impossible malformed grid; + # retain the nearest code as a defensive total fallback. + nearest = distance.argmin(dim=-1) + return torch.where(picked >= 0, picked, nearest).to(torch.uint8) + + +def _round_e4m3_positive(values: torch.Tensor) -> torch.Tensor: + """Encode finite non-negative values to E4M3 bytes with explicit bounds. + + E4M3's finite range is [0, 448]. Values above the finite range saturate + to 448 before the PyTorch cast (which otherwise produces the NaN sentinel), + and values below the representable subnormal range round to zero. This is + the documented, deterministic total policy for synthetic conversion. + """ + + if torch.any(~torch.isfinite(values)) or torch.any(values < 0): + raise ValueError("E4M3 scale encoding expects finite non-negative values") + bounded = values.to(torch.float32).clamp_(0.0, _E4M3_MAX) + # torch's CPU float8 conversion is round-to-nearest-even on the E4M3 grid. + return bounded.to(torch.float8_e4m3fn) + + +def encode_bf16_nvfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode a BF16/FP16/FP32 matrix into the native row-major NVFP4 triple. + + Args: + weight: ``[out_features, in_features]`` finite real matrix. The input + width must be divisible by 16, matching ``Nvfp4DenseLinear``. + + Returns: + ``(packed, block_scale, global_scale)`` where packed is uint8 + ``[N,K//2]``, block_scale is native ``torch.float8_e4m3fn`` + ``[N,K//16]`` and + global_scale is FP16 ``[N]``. The returned tensors are newly allocated + and no copy of ``weight`` is retained. + + Scale rule: + ``global = round_fp16(max_abs / 6)`` (clamped to the finite FP16 range, + with the smallest FP16 subnormal used when a positive value would round + to zero); ``block = round_e4m3(max_abs_block / (6*global))``. A zero + row uses global=1 and zero block scales. Quantization then rounds each + value to E2M1 after dividing by ``global*block``. Zero block scales + produce zero codes. These explicit bounds make conversion total for + every finite input, including under/overflow extrema. + """ + + if weight.ndim != 2: + raise ValueError(f"NVFP4 encoder expects a rank-2 matrix, got {tuple(weight.shape)}") + if weight.shape[1] % 16: + raise ValueError(f"NVFP4 input width must be divisible by 16, got {weight.shape[1]}") + if not weight.dtype.is_floating_point: + raise TypeError(f"NVFP4 encoder expects a floating tensor, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError("NVFP4 encoder rejects NaN and infinity inputs") + + source = weight.to(dtype=torch.float32) + n_rows, width = source.shape + abs_source = source.abs() + row_max = abs_source.amax(dim=1) + nonzero = row_max > 0 + + # Rounding through one float16 conversion is intentional. Explicitly clamp + # the target first because a large BF16 row otherwise converts to inf. + global_target = (row_max / 6.0).clamp(_FP16_MIN_SUBNORMAL, _FP16_MAX) + global_target = torch.where(nonzero, global_target, torch.ones_like(global_target)) + global_scale = global_target.to(torch.float16) + # A positive target below the FP16 subnormal can still become zero on some + # CPU implementations; repair it explicitly and deterministically. + global_scale = torch.where( + nonzero & (global_scale == 0), + torch.full_like(global_scale, _FP16_MIN_SUBNORMAL, dtype=torch.float16), + global_scale, + ) + + blocks = source.view(n_rows, width // 16, 16) + block_max = blocks.abs().amax(dim=-1) + denom = global_scale.float().unsqueeze(-1) * 6.0 + block_target = torch.where(block_max > 0, block_max / denom, torch.zeros_like(block_max)) + block_scale = _round_e4m3_positive(block_target) + block_real = block_scale.view(torch.float8_e4m3fn).float() + + # Quantize against the *rounded* scales consumed by the kernel. Saturating + # normalized values to +/-6 is the finite E2M1 endpoint policy. + scale_real = global_scale.float().unsqueeze(-1).unsqueeze(-1) * block_real.unsqueeze(-1) + normalized = torch.where(scale_real > 0, blocks / scale_real, torch.zeros_like(blocks)) + magnitude = normalized.abs().clamp_(0.0, 6.0) + mag_code = _round_e2m1_rne(magnitude.reshape(-1)).view(n_rows, width // 16, 16) + sign = (normalized < 0).to(torch.uint8) + code = mag_code | (sign << 3) + # Two values per byte, low nibble first, as required by the native kernels. + packed = code.reshape(n_rows, width // 2, 2) + packed = packed[..., 0] | (packed[..., 1] << 4) + return packed.contiguous(), block_scale.contiguous(), global_scale.contiguous() + + +def decode_nvfp4( + packed: torch.Tensor, + block_scale: torch.Tensor, + global_scale: torch.Tensor, + *, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Reference dequantization for the native row-major NVFP4 triple.""" + + if packed.dtype != torch.uint8 or block_scale.dtype not in (torch.uint8, torch.float8_e4m3fn): + raise TypeError("packed must be uint8 and block_scale must be uint8-view or float8_e4m3fn") + if packed.ndim != 2 or block_scale.ndim != 2 or global_scale.ndim != 1: + raise ValueError("NVFP4 tensors must be packed[N,K/2], scale[N,K/16], global[N]") + rows, packed_width = packed.shape + width = packed_width * 2 + if block_scale.shape != (rows, width // 16) or global_scale.shape != (rows,): + raise ValueError("NVFP4 tensor shapes do not agree") + lo = packed & 0x0F + hi = packed >> 4 + codes = torch.stack((lo, hi), dim=-1).reshape(rows, width).to(torch.long) + values = E2M1_SIGNED.to(device=packed.device)[codes] + scales = block_scale.view(torch.float8_e4m3fn).float().repeat_interleave(16, dim=1) + return (values * scales * global_scale.float().unsqueeze(-1)).to(dtype) + + +__all__ = ["E2M1_VALUES", "encode_bf16_nvfp4", "decode_nvfp4"] diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..b370f3cbe 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -53,7 +53,7 @@ class Qwen3_5GatedDeltaNet(BaseOP): def __init__( self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, rms_norm_eps, layer_id, expert_quant: str = "none", - attn_quant: str = "none", + attn_quant: str = "none", *, nvfp4_qkvz: bool = False, ): self.layer_id = layer_id # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as @@ -71,16 +71,25 @@ def __init__( self.value_dim = num_v_heads * head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim self.conv_kernel_size = conv_kernel_size - # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 - # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way - # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). + # qkv|z carry a weight scale (block-fp8 weight_scale_inv, per-tensor FP8 + # weight_scale, or native NVFP4 scales); b|a stay bf16. The explicit + # ``nvfp4_qkvz`` opt-in is used by Qwen4 only. Keeping it separate from + # ``attn_quant`` preserves Qwen3.5's historical NVFP4 behavior (where only + # out_proj is native FP4). self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" + self._nvfp4_qkvz = bool(nvfp4_qkvz) self._fp8 = self._block_fp8 or self._pertensor_fp8 + self._split_input = self._fp8 or self._nvfp4_qkvz self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._fp8: - ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged + if self._split_input: + if self._nvfp4_qkvz: + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged + + ColMerged = Nvfp4DenseColMerged + else: + ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False ) @@ -98,9 +107,9 @@ def __init__( self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32) self.A_log = torch.empty(num_v_heads, dtype=torch.float32) self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps) - # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors - # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors - # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. + # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / NVFP4 (W4A16) / + # bf16. In Qwen4's explicit ``nvfp4_qkvz`` mode, qkv|z is native NVFP4 while b|a + # remains BF16; Qwen3.5 callers retain the historical fused-BF16 input path. self.out_proj = make_replicated_quant( expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False ) @@ -161,7 +170,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._fp8: + if self._split_input: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 9f5ec6dc2..2e684ef60 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -27,8 +27,9 @@ from freetoken.models.blocks import BaseLLMModel from freetoken.models.config import ModelConfig from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention -from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet from freetoken.checkpoint.q3_ple import Q3PLEReader +from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet +from freetoken.models.quant_linear import make_col_merged_quant, make_replicated_quant from freetoken.utils import download_hf_weight, nvtx_annotate if TYPE_CHECKING: @@ -187,8 +188,16 @@ def __init__(self, config: ModelConfig, combine: bool = True): self.hidden_size = config.hidden_size hc_size = self.hc_count * self.hidden_size self.hc_norm = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) - self.input_mix_weight_down = LinearReplicated(hc_size, args.hc_lowrank, has_bias=False) - self.input_mix_weight_up = LinearReplicated(args.hc_lowrank, hc_size, has_bias=False) + # The frozen Qwen4 active map keeps mHC input-mix down/up native NVFP4 when + # ``dense_quant=nvfp4`` is explicitly selected. Block injection remains BF16. + self.input_mix_weight_down = make_replicated_quant( + "none", getattr(config, "dense_quant", "none"), hc_size, args.hc_lowrank, + has_bias=False, + ) + self.input_mix_weight_up = make_replicated_quant( + "none", getattr(config, "dense_quant", "none"), args.hc_lowrank, hc_size, + has_bias=False, + ) self.block_inject_weight = ( LinearReplicated(hc_size, self.hc_count, has_bias=False) if combine else None ) @@ -208,10 +217,14 @@ def forward(self, hyper_input: torch.Tensor): class _SharedExpert(BaseOP): def __init__(self, config: ModelConfig): width = config.shared_expert_intermediate_size - self.gate_up_proj = LinearColParallelMerged( - config.hidden_size, [width, width], has_bias=False + self.gate_up_proj = make_col_merged_quant( + "none", getattr(config, "dense_quant", "none"), config.hidden_size, + [width, width], has_bias=False, + ) + self.down_proj = make_replicated_quant( + "none", getattr(config, "dense_quant", "none"), width, config.hidden_size, + has_bias=False, ) - self.down_proj = LinearRowParallel(width, config.hidden_size, has_bias=False) def forward(self, hidden: torch.Tensor) -> torch.Tensor: return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(hidden))) @@ -670,7 +683,9 @@ class Qwen4ExpDecoderLayer(BaseOP): def __init__(self, config: ModelConfig, layer_id: int): self._layer_id = layer_id self._is_linear = config.is_linear_layer(layer_id) - dense_config = replace(config, expert_quant="none", attn_quant="none") + # Strip routed-expert quantization from the dense attention constructor, but + # preserve the explicit Qwen4 active ``attn_quant`` selection. + dense_config = replace(config, expert_quant="none") if self._is_linear: group = config.linear_attention_group() assert group is not None @@ -684,7 +699,8 @@ def __init__(self, config: ModelConfig, layer_id: int): rms_norm_eps=config.rms_norm_eps, layer_id=layer_id, expert_quant="none", - attn_quant="none", + attn_quant=config.attn_quant, + nvfp4_qkvz=(config.attn_quant == "nvfp4"), ) self.linear_attn.norm = _GatedRMSNorm( group.value_head_dim, diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 65793983d..6b60e1829 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -5,6 +5,7 @@ import safetensors import torch from freetoken.distributed import get_tp_info +from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 from freetoken.models.loader import iter_weight_files from tqdm import tqdm @@ -22,9 +23,14 @@ ".self_attn.k_proj.weight", ".self_attn.v_proj.weight", ), - ".linear_attn.in_proj.weight": ( + # Canonical runtime-state names match the explicit Qwen4 GDN split: native + # NVFP4 qkv|z and a separate BF16 b|a projection. Keeping these as two + # entries avoids a load-time dequant/re-fusion ambiguity. + ".linear_attn.in_proj_qkvz.weight": ( ".linear_attn.in_proj_qkv.weight", ".linear_attn.in_proj_z.weight", + ), + ".linear_attn.in_proj_ba.weight": ( ".linear_attn.in_proj_b.weight", ".linear_attn.in_proj_a.weight", ), @@ -107,6 +113,7 @@ def iter_weights( __all__ = [ "iter_weights", + "encode_bf16_nvfp4", "iter_weights_parallel", "load_nvfp4_expert_sources", "load_nvfp4_expert_sources_parallel", diff --git a/tests/kernels/test_qwen4_nvfp4_active.py b/tests/kernels/test_qwen4_nvfp4_active.py new file mode 100644 index 000000000..52667862b --- /dev/null +++ b/tests/kernels/test_qwen4_nvfp4_active.py @@ -0,0 +1,37 @@ +"""Bounded synthetic NVFP4 W4A16 differential tests (no model payloads).""" + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +@pytest.mark.parametrize("rows", [1, 2, 64, 65]) +def test_native_nvfp4_dense_matches_dequantized_reference(rows): + from freetoken.checkpoint.nvfp4 import decode_nvfp4, encode_bf16_nvfp4 + from freetoken.kernel.triton.nvfp4_linear import nvfp4_dense_linear + + torch.manual_seed(38038 + rows) + source = torch.randn(13, 32, dtype=torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + x = torch.randn(rows, 32, dtype=torch.bfloat16, device="cuda") + out = nvfp4_dense_linear(x, packed.cuda(), scales.cuda(), globals_.cuda()) + reference = x.float() @ decode_nvfp4(packed, scales, globals_).cuda().float().t() + torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) + assert torch.cuda.max_memory_allocated() < 6 * (1 << 30) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_native_nvfp4_state_dict_repackages_without_bf16_weight_copy(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseLinear + + torch.manual_seed(38039) + source = torch.randn(17, 32, dtype=torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + op = Nvfp4DenseLinear(32, 17) + state = {"weight": packed, "weight_scale": scales, "weight_global": globals_} + op.load_state_dict(state) + assert not state + assert op.weight.dtype == torch.int32 and op.weight.shape == (4, 17) + assert op.weight_scale.shape == (2, 17) + assert op.weight_global.dtype == torch.float16 diff --git a/tests/models/test_qwen4_exp_nvfp4_components.py b/tests/models/test_qwen4_exp_nvfp4_components.py new file mode 100644 index 000000000..c5cde438e --- /dev/null +++ b/tests/models/test_qwen4_exp_nvfp4_components.py @@ -0,0 +1,114 @@ +"""Synthetic Qwen4 active-weight NVFP4 component coverage. + +No model files are used. These tests exercise the deterministic host encoder, +canonical runtime fusion names, and the explicit GDN split without constructing +the full Qwen4 model. +""" + +from types import SimpleNamespace + +import pytest +import torch + + +def test_encoder_is_deterministic_and_round_trips_layout(): + from freetoken.checkpoint.nvfp4 import decode_nvfp4, encode_bf16_nvfp4 + + torch.manual_seed(38038) + source = torch.randn(3, 32, dtype=torch.bfloat16) + first = encode_bf16_nvfp4(source) + second = encode_bf16_nvfp4(source.clone()) + assert all(torch.equal(a, b) for a, b in zip(first, second)) + packed, scales, globals_ = first + assert packed.shape == (3, 16) and packed.dtype == torch.uint8 + assert scales.shape == (3, 2) and scales.dtype == torch.float8_e4m3fn + assert globals_.shape == (3,) and globals_.dtype == torch.float16 + assert decode_nvfp4(packed, scales, globals_).shape == source.shape + + +@pytest.mark.parametrize( + "value", + [ + 0.0, + 2.0**-133, # smallest finite BF16 subnormal + 2.0**-126, # smallest finite BF16 normal + 448.0, + -448.0, + 3.38953139e38, # largest finite BF16 (saturating policy is defined) + ], +) +def test_encoder_total_for_finite_extrema(value): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + + source = torch.tensor([[value] * 16], dtype=torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + assert torch.isfinite(globals_).all() + assert torch.isfinite(scales.view(torch.float8_e4m3fn).float()).all() + assert packed.numel() == 8 + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")]) +def test_encoder_rejects_nonfinite(bad): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + + with pytest.raises(ValueError, match="rejects NaN"): + encode_bf16_nvfp4(torch.full((1, 16), bad, dtype=torch.float32)) + + +def test_encoder_e2m1_midpoint_tie_to_even(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4, decode_nvfp4 + + # Select a single block with global=1 and block=1; row max=6 establishes + # that scale, then the midpoint pairs exercise each E2M1 tie. + values = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] + [6.0] * 9 + source = torch.tensor([values], dtype=torch.float32).to(torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + decoded = decode_nvfp4(packed, scales, globals_)[0] + # E2M1 even-code tie choices: 0, 1, 1, 2, 2, 4, 4. + assert decoded[:7].tolist() == pytest.approx([0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0]) + + +def test_qwen4_weight_fusions_use_runtime_state_names(): + from freetoken.models.qwen4_exp.weight import _try_fuse + + base = "model.layers.0.linear_attn." + buf = {} + assert _try_fuse(base + "in_proj_qkv.weight", torch.ones(2, 16), buf) == () + name, merged = _try_fuse(base + "in_proj_z.weight", torch.full((1, 16), 2.0), buf) + assert name == base + "in_proj_qkvz.weight" + assert merged.shape == (3, 16) + + buf = {} + assert _try_fuse(base + "in_proj_b.weight", torch.ones(1, 16), buf) == () + name, merged = _try_fuse(base + "in_proj_a.weight", torch.full((1, 16), 3.0), buf) + assert name == base + "in_proj_ba.weight" + assert merged[:, 0].tolist() == [1.0, 3.0] + + +def test_gdn_nvfp4_qkvz_is_explicit_opt_in(): + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear + from freetoken.layers import LinearColParallelMerged + from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + kwargs = dict( + hidden_size=32, + num_k_heads=2, + num_v_heads=2, + head_k_dim=8, + head_v_dim=8, + conv_kernel_size=4, + rms_norm_eps=1e-6, + layer_id=0, + expert_quant="none", + attn_quant="nvfp4", + ) + legacy = Qwen3_5GatedDeltaNet(**kwargs) + assert isinstance(legacy.in_proj, LinearColParallelMerged) + explicit = Qwen3_5GatedDeltaNet(**kwargs, nvfp4_qkvz=True) + assert isinstance(explicit.in_proj_qkvz, Nvfp4DenseColMerged) + assert isinstance(explicit.in_proj_ba, LinearColParallelMerged) + assert isinstance(explicit.out_proj, Nvfp4DenseLinear) From a123071d0b2b97a35a75cb334e1f740540f1b1df Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 06/17] feat(qwen4): add hardware-fit runtime components --- .../kernel/csrc/include/freetoken/tensor.h | 9 ++ .../kernel/csrc/jit/fast_index_copy.cuh | 10 +- python/freetoken/models/qwen4_exp/config.py | 14 +- python/freetoken/models/qwen4_exp/weight.py | 42 +++++ python/freetoken/moe/expert_source.py | 75 +++++---- python/freetoken/moe/offload_cache.py | 83 +++++++++- tests/checkpoint/test_q3_ple.py | 71 ++++++--- tests/kernels/test_qwen4_nvfp4_active.py | 38 ++++- tests/kernels/test_tensor_matcher.py | 150 ++++++++++++++++++ .../models/test_qwen4_exp_nvfp4_components.py | 127 +++++++++++++++ tests/models/test_qwen4_exp_raw_config.py | 24 +++ tests/moe/test_file_expert_source.py | 86 +++++++++- tests/moe/test_offload.py | 4 +- 13 files changed, 661 insertions(+), 72 deletions(-) create mode 100644 tests/kernels/test_tensor_matcher.py diff --git a/python/freetoken/kernel/csrc/include/freetoken/tensor.h b/python/freetoken/kernel/csrc/include/freetoken/tensor.h index 9b591a87d..2d0af9039 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/tensor.h +++ b/python/freetoken/kernel/csrc/include/freetoken/tensor.h @@ -397,6 +397,15 @@ struct TensorMatcher { } template + auto with_device(SymbolicDevice &device) && -> TensorMatcher && { + m_init_device(); + if constexpr (sizeof...(Codes) > 0) { + device.set_options(); + } + m_device.rebind(device); + return std::move(*this); + } + auto with_device(DeviceRef &&device) && -> TensorMatcher && { m_init_device(); m_device.rebind(*device); diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23ed..c6a471365 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -354,7 +354,7 @@ struct FastIndexCopyKernel { TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .template with_device(device) .verify(src_indices) .verify(dst_indices); @@ -363,7 +363,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .template with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -529,14 +529,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).template with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).template with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).template with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index f57f461d6..2846f1ba2 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -143,6 +143,11 @@ def parse_config(hf_config: Any) -> ModelConfig: f"detected {detected_quant!r}" ) + active_quant = getattr(hf_config, "freetoken_active_quant", None) + if active_quant not in (None, "nvfp4_w4a16_v1"): + raise ValueError(f"unsupported Qwen4 active-weight format: {active_quant}") + active_linear_quant = "nvfp4" if active_quant == "nvfp4_w4a16_v1" else "none" + return ModelConfig( num_layers=int(text.num_hidden_layers), num_qo_heads=int(text.num_attention_heads), @@ -167,10 +172,11 @@ def parse_config(hf_config: Any) -> ModelConfig: moe_enabled=True, expert_quant=expert_quant, weight_block_size=block_size, - # Only routed experts and PLE are FP8 in the official checkpoint. All - # attention, hyper-connection, and shared-expert projections stay BF16. - attn_quant="none", - dense_quant="none", + # The published experts-only checkpoint has no active-weight marker and + # therefore retains BF16 operators. Only the canonical preconverted + # FTW artifact may opt into the frozen native W4A16 map. + attn_quant=active_linear_quant, + dense_quant=active_linear_quant, lm_head_quant="none", use_qk_norm=True, # Qwen3.8-Flash-Next is a VL checkpoint. Vision is part of this model, diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 6b60e1829..7f754eb7e 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Iterable from typing import Iterator import safetensors @@ -40,6 +41,44 @@ ), } +ACTIVE_NVFP4_FORMAT = "nvfp4_w4a16_v1" +_ACTIVE_NVFP4_WEIGHT_SUFFIXES = ( + ".self_attn.qkv_proj.weight", + ".self_attn.o_proj.weight", + ".linear_attn.in_proj_qkvz.weight", + ".linear_attn.out_proj.weight", + ".attn_hyper_connection.input_mix_weight_down.weight", + ".attn_hyper_connection.input_mix_weight_up.weight", + ".mlp_hyper_connection.input_mix_weight_down.weight", + ".mlp_hyper_connection.input_mix_weight_up.weight", + ".hyper_connection_mixer.input_mix_weight_down.weight", + ".hyper_connection_mixer.input_mix_weight_up.weight", + ".mlp.shared_expert.gate_up_proj.weight", + ".mlp.shared_expert.down_proj.weight", +) + + +def is_active_nvfp4_weight(name: str) -> bool: + """Whether a fused runtime-state weight belongs to the frozen Qwen4 map.""" + + return name.endswith(_ACTIVE_NVFP4_WEIGHT_SUFFIXES) + + +def iter_active_nvfp4_runtime_entries( + entries: Iterable[tuple[str, torch.Tensor]], +) -> Iterator[tuple[str, torch.Tensor]]: + """Stream fused BF16 state into canonical native NVFP4 FTW entries.""" + + for name, tensor in entries: + if not is_active_nvfp4_weight(name): + yield name, tensor + continue + packed, scale, global_scale = encode_bf16_nvfp4(tensor) + prefix = name.removesuffix(".weight") + yield name, packed + yield prefix + ".weight_scale", scale + yield prefix + ".weight_global", global_scale + def _rename(raw_name: str) -> str | None: if raw_name.startswith("mtp."): @@ -114,6 +153,9 @@ def iter_weights( __all__ = [ "iter_weights", "encode_bf16_nvfp4", + "ACTIVE_NVFP4_FORMAT", + "is_active_nvfp4_weight", + "iter_active_nvfp4_runtime_entries", "iter_weights_parallel", "load_nvfp4_expert_sources", "load_nvfp4_expert_sources_parallel", diff --git a/python/freetoken/moe/expert_source.py b/python/freetoken/moe/expert_source.py index 2ddd311e4..9ac562fef 100644 --- a/python/freetoken/moe/expert_source.py +++ b/python/freetoken/moe/expert_source.py @@ -44,7 +44,15 @@ _cursor += _size assert _cursor == RAW_RECORD_BYTES -_HEADER_STRUCT = struct.Struct("<8sIIIIQQ16s32s32s") +_HEADER_STRUCT = struct.Struct("<8sIIIIIQQ16s32s32s") + + +def _sha256_path(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(8 << 20): + digest.update(chunk) + return digest.hexdigest() class ExpertSourceError(RuntimeError): @@ -70,6 +78,17 @@ def _dtype(name: str) -> torch.dtype: }[name] +def _plane_shape(name: str) -> tuple[int, ...]: + return { + "gate_up_packed": (1280, 1280), + "gate_up_scale": (1280, 160), + "gate_up_global": (1280,), + "down_packed": (2560, 320), + "down_scale": (2560, 40), + "down_global": (2560,), + }[name] + + class FileExpertSource: """Read fixed NVFP4 expert records from one layer sidecar. @@ -86,6 +105,10 @@ class FileExpertSource: raw_record_bytes = RAW_RECORD_BYTES num_experts = NUM_EXPERTS max_queue_depth = 16 + plane_specs = { + name: (_plane_shape(name), _dtype(dtype_name)) + for name, _size, dtype_name in PLANE_LAYOUT + } def __init__( self, @@ -93,6 +116,7 @@ def __init__( *, expected_sha256: str | None = None, expected_source_fingerprint: str | bytes | None = None, + expected_layer_id: int | None = None, num_experts: int = NUM_EXPERTS, max_queue_depth: int = 1, verify_hash: bool = True, @@ -112,7 +136,7 @@ def __init__( self._closed = False self._fd = os.open(str(self.path), os.O_RDONLY | getattr(os, "O_BINARY", 0)) try: - self._validate_file(expected_source_fingerprint) + self._validate_file(expected_source_fingerprint, expected_layer_id) self.sha256 = self._hash_file() if verify_hash else None if expected_sha256 is not None: expected_sha256 = expected_sha256.lower() @@ -138,6 +162,7 @@ def create_synthetic( num_experts: int = NUM_EXPERTS, records: Iterable[bytes] | None = None, source_fingerprint: bytes | None = None, + layer_id: int = 0, ) -> str: """Create a tiny deterministic sidecar for tests (never model data).""" path = _z_path(path) @@ -158,6 +183,7 @@ def create_synthetic( HEADER_BYTES, num_experts, len(PLANE_LAYOUT), + int(layer_id), RAW_RECORD_BYTES, RECORD_BYTES, b"nvfp4-qwen4-v1\0\0", # 16-byte layout tag @@ -176,25 +202,35 @@ def create_synthetic( with path.open("r+b") as fh: fh.seek(0) header = bytearray(fh.read(HEADER_BYTES)) - header[struct.calcsize("<8sIIIIQQ16s32s") : _HEADER_STRUCT.size] = digest + header[_HEADER_STRUCT.size - 32 : _HEADER_STRUCT.size] = digest fh.seek(0) fh.write(header) - return hashlib.sha256(path.read_bytes()).hexdigest() + # Keep synthetic native-geometry fixtures bounded: the real 512-record + # shape is roughly 1.42 GB and must never be read into one Python bytes + # object merely to produce its verification hash. + return _sha256_path(path) - def _validate_file(self, expected_source_fingerprint: str | bytes | None) -> None: + def _validate_file( + self, expected_source_fingerprint: str | bytes | None, expected_layer_id: int | None + ) -> None: size = self.path.stat().st_size expected_size = HEADER_BYTES + self.num_experts * self.record_bytes header = self._read_exact(HEADER_BYTES, 0) if len(header) != HEADER_BYTES: raise ExpertSourceError("truncated expert tier header") try: - magic, version, hbytes, experts, planes, raw_bytes, rec_bytes, tag, fingerprint, payload_hash = _HEADER_STRUCT.unpack_from(header) + magic, version, hbytes, experts, planes, layer_id, raw_bytes, rec_bytes, tag, fingerprint, payload_hash = _HEADER_STRUCT.unpack_from(header) except struct.error as exc: raise ExpertSourceError("malformed expert tier header") from exc if magic != MAGIC or version != VERSION or hbytes != HEADER_BYTES: raise ExpertSourceError("unsupported expert tier magic/version/header") if experts != self.num_experts or planes != len(PLANE_LAYOUT): raise ExpertSourceError("expert tier geometry mismatch") + self.layer_id = int(layer_id) + if expected_layer_id is not None and self.layer_id != int(expected_layer_id): + raise ExpertSourceError( + f"expert tier layer mismatch: {self.layer_id} != {int(expected_layer_id)}" + ) if raw_bytes != RAW_RECORD_BYTES or rec_bytes != RECORD_BYTES or tag.rstrip(b"\0") != b"nvfp4-qwen4-v1": raise ExpertSourceError("expert tier layout mismatch") if size != expected_size: @@ -231,14 +267,7 @@ def _read_exact(self, size: int, offset: int) -> bytes: return data def _hash_file(self) -> str: - h = hashlib.sha256() - with self.path.open("rb") as fh: - while True: - block = fh.read(8 << 20) - if not block: - break - h.update(block) - return h.hexdigest() + return _sha256_path(self.path) def _record_bytes(self, expert_id: int) -> bytes: if self._closed: @@ -265,19 +294,11 @@ def read_record(self, expert_id: int) -> dict[str, torch.Tensor]: for name, size, dtype_name in PLANE_LAYOUT: offset = _PLANE_OFFSETS[name] dtype = _dtype(dtype_name) - if name == "gate_up_packed": - shape = (1280, 1280) - elif name == "down_packed": - shape = (2560, 320) - elif name == "gate_up_scale": - shape = (1280, 160) - elif name == "gate_up_global": - shape = (1280,) - elif name == "down_scale": - shape = (2560, 40) - else: - shape = (2560,) - out[name] = torch.frombuffer(memoryview(raw)[offset : offset + size], dtype=dtype).clone().reshape(shape) + shape = _plane_shape(name) + # bytearray supplies a writable, record-local staging buffer and + # avoids exposing Python's immutable bytes through a writable tensor. + staging = bytearray(raw[offset : offset + size]) + out[name] = torch.frombuffer(staging, dtype=dtype).reshape(shape) return out def read_records(self, expert_ids: Iterable[int], *, max_concurrency: int = 1) -> list[dict[str, torch.Tensor]]: diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index aa588c8b6..35c536432 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -203,6 +203,7 @@ def __post_init__(self) -> None: self.file_sources = {} self._pending_file_fetches: list[tuple[int, int]] = [] self._pending_file_materialize = False + self._pending_file_rollback = None # per-layer host residency: the GPU movement paths require "pinned"; LOCKED/PAGEABLE layers decode on the CPU executor and prefill via copy_missing's pageable branch # _unpinned_layers is the derived id set the hot paths test against self.layer_residency: list[str] = [] @@ -368,8 +369,28 @@ def set_file_sources(self, sources: dict[int, object]) -> None: raise ValueError("file-backed expert tiers are GPU-only; use resident HostBanks for CPU/hybrid") if self.prefill_overlap: raise ValueError("file-backed expert tiers require prefill_overlap=False") + if not sources: + return if not self.bank_sources: - raise ValueError("set_bank_sources must allocate cache planes before file tiers") + # A pure file-tier fixture (and a future all-file policy) has no + # resident HostBank from which to infer cache-plane geometry. The + # source contract supplies the exact per-expert native plane specs. + representative = next(iter(sources.values())) + if tuple(getattr(representative, "bank_schema", ())) != tuple(self.bank_schema): + raise ValueError("file tier bank schema does not match cache quant_format") + specs = getattr(representative, "plane_specs", None) + if not isinstance(specs, dict): + raise ValueError("file source does not declare native plane_specs") + for name in self.bank_schema: + shape, dtype = specs[name] + self.bank_sources[name] = [None] * self.num_layers + self.bank_caches[name] = torch.empty( + (self.cache_size, *shape), dtype=dtype, device=self.device + ) + self.banks = [ + (self.bank_sources[name], self.bank_caches[name]) for name in self.bank_schema + ] + self._build_copy_plan() for layer_id, source in sources.items(): layer_id = int(layer_id) if not 0 <= layer_id < self.num_layers: @@ -378,6 +399,10 @@ def set_file_sources(self, sources: dict[int, object]) -> None: raise ValueError("file tier bank schema does not match cache quant_format") if int(source.num_experts) != self.num_experts: raise ValueError("file tier expert count does not match cache geometry") + if int(getattr(source, "layer_id", -1)) != layer_id: + raise ValueError( + f"file tier declares layer {getattr(source, 'layer_id', None)}, registered as {layer_id}" + ) for name in self.bank_schema: if self.bank_sources[name][layer_id] is not None: raise ValueError(f"layer {layer_id} already has a resident HostBank for {name}") @@ -505,9 +530,14 @@ def rebuild(self, cache_size: int) -> None: (source for source in self.bank_sources[name] if source is not None), None ) if head is None: - raise ValueError(f"bank {name!r} has no resident shape for rebuild") + representative = next(iter(self.file_sources.values()), None) + if representative is None: + raise ValueError(f"bank {name!r} has no source shape for rebuild") + shape, dtype = representative.plane_specs[name] + else: + shape, dtype = head.shape[1:], head.dtype self.bank_caches[name] = torch.empty( - (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device + (cache_size, *shape), dtype=dtype, device=self.device ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs @@ -881,6 +911,15 @@ def ensure_file_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: if self.decode_target != "gpu": raise RuntimeError("file-backed experts cannot serve CPU/hybrid decode") original_shape = tuple(expert_ids.shape) + # File reads happen after the host-side slot plan is installed. Retain + # the small bookkeeping tensors so an I/O error cannot expose a cache + # slot whose six planes were never completely populated. + rollback = ( + self.slot_for_id.clone(), + self.id_of_slot.clone(), + self.usage.clone(), + self.step.clone(), + ) raw_ids = [int(value) for value in expert_ids.detach().cpu().reshape(-1).tolist()] self.step += 1 step = int(self.step.item()) @@ -912,6 +951,7 @@ def ensure_file_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: # keep num_indices truthful for diagnostics and callers that inspect it. self.num_indices[0] = len(pending) self._pending_file_fetches = pending + self._pending_file_rollback = rollback if pending else None replacement = torch.tensor(mapped, dtype=expert_ids.dtype, device=expert_ids.device).reshape(original_shape) expert_ids.copy_(replacement) @@ -938,6 +978,25 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None def materialize_layer(self, layer_id: int) -> None: if layer_id in self.file_sources: + if self.cache_size < self.num_experts: + raise RuntimeError("file-tier prefill requires one slot per expert") + self._pending_file_rollback = ( + self.slot_for_id.clone(), + self.id_of_slot.clone(), + self.usage.clone(), + self.step.clone(), + ) + self.step += 1 + step = int(self.step.item()) + self.slot_for_id[layer_id].fill_(-1) + for expert_id in range(self.num_experts): + slot = expert_id + old = int(self.id_of_slot[slot].item()) + if old >= 0: + self.slot_for_id.view(-1)[old] = -1 + self.slot_for_id[layer_id, expert_id] = slot + self.id_of_slot[slot] = layer_id * self.num_experts + expert_id + self.usage[slot] = step self._pending_src_layer = layer_id self._pending_whole_layer = True self._pending_file_materialize = True @@ -968,6 +1027,7 @@ def reset(self) -> None: self.expert_recency.fill_(-1) self._pending_file_fetches = [] self._pending_file_materialize = False + self._pending_file_rollback = None def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -1102,10 +1162,23 @@ def copy_missing(self) -> None: destinations = { name: cache for name, (_, cache) in zip(self.bank_schema, self.banks) } - for slot, expert_id in pairs: - source.read_into(expert_id, destinations, slot) + try: + for slot, expert_id in pairs: + source.read_into(expert_id, destinations, slot) + except Exception: + if self._pending_file_rollback is not None: + slot_for_id, id_of_slot, usage, step = self._pending_file_rollback + self.slot_for_id.copy_(slot_for_id) + self.id_of_slot.copy_(id_of_slot) + self.usage.copy_(usage) + self.step.copy_(step) + self._pending_file_fetches = [] + self._pending_file_materialize = False + self._pending_file_rollback = None + raise self._pending_file_fetches = [] self._pending_file_materialize = False + self._pending_file_rollback = None return if layer_id in self._unpinned_layers: if not self._pending_whole_layer: diff --git a/tests/checkpoint/test_q3_ple.py b/tests/checkpoint/test_q3_ple.py index 76e6d88a9..004bd94b0 100644 --- a/tests/checkpoint/test_q3_ple.py +++ b/tests/checkpoint/test_q3_ple.py @@ -3,8 +3,10 @@ import hashlib import json import os +import shutil import sys from pathlib import Path +from uuid import uuid4 import pytest import torch @@ -22,23 +24,23 @@ from q3_ple_32_reference import dequantize_row, encode_table -FIXTURE_ROOT = ROOT / "artifacts" / "stage6_q3ple_fixture" -MANIFEST_PATH = FIXTURE_ROOT / "ple-q3.json" -DATA_PATH = FIXTURE_ROOT / "ple-q3-000.bin" - - -def _write_fixture() -> tuple[Path, Path, int]: - FIXTURE_ROOT.mkdir(parents=True, exist_ok=True) +def _write_fixture(fixture_root: Path) -> tuple[Path, Path, int]: + fixture_root.mkdir(parents=True, exist_ok=True) + manifest_path = fixture_root / "ple-q3.json" + data_path = fixture_root / "ple-q3-000.bin" rows = [] - for row in range(12): + for row in range(20): rows.append([((row + 1) * 0.125) * ((i % 17) - 8) for i in range(ROW_VALUES)]) + # Explicit zero/extreme rows exercise every one of the five block decoders. + rows[0] = [0.0] * ROW_VALUES + rows[-1] = [(-1.0 if i & 1 else 1.0) * 448.0 for i in range(ROW_VALUES)] encoded = encode_table(rows, refinement_passes=2, scale_dtype="bf16") - split = 5 * ROW_BYTES + split = 9 * ROW_BYTES second_offset = ALIGN payload = encoded[:split] + bytes(second_offset - split) + encoded[split:] - DATA_PATH.write_bytes(payload) + data_path.write_bytes(payload) segments = [] - for first, end, offset in ((0, 5, 0), (5, 12, second_offset)): + for first, end, offset in ((0, 9, 0), (9, 20, second_offset)): segment_bytes = encoded[first * ROW_BYTES : end * ROW_BYTES] segments.append( { @@ -60,18 +62,23 @@ def _write_fixture() -> tuple[Path, Path, int]: "rows": len(rows), "payload_bytes": len(encoded), "file_bytes": len(payload), - "data_file": DATA_PATH.name, + "data_file": data_path.name, "weight_scale": 1.25, "sha256": hashlib.sha256(payload).hexdigest(), "segments": segments, } - MANIFEST_PATH.write_text(json.dumps(manifest, indent=2), encoding="utf-8") - return MANIFEST_PATH, DATA_PATH, len(rows) + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return manifest_path, data_path, len(rows) @pytest.fixture(scope="module") def fixture_paths(): - return _write_fixture() + root = ROOT / ".stage6-test-fixtures" / uuid4().hex + assert root.drive.upper() == "Z:" + try: + yield _write_fixture(root) + finally: + shutil.rmtree(root, ignore_errors=True) def test_q3_reader_constants_and_ordered_gather(fixture_paths): @@ -80,7 +87,7 @@ def test_q3_reader_constants_and_ordered_gather(fixture_paths): assert reader.row_count == row_count assert reader.total_payload_bytes == row_count * 70 assert reader.gather([11, 0, 11]).shape == (3, ROW_VALUES) - scaled = reader.gather16(list(range(16)) if row_count >= 16 else [0] * 16, apply_weight_scale=True) + scaled = reader.gather16(list(range(16)), apply_weight_scale=True) assert scaled.dtype == torch.bfloat16 assert scaled.shape == (16, ROW_VALUES) assert torch.equal(scaled[0], reader.gather([0], apply_weight_scale=True)[0]) @@ -91,9 +98,9 @@ def test_q3_reader_matches_authoritative_codec(fixture_paths): # The fixture uses two segments and includes alignment padding between them. encoded = data.read_bytes() with Q3PLEReader(manifest) as reader: - for row in (0, 1, 4, 5, 6, 11): + for row in (0, 1, 8, 9, 10, 19): raw_row = (encoded[row * ROW_BYTES : (row + 1) * ROW_BYTES] - if row < 5 else encoded[ALIGN + (row - 5) * ROW_BYTES : ALIGN + (row - 4) * ROW_BYTES]) + if row < 9 else encoded[ALIGN + (row - 9) * ROW_BYTES : ALIGN + (row - 8) * ROW_BYTES]) expected = torch.tensor( dequantize_row(raw_row), dtype=torch.bfloat16, @@ -101,6 +108,21 @@ def test_q3_reader_matches_authoritative_codec(fixture_paths): assert torch.equal(reader.gather([row])[0], expected) +def test_q3_reader_all_block_boundaries_and_random_order(fixture_paths): + manifest, data, _ = fixture_paths + encoded = data.read_bytes() + with Q3PLEReader(manifest) as reader: + rows = reader.gather([19, 9, 0, 19, 8]) + assert rows.shape == (5, ROW_VALUES) + assert torch.equal(rows[0], rows[3]) + for row_index, row in zip((19, 9, 0, 19, 8), rows): + offset = row_index * ROW_BYTES if row_index < 9 else ALIGN + (row_index - 9) * ROW_BYTES + expected = torch.tensor(dequantize_row(encoded[offset : offset + ROW_BYTES]), dtype=torch.bfloat16) + assert torch.equal(row, expected) + for boundary in (0, 31, 32, 63, 64, 95, 96, 127, 128, 159): + assert row[boundary] == expected[boundary] + + @pytest.mark.parametrize( "field,value", [("version", 2), ("endianness", "big"), ("row_bytes", 69), ("sha256", "0" * 64)], @@ -145,6 +167,19 @@ def test_q3_reader_rejects_gap_overlap_and_truncation(fixture_paths): truncated.unlink(missing_ok=True) +def test_q3_reader_rejects_corrupt_segment_hash(fixture_paths): + manifest, _, _ = fixture_paths + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw["segments"][1]["sha256"] = "f" * 64 + bad = manifest.with_name("bad-segment-hash.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError, match="segment hash mismatch"): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + + def test_q3_reader_requires_z_backing(): with pytest.raises((ValueError, FileNotFoundError)): Q3PLEReader("C:\\q3-ple\\ple-q3.json") diff --git a/tests/kernels/test_qwen4_nvfp4_active.py b/tests/kernels/test_qwen4_nvfp4_active.py index 52667862b..350481e81 100644 --- a/tests/kernels/test_qwen4_nvfp4_active.py +++ b/tests/kernels/test_qwen4_nvfp4_active.py @@ -5,18 +5,40 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +@pytest.mark.parametrize( + "out_features,in_features", + [ + (320, 10240), # mHC down + (10240, 320), # mHC up + (1280, 2560), # shared gate|up + (2560, 640), # shared down + (13312, 2560), # QSA q|k|v + (2560, 6144), # QSA/GDN output + (16384, 2560), # GDN qkv|z + ], +) @pytest.mark.parametrize("rows", [1, 2, 64, 65]) -def test_native_nvfp4_dense_matches_dequantized_reference(rows): - from freetoken.checkpoint.nvfp4 import decode_nvfp4, encode_bf16_nvfp4 +def test_native_nvfp4_dense_matches_dequantized_reference(rows, out_features, in_features): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 from freetoken.kernel.triton.nvfp4_linear import nvfp4_dense_linear + from freetoken.kernel.triton.nvfp4_dequant import dequant_nvfp4 - torch.manual_seed(38038 + rows) - source = torch.randn(13, 32, dtype=torch.bfloat16) + torch.manual_seed(38038 + rows + out_features) + source = torch.randn(out_features, in_features, dtype=torch.bfloat16) packed, scales, globals_ = encode_bf16_nvfp4(source) - x = torch.randn(rows, 32, dtype=torch.bfloat16, device="cuda") - out = nvfp4_dense_linear(x, packed.cuda(), scales.cuda(), globals_.cuda()) - reference = x.float() @ decode_nvfp4(packed, scales, globals_).cuda().float().t() - torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) + x = torch.randn(rows, in_features, dtype=torch.bfloat16, device="cuda") + packed_cuda, scales_cuda, globals_cuda = packed.cuda(), scales.cuda(), globals_.cuda() + out = nvfp4_dense_linear(x, packed_cuda, scales_cuda, globals_cuda) + weight = dequant_nvfp4( + packed_cuda.unsqueeze(0), scales_cuda.unsqueeze(0), globals_cuda.unsqueeze(0), + torch.zeros(1, dtype=torch.int32, device="cuda"), dtype=torch.bfloat16, + )[0] + reference = x @ weight.t() + # Reuse the pre-existing native NVFP4 backend tolerance verbatim: BF16 + # grouped/dense GEMMs accumulate large reductions, so the absolute bound is + # relative to this fixture's output magnitude. + atol = 0.03 * float(reference.abs().max()) + torch.testing.assert_close(out.float(), reference.float(), rtol=3e-2, atol=atol) assert torch.cuda.max_memory_allocated() < 6 * (1 << 30) diff --git a/tests/kernels/test_tensor_matcher.py b/tests/kernels/test_tensor_matcher.py new file mode 100644 index 000000000..a245aa0d2 --- /dev/null +++ b/tests/kernels/test_tensor_matcher.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +from functools import lru_cache +import os +from pathlib import Path +import sys + +import pytest +import torch + + +_SOURCE = r""" +#include +#include + +int symbolic_cuda_same(tvm::ffi::TensorView first, tvm::ffi::TensorView second) { + auto device = host::SymbolicDevice{}; + host::TensorMatcher({-1}) + .with_device(device) + .verify(first) + .verify(second); + return device.unwrap().device_id; +} + +int symbolic_cuda_same_typed(tvm::ffi::TensorView first, + tvm::ffi::TensorView second) { + auto length = host::SymbolicSize{"length"}; + auto dtype = host::SymbolicDType{}; + auto device = host::SymbolicDevice{}; + host::TensorMatcher({length}) + .with_dtype(dtype) + .with_device(device) + .verify(first) + .verify(second); + return device.unwrap().device_id; +} + +template struct SymbolicCudaTemplate { + static int run(tvm::ffi::TensorView first, tvm::ffi::TensorView second) { + auto length = host::SymbolicSize{"length"}; + auto dtype = host::SymbolicDType{}; + auto device = host::SymbolicDevice{}; + host::TensorMatcher({length}) + .with_dtype(dtype) + .template with_device(device) + .verify(first) + .verify(second); + return device.unwrap().device_id + Tag - Tag; + } +}; + +int symbolic_cuda_same_templated(tvm::ffi::TensorView first, + tvm::ffi::TensorView second) { + return SymbolicCudaTemplate<1>::run(first, second); +} + +int symbolic_unrestricted(tvm::ffi::TensorView value) { + auto device = host::SymbolicDevice{}; + host::TensorMatcher({-1}).with_device(device).verify(value); + return static_cast(device.unwrap().device_type); +} + +void fixed_cpu(tvm::ffi::TensorView value) { + host::TensorMatcher({-1}).with_device().verify(value); +} + +void explicit_cpu(tvm::ffi::TensorView value) { + host::TensorMatcher({-1}).with_device({{kDLCPU, 0}}).verify(value); +} + +void reject_different_cuda_device_ids() { + auto device = host::SymbolicDevice{}; + device.set_options(); + device.verify({kDLCUDA, 0}); + device.verify({kDLCUDA, 1}); +} +""" + +_FUNCTIONS = [ + "symbolic_cuda_same", + "symbolic_cuda_same_typed", + "symbolic_cuda_same_templated", + "symbolic_unrestricted", + "fixed_cpu", + "explicit_cpu", + "reject_different_cuda_device_ids", +] + + +@lru_cache(maxsize=1) +def _cpu_module(): + from freetoken.kernel.utils import DEFAULT_CFLAGS, DEFAULT_INCLUDE + from tvm_ffi.cpp import load_inline + + return load_inline( + "freetoken_tensor_matcher_cpu_test_v7", + cpp_sources=_SOURCE, + functions=_FUNCTIONS, + extra_cflags=DEFAULT_CFLAGS, + extra_include_paths=DEFAULT_INCLUDE, + ) + + +@lru_cache(maxsize=1) +def _cuda_module(): + from freetoken.kernel.utils import DEFAULT_INCLUDE, _cuda_cflags + from tvm_ffi.cpp import load_inline + + extra_ldflags = [] + if sys.platform == "win32": + cuda_home = Path(os.environ["CUDA_HOME"]) + extra_ldflags = [f"/LIBPATH:{cuda_home / 'lib' / 'x64'}", "cudart.lib"] + + return load_inline( + "freetoken_tensor_matcher_cuda_test_v7", + cuda_sources=_SOURCE, + functions=_FUNCTIONS, + extra_cuda_cflags=_cuda_cflags([]), + extra_ldflags=extra_ldflags, + extra_include_paths=DEFAULT_INCLUDE, + backend="cuda", + ) + + +def test_symbolic_device_restrictions_and_existing_cpu_paths(): + module = _cpu_module() + value = torch.ones(4) + + assert module.symbolic_unrestricted(value) == 1 # DLPack kDLCPU + module.fixed_cpu(value) + module.explicit_cpu(value) + with pytest.raises(Exception, match="Device"): + module.symbolic_cuda_same(value, value) + with pytest.raises(Exception, match="Device"): + module.reject_different_cuda_device_ids() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_symbolic_cuda_binding_compiles_and_checks_same_device(): + module = _cuda_module() + first = torch.ones(4, device="cuda") + second = torch.zeros(4, device="cuda") + + assert module.symbolic_cuda_same(first, second) == torch.cuda.current_device() + first_int = torch.ones(4, dtype=torch.int32, device="cuda") + second_int = torch.zeros(4, dtype=torch.int32, device="cuda") + assert module.symbolic_cuda_same_typed(first_int, second_int) == torch.cuda.current_device() + assert module.symbolic_cuda_same_templated(first_int, second_int) == torch.cuda.current_device() + with pytest.raises(Exception, match="Device"): + module.symbolic_cuda_same(torch.ones(4), torch.zeros(4)) diff --git a/tests/models/test_qwen4_exp_nvfp4_components.py b/tests/models/test_qwen4_exp_nvfp4_components.py index c5cde438e..3efc69d95 100644 --- a/tests/models/test_qwen4_exp_nvfp4_components.py +++ b/tests/models/test_qwen4_exp_nvfp4_components.py @@ -6,11 +6,48 @@ """ from types import SimpleNamespace +import math import pytest import torch +def _independent_nvfp4_reference(source: torch.Tensor): + """Tiny scalar oracle implementing the documented format, not FreeToken helpers.""" + + grid = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + packed_rows, scale_rows, globals_out = [], [], [] + for row in source.float().tolist(): + row_max = max(abs(value) for value in row) + if row_max == 0: + global_scale = 1.0 + else: + target = min(max(row_max / 6.0, 2.0**-24), 65504.0) + global_scale = float(torch.tensor(target, dtype=torch.float16)) + if global_scale == 0: + global_scale = 2.0**-24 + globals_out.append(global_scale) + codes, block_scales = [], [] + for start in range(0, len(row), 16): + block = row[start : start + 16] + block_max = max(abs(value) for value in block) + target = 0.0 if block_max == 0 else min(block_max / (6.0 * global_scale), 448.0) + scale = float(torch.tensor(target, dtype=torch.float8_e4m3fn)) + block_scales.append(scale) + for value in block: + normalized = 0.0 if scale == 0 else max(-6.0, min(6.0, value / (scale * global_scale))) + magnitude = abs(normalized) + code = min(range(8), key=lambda item: (abs(grid[item] - magnitude), item & 1)) + codes.append(code | (8 if normalized < 0 else 0)) + packed_rows.append([codes[i] | (codes[i + 1] << 4) for i in range(0, len(codes), 2)]) + scale_rows.append(block_scales) + return ( + torch.tensor(packed_rows, dtype=torch.uint8), + torch.tensor(scale_rows, dtype=torch.float8_e4m3fn), + torch.tensor(globals_out, dtype=torch.float16), + ) + + def test_encoder_is_deterministic_and_round_trips_layout(): from freetoken.checkpoint.nvfp4 import decode_nvfp4, encode_bf16_nvfp4 @@ -26,6 +63,16 @@ def test_encoder_is_deterministic_and_round_trips_layout(): assert decode_nvfp4(packed, scales, globals_).shape == source.shape +def test_encoder_bytes_match_independent_scalar_oracle(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + + torch.manual_seed(38038) + source = torch.randn(4, 32, dtype=torch.bfloat16) + actual = encode_bf16_nvfp4(source) + expected = _independent_nvfp4_reference(source) + assert all(torch.equal(a, b) for a, b in zip(actual, expected)) + + @pytest.mark.parametrize( "value", [ @@ -85,6 +132,39 @@ def test_qwen4_weight_fusions_use_runtime_state_names(): assert merged[:, 0].tolist() == [1.0, 3.0] +def test_active_converter_emits_canonical_runtime_triple_and_preserves_slices(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + q = torch.full((2, 16), 1.0, dtype=torch.bfloat16) + k = torch.full((1, 16), 2.0, dtype=torch.bfloat16) + v = torch.full((1, 16), -3.0, dtype=torch.bfloat16) + name = "model.layers.0.self_attn.qkv_proj.weight" + emitted = dict(iter_active_nvfp4_runtime_entries([(name, torch.cat((q, k, v), dim=0))])) + prefix = name.removesuffix(".weight") + assert list(emitted) == [name, prefix + ".weight_scale", prefix + ".weight_global"] + packed, scales, globals_ = emitted[name], emitted[prefix + ".weight_scale"], emitted[prefix + ".weight_global"] + cursor = 0 + for constituent in (q, k, v): + expected = encode_bf16_nvfp4(constituent) + end = cursor + constituent.shape[0] + assert torch.equal(packed[cursor:end], expected[0]) + assert torch.equal(scales[cursor:end], expected[1]) + assert torch.equal(globals_[cursor:end], expected[2]) + cursor = end + + +def test_active_converter_protects_non_map_tensors(): + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + protected = torch.randn(3, 16, dtype=torch.bfloat16) + rows = list(iter_active_nvfp4_runtime_entries([ + ("model.layers.0.self_attn.index_qk_proj.weight", protected), + ])) + assert len(rows) == 1 and rows[0][0].endswith("index_qk_proj.weight") + assert rows[0][1] is protected + + def test_gdn_nvfp4_qkvz_is_explicit_opt_in(): from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear from freetoken.layers import LinearColParallelMerged @@ -112,3 +192,50 @@ def test_gdn_nvfp4_qkvz_is_explicit_opt_in(): assert isinstance(explicit.in_proj_qkvz, Nvfp4DenseColMerged) assert isinstance(explicit.in_proj_ba, LinearColParallelMerged) assert isinstance(explicit.out_proj, Nvfp4DenseLinear) + + +def test_qwen4_frozen_operator_map_and_protected_linears(): + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear + from freetoken.layers import LinearReplicated + from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention + from freetoken.models.qwen4_exp.model import _GatedResidual, _SharedExpert + + rotary = SimpleNamespace(rotary_dim=128, max_position=128, base=10_000.0, scaling=None) + attention_config = SimpleNamespace( + head_dim=128, num_qo_heads=2, num_kv_heads=1, hidden_size=256, + rms_norm_eps=1e-6, rotary_config=rotary, expert_quant="none", attn_quant="nvfp4", + ) + attention = Qwen3_5Attention(attention_config, 0) + assert isinstance(attention.qkv_proj, Nvfp4DenseColMerged) + assert isinstance(attention.o_proj, Nvfp4DenseLinear) + + qwen_config = SimpleNamespace( + hidden_size=32, rms_norm_eps=1e-6, dense_quant="nvfp4", + shared_expert_intermediate_size=16, + qwen4_args=SimpleNamespace(hc_count=4, hc_lowrank=16), + ) + residual = _GatedResidual(qwen_config, combine=True) + assert isinstance(residual.input_mix_weight_down, Nvfp4DenseLinear) + assert isinstance(residual.input_mix_weight_up, Nvfp4DenseLinear) + assert isinstance(residual.block_inject_weight, LinearReplicated) + shared = _SharedExpert(qwen_config) + assert isinstance(shared.gate_up_proj, Nvfp4DenseColMerged) + assert isinstance(shared.down_proj, Nvfp4DenseLinear) + + +def test_legacy_gdn_quant_modes_keep_their_original_dispatch(): + from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged + from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged + from freetoken.layers import LinearColParallelMerged + from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet + + base = dict( + hidden_size=256, num_k_heads=1, num_v_heads=1, head_k_dim=128, + head_v_dim=128, conv_kernel_size=4, rms_norm_eps=1e-6, layer_id=0, + ) + bf16 = Qwen3_5GatedDeltaNet(**base, expert_quant="none", attn_quant="none") + assert isinstance(bf16.in_proj, LinearColParallelMerged) + block = Qwen3_5GatedDeltaNet(**base, expert_quant="fp8_block", attn_quant="none") + assert isinstance(block.in_proj_qkvz, Fp8BlockColMerged) + pertensor = Qwen3_5GatedDeltaNet(**base, expert_quant="none", attn_quant="fp8_pertensor") + assert isinstance(pertensor.in_proj_qkvz, Fp8PerTensorColMerged) diff --git a/tests/models/test_qwen4_exp_raw_config.py b/tests/models/test_qwen4_exp_raw_config.py index 81745aa9d..966917cd0 100644 --- a/tests/models/test_qwen4_exp_raw_config.py +++ b/tests/models/test_qwen4_exp_raw_config.py @@ -76,3 +76,27 @@ def test_qwen4_raw_config_uses_official_topk_normalization_default(): assert config.norm_topk_prob is True assert config.rotary_config.max_position == 262_144 assert config.attn_type_for_layer(3).value == "qsa" + + +def test_qwen4_active_nvfp4_requires_explicit_artifact_marker(): + published = _raw_checkpoint_config() + config = parse_config(published) + assert config.attn_quant == "none" + assert config.dense_quant == "none" + + converted = _raw_checkpoint_config() + converted.freetoken_active_quant = "nvfp4_w4a16_v1" + config = parse_config(converted) + assert config.attn_quant == "nvfp4" + assert config.dense_quant == "nvfp4" + + +def test_qwen4_rejects_unknown_active_weight_marker(): + config = _raw_checkpoint_config() + config.freetoken_active_quant = "ambiguous-q8" + try: + parse_config(config) + except ValueError as exc: + assert "unsupported Qwen4 active-weight format" in str(exc) + else: + raise AssertionError("unknown active-weight marker was accepted") diff --git a/tests/moe/test_file_expert_source.py b/tests/moe/test_file_expert_source.py index 39b9e9130..226e61b6e 100644 --- a/tests/moe/test_file_expert_source.py +++ b/tests/moe/test_file_expert_source.py @@ -64,8 +64,8 @@ def test_file_expert_source_reads_exact_planes_and_rejects_tamper(z_fixture_dir) def test_file_expert_source_cache_miss_fills_slots_without_host_layer(z_fixture_dir): path = z_fixture_dir / "experts-L01.nvfp4" - digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 11) for i in range(3)]) - src = FileExpertSource(path, num_experts=3, expected_sha256=digest) + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 11) for i in range(3)], layer_id=1) + src = FileExpertSource(path, num_experts=3, expected_sha256=digest, expected_layer_id=1) try: from freetoken.moe.offload_cache import OffloadMoeCache @@ -97,8 +97,8 @@ def test_file_expert_source_cache_miss_fills_slots_without_host_layer(z_fixture_ def test_file_tier_materialize_streams_complete_layer(z_fixture_dir): path = z_fixture_dir / "experts-L01.nvfp4" - digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 21) for i in range(3)]) - src = FileExpertSource(path, num_experts=3, expected_sha256=digest) + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 21) for i in range(3)], layer_id=1) + src = FileExpertSource(path, num_experts=3, expected_sha256=digest, expected_layer_id=1) try: from freetoken.moe.offload_cache import OffloadMoeCache @@ -110,6 +110,84 @@ def test_file_tier_materialize_streams_complete_layer(z_fixture_dir): values = cache.bank_caches["gate_up_packed"][:, 0, 0].tolist() assert values == [21, 22, 23] assert src.read_count == 3 + assert cache.slot_for_id[1].tolist() == [0, 1, 2] + assert cache.id_of_slot.tolist() == [3, 4, 5] + finally: + src.close() + + +def test_file_only_cache_derives_planes_without_any_hostbank(z_fixture_dir): + path = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=2, records=[_record(41), _record(42)]) + src = FileExpertSource(path, num_experts=2, expected_sha256=digest) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(1, 2, 2, torch.device("cpu"), quant_format="nvfp4") + cache.set_file_sources({0: src}) + assert all(all(item is None for item in layers) for layers in cache.bank_sources.values()) + ids = torch.tensor([1], dtype=torch.int32) + cache.ensure_experts(0, ids) + cache.copy_missing() + assert cache.bank_caches["gate_up_packed"][0, 0, 0].item() == 42 + finally: + src.close() + + +def test_file_only_multilayer_identity_eviction_and_reset(z_fixture_dir): + path0 = z_fixture_dir / "experts-L00.nvfp4" + path1 = z_fixture_dir / "experts-L01.nvfp4" + digest0 = FileExpertSource.create_synthetic(path0, num_experts=2, records=[_record(51), _record(52)], layer_id=0) + digest1 = FileExpertSource.create_synthetic(path1, num_experts=2, records=[_record(61), _record(62)], layer_id=1) + src0 = FileExpertSource(path0, num_experts=2, expected_sha256=digest0, expected_layer_id=0) + src1 = FileExpertSource(path1, num_experts=2, expected_sha256=digest1, expected_layer_id=1) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(2, 2, 2, torch.device("cpu"), quant_format="nvfp4") + cache.set_file_sources({0: src0, 1: src1}) + ids0 = torch.tensor([0, 1], dtype=torch.int32) + cache.ensure_experts(0, ids0) + cache.copy_missing() + ids1 = torch.tensor([0], dtype=torch.int32) + cache.ensure_experts(1, ids1) + cache.copy_missing() + slot = int(ids1.item()) + assert cache.slot_for_id[1, 0].item() == slot + assert cache.slot_for_id[0, slot].item() == -1 + assert cache.bank_caches["gate_up_packed"][slot, 0, 0].item() == 61 + cache.reset() + assert (cache.slot_for_id == -1).all() + assert (cache.id_of_slot == -1).all() + finally: + src0.close() + src1.close() + + +def test_file_tier_read_failure_rolls_back_slot_bookkeeping(z_fixture_dir, monkeypatch): + path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=2, records=[_record(31), _record(32)], layer_id=1) + src = FileExpertSource(path, num_experts=2, expected_sha256=digest, expected_layer_id=1) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(2, 2, 2, torch.device("cpu"), quant_format="nvfp4") + cache.set_bank_sources(_resident_banks(2, 2)) + cache.set_file_sources({1: src}) + before_slots = cache.slot_for_id.clone() + before_ids = cache.id_of_slot.clone() + ids = torch.tensor([1], dtype=torch.int32) + cache.ensure_experts(1, ids) + + def fail_read(*_args, **_kwargs): + raise ExpertSourceError("synthetic short read") + + monkeypatch.setattr(src, "read_into", fail_read) + with pytest.raises(ExpertSourceError, match="short read"): + cache.copy_missing() + assert torch.equal(cache.slot_for_id, before_slots) + assert torch.equal(cache.id_of_slot, before_ids) + assert cache._pending_file_fetches == [] finally: src.close() diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 8b62a9076..e573c929b 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -415,7 +415,9 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(): model_path="/tmp/freetoken-test-model", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16, - attention_backend="fi", + # This test exercises MoE cache-rate conversion, not FlashInfer package + # discovery; keep it runnable in the project-local component-test env. + attention_backend="triton", moe_cache_rate=0.3, ) object.__setattr__( From 82533b81a55a7cad1ff932b62021372e64ede040 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 07/17] feat(qwen4): add streaming Q3 PLE sidecar writer --- python/freetoken/checkpoint/__init__.py | 4 +- python/freetoken/checkpoint/q3_ple.py | 380 +++++++++++++++++++++++- tests/checkpoint/test_q3_ple_writer.py | 164 ++++++++++ 3 files changed, 545 insertions(+), 3 deletions(-) create mode 100644 tests/checkpoint/test_q3_ple_writer.py diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index d11cb7fcd..f004cb5ce 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -12,10 +12,10 @@ load_ftw_banks, ) from .convert import convert_checkpoint -from .q3_ple import Q3PLEReader, Q3PLESegment +from .q3_ple import Q3PLEReader, Q3PLESegment, write_q3_ple_sidecar __all__ = [ "FTWReader", "FTWWriter", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", - "Q3PLEReader", "Q3PLESegment", + "Q3PLEReader", "Q3PLESegment", "write_q3_ple_sidecar", ] diff --git a/python/freetoken/checkpoint/q3_ple.py b/python/freetoken/checkpoint/q3_ple.py index 0e968fb4a..b302b40a6 100644 --- a/python/freetoken/checkpoint/q3_ple.py +++ b/python/freetoken/checkpoint/q3_ple.py @@ -13,11 +13,13 @@ import json import math import os +import operator import struct import threading +import uuid from dataclasses import dataclass from pathlib import Path -from typing import Sequence +from typing import Iterable, Sequence import torch @@ -30,6 +32,8 @@ FORMAT = "q3_ple_32" VERSION = 1 ALIGN = 4096 +REFINEMENT_PASSES = 2 +DEFAULT_SEGMENT_ROWS = 128 def _z_path(path: str | os.PathLike[str]) -> Path: @@ -45,6 +49,133 @@ def _z_path(path: str | os.PathLike[str]) -> Path: return resolved +def _z_output_path(path: str | os.PathLike[str]) -> Path: + """Resolve an output path and require an existing ``Z:`` parent directory. + + Unlike :func:`_z_path`, this helper permits the leaf file not to exist. The + writer intentionally does not create arbitrary parent directories: callers + must choose an already-created, Z-backed fixture or checkpoint directory. + """ + + candidate = Path(path).expanduser() + if not candidate.is_absolute(): + raise ValueError(f"Q3_PLE_32 output path must be absolute: {path}") + parent = candidate.parent.resolve(strict=True) + resolved = parent / candidate.name + drive, _ = os.path.splitdrive(str(resolved)) + if drive.upper() != "Z:" and not str(resolved).lower().startswith("/z/"): + raise ValueError(f"Q3_PLE_32 output must resolve to Z:, got {resolved}") + return resolved + + +def _pack_codes(codes: Sequence[int]) -> bytes: + if len(codes) != BLOCK_VALUES: + raise ValueError(f"expected {BLOCK_VALUES} codes, got {len(codes)}") + packed = 0 + for index, code in enumerate(codes): + if not 0 <= code <= 7: + raise ValueError(f"code {index} is outside 0..7: {code}") + packed |= int(code) << (3 * index) + return packed.to_bytes(12, "little") + + +def _bf16_bits(value: float) -> int: + """Round a finite Python float to an IEEE BF16 bit pattern.""" + + try: + bits = struct.unpack("> 16) & 1) + return (rounded >> 16) & 0xFFFF + + +def _bf16_from_bits(bits: int) -> float: + return struct.unpack(" tuple[bytes, float]: + if not math.isfinite(value) or value < 0.0: + raise ValueError(f"Q3_PLE_32 scale must be finite and non-negative: {value!r}") + bits = _bf16_bits(value) + # A positive source scale must remain representable after BF16 storage. A + # zero source scale is reserved for an all-zero block. + if value > 0.0 and bits == 0: + bits = 1 + stored = _bf16_from_bits(bits) + if not math.isfinite(stored): + raise ValueError("Q3_PLE_32 stored BF16 scale is not finite") + return struct.pack(" list[int]: + if scale == 0.0: + return [4] * BLOCK_VALUES + return [max(-4, min(3, int(round(value / scale)))) + 4 for value in values] + + +def quantize_block(values: Sequence[float], *, refinement_passes: int = REFINEMENT_PASSES) -> bytes: + """Encode one 32-value block using the canonical Q3_PLE_32 recipe. + + This mirrors ``scripts/q3_ple_32_reference.py`` while keeping production + conversion independent of the repository's executable reference script. + """ + + if len(values) != BLOCK_VALUES: + raise ValueError(f"expected {BLOCK_VALUES} values, got {len(values)}") + if refinement_passes < 0: + raise ValueError("refinement_passes must be non-negative") + try: + source = [float(value) for value in values] + except (TypeError, ValueError) as exc: + raise ValueError("Q3_PLE_32 values must be numeric") from exc + if not all(math.isfinite(value) for value in source): + raise ValueError("Q3_PLE_32 cannot encode non-finite values") + + minimum = min(source) + maximum = max(source) + scale = max(-minimum / 4.0, maximum / 3.0) + if scale == 0.0: + scale_bytes, _ = _store_bf16_scale(0.0) + return scale_bytes + _pack_codes([4] * BLOCK_VALUES) + + codes = _codes_for_scale(source, scale) + for _ in range(refinement_passes): + quants = [code - 4 for code in codes] + denominator = sum(quant * quant for quant in quants) + if denominator == 0: + break + refined = sum(value * quant for value, quant in zip(source, quants)) / denominator + if refined <= 0.0 or not math.isfinite(refined): + break + new_codes = _codes_for_scale(source, refined) + scale = refined + if new_codes == codes: + codes = new_codes + break + codes = new_codes + + scale_bytes, stored_scale = _store_bf16_scale(scale) + # Requantize once against the stored BF16 value so decoding exactly follows + # runtime behavior rather than the pre-rounded Python scale. + codes = _codes_for_scale(source, stored_scale) + block = scale_bytes + _pack_codes(codes) + if len(block) != BLOCK_BYTES: + raise AssertionError(f"Q3_PLE_32 block has wrong size: {len(block)}") + return block + + +def quantize_row(values: Sequence[float], *, refinement_passes: int = REFINEMENT_PASSES) -> bytes: + """Encode a 160-value row as five canonical Q3_PLE_32 blocks.""" + + if len(values) != ROW_VALUES: + raise ValueError(f"expected {ROW_VALUES} values, got {len(values)}") + return b"".join( + quantize_block(values[offset : offset + BLOCK_VALUES], refinement_passes=refinement_passes) + for offset in range(0, ROW_VALUES, BLOCK_VALUES) + ) + + def _unpack_codes(payload: bytes) -> list[int]: if len(payload) != 12: raise ValueError(f"Q3_PLE_32 code payload must be 12 bytes, got {len(payload)}") @@ -105,6 +236,11 @@ def __init__(self, manifest_path: str | os.PathLike[str], *, data_path: str | os raise ValueError("Q3_PLE_32 row_values mismatch") if int(manifest.get("row_bytes", ROW_BYTES)) != ROW_BYTES: raise ValueError("Q3_PLE_32 row_bytes mismatch") + # Older Stage 6 fixtures predate this field, so absence remains + # readable. A present fingerprint is always canonical SHA-256 hex; + # malformed provenance must fail closed rather than being ignored. + if "source_fingerprint" in manifest: + _validate_source_fingerprint(manifest["source_fingerprint"]) candidate = data_path if candidate is None: @@ -263,14 +399,256 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: self.close() +def _align_up(value: int, alignment: int = ALIGN) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _validate_source_fingerprint(source_fingerprint: str) -> str: + if not isinstance(source_fingerprint, str): + raise ValueError("source_fingerprint must be a 64-character SHA-256 hex string") + fingerprint = source_fingerprint.lower() + if len(fingerprint) != 64 or any(char not in "0123456789abcdef" for char in fingerprint): + raise ValueError("source_fingerprint must be a 64-character SHA-256 hex string") + return fingerprint + + +def _materialize_row(row: object) -> list[float]: + """Materialize one bounded row without retaining any other source rows.""" + + if isinstance(row, torch.Tensor): + if row.ndim != 1 or row.numel() != ROW_VALUES: + raise ValueError(f"Q3_PLE_32 row must contain exactly {ROW_VALUES} values") + try: + values = row.detach().cpu().tolist() + except Exception as exc: + raise ValueError("Q3_PLE_32 row tensor could not be copied to CPU") from exc + else: + try: + values = list(row) # type: ignore[arg-type] + except (TypeError, ValueError) as exc: + raise ValueError(f"Q3_PLE_32 row must contain exactly {ROW_VALUES} values") from exc + if len(values) != ROW_VALUES: + raise ValueError(f"Q3_PLE_32 row must contain exactly {ROW_VALUES} values, got {len(values)}") + return values + + +def _partial_path(path: Path, token: str) -> Path: + return path.with_name(f".{path.name}.partial-{os.getpid()}-{threading.get_ident()}-{token}") + + +def _validate_segment_directory( + segments: Sequence[dict[str, int | str]], row_count: int, file_bytes: int +) -> None: + expected_row = 0 + previous_end = 0 + for segment in segments: + first_row = int(segment["first_row"]) + end_row = int(segment["end_row"]) + data_offset = int(segment["data_offset"]) + byte_length = int(segment["byte_length"]) + digest = str(segment["sha256"]) + if first_row != expected_row or end_row <= first_row: + raise ValueError("Q3_PLE_32 writer generated a malformed segment directory") + if data_offset < 0 or data_offset % ALIGN: + raise ValueError("Q3_PLE_32 writer generated an unaligned segment") + if byte_length != (end_row - first_row) * ROW_BYTES: + raise ValueError("Q3_PLE_32 writer generated a segment length mismatch") + if data_offset < previous_end or data_offset + byte_length > file_bytes: + raise ValueError("Q3_PLE_32 writer generated overlapping/out-of-range segments") + if len(digest) != 64 or any(char not in "0123456789abcdef" for char in digest): + raise ValueError("Q3_PLE_32 writer generated a malformed segment hash") + expected_row = end_row + previous_end = data_offset + byte_length + if not segments or expected_row != row_count: + raise ValueError("Q3_PLE_32 writer generated incomplete segment coverage") + + +def _fsync(handle: object) -> None: + # This helper exists to keep the finalize path explicit and easy to audit; + # the writer only passes ordinary binary file handles here. + file_handle = handle # type narrowing for type checkers without a runtime dependency + file_handle.flush() # type: ignore[attr-defined] + os.fsync(file_handle.fileno()) # type: ignore[attr-defined] + + +def write_q3_ple_sidecar( + rows: Iterable[Sequence[float] | torch.Tensor], + data_path: str | os.PathLike[str], + manifest_path: str | os.PathLike[str], + *, + source_fingerprint: str, + weight_scale: float, + segment_rows: int = DEFAULT_SEGMENT_ROWS, +) -> dict: + """Stream rows into an atomic, reader-compatible Q3_PLE_32 sidecar. + + ``rows`` is consumed exactly once and only the current 160-value row is held + in memory. Segments are split in source order and begin at 4 KiB-aligned + offsets; alignment bytes are included in the whole-file hash but never in a + segment's logical length/hash. Both files are written under unique partial + names, fsynced, and atomically renamed into place on successful completion. + + The converter rule is intentionally fixed at two least-squares refinement + passes (the provisional Q3_PLE_32 recipe), and ``weight_scale`` is metadata + applied by the runtime after row dequantization rather than folded into the + per-block scales. + """ + + data_final = _z_output_path(data_path) + manifest_final = _z_output_path(manifest_path) + if data_final == manifest_final: + raise ValueError("Q3_PLE_32 data_path and manifest_path must differ") + source_digest = _validate_source_fingerprint(source_fingerprint) + if isinstance(segment_rows, bool): + raise ValueError("segment_rows must be a positive integer") + try: + segment_size = operator.index(segment_rows) + except TypeError as exc: + raise ValueError("segment_rows must be a positive integer") from exc + if segment_size <= 0: + raise ValueError("segment_rows must be a positive integer") + try: + global_scale = float(weight_scale) + except (TypeError, ValueError) as exc: + raise ValueError("weight_scale must be finite") from exc + if not math.isfinite(global_scale): + raise ValueError("weight_scale must be finite") + + # Unique tokens make concurrent conversion attempts independent and avoid + # ever truncating a stale partial file left by an interrupted process. + token = uuid.uuid4().hex + data_partial = _partial_path(data_final, token) + manifest_partial = _partial_path(manifest_final, token) + segments: list[dict[str, int | str]] = [] + whole_digest = hashlib.sha256() + payload_digest = hashlib.sha256() + rows_written = 0 + file_offset = 0 + current_segment: dict[str, int | str] | None = None + segment_digest: hashlib._Hash | None = None + + def write_padding(handle: object, count: int) -> None: + if count <= 0: + return + padding = bytes(min(1 << 20, count)) + remaining = count + while remaining: + take = min(remaining, len(padding)) + chunk = padding[:take] + handle.write(chunk) # type: ignore[attr-defined] + whole_digest.update(chunk) + remaining -= take + + try: + with data_partial.open("wb") as output: + for source_row in rows: + row_values = _materialize_row(source_row) + encoded_row = quantize_row(row_values, refinement_passes=REFINEMENT_PASSES) + if len(encoded_row) != ROW_BYTES: + raise AssertionError(f"Q3_PLE_32 row has wrong size: {len(encoded_row)}") + + if rows_written % segment_size == 0: + if current_segment is not None: + assert segment_digest is not None + current_segment["end_row"] = rows_written + current_segment["byte_length"] = ( + rows_written * ROW_BYTES - int(current_segment["first_row"]) * ROW_BYTES + ) + current_segment["sha256"] = segment_digest.hexdigest() + segments.append(current_segment) + aligned_offset = _align_up(file_offset) + write_padding(output, aligned_offset - file_offset) + file_offset = aligned_offset + current_segment = { + "first_row": rows_written, + "end_row": rows_written, + "data_offset": file_offset, + "byte_length": 0, + "sha256": "", + } + segment_digest = hashlib.sha256() + + assert current_segment is not None and segment_digest is not None + output.write(encoded_row) + whole_digest.update(encoded_row) + payload_digest.update(encoded_row) + segment_digest.update(encoded_row) + file_offset += len(encoded_row) + rows_written += 1 + + if current_segment is None: + raise ValueError("Q3_PLE_32 rows must contain at least one row") + assert segment_digest is not None + current_segment["end_row"] = rows_written + current_segment["byte_length"] = ( + rows_written * ROW_BYTES - int(current_segment["first_row"]) * ROW_BYTES + ) + current_segment["sha256"] = segment_digest.hexdigest() + segments.append(current_segment) + _fsync(output) + except Exception: + data_partial.unlink(missing_ok=True) + manifest_partial.unlink(missing_ok=True) + raise + + try: + file_bytes = data_partial.stat().st_size + except Exception: + data_partial.unlink(missing_ok=True) + raise + if file_bytes != file_offset: + data_partial.unlink(missing_ok=True) + raise OSError(f"Q3_PLE_32 partial length mismatch: {file_bytes} != {file_offset}") + try: + _validate_segment_directory(segments, rows_written, file_bytes) + except Exception: + data_partial.unlink(missing_ok=True) + raise + manifest = { + "format": FORMAT, + "version": VERSION, + "endianness": "little", + "block_values": BLOCK_VALUES, + "block_bytes": BLOCK_BYTES, + "row_values": ROW_VALUES, + "row_bytes": ROW_BYTES, + "rows": rows_written, + "payload_bytes": rows_written * ROW_BYTES, + "file_bytes": file_bytes, + "data_file": os.path.relpath(data_final, manifest_final.parent), + "weight_scale": global_scale, + "source_fingerprint": source_digest, + "sha256": whole_digest.hexdigest(), + "payload_sha256": payload_digest.hexdigest(), + "segments": segments, + } + try: + with manifest_partial.open("w", encoding="utf-8", newline="\n") as manifest_handle: + json.dump(manifest, manifest_handle, ensure_ascii=False, indent=2, sort_keys=True) + manifest_handle.write("\n") + _fsync(manifest_handle) + os.replace(data_partial, data_final) + os.replace(manifest_partial, manifest_final) + except Exception: + data_partial.unlink(missing_ok=True) + manifest_partial.unlink(missing_ok=True) + raise + return manifest + + __all__ = [ "ALIGN", "BLOCK_BYTES", "BLOCK_VALUES", + "DEFAULT_SEGMENT_ROWS", "FORMAT", + "REFINEMENT_PASSES", "Q3PLEReader", "Q3PLESegment", "ROW_BYTES", "ROW_VALUES", "VERSION", + "quantize_block", + "quantize_row", + "write_q3_ple_sidecar", ] diff --git a/tests/checkpoint/test_q3_ple_writer.py b/tests/checkpoint/test_q3_ple_writer.py new file mode 100644 index 000000000..b101f9ec9 --- /dev/null +++ b/tests/checkpoint/test_q3_ple_writer.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +import hashlib +import importlib.util +import json +import shutil +import sys +from uuid import uuid4 +from pathlib import Path + +import pytest + +from freetoken.checkpoint.q3_ple import ( + ALIGN, + BLOCK_BYTES, + ROW_BYTES, + ROW_VALUES, + Q3PLEReader, + quantize_block, + write_q3_ple_sidecar, +) + + +def _rows(count: int): + for row_index in range(count): + yield [((row_index + 1) * 0.125) * ((column % 17) - 8) for column in range(ROW_VALUES)] + + +def _load_authoritative_reference(): + reference_path = Path(__file__).resolve().parents[4] / "scripts" / "q3_ple_32_reference.py" + if not reference_path.exists(): + pytest.skip("authoritative Q3_PLE_32 reference script is not present in this checkout") + spec = importlib.util.spec_from_file_location("q3_ple_32_authoritative_reference", reference_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture() +def z_fixture_dir() -> Path: + root = Path(__file__).resolve().parents[2] / ".stage7-q3-writer-fixtures" / uuid4().hex + assert root.drive.upper() == "Z:" + root.mkdir(parents=True) + try: + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +def test_writer_matches_reference_vectors_and_reader(z_fixture_dir: Path) -> None: + # These vectors are the byte-for-byte values from the authoritative + # scripts/q3_ple_32_reference.py codec (two refinement passes, BF16 scale). + assert quantize_block([0.0] * 32).hex() == "0000244992244992244992244992" + assert quantize_block([(index - 16) / 4.0 for index in range(32)]).hex() == ( + "9c3f492249dab69124dbb6b6edff" + ) + + data_path = z_fixture_dir / "ple-q3.bin" + manifest_path = z_fixture_dir / "ple-q3.json" + manifest = write_q3_ple_sidecar( + _rows(5), + data_path, + manifest_path, + source_fingerprint="A" * 64, + weight_scale=1.25, + segment_rows=2, + ) + assert manifest["source_fingerprint"] == "a" * 64 + assert manifest["weight_scale"] == 1.25 + assert manifest["payload_bytes"] == 5 * ROW_BYTES + assert manifest["file_bytes"] == data_path.stat().st_size + assert [segment["first_row"] for segment in manifest["segments"]] == [0, 2, 4] + assert [segment["data_offset"] for segment in manifest["segments"]] == [0, ALIGN, ALIGN * 2] + assert all( + segment["byte_length"] == (2 if segment["first_row"] < 4 else 1) * ROW_BYTES + for segment in manifest["segments"] + ) + + raw = data_path.read_bytes() + assert hashlib.sha256(raw).hexdigest() == manifest["sha256"] + loaded = json.loads(manifest_path.read_text(encoding="utf-8")) + assert loaded == manifest + with Q3PLEReader(manifest_path) as reader: + gathered = reader.gather([4, 0, 4]) + assert gathered.shape == (3, ROW_VALUES) + assert gathered[0].equal(gathered[2]) + scaled = reader.gather([0], apply_weight_scale=True) + assert scaled.equal(reader.gather([0]) * 1.25) + + +def test_writer_payload_is_byte_identical_to_authoritative_reference(z_fixture_dir: Path) -> None: + reference = _load_authoritative_reference() + rows = list(_rows(3)) + expected = reference.encode_table(rows, refinement_passes=2, scale_dtype="bf16") + manifest = write_q3_ple_sidecar( + iter(rows), + z_fixture_dir / "reference.bin", + z_fixture_dir / "reference.json", + source_fingerprint="e" * 64, + weight_scale=1.0, + segment_rows=128, + ) + assert (z_fixture_dir / "reference.bin").read_bytes() == expected + assert manifest["file_bytes"] == len(expected) + + +def test_writer_consumes_rows_once_and_rejects_nonfinite_without_finalizing(z_fixture_dir: Path) -> None: + data_path = z_fixture_dir / "ple-q3.bin" + manifest_path = z_fixture_dir / "ple-q3.json" + consumed = 0 + + def source_rows(): + nonlocal consumed + consumed += 1 + yield [0.0] * ROW_VALUES + consumed += 1 + bad = [0.0] * ROW_VALUES + bad[31] = float("nan") + yield bad + + with pytest.raises(ValueError, match="non-finite"): + write_q3_ple_sidecar( + source_rows(), + data_path, + manifest_path, + source_fingerprint="b" * 64, + weight_scale=1.0, + segment_rows=1, + ) + assert consumed == 2 + assert not data_path.exists() + assert not manifest_path.exists() + assert not list(z_fixture_dir.glob(".*.partial-*")) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"source_fingerprint": "short", "weight_scale": 1.0}, + {"source_fingerprint": "c" * 64, "weight_scale": float("inf")}, + {"source_fingerprint": "c" * 64, "weight_scale": 1.0, "segment_rows": 0}, + ], +) +def test_writer_rejects_bad_integrity_metadata(z_fixture_dir: Path, kwargs: dict) -> None: + with pytest.raises(ValueError): + write_q3_ple_sidecar( + _rows(1), + z_fixture_dir / "ple-q3.bin", + z_fixture_dir / "ple-q3.json", + **kwargs, + ) + + +def test_writer_requires_z_backing(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Z:"): + write_q3_ple_sidecar( + _rows(1), + tmp_path / "ple-q3.bin", + tmp_path / "ple-q3.json", + source_fingerprint="d" * 64, + weight_scale=1.0, + ) From 5d1102d60c54f47d04fd06f86e5c6d2d93b11491 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 08/17] feat(qwen4): add FTEXPERT1 sidecar writer --- python/freetoken/moe/__init__.py | 9 +- python/freetoken/moe/expert_source.py | 687 +++++++++++++++++++++--- tests/moe/test_expert_sidecar_writer.py | 142 +++++ 3 files changed, 770 insertions(+), 68 deletions(-) create mode 100644 tests/moe/test_expert_sidecar_writer.py diff --git a/python/freetoken/moe/__init__.py b/python/freetoken/moe/__init__.py index 9e0dc1706..b1edb9c1d 100644 --- a/python/freetoken/moe/__init__.py +++ b/python/freetoken/moe/__init__.py @@ -73,8 +73,15 @@ def create_moe_backend(backend: str) -> BaseMoeBackend: "is_offload_moe_backend", "FileExpertSource", "ExpertSourceError", + "write_expert_sidecar", + "adapt_expert_tensor_record", ] # Kept at module bottom to avoid importing torch/file-I/O helpers while the # backend registry is initialized by lightweight callers. -from .expert_source import ExpertSourceError, FileExpertSource # noqa: E402 +from .expert_source import ( # noqa: E402 + ExpertSourceError, + FileExpertSource, + adapt_expert_tensor_record, + write_expert_sidecar, +) diff --git a/python/freetoken/moe/expert_source.py b/python/freetoken/moe/expert_source.py index 9ac562fef..a7b662b55 100644 --- a/python/freetoken/moe/expert_source.py +++ b/python/freetoken/moe/expert_source.py @@ -14,12 +14,16 @@ import struct import threading from pathlib import Path -from typing import Iterable +from collections.abc import Iterable, Mapping, Sequence +from typing import Any import torch -MAGIC = b"FTEXNV4\0" +MAGIC = b"FTEXPERT1" +# ``FTEXNV4`` was the private Stage 6 fixture format. Keep the reader able to +# reopen those fixtures while making every new artifact unambiguously FTEXPERT1. +LEGACY_MAGIC = b"FTEXNV4\0" VERSION = 1 HEADER_BYTES = 4096 RECORD_BYTES = 2_772_992 @@ -44,15 +48,38 @@ _cursor += _size assert _cursor == RAW_RECORD_BYTES -_HEADER_STRUCT = struct.Struct("<8sIIIIIQQ16s32s32s") - - -def _sha256_path(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - while chunk := handle.read(8 << 20): - digest.update(chunk) - return digest.hexdigest() +_HEADER_STRUCT = struct.Struct("<9sIIIIIQQ16s32s32s32s") +_LEGACY_HEADER_STRUCT = struct.Struct("<8sIIIIIQQ16s32s32s") + +# A fixed descriptor in the otherwise-reserved header makes reduced synthetic +# geometry self-describing. Native artifacts retain the exact six-plane +# ModelOpt layout; tests can use tiny tensors without teaching the reader a +# second out-of-band schema. The descriptor is deliberately binary and fixed +# width so two writes of the same inputs are byte-for-byte identical. +_GEOMETRY_MAGIC = b"GEO1" +_GEOMETRY_ENTRY = struct.Struct(" str: + return { + torch.uint8: "uint8", + torch.float8_e4m3fn: "float8_e4m3fn", + torch.float16: "float16", + torch.float32: "float32", + torch.bfloat16: "bfloat16", + }[dtype] class ExpertSourceError(RuntimeError): @@ -89,6 +116,483 @@ def _plane_shape(name: str) -> tuple[int, ...]: }[name] +def _align_up(value: int, alignment: int = ALIGNMENT) -> int: + return ((int(value) + alignment - 1) // alignment) * alignment + + +def _normalise_fingerprint(value: str | bytes | bytearray | None) -> bytes: + """Return the fixed 32-byte source fingerprint stored in the header. + + Fingerprints are normally SHA-256 bytes or their 64-character hexadecimal + spelling. Short byte strings are accepted for deterministic synthetic + fixtures and zero padded; this mirrors the Stage 6 fixture contract while + still rejecting an accidentally over-wide digest. + """ + + if value is None: + return b"\0" * 32 + if isinstance(value, str): + try: + value = bytes.fromhex(value) + except ValueError: + # Human-readable fixture labels are useful in bounded tests; keep + # them deterministic while documenting that production callers + # should pass SHA-256 bytes/hex. + value = value.encode("utf-8") + raw = bytes(value) + if len(raw) > 32: + raise ValueError("source_fingerprint must be at most 32 bytes") + return raw.ljust(32, b"\0") + + +def _normalise_geometry( + geometry: Mapping[str, Any] | Sequence[tuple[str, Any, Any]] | None, +) -> tuple[tuple[str, tuple[int, ...], torch.dtype], ...]: + """Validate/normalise a six-plane geometry declaration. + + ``geometry`` may map plane names to ``(shape, dtype)`` pairs, or be a + sequence of ``(name, shape, dtype)`` entries. Names must appear exactly in + :data:`PLANE_LAYOUT` order; accepting a mapping is convenient for callers, + but serialization remains ordered and deterministic. + """ + + if geometry is None: + return tuple((name, _plane_shape(name), _dtype(dtype_name)) for name, _size, dtype_name in PLANE_LAYOUT) + if isinstance(geometry, Mapping): + unknown = set(geometry) - set(name for name, _size, _dtype_name in PLANE_LAYOUT) + missing = set(name for name, _size, _dtype_name in PLANE_LAYOUT) - set(geometry) + if unknown or missing: + raise ValueError(f"geometry must contain exactly six planes; missing={sorted(missing)}, unknown={sorted(unknown)}") + entries = [] + for name, _size, _dtype_name in PLANE_LAYOUT: + value = geometry[name] + if not isinstance(value, (tuple, list)) or len(value) != 2: + raise TypeError(f"geometry[{name!r}] must be (shape, dtype)") + shape, dtype = value + entries.append((name, tuple(int(dim) for dim in shape), dtype)) + else: + entries = [] + for item in geometry: + if not isinstance(item, (tuple, list)) or len(item) != 3: + raise TypeError("geometry entries must be (name, shape, dtype)") + name, shape, dtype = item + entries.append((str(name), tuple(int(dim) for dim in shape), dtype)) + expected_names = tuple(name for name, _size, _dtype_name in PLANE_LAYOUT) + actual_names = tuple(name for name, _shape, _dtype in entries) + if actual_names != expected_names: + raise ValueError(f"geometry plane order must be {expected_names}, got {actual_names}") + normalised = [] + for name, shape, dtype in entries: + if not shape or any(dim <= 0 for dim in shape) or len(shape) > 4: + raise ValueError(f"{name} shape must have 1-4 positive dimensions") + if not isinstance(dtype, torch.dtype) or dtype not in _DTYPE_CODES: + raise TypeError(f"unsupported dtype for {name}: {dtype!r}") + normalised.append((name, shape, dtype)) + return tuple(normalised) + + +def _geometry_descriptor( + specs: tuple[tuple[str, tuple[int, ...], torch.dtype], ...], +) -> bytes: + descriptor = bytearray(_GEOMETRY_BYTES) + descriptor[:4] = _GEOMETRY_MAGIC + cursor = 4 + for _name, shape, dtype in specs: + _GEOMETRY_ENTRY.pack_into( + descriptor, + cursor, + _DTYPE_CODES[dtype], + len(shape), + 0, + *(tuple(shape) + (0,) * (4 - len(shape))), + ) + cursor += _GEOMETRY_ENTRY.size + return bytes(descriptor) + + +def _parse_geometry_descriptor(header: bytes) -> tuple[tuple[str, tuple[int, ...], torch.dtype], ...] | None: + if len(header) < _GEOMETRY_BYTES: + return None + offset = _HEADER_STRUCT.size + if header[offset : offset + 4] != _GEOMETRY_MAGIC: + return None + cursor = offset + 4 + specs = [] + try: + for name, _size, _dtype_name in PLANE_LAYOUT: + code, rank, _reserved, d0, d1, d2, d3 = _GEOMETRY_ENTRY.unpack_from(header, cursor) + dtype = _CODE_DTYPES[code] + if not 1 <= rank <= 4: + return None + dims = (d0, d1, d2, d3)[:rank] + if any(dim <= 0 for dim in dims): + return None + specs.append((name, dims, dtype)) + cursor += _GEOMETRY_ENTRY.size + except (KeyError, struct.error): + return None + return tuple(specs) + + +def _canonical_whole_sha256(path: Path) -> str: + """Hash a finalized sidecar with the stored whole-hash field zeroed.""" + + digest = hashlib.sha256() + with path.open("rb") as handle: + offset = 0 + while chunk := handle.read(8 << 20): + if offset <= _WHOLE_HASH_OFFSET < offset + len(chunk): + begin = _WHOLE_HASH_OFFSET - offset + chunk = chunk[:begin] + b"\0" * 32 + chunk[begin + 32 :] + digest.update(chunk) + offset += len(chunk) + return digest.hexdigest() + + +def _sha256_path(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(8 << 20): + digest.update(chunk) + return digest.hexdigest() + + +def _tensor_bytes(value: Any, *, name: str, shape: tuple[int, ...], dtype: torch.dtype) -> bytes: + """Validate one plane and return its native little-endian bytes.""" + + if isinstance(value, (bytes, bytearray, memoryview)): + raw = bytes(value) + expected = int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() + if len(raw) != expected: + raise ValueError(f"{name} bytes length {len(raw)} != expected {expected}") + return raw + try: + tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + except Exception as exc: + raise TypeError(f"{name} must be a torch tensor or bytes-like value") from exc + # ModelOpt emits a scalar ``weight_scale_2`` for each source projection. + # Adapters may pass that scalar directly for a global plane; expand only in + # this explicit case and require the declared FP16 dtype afterward. + if name.endswith("_global") and tensor.numel() == 1 and tuple(tensor.shape) != shape: + if not tensor.dtype.is_floating_point: + raise TypeError(f"{name} scalar expansion requires a floating source") + tensor = tensor.to(dtype=dtype).expand(shape) + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} != expected {shape}") + if tensor.dtype != dtype: + raise TypeError(f"{name} dtype {tensor.dtype} != expected {dtype}") + tensor = tensor.detach().to(device="cpu").contiguous() + return tensor.view(torch.uint8).numpy().tobytes() + + +def _record_from_planes( + record: Mapping[str, Any], + specs: tuple[tuple[str, tuple[int, ...], torch.dtype], ...], +) -> bytes: + expected = tuple(name for name, _shape, _dtype in specs) + keys = tuple(key for key in record if key not in {"expert_id", "id"}) + if set(keys) != set(expected): + missing = sorted(set(expected) - set(keys)) + unknown = sorted(set(keys) - set(expected)) + raise ValueError(f"expert record planes mismatch; missing={missing}, unknown={unknown}") + return b"".join( + _tensor_bytes(record[name], name=name, shape=shape, dtype=dtype) + for name, shape, dtype in specs + ) + + +_SOURCE_TENSOR_NAMES = tuple( + f"{projection}.{suffix}" + for projection in ("gate_proj", "up_proj", "down_proj") + for suffix in ("weight", "weight_scale", "weight_scale_2", "input_scale") +) + + +def _source_scalar(value: Any, *, name: str) -> torch.Tensor: + """Validate one ModelOpt source scalar (F32, rank-zero or one element).""" + + tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + if tensor.dtype != torch.float32: + raise TypeError(f"{name} dtype {tensor.dtype} != expected torch.float32") + if tensor.numel() != 1: + raise ValueError(f"{name} must be a scalar, got shape {tuple(tensor.shape)}") + return tensor.detach().reshape(()) + + +def _as_named_mapping(record: Any) -> dict[str, Any]: + if isinstance(record, Mapping): + return dict(record) + try: + items = list(record) + except TypeError as exc: + raise TypeError("expert tensor record must be a mapping or (name, tensor) iterable") from exc + mapped: dict[str, Any] = {} + for item in items: + if not isinstance(item, (tuple, list)) or len(item) != 2: + raise TypeError("expert tensor entries must be (name, tensor) pairs") + name, value = item + if name in mapped: + raise ValueError(f"duplicate expert tensor name {name!r}") + mapped[str(name)] = value + return mapped + + +def adapt_expert_tensor_record(record: Mapping[str, Any] | Iterable[tuple[str, Any]]) -> dict[str, Any]: + """Adapt one ModelOpt expert's twelve source tensors to six native planes. + + ModelOpt stores ``gate_proj``, ``up_proj`` and ``down_proj`` independently, + each with ``weight``, ``weight_scale``, ``weight_scale_2`` and + ``input_scale``. The native runtime sidecar fuses gate/up along the output + axis, ignores the activation ``input_scale`` (W4A16), and expands each F32 + ``weight_scale_2`` scalar to the FP16 per-output-row global plane. + + A six-plane mapping is returned unchanged (but copied), allowing callers + that have already performed the adaptation to use the same writer API. + """ + + # ``_iter_record_items`` permits an explicit ID alongside the planes; the + # ID is routing metadata, not one of the twelve source tensors. + named_record = _as_named_mapping(record) + clean_record = {key: value for key, value in named_record.items() if key not in {"expert_id", "id"}} + keys = tuple(clean_record) + native_names = tuple(name for name, _shape, _dtype in _normalise_geometry(None)) + if set(keys) == set(native_names): + return {name: clean_record[name] for name in native_names} + if set(keys) != set(_SOURCE_TENSOR_NAMES): + missing = sorted(set(_SOURCE_TENSOR_NAMES) - set(keys)) + unknown = sorted(set(keys) - set(_SOURCE_TENSOR_NAMES)) + raise ValueError(f"expert source tensors mismatch; missing={missing}, unknown={unknown}") + # Validate all twelve source names, including both metadata scalar kinds. + for projection in ("gate_proj", "up_proj", "down_proj"): + _source_scalar(clean_record[f"{projection}.input_scale"], name=f"{projection}.input_scale") + _source_scalar(clean_record[f"{projection}.weight_scale_2"], name=f"{projection}.weight_scale_2") + gate_weight = clean_record["gate_proj.weight"] + up_weight = clean_record["up_proj.weight"] + gate_scale = clean_record["gate_proj.weight_scale"] + up_scale = clean_record["up_proj.weight_scale"] + down_weight = clean_record["down_proj.weight"] + down_scale = clean_record["down_proj.weight_scale"] + # Let the regular six-plane validator provide exact dtype/shape diagnostics + # after these bounded concatenations. Concatenation happens per expert and + # therefore never materialises a layer or model-sized tensor. + gate_weight_t = gate_weight if isinstance(gate_weight, torch.Tensor) else torch.as_tensor(gate_weight) + up_weight_t = up_weight if isinstance(up_weight, torch.Tensor) else torch.as_tensor(up_weight) + gate_scale_t = gate_scale if isinstance(gate_scale, torch.Tensor) else torch.as_tensor(gate_scale) + up_scale_t = up_scale if isinstance(up_scale, torch.Tensor) else torch.as_tensor(up_scale) + down_weight_t = down_weight if isinstance(down_weight, torch.Tensor) else torch.as_tensor(down_weight) + down_scale_t = down_scale if isinstance(down_scale, torch.Tensor) else torch.as_tensor(down_scale) + if gate_weight_t.ndim != up_weight_t.ndim or gate_weight_t.ndim < 1: + raise ValueError("gate_proj.weight and up_proj.weight must have matching rank") + if gate_scale_t.ndim != up_scale_t.ndim or gate_scale_t.ndim < 1: + raise ValueError("gate_proj.weight_scale and up_proj.weight_scale must have matching rank") + gate_rows = int(gate_weight_t.shape[0]) + up_rows = int(up_weight_t.shape[0]) + gate_global = _source_scalar(clean_record["gate_proj.weight_scale_2"], name="gate_proj.weight_scale_2").to(torch.float16).expand(gate_rows) + up_global = _source_scalar(clean_record["up_proj.weight_scale_2"], name="up_proj.weight_scale_2").to(torch.float16).expand(up_rows) + down_global = _source_scalar(clean_record["down_proj.weight_scale_2"], name="down_proj.weight_scale_2").to(torch.float16).expand(int(down_weight_t.shape[0])) + return { + "gate_up_packed": torch.cat((gate_weight_t, up_weight_t), dim=0), + "gate_up_scale": torch.cat((gate_scale_t, up_scale_t), dim=0), + "gate_up_global": torch.cat((gate_global, up_global), dim=0), + "down_packed": down_weight_t, + "down_scale": down_scale_t, + "down_global": down_global, + } + + +def _iter_record_items( + records_or_planes: Any, + *, + num_experts: int, +) -> Iterable[tuple[int, Any]]: + """Yield ``(expert_id, record)`` without materialising the expert bank.""" + + expected_names = {name for name, _size, _dtype_name in PLANE_LAYOUT} + if isinstance(records_or_planes, Mapping): + keys = set(records_or_planes) + source_names = set(_SOURCE_TENSOR_NAMES) + if keys & (expected_names | source_names): + if keys.issubset(source_names): + bank_names = tuple(_SOURCE_TENSOR_NAMES) + elif keys == expected_names: + bank_names = tuple(name for name, _s, _d in PLANE_LAYOUT) + else: + raise ValueError("plane-bank input must contain exactly the six native or twelve source tensor names") + # Stacked tensors/sequences are indexed lazily one expert at a time. + banks = records_or_planes + for eid in range(num_experts): + row = {} + for name in bank_names: + bank = banks[name] + if isinstance(bank, Mapping): + bank_ids = set(bank) + expected_ids = set(range(num_experts)) + if bank_ids != expected_ids: + raise ValueError( + f"plane bank {name} IDs mismatch; missing={sorted(expected_ids - bank_ids)}, " + f"unknown={sorted(bank_ids - expected_ids)}" + ) + if eid not in bank: + raise ValueError(f"missing expert id {eid} in plane bank {name}") + row[name] = bank[eid] + else: + try: + bank_len = len(bank) + except TypeError: + bank_len = None + if bank_len is not None and bank_len != num_experts: + raise ValueError(f"plane bank {name} length {bank_len} != num_experts {num_experts}") + try: + row[name] = bank[eid] + except (IndexError, KeyError, TypeError) as exc: + raise ValueError(f"missing expert id {eid} in plane bank {name}") from exc + yield eid, row + return + for key in sorted(records_or_planes): + if not isinstance(key, int): + raise TypeError("expert mapping keys must be integer expert IDs") + yield key, records_or_planes[key] + return + + for index, item in enumerate(records_or_planes): + expert_id = index + record = item + if isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], int): + expert_id, record = item + elif isinstance(item, Mapping): + explicit = item.get("expert_id", item.get("id", None)) + if explicit is not None: + expert_id = explicit + yield int(expert_id), record + + +def write_expert_sidecar( + path: str | os.PathLike[str], + records_or_planes: Any, + *, + layer_id: int, + source_fingerprint: str | bytes, + num_experts: int = NUM_EXPERTS, + geometry: Mapping[str, Any] | Sequence[tuple[str, Any, Any]] | None = None, + overwrite: bool = True, +) -> dict[str, Any]: + """Stream a deterministic FTEXPERT1 routed-expert sidecar to ``path``. + + ``records_or_planes`` can be an iterable of six-plane mappings (or the + twelve ModelOpt source-tensor mappings, optionally ``(expert_id, mapping)``), + an ``{expert_id: mapping}`` mapping, or a mapping of six/twelve plane names + to stacked tensors/sequences. Exactly one record for every ID + ``0..num_experts-1`` is required. The destination is written to a + ``.partial`` sibling and atomically replaced only after all validation, + payload hashing, and header hashes complete. + """ + + destination = _z_path(path) + if int(num_experts) < 1 or int(num_experts) > NUM_EXPERTS: + raise ValueError(f"num_experts must be in [1, {NUM_EXPERTS}]") + num_experts = int(num_experts) + if not 0 <= int(layer_id) <= 0xFFFFFFFF: + raise ValueError("layer_id must fit an unsigned 32-bit field") + layer_id = int(layer_id) + fingerprint = _normalise_fingerprint(source_fingerprint) + specs = _normalise_geometry(geometry) + sizes = tuple(int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() for _name, shape, dtype in specs) + raw_record_bytes = sum(sizes) + record_bytes = _align_up(raw_record_bytes) + descriptor = _geometry_descriptor(specs) + destination.parent.mkdir(parents=True, exist_ok=True) + partial = Path(str(destination) + ".partial") + if partial.exists(): + partial.unlink() + payload_hash = hashlib.sha256() + seen: set[int] = set() + try: + with partial.open("wb") as handle: + header = bytearray(HEADER_BYTES) + _HEADER_STRUCT.pack_into( + header, + 0, + MAGIC, + VERSION, + HEADER_BYTES, + num_experts, + len(specs), + layer_id, + raw_record_bytes, + record_bytes, + b"nvfp4-qwen4-v1\0\0", + fingerprint, + b"\0" * 32, + b"\0" * 32, + ) + header[_HEADER_STRUCT.size : _HEADER_STRUCT.size + len(descriptor)] = descriptor + handle.write(header) + for expert_id, record in _iter_record_items(records_or_planes, num_experts=num_experts): + if not 0 <= expert_id < num_experts: + raise ValueError(f"expert_id {expert_id} outside [0, {num_experts})") + if expert_id in seen: + raise ValueError(f"duplicate expert_id {expert_id}") + seen.add(expert_id) + if isinstance(record, (bytes, bytearray, memoryview)): + if len(record) != raw_record_bytes: + raise ValueError(f"expert {expert_id} raw bytes length {len(record)} != {raw_record_bytes}") + raw = bytes(record) + elif isinstance(record, Mapping) or isinstance(record, Iterable): + raw = _record_from_planes(adapt_expert_tensor_record(record), specs) + else: + raise TypeError(f"expert {expert_id} must be a plane mapping or raw bytes") + if len(raw) != raw_record_bytes: + raise ValueError(f"expert {expert_id} serialized length {len(raw)} != {raw_record_bytes}") + padded = raw + b"\0" * (record_bytes - raw_record_bytes) + payload_hash.update(padded) + handle.write(padded) + missing = sorted(set(range(num_experts)) - seen) + if missing: + raise ValueError(f"missing expert IDs: {missing}") + handle.flush() + os.fsync(handle.fileno()) + digest = payload_hash.digest() + # Patch payload hash first, then derive a canonical whole hash over the + # finalized header with only the whole-hash field zeroed. + with partial.open("r+b") as handle: + handle.seek(_PAYLOAD_HASH_OFFSET) + handle.write(digest) + handle.flush() + os.fsync(handle.fileno()) + canonical = _canonical_whole_sha256(partial) + with partial.open("r+b") as handle: + handle.seek(_WHOLE_HASH_OFFSET) + handle.write(bytes.fromhex(canonical)) + handle.flush() + os.fsync(handle.fileno()) + if not overwrite and destination.exists(): + raise FileExistsError(destination) + os.replace(partial, destination) + except Exception: + try: + partial.unlink() + except FileNotFoundError: + pass + raise + whole = _sha256_path(destination) + return { + "path": str(destination), + "format": "FTEXPERT1", + "version": VERSION, + "layer_id": layer_id, + "num_experts": num_experts, + "planes": tuple(name for name, _shape, _dtype in specs), + "raw_record_bytes": raw_record_bytes, + "record_bytes": record_bytes, + "source_fingerprint": fingerprint.hex(), + "payload_sha256": digest.hex(), + "canonical_sha256": canonical, + "whole_sha256": whole, + "sha256": whole, + "sample_ids": (0, num_experts - 1), + } + + class FileExpertSource: """Read fixed NVFP4 expert records from one layer sidecar. @@ -137,6 +641,7 @@ def __init__( self._fd = os.open(str(self.path), os.O_RDONLY | getattr(os, "O_BINARY", 0)) try: self._validate_file(expected_source_fingerprint, expected_layer_id) + self.staging_record_bytes = self.record_bytes self.sha256 = self._hash_file() if verify_hash else None if expected_sha256 is not None: expected_sha256 = expected_sha256.lower() @@ -169,89 +674,128 @@ def create_synthetic( if num_experts < 1: raise ValueError("num_experts must be positive") source_fingerprint = source_fingerprint or hashlib.sha256(b"synthetic").digest() - source_fingerprint = bytes(source_fingerprint[:32]).ljust(32, b"\0") rows = iter(records) if records is not None else None - payload_hash = hashlib.sha256() - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("wb") as fh: - header = bytearray(HEADER_BYTES) - _HEADER_STRUCT.pack_into( - header, - 0, - MAGIC, - VERSION, - HEADER_BYTES, - num_experts, - len(PLANE_LAYOUT), - int(layer_id), - RAW_RECORD_BYTES, - RECORD_BYTES, - b"nvfp4-qwen4-v1\0\0", # 16-byte layout tag - source_fingerprint, - b"\0" * 32, - ) - fh.write(header) + + def _rows() -> Iterable[bytes]: for expert_id in range(num_experts): raw = next(rows) if rows is not None else bytes([expert_id & 0xFF]) * RAW_RECORD_BYTES if len(raw) != RAW_RECORD_BYTES: raise ValueError("synthetic record must contain exactly RAW_RECORD_BYTES") - record = raw + bytes(RECORD_BYTES - RAW_RECORD_BYTES) - payload_hash.update(record) - fh.write(record) - digest = payload_hash.digest() - with path.open("r+b") as fh: - fh.seek(0) - header = bytearray(fh.read(HEADER_BYTES)) - header[_HEADER_STRUCT.size - 32 : _HEADER_STRUCT.size] = digest - fh.seek(0) - fh.write(header) + yield raw + + result = write_expert_sidecar( + path, + _rows(), + layer_id=layer_id, + source_fingerprint=source_fingerprint, + num_experts=num_experts, + ) # Keep synthetic native-geometry fixtures bounded: the real 512-record # shape is roughly 1.42 GB and must never be read into one Python bytes # object merely to produce its verification hash. - return _sha256_path(path) + return str(result["sha256"]) def _validate_file( self, expected_source_fingerprint: str | bytes | None, expected_layer_id: int | None ) -> None: size = self.path.stat().st_size - expected_size = HEADER_BYTES + self.num_experts * self.record_bytes header = self._read_exact(HEADER_BYTES, 0) if len(header) != HEADER_BYTES: raise ExpertSourceError("truncated expert tier header") try: - magic, version, hbytes, experts, planes, layer_id, raw_bytes, rec_bytes, tag, fingerprint, payload_hash = _HEADER_STRUCT.unpack_from(header) + if header[: len(MAGIC)] == MAGIC: + ( + magic, + version, + hbytes, + experts, + planes, + layer_id, + raw_bytes, + rec_bytes, + tag, + fingerprint, + payload_hash, + whole_hash, + ) = _HEADER_STRUCT.unpack_from(header) + whole_offset = _WHOLE_HASH_OFFSET + specs = _parse_geometry_descriptor(header) + if specs is None: + specs = tuple((name, _plane_shape(name), _dtype(dtype_name)) for name, _size, dtype_name in PLANE_LAYOUT) + elif header[: len(LEGACY_MAGIC)] == LEGACY_MAGIC: + ( + magic, + version, + hbytes, + experts, + planes, + layer_id, + raw_bytes, + rec_bytes, + tag, + fingerprint, + payload_hash, + ) = _LEGACY_HEADER_STRUCT.unpack_from(header) + whole_hash = b"\0" * 32 + whole_offset = None + specs = tuple((name, _plane_shape(name), _dtype(dtype_name)) for name, _size, dtype_name in PLANE_LAYOUT) + else: + raise ExpertSourceError("unsupported expert tier magic/version/header") except struct.error as exc: raise ExpertSourceError("malformed expert tier header") from exc - if magic != MAGIC or version != VERSION or hbytes != HEADER_BYTES: + if magic not in (MAGIC, LEGACY_MAGIC) or version != VERSION or hbytes != HEADER_BYTES: raise ExpertSourceError("unsupported expert tier magic/version/header") if experts != self.num_experts or planes != len(PLANE_LAYOUT): raise ExpertSourceError("expert tier geometry mismatch") + if tag.rstrip(b"\0") != b"nvfp4-qwen4-v1": + raise ExpertSourceError("expert tier layout mismatch") self.layer_id = int(layer_id) if expected_layer_id is not None and self.layer_id != int(expected_layer_id): - raise ExpertSourceError( - f"expert tier layer mismatch: {self.layer_id} != {int(expected_layer_id)}" - ) - if raw_bytes != RAW_RECORD_BYTES or rec_bytes != RECORD_BYTES or tag.rstrip(b"\0") != b"nvfp4-qwen4-v1": + raise ExpertSourceError(f"expert tier layer mismatch: {self.layer_id} != {int(expected_layer_id)}") + self.plane_layout = tuple( + (name, int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size(), _dtype_name(dtype)) + for name, shape, dtype in specs + ) + self.plane_specs = {name: (shape, dtype) for name, shape, dtype in specs} + self._plane_offsets = {} + cursor = 0 + for name, shape, dtype in specs: + self._plane_offsets[name] = cursor + cursor += int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() + if cursor != int(raw_bytes) or int(rec_bytes) != _align_up(cursor): raise ExpertSourceError("expert tier layout mismatch") + self.raw_record_bytes = int(raw_bytes) + self.record_bytes = int(rec_bytes) + expected_size = HEADER_BYTES + self.num_experts * self.record_bytes if size != expected_size: raise ExpertSourceError(f"expert tier length mismatch: {size} != {expected_size}") if expected_source_fingerprint is not None: - expected = bytes.fromhex(expected_source_fingerprint) if isinstance(expected_source_fingerprint, str) else bytes(expected_source_fingerprint) - if fingerprint != expected[:32].ljust(32, b"\0"): + try: + expected = _normalise_fingerprint(expected_source_fingerprint) + except ValueError as exc: + raise ExpertSourceError("invalid expected source fingerprint") from exc + if fingerprint != expected: raise ExpertSourceError("expert tier source fingerprint mismatch") + self.source_fingerprint = fingerprint.hex() self.payload_sha256 = payload_hash.hex() + if magic == MAGIC and payload_hash == b"\0" * 32: + raise ExpertSourceError("expert tier missing payload hash") if self.payload_sha256 != "00" * 32: h = hashlib.sha256() - offset = HEADER_BYTES with self.path.open("rb") as fh: - fh.seek(offset) - while True: - block = fh.read(8 << 20) - if not block: - break + fh.seek(HEADER_BYTES) + while block := fh.read(8 << 20): h.update(block) if h.digest() != payload_hash: raise ExpertSourceError("expert tier payload hash mismatch") + self.whole_sha256 = _sha256_path(self.path) + self.canonical_sha256 = None + if magic == MAGIC and whole_hash == b"\0" * 32: + raise ExpertSourceError("expert tier missing whole hash") + if whole_offset is not None and whole_hash != b"\0" * 32: + self.canonical_sha256 = _canonical_whole_sha256(self.path) + if self.canonical_sha256 != whole_hash.hex(): + raise ExpertSourceError("expert tier whole hash mismatch") def _read_exact(self, size: int, offset: int) -> bytes: if self._closed: @@ -291,14 +835,13 @@ def _record_bytes(self, expert_id: int) -> bytes: def read_record(self, expert_id: int) -> dict[str, torch.Tensor]: raw = self._record_bytes(expert_id) out: dict[str, torch.Tensor] = {} - for name, size, dtype_name in PLANE_LAYOUT: - offset = _PLANE_OFFSETS[name] - dtype = _dtype(dtype_name) - shape = _plane_shape(name) - # bytearray supplies a writable, record-local staging buffer and - # avoids exposing Python's immutable bytes through a writable tensor. - staging = bytearray(raw[offset : offset + size]) - out[name] = torch.frombuffer(staging, dtype=dtype).reshape(shape) + for name, (shape, dtype) in self.plane_specs.items(): + size = int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() + offset = self._plane_offsets[name] + # Clone the record-local staging slice so returned tensors remain + # independent after this bounded read buffer is released. + staging = raw[offset : offset + size] + out[name] = torch.frombuffer(bytearray(staging), dtype=dtype).clone().reshape(shape) return out def read_records(self, expert_ids: Iterable[int], *, max_concurrency: int = 1) -> list[dict[str, torch.Tensor]]: @@ -332,4 +875,14 @@ def __exit__(self, *_exc) -> None: self.close() -__all__ = ["FileExpertSource", "ExpertSourceError", "PLANE_LAYOUT", "HEADER_BYTES", "RECORD_BYTES", "RAW_RECORD_BYTES"] +__all__ = [ + "FileExpertSource", + "ExpertSourceError", + "PLANE_LAYOUT", + "HEADER_BYTES", + "RECORD_BYTES", + "RAW_RECORD_BYTES", + "MAGIC", + "write_expert_sidecar", + "adapt_expert_tensor_record", +] diff --git a/tests/moe/test_expert_sidecar_writer.py b/tests/moe/test_expert_sidecar_writer.py new file mode 100644 index 000000000..7d9b45a32 --- /dev/null +++ b/tests/moe/test_expert_sidecar_writer.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch + +from freetoken.moe.expert_source import ( + ExpertSourceError, + FileExpertSource, + MAGIC, + adapt_expert_tensor_record, + write_expert_sidecar, +) + + +@pytest.fixture +def z_dir(): + root = Path.cwd() / ".stage7-test-fixtures" / uuid4().hex + root.mkdir(parents=True) + assert (root.drive or "").upper() == "Z:" + try: + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +def _geometry(): + return { + "gate_up_packed": ((4, 4), torch.uint8), + "gate_up_scale": ((4, 1), torch.float8_e4m3fn), + "gate_up_global": ((4,), torch.float16), + "down_packed": ((2, 16), torch.uint8), + "down_scale": ((2, 1), torch.float8_e4m3fn), + "down_global": ((2,), torch.float16), + } + + +def _planes(value: int): + out = {} + for name, (shape, dtype) in _geometry().items(): + if dtype == torch.float8_e4m3fn: + out[name] = torch.full(shape, 1, dtype=dtype) + else: + out[name] = torch.full(shape, value, dtype=dtype) + return out + + +def _source_record(value: int): + out = {} + for projection, rows, width in (("gate_proj", 2, 4), ("up_proj", 2, 4), ("down_proj", 2, 16)): + out[f"{projection}.weight"] = torch.full((rows, width), value, dtype=torch.uint8) + out[f"{projection}.weight_scale"] = torch.full((rows, max(1, width // 16)), 1, dtype=torch.float8_e4m3fn) + out[f"{projection}.weight_scale_2"] = torch.tensor(1.5, dtype=torch.float32) + out[f"{projection}.input_scale"] = torch.tensor(1.0, dtype=torch.float32) + return out + + +def test_writer_reduced_geometry_reopens_and_is_deterministic(z_dir): + first = z_dir / "layer-00.ftex" + second = z_dir / "layer-00-copy.ftex" + kwargs = dict(layer_id=7, source_fingerprint=b"source", num_experts=3, geometry=_geometry()) + result = write_expert_sidecar(first, [_planes(i) for i in range(3)], **kwargs) + result_copy = write_expert_sidecar(second, [_planes(i) for i in range(3)], **kwargs) + assert result["format"] == "FTEXPERT1" + assert result["raw_record_bytes"] == 66 + assert result["record_bytes"] == 4096 + assert result["sample_ids"] == (0, 2) + assert first.read_bytes() == second.read_bytes() + assert result["sha256"] == hashlib.sha256(first.read_bytes()).hexdigest() + assert result["sha256"] == result_copy["sha256"] + with FileExpertSource(first, num_experts=3, expected_sha256=result["sha256"], expected_layer_id=7) as source: + assert source.read_record(0)["gate_up_packed"].flatten()[0].item() == 0 + assert source.read_record(2)["gate_up_packed"].flatten()[0].item() == 2 + assert source.record_bytes == 4096 + + +def test_source_tensor_adapter_validates_twelve_names_and_expands_globals(z_dir): + source_record = _source_record(9) + adapted = adapt_expert_tensor_record(source_record) + assert adapted["gate_up_packed"].shape == (4, 4) + assert adapted["gate_up_scale"].shape == (4, 1) + assert adapted["gate_up_global"].dtype == torch.float16 + assert adapted["gate_up_global"].tolist() == [1.5] * 4 + assert adapted["down_global"].tolist() == [1.5] * 2 + result = write_expert_sidecar( + z_dir / "source.ftex", [source_record], layer_id=0, source_fingerprint=b"source", num_experts=1, geometry=_geometry() + ) + with FileExpertSource(result["path"], num_experts=1) as source: + assert source.read_record(0)["gate_up_packed"].flatten()[0].item() == 9 + + +def test_writer_accepts_explicit_id_and_named_pairs(z_dir): + named = [(name, value) for name, value in _planes(4).items()] + path = z_dir / "explicit.ftex" + write_expert_sidecar( + path, + [(0, named)], + layer_id=0, + source_fingerprint="fixture", + num_experts=1, + geometry=_geometry(), + ) + with FileExpertSource(path, num_experts=1) as source: + assert source.read_record(0)["down_global"].flatten()[0].item() == 4 + + +@pytest.mark.parametrize("records", [ + [(0, _planes(1)), (0, _planes(2))], + [(0, _planes(1))], + [(0, _planes(1)), (2, _planes(2))], +]) +def test_writer_rejects_duplicate_or_missing_ids_without_publishing(z_dir, records): + path = z_dir / "bad.ftex" + with pytest.raises(ValueError, match="(duplicate|missing|outside)"): + write_expert_sidecar(path, records, layer_id=0, source_fingerprint=b"x", num_experts=2, geometry=_geometry()) + assert not path.exists() + assert not Path(str(path) + ".partial").exists() + + +def test_writer_partial_and_payload_corruption_fail_closed(z_dir): + path = z_dir / "corrupt.ftex" + result = write_expert_sidecar(path, [_planes(1)], layer_id=0, source_fingerprint=b"x", num_experts=1, geometry=_geometry()) + with path.open("r+b") as handle: + handle.seek(4096 + 1) + handle.write(b"x") + with pytest.raises(ExpertSourceError, match="payload hash mismatch"): + FileExpertSource(path, num_experts=1) + path.unlink() + partial = Path(str(path) + ".partial") + partial.write_bytes(b"partial") + with pytest.raises(ExpertSourceError): + FileExpertSource(partial, num_experts=1) + + +def test_writer_emits_ft_expert_magic(z_dir): + path = z_dir / "magic.ftex" + write_expert_sidecar(path, [_planes(1)], layer_id=0, source_fingerprint=b"x", num_experts=1, geometry=_geometry()) + assert path.read_bytes()[: len(MAGIC)] == MAGIC From 30ea4f13e8403a9edceef00e27864d00ca969261 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 09/17] feat(qwen4): wire modular artifact manifest --- python/freetoken/checkpoint/convert.py | 48 +- python/freetoken/checkpoint/qwen4_artifact.py | 457 ++++++++++++++++++ python/freetoken/engine/engine.py | 59 +++ python/freetoken/models/qwen4_exp/config.py | 9 +- python/freetoken/models/qwen4_exp/model.py | 30 ++ python/freetoken/models/weight.py | 20 + tests/checkpoint/test_qwen4_artifact.py | 145 ++++++ 7 files changed, 764 insertions(+), 4 deletions(-) create mode 100644 python/freetoken/checkpoint/qwen4_artifact.py create mode 100644 tests/checkpoint/test_qwen4_artifact.py diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 2f643bca9..223d61b12 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -217,6 +217,7 @@ def convert_checkpoint( moe_backend: str = "offload", shard_limit: int = DEFAULT_SHARD_LIMIT, device: str | None = None, + artifact_format: str | None = None, ) -> dict: """Write ``model_path`` as an FTW checkpoint at ``out_dir``. Returns the index dict. @@ -246,6 +247,13 @@ def convert_checkpoint( cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), dtype=dtype, moe_backend=moe_backend) mc = cfg.model_config + if artifact_format not in (None, "qwen4_modular_v1"): + raise ValueError( + f"unsupported artifact_format {artifact_format!r}; expected None or 'qwen4_modular_v1'" + ) + is_qwen4 = any("Qwen4" in str(arch) for arch in getattr(mc, "architectures", ())) + if artifact_format is not None and not is_qwen4: + raise ValueError("artifact_format='qwen4_modular_v1' requires a Qwen4 checkpoint") offload = moe_backend == "offload" and getattr(mc, "is_moe", False) include_moe_experts = not offload @@ -257,8 +265,31 @@ def convert_checkpoint( # 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) dense_bytes = 0 - for name, tensor in count_bar(load_weight(model_path, torch.device("cpu"), - include_moe_experts=include_moe_experts), + dense_entries = load_weight( + model_path, + torch.device("cpu"), + include_moe_experts=include_moe_experts, + ) + if artifact_format == "qwen4_modular_v1": + # Quantization is an explicit artifact-build policy, never a generic runtime + # fallback. The Qwen4 iterator has already fused canonical projections; this + # wrapper only converts the frozen active map while leaving routers, PLE, and + # all non-active entries untouched. + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + def _modular_dense_entries(): + from freetoken.models.config import VISION_KEY_PREFIXES + + for name, tensor in iter_active_nvfp4_runtime_entries(dense_entries): + # The modular target is text-only by contract. Qwen4's canonical + # rename uses ``visual.*`` while other wrappers retain the generic + # vision prefixes; drop both explicitly at conversion time. + if name.startswith(("visual.",) + VISION_KEY_PREFIXES): + continue + yield name, tensor + + dense_entries = _modular_dense_entries() + for name, tensor in count_bar(dense_entries, "Converting dense weights"): writer.add_tensor(name, tensor, kind="weight") n_weight += 1 @@ -325,6 +356,19 @@ def convert_checkpoint( _progress("finalize") # writing shard index + copying config/tokenizer copied = _copy_metadata(model_path, out_dir) + if artifact_format == "qwen4_modular_v1": + config_path = os.path.join(out_dir, "config.json") + if not os.path.isfile(config_path): + raise ValueError("Qwen4 modular conversion requires a copied config.json") + with open(config_path, "r", encoding="utf-8") as handle: + config_data = json.load(handle) + config_data["freetoken_text_only"] = "qwen4_text_only_v1" + config_data["freetoken_active_quant"] = "nvfp4_w4a16_v1" + tmp_config = config_path + ".tmp" + with open(tmp_config, "w", encoding="utf-8") as handle: + json.dump(config_data, handle, indent=2, sort_keys=True) + handle.write("\n") + os.replace(tmp_config, config_path) try: fingerprint = _source_fingerprint(model_path, mc, device=dev) diff --git a/python/freetoken/checkpoint/qwen4_artifact.py b/python/freetoken/checkpoint/qwen4_artifact.py new file mode 100644 index 000000000..ea15897d2 --- /dev/null +++ b/python/freetoken/checkpoint/qwen4_artifact.py @@ -0,0 +1,457 @@ +"""Validation and runtime wiring for the Qwen4 modular artifact. + +The modular artifact is intentionally a small manifest around three independently +validated pieces: native active weights, a Q3 PLE sidecar, and a mixed resident/file +expert tier. This module owns only the manifest contract and the wiring seam. The +large writers and the byte-level readers live in their existing modules. + +An absent ``manifest.json`` is not an error and keeps the normal Qwen4 checkpoint path +unchanged. Once the marker is present, malformed or unknown values fail closed rather +than silently falling back to a different representation. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + + +FORMAT = "freetoken-qwen4-modular-v1" +VERSION = 1 +TEXT_ONLY_MARKER = "qwen4_text_only_v1" +ARTIFACT_FORMAT = "qwen4_modular_v1" +ACTIVE_FORMAT = "nvfp4_w4a16_v1" +PLE_FORMAT = "q3_ple_32" +EXPERT_FORMAT = "ftexpert1_nvfp4_v1" +REQUIRED_VOLUME = "Z:" +MANIFEST_NAME = "manifest.json" + + +class Qwen4ArtifactError(ValueError): + """Raised when a Qwen4 modular manifest cannot be trusted.""" + + +def _resolve_z(path: str | os.PathLike[str], *, label: str) -> Path: + resolved = Path(path).expanduser().resolve() + drive = (resolved.drive or os.path.splitdrive(str(resolved))[0]).upper() + if drive != REQUIRED_VOLUME: + raise Qwen4ArtifactError(f"{label} must resolve to Z:, got {resolved}") + return resolved + + +def _path_from(root: Path, value: object, *, label: str) -> Path: + if not isinstance(value, str) or not value.strip(): + raise Qwen4ArtifactError(f"{label} must be a non-empty path") + candidate = Path(value) + if not candidate.is_absolute(): + candidate = root / candidate + return _resolve_z(candidate, label=label) + + +def _require_mapping(value: object, *, label: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping): + raise Qwen4ArtifactError(f"{label} must be an object") + return value + + +def _require_sha(value: object, *, label: str) -> str: + text = str(value).lower() + if len(text) != 64 or any(char not in "0123456789abcdef" for char in text): + raise Qwen4ArtifactError(f"{label} must be a SHA-256 hex digest") + return text + + +@dataclass(frozen=True) +class ExpertFile: + layer: int + path: Path + bytes: int + sha256: str + + +@dataclass(frozen=True) +class Qwen4ArtifactManifest: + """Validated view of ``manifest.json``. + + ``raw`` is retained for provenance and future additive fields. Paths are absolute + Z: paths so callers never accidentally resolve a relative sidecar against cwd. + """ + + path: Path + raw: Mapping[str, Any] + active_path: Path + ple_manifest_path: Path + ple_data_bytes: int + ple_sha256: str + expert_files: tuple[ExpertFile, ...] + file_tier_layers: tuple[int, ...] + resident_layers: tuple[int, ...] + + @property + def root(self) -> Path: + return self.path.parent + + @property + def manifest_path(self) -> Path: + return self.path + + def __getitem__(self, key: str): + return self.raw[key] + + def get(self, key: str, default=None): + return self.raw.get(key, default) + + @property + def text_only(self) -> bool: + return bool(self.raw.get("text_only", False)) + + @property + def active_format(self) -> str: + return str(_require_mapping(self.raw.get("active"), label="active")["format"]) + + @property + def source(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("source", {}), label="source") + + @property + def active(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("active"), label="active") + + @property + def ple(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("ple"), label="ple") + + @property + def experts(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("experts"), label="experts") + + @property + def expert_format(self) -> str: + return str(_require_mapping(self.raw.get("experts"), label="experts")["format"]) + + @property + def num_layers(self) -> int: + layers = set(self.file_tier_layers) | set(self.resident_layers) + return max(layers) + 1 if layers else 0 + + def file_for_layer(self, layer_id: int) -> ExpertFile: + for entry in self.expert_files: + if entry.layer == int(layer_id): + return entry + raise Qwen4ArtifactError(f"manifest has no expert file for layer {layer_id}") + + +def _validate_layers(value: object, *, label: str) -> tuple[int, ...]: + if not isinstance(value, list): + raise Qwen4ArtifactError(f"experts.{label} must be a list") + result = [] + for item in value: + if isinstance(item, bool): + raise Qwen4ArtifactError(f"experts.{label} contains a non-integer layer") + try: + layer = int(item) + except (TypeError, ValueError) as exc: + raise Qwen4ArtifactError(f"experts.{label} contains a non-integer layer") from exc + if layer < 0 or layer in result: + raise Qwen4ArtifactError(f"experts.{label} contains an invalid/duplicate layer") + result.append(layer) + return tuple(result) + + +def _read_manifest_path(model_path: str | os.PathLike[str]) -> Path | None: + candidate = Path(model_path) + if candidate.name == MANIFEST_NAME and candidate.is_file(): + return _resolve_z(candidate, label="Qwen4 modular manifest") + if candidate.is_dir(): + path = candidate / MANIFEST_NAME + if path.is_file(): + return _resolve_z(path, label="Qwen4 modular manifest") + return None + + +def load_qwen4_artifact_manifest( + model_path: str | os.PathLike[str], *, require: bool = False +) -> Qwen4ArtifactManifest | None: + """Load and validate a Qwen4 modular manifest. + + ``None`` means no manifest is present. This is the compatibility path for all + unmarked checkpoints. ``require=True`` is useful at a marked call site where a + missing manifest must not silently fall back to source weights. + """ + + path = _read_manifest_path(model_path) + if path is None: + if require: + raise Qwen4ArtifactError(f"Qwen4 modular manifest missing under {model_path}") + return None + try: + with path.open("r", encoding="utf-8") as handle: + raw = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + raise Qwen4ArtifactError(f"cannot read Qwen4 modular manifest {path}") from exc + root = _require_mapping(raw, label="Qwen4 modular manifest") + if root.get("format") != FORMAT or int(root.get("version", -1)) != VERSION: + raise Qwen4ArtifactError("unsupported Qwen4 modular manifest format/version") + if root.get("text_only") is not True: + raise Qwen4ArtifactError("Qwen4 modular manifest must declare text_only=true") + + active = _require_mapping(root.get("active"), label="active") + if active.get("format") != ACTIVE_FORMAT: + raise Qwen4ArtifactError(f"unsupported active format {active.get('format')!r}") + active_path = _path_from(path.parent, active.get("path"), label="active.path") + if "bytes" in active and int(active["bytes"]) < 0: + raise Qwen4ArtifactError("active.bytes must be non-negative") + if "sha256" in active: + _require_sha(active["sha256"], label="active.sha256") + + ple = _require_mapping(root.get("ple"), label="ple") + if ple.get("format") != PLE_FORMAT: + raise Qwen4ArtifactError(f"unsupported PLE format {ple.get('format')!r}") + if str(ple.get("required_volume", REQUIRED_VOLUME)).upper() != REQUIRED_VOLUME: + raise Qwen4ArtifactError("Qwen4 PLE sidecar must reside on Z:") + ple_manifest_path = _path_from(path.parent, ple.get("manifest"), label="ple.manifest") + ple_data_bytes = int(ple.get("data_bytes", 0)) + if ple_data_bytes < 0: + raise Qwen4ArtifactError("ple.data_bytes must be non-negative") + ple_sha256 = _require_sha(ple.get("sha256"), label="ple.sha256") + + experts = _require_mapping(root.get("experts"), label="experts") + if experts.get("format") != EXPERT_FORMAT: + raise Qwen4ArtifactError(f"unsupported expert format {experts.get('format')!r}") + if str(experts.get("required_volume", REQUIRED_VOLUME)).upper() != REQUIRED_VOLUME: + raise Qwen4ArtifactError("Qwen4 expert sidecars must reside on Z:") + file_layers = _validate_layers(experts.get("file_tier_layers"), label="file_tier_layers") + resident_layers = _validate_layers(experts.get("resident_layers"), label="resident_layers") + if set(file_layers) & set(resident_layers): + raise Qwen4ArtifactError("experts file_tier_layers and resident_layers overlap") + files_raw = experts.get("files") + if not isinstance(files_raw, list) or not files_raw: + raise Qwen4ArtifactError("experts.files must be a non-empty list") + files: list[ExpertFile] = [] + seen: set[int] = set() + for item in files_raw: + entry = _require_mapping(item, label="experts.files[]") + try: + layer = int(entry["layer"]) + size = int(entry["bytes"]) + except (KeyError, TypeError, ValueError) as exc: + raise Qwen4ArtifactError("expert file requires integer layer and bytes") from exc + if layer < 0 or layer in seen or size < 0: + raise Qwen4ArtifactError("expert file has invalid/duplicate layer or bytes") + seen.add(layer) + files.append( + ExpertFile( + layer=layer, + path=_path_from(path.parent, entry.get("path"), label=f"experts.files[{layer}].path"), + bytes=size, + sha256=_require_sha(entry.get("sha256"), label=f"experts.files[{layer}].sha256"), + ) + ) + declared_layers = set(file_layers) | set(resident_layers) + if declared_layers != seen: + raise Qwen4ArtifactError( + "experts.files layers must exactly match file_tier_layers + resident_layers" + ) + return Qwen4ArtifactManifest( + path=path, + raw=root, + active_path=active_path, + ple_manifest_path=ple_manifest_path, + ple_data_bytes=ple_data_bytes, + ple_sha256=ple_sha256, + expert_files=tuple(sorted(files, key=lambda item: item.layer)), + file_tier_layers=file_layers, + resident_layers=resident_layers, + ) + + +def qwen4_text_only_marker(config: Any) -> bool: + """Validate and return the explicit target config marker. + + The marker is deliberately target-specific. A typo or future marker is an error, + not permission to disable vision under an unknown policy. + """ + + marker = getattr(config, "freetoken_text_only", None) + if marker is None: + return False + if marker != TEXT_ONLY_MARKER: + raise Qwen4ArtifactError(f"unsupported freetoken_text_only marker {marker!r}") + return True + + +def build_mixed_expert_sources( + manifest: Qwen4ArtifactManifest, + *, + num_experts: int = 512, + resident_residency: list[str] | None = None, + allocator=None, + verify_hash: bool = True, +): + """Materialize only resident layers from bounded expert sidecars. + + Each resident layer is allocated as six independent :class:`HostBank` buffers and + filled one record at a time through ``FileExpertSource.read_record``. File-tier + layers remain ``None`` in the returned bank lists and retain an open + ``FileExpertSource`` for demand paging. ``allocator`` is an injectable + ``(shape, dtype) -> buffer`` callback for CPU-only tests; a production call uses + the normal HostBank allocator and settles each completed layer to the requested + residency class. + """ + + if not isinstance(manifest, Qwen4ArtifactManifest): + raise TypeError("manifest must be a validated Qwen4ArtifactManifest") + num_experts = int(num_experts) + from freetoken.moe.expert_source import FileExpertSource + from freetoken.moe.host_banks import HostBank, HostResidency + + if not 1 <= num_experts <= FileExpertSource.num_experts: + raise ValueError( + f"Qwen4 modular expert sidecars support num_experts in [1, {FileExpertSource.num_experts}]" + ) + + layers = manifest.num_layers + residency = resident_residency or [HostResidency.PINNED.value] * layers + if len(residency) != layers: + raise ValueError(f"resident_residency has {len(residency)} layers, expected {layers}") + unknown_residency = set(residency) - {item.value for item in HostResidency} + if unknown_residency: + raise ValueError(f"unknown host residency values: {sorted(unknown_residency)}") + if allocator is None: + allocator = lambda shape, dtype: HostBank(shape, dtype) + + sources = {name: [None] * layers for name in FileExpertSource.bank_schema} + file_sources = {} + for layer in sorted(manifest.file_tier_layers): + entry = manifest.file_for_layer(layer) + file_sources[layer] = FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_layer_id=layer, + num_experts=num_experts, + verify_hash=verify_hash, + ) + + # Resident layers are bounded by one six-plane record at a time. We do not + # retain a second full-layer staging tensor and close each source after fill. + for layer in sorted(manifest.resident_layers): + entry = manifest.file_for_layer(layer) + with FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_layer_id=layer, + num_experts=num_experts, + verify_hash=verify_hash, + ) as source: + buffers = { + name: allocator(shape=(num_experts, *shape), dtype=dtype) + for name, (shape, dtype) in FileExpertSource.plane_specs.items() + } + for expert_id in range(num_experts): + row = source.read_record(expert_id) + for name, buffer in buffers.items(): + destination = getattr(buffer, "tensor", buffer) + destination[expert_id].copy_(row[name]) + settle = residency[layer] + for buffer in buffers.values(): + if settle == HostResidency.PINNED.value and hasattr(buffer, "pin"): + buffer.pin() + elif settle == HostResidency.LOCKED.value and hasattr(buffer, "lock"): + buffer.lock() + for name, buffer in buffers.items(): + sources[name][layer] = getattr(buffer, "tensor", buffer) + return sources, file_sources + + +def configure_mixed_expert_sources( + cache, + manifest: Qwen4ArtifactManifest | Mapping[str, Any], + resident_sources, + *, + file_sources: Mapping[int, object] | None = None, +): + """Wire resident bank entries and file-backed layers into an offload cache. + + ``resident_sources`` is injected by the caller (normally the existing FTW bank + loader). This keeps the helper unit-testable with tiny synthetic tensors and avoids + implementing another expert writer here. File tiers use only the public + :class:`FileExpertSource` reader API. + """ + + if not isinstance(manifest, Qwen4ArtifactManifest): + raise TypeError("manifest must be a validated Qwen4ArtifactManifest") + if cache.decode_target != "gpu": + raise ValueError("Qwen4 modular file tiers are GPU-only") + if cache.prefill_overlap: + raise ValueError("Qwen4 modular file tiers require prefill_overlap=False") + # The artifact's native expert geometry is 512 slots. Enforce this independently + # of a malformed/synthetic model config so rebuilds cannot create an undersized cache. + if int(cache.cache_size) < 512: + raise ValueError("Qwen4 modular expert cache requires at least 512 slots") + if not isinstance(resident_sources, Mapping): + raise TypeError("resident_sources must be a bank-name -> per-layer mapping") + layers = manifest.num_layers + if layers <= 0 or int(cache.num_layers) != layers: + raise ValueError( + f"manifest layer geometry ({layers}) does not match cache ({cache.num_layers})" + ) + file_layers = set(manifest.file_tier_layers) + resident_layers = set(manifest.resident_layers) + if file_layers | resident_layers != set(range(layers)): + raise ValueError("manifest expert layer sets must cover a contiguous model") + if set(resident_sources) != set(cache.bank_schema): + raise ValueError("resident bank schema does not match cache quant_format") + per_layer: dict[str, list[Any]] = {} + for name in cache.bank_schema: + values = list(resident_sources[name]) + if len(values) != layers: + raise ValueError(f"resident bank {name!r} has {len(values)} layers, expected {layers}") + for layer in file_layers: + if values[layer] is not None: + # A file tier must be represented explicitly as a None resident source. + values[layer] = None + for layer in resident_layers: + if values[layer] is None: + raise ValueError(f"resident layer {layer} has no resident source for {name}") + per_layer[name] = values + cache.set_bank_sources(per_layer) + from freetoken.moe.expert_source import FileExpertSource + + if file_sources is None: + file_sources = {} + for layer in sorted(file_layers): + entry = manifest.file_for_layer(layer) + source = FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_layer_id=layer, + num_experts=cache.num_experts, + ) + file_sources[layer] = source + else: + file_sources = {int(layer): source for layer, source in file_sources.items()} + if set(file_sources) != file_layers: + raise ValueError("file_sources keys do not match manifest file_tier_layers") + cache.set_file_sources(dict(file_sources)) + return dict(file_sources) + + +__all__ = [ + "ACTIVE_FORMAT", + "ARTIFACT_FORMAT", + "EXPERT_FORMAT", + "FORMAT", + "MANIFEST_NAME", + "PLE_FORMAT", + "Qwen4ArtifactError", + "Qwen4ArtifactManifest", + "TEXT_ONLY_MARKER", + "VERSION", + "configure_mixed_expert_sources", + "build_mixed_expert_sources", + "load_qwen4_artifact_manifest", + "qwen4_text_only_marker", +] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index dd9a8499b..bb60c05f0 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -506,6 +506,65 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: + # Qwen4 modular artifacts carry their own mixed expert tier. The ordinary + # loader cannot infer file-backed layers from a generic checkpoint path, so + # resolve the manifest first and route through the explicit wiring seam. + from freetoken.checkpoint.qwen4_artifact import ( + build_mixed_expert_sources, + configure_mixed_expert_sources, + load_qwen4_artifact_manifest, + ) + + artifact = load_qwen4_artifact_manifest(config.model_path) + if artifact is not None: + if config.moe_backend != "offload": + raise ValueError( + "Qwen4 modular expert tiers are GPU-only; use --moe-backend offload" + ) + if config.moe_prefill_overlap: + object.__setattr__(config, "moe_prefill_overlap", False) + # Expert sidecars are authoritative for both tiers. Resident layers are + # streamed one record at a time into HostBanks; file layers remain open + # FileExpertSource readers and never materialize a full layer. + resident_sources, file_sources = build_mixed_expert_sources( + artifact, + num_experts=config.model_config.num_experts, + resident_residency=["pinned"] * artifact.num_layers, + verify_hash=not config.use_dummy_weight, + ) + from freetoken.moe.expert_banks import ExpertBanks + + banks = ExpertBanks("nvfp4", resident_sources) + if config.moe_cache_auto: + size, pages, _overlap = self._resolve_auto_moe_cache_size(config, banks) + object.__setattr__(config, "moe_cache_size", max(size, 512)) + if config.num_page_override is None: + object.__setattr__(config, "num_page_override", pages) + _require_offload_cache_size(config.moe_cache_size, 512) + cache = OffloadMoeCache( + num_layers=config.model_config.num_moe_layers, + num_experts=config.model_config.num_experts, + cache_size=config.moe_cache_size, + device=self.device, + cache_policy=config.moe_cache_policy, + prefill_overlap=False, + prefill_hit_d2d=False, + quant_format=banks.quant_format, + decode_target="gpu", + hybrid_max_fetch=config.moe_hybrid_max_fetch, + ) + cache.cpu_layer_ids = frozenset() + configure_mixed_expert_sources( + cache, artifact, banks.sources, file_sources=file_sources + ) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = config.moe_collect_stats + layers = attach_offload_moe_cache(self.model, cache) + assert len(layers) == config.model_config.num_moe_layers + self.ctx.moe_offload_cache = cache + self.moe_offload_cache = cache + return cache + # A model may fully own cache construction via make_offload_moe_cache. # Otherwise load_expert_banks gives the model module a setup hook first, then # falls back to per-quant providers, and the engine wires the banks into cache. diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 2846f1ba2..ccd0a319f 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -9,6 +9,7 @@ RotaryConfig, detect_expert_quant, ) +from freetoken.checkpoint.qwen4_artifact import qwen4_text_only_marker from .args import Qwen4ExpArgs, Qwen4VisionConfig @@ -94,7 +95,11 @@ def parse_config(hf_config: Any) -> ModelConfig: "Qwen4-Exp mrope_section must cover the rotary dimension: " f"{qwen4_args.mrope_section} vs {rotary_dim}" ) - raw_vision = getattr(hf_config, "vision_config", None) + # A modular Qwen4 artifact is explicitly text-only. Keep the ordinary source + # checkpoint behavior untouched (vision remains part of the parsed model) and + # fail closed on an unknown target marker in qwen4_text_only_marker(). + text_only = qwen4_text_only_marker(hf_config) + raw_vision = None if text_only else getattr(hf_config, "vision_config", None) vision_config = None if raw_vision is not None: vision_config = Qwen4VisionConfig( @@ -182,7 +187,7 @@ def parse_config(hf_config: Any) -> ModelConfig: # Qwen3.8-Flash-Next is a VL checkpoint. Vision is part of this model, # not an optional text-only add-on. vision_config=vision_config, - image_token_id=getattr(hf_config, "image_token_id", None), + image_token_id=None if text_only else getattr(hf_config, "image_token_id", None), attention_groups=groups, qwen4_args=qwen4_args, # PLE keeps per-request dilated-convolution state outside the generic diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 2e684ef60..84e6bb27c 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -745,6 +745,36 @@ def __init__(self, config: ModelConfig): self._image_token_id = config.image_token_id def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + # Modular Qwen4 artifacts carry an explicit Q3_PLE sidecar. Once the + # target marker is present, selecting that sidecar is mandatory: silently + # falling back to the 51-GiB FP8 safetensors table would defeat the artifact's + # bounded host-memory contract. Unmarked source checkpoints keep the original + # FP8 path unchanged. + if not dummy: + from freetoken.checkpoint.qwen4_artifact import ( + load_qwen4_artifact_manifest, + qwen4_text_only_marker, + ) + + artifact = load_qwen4_artifact_manifest(model_path) + if artifact is not None: + self.load_q3_ple_weights(str(artifact.ple_manifest_path)) + return + # A known text-only target marker without its modular manifest is + # incomplete. Do not silently reopen the source FP8 PLE table. + try: + from freetoken.utils import cached_load_hf_config + + if qwen4_text_only_marker(cached_load_hf_config(model_path)): + raise ValueError( + "Qwen4 text-only target is marked but manifest.json is missing" + ) + except ValueError: + raise + except Exception: + # Unmarked hub/source paths retain the historical loader behavior; + # parse_config remains the authoritative marker validator. + pass for layer in self.layers.op_list: if layer.ple is not None: layer.ple.load_host_weights(model_path, dummy=dummy) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9b..c8cf2c494 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -232,8 +232,28 @@ def load_weight( # fails loudly in load_state_dict (strict missing/unexpected expert keys), so the reader # just yields the stored weight tensors regardless of the include_moe_experts flag. from freetoken.checkpoint.ftw import is_ftw_checkpoint, iter_ftw_weights + from freetoken.checkpoint.qwen4_artifact import load_qwen4_artifact_manifest from freetoken.models.config import VISION_KEY_PREFIXES, vision_load_enabled + # A modular Qwen4 artifact keeps active weights in a nested FTW target while + # the root directory owns the manifest, Q3 PLE sidecar, and mixed expert tiers. + # Resolve that target before the ordinary FTW/source branches. The resolver is + # deliberately optional so every unmarked checkpoint follows the historical path. + artifact = load_qwen4_artifact_manifest(model_path) + if artifact is not None: + active_path = artifact.active_path + if active_path.is_file() and active_path.name == "freetoken_weight.json": + active_path = active_path.parent + if not is_ftw_checkpoint(str(active_path)): + raise ValueError(f"Qwen4 modular active target is not an FTW checkpoint: {active_path}") + for name, tensor in iter_ftw_weights(str(active_path)): + # Qwen4's loader uses ``visual.*`` after canonical renaming; the generic + # prefixes cover model variants that retain ``vision_tower``/``embed_vision``. + if name.startswith(("visual.",) + VISION_KEY_PREFIXES): + continue + yield name, tensor + return + if is_ftw_checkpoint(model_path): # The FTW dense shard stores whatever existed at conversion, including the vision # stack. Vision is opt-in (default OFF, see vision_load_enabled): when it is off the diff --git a/tests/checkpoint/test_qwen4_artifact.py b/tests/checkpoint/test_qwen4_artifact.py new file mode 100644 index 000000000..7a7778866 --- /dev/null +++ b/tests/checkpoint/test_qwen4_artifact.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch + +from freetoken.checkpoint.qwen4_artifact import ( + FORMAT, + TEXT_ONLY_MARKER, + configure_mixed_expert_sources, + build_mixed_expert_sources, + load_qwen4_artifact_manifest, + qwen4_text_only_marker, +) +from freetoken.moe.expert_source import FileExpertSource, RAW_RECORD_BYTES + + +@pytest.fixture +def z_fixture_dir(): + root = Path.cwd() / ".stage7-test-fixtures" / uuid4().hex + root.mkdir(parents=True, exist_ok=False) + try: + assert (root.drive or "").upper() == "Z:", root + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +class _FakeCache: + bank_schema = FileExpertSource.bank_schema + decode_target = "gpu" + prefill_overlap = False + cache_size = 512 + num_layers = 2 + num_experts = 1 + + def set_bank_sources(self, sources): + self.bank_sources = sources + + def set_file_sources(self, sources): + self.file_sources = sources + + +def _manifest(root: Path, sidecar: Path, digest: str) -> Path: + data = { + "format": FORMAT, + "version": 1, + "text_only": True, + "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": "0" * 64}, + "minimum_freetoken_commit": "0" * 40, + "tvm_ffi_patch_sha256": "0" * 64, + "active": {"format": "nvfp4_w4a16_v1", "path": ".", "bytes": 0, "sha256": "0" * 64}, + "ple": { + "format": "q3_ple_32", + "manifest": "ple-q3.json", + "data_bytes": 0, + "sha256": "0" * 64, + "required_volume": "Z:", + }, + "experts": { + "format": "ftexpert1_nvfp4_v1", + "files": [{"layer": 0, "path": sidecar.name, "bytes": sidecar.stat().st_size, "sha256": digest}], + "file_tier_layers": [0], + "resident_layers": [1], + "required_volume": "Z:", + }, + "metadata": {"config_sha256": "0" * 64}, + "complete_artifact_fingerprint": "0" * 64, + } + path = root / "manifest.json" + path.write_text(json.dumps(data), encoding="utf-8") + return path + + +def test_qwen4_marker_is_explicit_and_unknown_fails_closed(): + class Config: + freetoken_text_only = TEXT_ONLY_MARKER + + assert qwen4_text_only_marker(Config()) is True + Config.freetoken_text_only = "future_policy" + with pytest.raises(ValueError, match="unsupported freetoken_text_only"): + qwen4_text_only_marker(Config()) + + +def test_manifest_reopens_and_wires_mixed_sources(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + manifest_path = _manifest(z_fixture_dir, sidecar, digest) + manifest = load_qwen4_artifact_manifest(manifest_path, require=True) + assert manifest is not None + assert manifest.file_tier_layers == (0,) + assert manifest.resident_layers == (1,) + + resident = {name: [None, object()] for name in FileExpertSource.bank_schema} + cache = _FakeCache() + sources = configure_mixed_expert_sources(cache, manifest, resident) + assert sorted(sources) == [0] + assert all(cache.bank_sources[name][0] is None for name in cache.bank_schema) + assert all(cache.bank_sources[name][1] is not None for name in cache.bank_schema) + assert cache.file_sources[0].layer_id == 0 + cache.file_sources[0].close() + + +def test_resident_builder_streams_one_record_without_full_layer(z_fixture_dir): + first_path = z_fixture_dir / "experts-L00.nvfp4" + first_digest = FileExpertSource.create_synthetic( + first_path, num_experts=1, records=[bytes([8]) * RAW_RECORD_BYTES], layer_id=0 + ) + resident_path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic( + resident_path, num_experts=1, records=[bytes([9]) * RAW_RECORD_BYTES], layer_id=1 + ) + manifest_data = { + "format": FORMAT, + "version": 1, + "text_only": True, + "active": {"format": "nvfp4_w4a16_v1", "path": "."}, + "ple": {"format": "q3_ple_32", "manifest": "ple-q3.json", "sha256": "0" * 64}, + "experts": { + "format": "ftexpert1_nvfp4_v1", + "files": [{"layer": 0, "path": first_path.name, "bytes": first_path.stat().st_size, "sha256": first_digest}, + {"layer": 1, "path": resident_path.name, "bytes": resident_path.stat().st_size, "sha256": digest}], + "file_tier_layers": [], + "resident_layers": [0, 1], + "required_volume": "Z:", + }, + } + path = z_fixture_dir / "manifest.json" + path.write_text(json.dumps(manifest_data), encoding="utf-8") + manifest = load_qwen4_artifact_manifest(path, require=True) + resident_sources, file_sources = build_mixed_expert_sources( + manifest, + num_experts=1, + resident_residency=["pageable", "pageable"], + allocator=lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + assert file_sources == {} + assert resident_sources["gate_up_packed"][0].shape == (1, 1280, 1280) + assert int(resident_sources["gate_up_packed"][1][0, 0, 0]) == 9 From b6c53cb05a933ca441eb3d4f293e2f873291b2e9 Mon Sep 17 00:00:00 2001 From: Nicholas Mattteo <130715692+nickmatteo@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:21:57 -0400 Subject: [PATCH 10/17] feat(qwen4): orchestrate modular artifact conversion --- python/freetoken/checkpoint/__init__.py | 8 +- python/freetoken/checkpoint/convert.py | 80 ++- python/freetoken/checkpoint/q3_ple.py | 108 +++- python/freetoken/checkpoint/qwen4_artifact.py | 511 ++++++++++++++++-- python/freetoken/engine/cache_budget.py | 8 +- python/freetoken/engine/engine.py | 56 +- python/freetoken/models/qwen4_exp/model.py | 6 +- python/freetoken/models/weight.py | 1 + python/freetoken/moe/__init__.py | 2 + python/freetoken/moe/expert_source.py | 100 +++- python/freetoken/moe/offload_cache.py | 11 +- tests/checkpoint/test_convert_metadata.py | 108 +++- tests/checkpoint/test_q3_ple_writer.py | 54 +- tests/checkpoint/test_qwen4_artifact.py | 369 ++++++++++++- tests/engine/test_cache_budget.py | 8 + tests/models/test_qwen4_exp_raw_config.py | 45 ++ tests/moe/test_expert_sidecar_writer.py | 34 +- tests/moe/test_file_expert_source.py | 8 +- 18 files changed, 1402 insertions(+), 115 deletions(-) diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index f004cb5ce..f78d90ca9 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -12,10 +12,16 @@ load_ftw_banks, ) from .convert import convert_checkpoint -from .q3_ple import Q3PLEReader, Q3PLESegment, write_q3_ple_sidecar +from .q3_ple import ( + Q3PLEReader, + Q3PLESegment, + write_q3_ple_from_safetensors, + write_q3_ple_sidecar, +) __all__ = [ "FTWReader", "FTWWriter", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", "Q3PLEReader", "Q3PLESegment", "write_q3_ple_sidecar", + "write_q3_ple_from_safetensors", ] diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 223d61b12..106cbdb22 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -114,7 +114,9 @@ def _copy_host_mapped_weights(model_path: str, out_dir: str) -> list[str]: return copied -def _copy_metadata(model_path: str, out_dir: str) -> list[str]: +def _copy_metadata( + model_path: str, out_dir: str, *, include_host_mapped_weights: bool = True +) -> list[str]: """Copy all non-weight files (config, tokenizer, remote-code, nested model configs) preserving directory structure, so the FTW dir is a self-contained checkpoint.""" if os.path.isfile(model_path): @@ -148,10 +150,22 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]: os.makedirs(os.path.dirname(dst), exist_ok=True) shutil.copy2(src, dst) copied.append(rel) - copied.extend(_copy_host_mapped_weights(model_path, out_dir)) + if include_host_mapped_weights: + copied.extend(_copy_host_mapped_weights(model_path, out_dir)) return copied +def _iter_qwen4_modular_dense_entries(entries): + """Apply the frozen active map and text-only filtering to a source stream.""" + from freetoken.models.config import VISION_KEY_PREFIXES + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + for name, tensor in iter_active_nvfp4_runtime_entries(entries): + if name.startswith(("visual.",) + VISION_KEY_PREFIXES): + continue + yield name, tensor + + class _ConvertSink: """Layer-completion sink for ``load_expert_banks(layer_sink=...)``: writes each completed layer's banks as their own FTW entries immediately (name @@ -218,6 +232,7 @@ def convert_checkpoint( shard_limit: int = DEFAULT_SHARD_LIMIT, device: str | None = None, artifact_format: str | None = None, + source_inventory_sha256: str | None = None, ) -> dict: """Write ``model_path`` as an FTW checkpoint at ``out_dir``. Returns the index dict. @@ -240,26 +255,45 @@ def convert_checkpoint( f"FTW conversion runs single-process and the format records no TP layout, " f"but TP is already set to size={tp.size}" ) - dev = torch.device(device or "cuda:0") - torch.cuda.set_device(dev) - torch.zeros(1, device=dev) # init CUDA context (needed by nvfp4 backend pick / pinning) - - cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), - dtype=dtype, moe_backend=moe_backend) - mc = cfg.model_config if artifact_format not in (None, "qwen4_modular_v1"): raise ValueError( f"unsupported artifact_format {artifact_format!r}; expected None or 'qwen4_modular_v1'" ) + # The modular target is pre-encoded entirely on CPU. It deliberately omits + # expert banks from this active FTW component, so initializing CUDA here would + # add an unnecessary conversion dependency and obscure the zero-VRAM envelope. + dev = torch.device("cpu" if artifact_format == "qwen4_modular_v1" else (device or "cuda:0")) + if dev.type == "cuda": + torch.cuda.set_device(dev) + torch.zeros(1, device=dev) # needed by legacy expert backend selection / pinning + + cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), + dtype=dtype, moe_backend=moe_backend) + mc = cfg.model_config is_qwen4 = any("Qwen4" in str(arch) for arch in getattr(mc, "architectures", ())) if artifact_format is not None and not is_qwen4: raise ValueError("artifact_format='qwen4_modular_v1' requires a Qwen4 checkpoint") - offload = moe_backend == "offload" and getattr(mc, "is_moe", False) - include_moe_experts = not offload + modular = artifact_format == "qwen4_modular_v1" + if modular: + source_inventory_sha256 = str(source_inventory_sha256 or "").lower() + if len(source_inventory_sha256) != 64 or any( + char not in "0123456789abcdef" for char in source_inventory_sha256 + ): + raise ValueError( + "qwen4_modular_v1 conversion requires source_inventory_sha256" + ) + offload = not modular and moe_backend == "offload" and getattr(mc, "is_moe", False) + include_moe_experts = False if modular else not offload from freetoken.utils.progress import byte_bar, count_bar - writer = FTWWriter(out_dir, shard_limit=shard_limit) + # For the modular target ``out_dir`` is the artifact root; active FTW bytes + # live in their own component directory while config/tokenizer metadata stays + # at the root used by normal Engine startup. + active_out_dir = ( + os.path.join(out_dir, "qwen4-active-v1.ftw") if modular else out_dir + ) + writer = FTWWriter(active_out_dir, shard_limit=shard_limit) n_weight = n_bank = n_alpha = 0 # 1) dense weights (host tensors; load straight to CPU to avoid GPU pressure) @@ -275,20 +309,7 @@ def convert_checkpoint( # fallback. The Qwen4 iterator has already fused canonical projections; this # wrapper only converts the frozen active map while leaving routers, PLE, and # all non-active entries untouched. - from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries - - def _modular_dense_entries(): - from freetoken.models.config import VISION_KEY_PREFIXES - - for name, tensor in iter_active_nvfp4_runtime_entries(dense_entries): - # The modular target is text-only by contract. Qwen4's canonical - # rename uses ``visual.*`` while other wrappers retain the generic - # vision prefixes; drop both explicitly at conversion time. - if name.startswith(("visual.",) + VISION_KEY_PREFIXES): - continue - yield name, tensor - - dense_entries = _modular_dense_entries() + dense_entries = _iter_qwen4_modular_dense_entries(dense_entries) for name, tensor in count_bar(dense_entries, "Converting dense weights"): writer.add_tensor(name, tensor, kind="weight") @@ -355,7 +376,11 @@ def _modular_dense_entries(): bar.close() _progress("finalize") # writing shard index + copying config/tokenizer - copied = _copy_metadata(model_path, out_dir) + copied = _copy_metadata( + model_path, + out_dir, + include_host_mapped_weights=artifact_format != "qwen4_modular_v1", + ) if artifact_format == "qwen4_modular_v1": config_path = os.path.join(out_dir, "config.json") if not os.path.isfile(config_path): @@ -378,6 +403,7 @@ def _modular_dense_entries(): index = writer.finalize({ "source_model_path": os.path.abspath(model_path), "fingerprint": fingerprint, + "source_inventory_sha256": source_inventory_sha256 if modular else None, # quant_format records the actual on-disk bank layout (e.g. nvfp4_marlin vs # nvfp4_b12x): the suffix is a runtime backend pick (GPU capability / env), NOT in # config, and the stored bytes are physically repacked into it -- so it's kept and diff --git a/python/freetoken/checkpoint/q3_ple.py b/python/freetoken/checkpoint/q3_ple.py index b302b40a6..cd62cc5b2 100644 --- a/python/freetoken/checkpoint/q3_ple.py +++ b/python/freetoken/checkpoint/q3_ple.py @@ -60,6 +60,9 @@ def _z_output_path(path: str | os.PathLike[str]) -> Path: candidate = Path(path).expanduser() if not candidate.is_absolute(): raise ValueError(f"Q3_PLE_32 output path must be absolute: {path}") + lexical_drive, _ = os.path.splitdrive(str(candidate)) + if lexical_drive.upper() != "Z:" and not str(candidate).lower().startswith("/z/"): + raise ValueError(f"Q3_PLE_32 output must resolve to Z:, got {candidate}") parent = candidate.parent.resolve(strict=True) resolved = parent / candidate.name drive, _ = os.path.splitdrive(str(resolved)) @@ -306,10 +309,15 @@ def _parse_segment(self, item: object) -> Q3PLESegment: def _validate_segments(self) -> None: expected_row = 0 previous_end = 0 + contiguous = self.manifest.get("storage_layout") == "contiguous_rows_v1" for segment in self.segments: if segment.first_row != expected_row or segment.end_row <= segment.first_row: raise ValueError("Q3_PLE_32 segment rows have a gap, overlap, or bad order") - if segment.data_offset < 0 or segment.data_offset % ALIGN: + if segment.data_offset < 0: + raise ValueError("Q3_PLE_32 segment data offset is negative") + if contiguous and segment.data_offset != previous_end: + raise ValueError("Q3_PLE_32 contiguous segment directory has a gap or overlap") + if not contiguous and segment.data_offset % ALIGN: raise ValueError("Q3_PLE_32 segment data offset is not 4 KiB aligned") if segment.byte_length != segment.rows * ROW_BYTES: raise ValueError("Q3_PLE_32 segment byte length does not match rows") @@ -449,11 +457,11 @@ def _validate_segment_directory( digest = str(segment["sha256"]) if first_row != expected_row or end_row <= first_row: raise ValueError("Q3_PLE_32 writer generated a malformed segment directory") - if data_offset < 0 or data_offset % ALIGN: - raise ValueError("Q3_PLE_32 writer generated an unaligned segment") + if data_offset != previous_end: + raise ValueError("Q3_PLE_32 writer generated a non-contiguous segment") if byte_length != (end_row - first_row) * ROW_BYTES: raise ValueError("Q3_PLE_32 writer generated a segment length mismatch") - if data_offset < previous_end or data_offset + byte_length > file_bytes: + if data_offset + byte_length > file_bytes: raise ValueError("Q3_PLE_32 writer generated overlapping/out-of-range segments") if len(digest) != 64 or any(char not in "0123456789abcdef" for char in digest): raise ValueError("Q3_PLE_32 writer generated a malformed segment hash") @@ -483,9 +491,9 @@ def write_q3_ple_sidecar( """Stream rows into an atomic, reader-compatible Q3_PLE_32 sidecar. ``rows`` is consumed exactly once and only the current 160-value row is held - in memory. Segments are split in source order and begin at 4 KiB-aligned - offsets; alignment bytes are included in the whole-file hash but never in a - segment's logical length/hash. Both files are written under unique partial + in memory. Segments are logical hash/addressing ranges over one contiguous + row stream. There is no inter-row or inter-segment padding: the production + table is exactly ``rows * 70`` bytes. Both files are written under unique partial names, fsynced, and atomically renamed into place on successful completion. The converter rule is intentionally fixed at two least-squares refinement @@ -527,18 +535,6 @@ def write_q3_ple_sidecar( current_segment: dict[str, int | str] | None = None segment_digest: hashlib._Hash | None = None - def write_padding(handle: object, count: int) -> None: - if count <= 0: - return - padding = bytes(min(1 << 20, count)) - remaining = count - while remaining: - take = min(remaining, len(padding)) - chunk = padding[:take] - handle.write(chunk) # type: ignore[attr-defined] - whole_digest.update(chunk) - remaining -= take - try: with data_partial.open("wb") as output: for source_row in rows: @@ -556,9 +552,6 @@ def write_padding(handle: object, count: int) -> None: ) current_segment["sha256"] = segment_digest.hexdigest() segments.append(current_segment) - aligned_offset = _align_up(file_offset) - write_padding(output, aligned_offset - file_offset) - file_offset = aligned_offset current_segment = { "first_row": rows_written, "end_row": rows_written, @@ -615,6 +608,7 @@ def write_padding(handle: object, count: int) -> None: "rows": rows_written, "payload_bytes": rows_written * ROW_BYTES, "file_bytes": file_bytes, + "storage_layout": "contiguous_rows_v1", "data_file": os.path.relpath(data_final, manifest_final.parent), "weight_scale": global_scale, "source_fingerprint": source_digest, @@ -636,6 +630,75 @@ def write_padding(handle: object, count: int) -> None: return manifest +def write_q3_ple_from_safetensors( + model_path: str | os.PathLike[str], + data_path: str | os.PathLike[str], + manifest_path: str | os.PathLike[str], + *, + layer_id: int, + split_parts: int, + source_fingerprint: str, + rows_per_chunk: int = 8192, + segment_rows: int = DEFAULT_SEGMENT_ROWS, +) -> dict: + """Stream the official FP8 PLE shards into the native Q3 sidecar. + + Shards and rows are consumed in exact ``shard_0..shard_N`` order. A + Safetensors slice is read in bounded row chunks; the full 51.2-GiB table is + never materialized. The source per-model ``weight_scale`` remains a separate + scalar in the Q3 manifest and is not folded into block scales. + """ + + folder = Path(model_path).expanduser().resolve() + if not folder.is_dir(): + raise ValueError(f"Q3 PLE source must be a local checkpoint directory: {folder}") + if rows_per_chunk <= 0 or split_parts <= 0: + raise ValueError("rows_per_chunk and split_parts must be positive") + index_path = folder / "model.safetensors.index.json" + with index_path.open("r", encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + prefix = ( + f"model.language_model.layers.{int(layer_id)}.ple.ple_embedding." + "ngram_embedding" + ) + shard_keys = [f"{prefix}.shard_{part}.weight" for part in range(int(split_parts))] + missing = [key for key in shard_keys if key not in weight_map] + scale_key = prefix + ".weight_scale" + if missing or scale_key not in weight_map: + raise ValueError(f"incomplete PLE source mapping under {prefix}") + + import safetensors + + scale_file = folder / weight_map[scale_key] + with safetensors.safe_open(scale_file, framework="pt", device="cpu") as handle: + scale = handle.get_tensor(scale_key).reshape(()) + weight_scale = float(scale.float().item()) + + def iter_rows(): + for key in shard_keys: + source_file = folder / weight_map[key] + with safetensors.safe_open(source_file, framework="pt", device="cpu") as handle: + sliced = handle.get_slice(key) + shape = tuple(int(value) for value in sliced.get_shape()) + if len(shape) != 2 or shape[1] != ROW_VALUES: + raise ValueError(f"unexpected PLE source shape for {key}: {shape}") + for start in range(0, shape[0], int(rows_per_chunk)): + chunk = sliced[start : min(start + int(rows_per_chunk), shape[0])] + if chunk.dtype != torch.float8_e4m3fn: + raise ValueError(f"unexpected PLE source dtype for {key}: {chunk.dtype}") + for row in chunk.float(): + yield row + + return write_q3_ple_sidecar( + iter_rows(), + data_path, + manifest_path, + source_fingerprint=source_fingerprint, + weight_scale=weight_scale, + segment_rows=segment_rows, + ) + + __all__ = [ "ALIGN", "BLOCK_BYTES", @@ -651,4 +714,5 @@ def write_padding(handle: object, count: int) -> None: "quantize_block", "quantize_row", "write_q3_ple_sidecar", + "write_q3_ple_from_safetensors", ] diff --git a/python/freetoken/checkpoint/qwen4_artifact.py b/python/freetoken/checkpoint/qwen4_artifact.py index ea15897d2..4e2096286 100644 --- a/python/freetoken/checkpoint/qwen4_artifact.py +++ b/python/freetoken/checkpoint/qwen4_artifact.py @@ -13,6 +13,7 @@ from __future__ import annotations import json +import hashlib import os from dataclasses import dataclass from pathlib import Path @@ -28,6 +29,16 @@ EXPERT_FORMAT = "ftexpert1_nvfp4_v1" REQUIRED_VOLUME = "Z:" MANIFEST_NAME = "manifest.json" +ACTIVE_TARGET_BYTES = 4_804_403_200 +PLE_TARGET_BYTES = 22_400_107_520 +EXPERT_FILE_BYTES = 1_419_776_000 +EXPERT_LAYERS = 48 +EXPERT_NUM_EXPERTS = 512 +KNOWN_TARGET_BYTES = 95_353_758_720 +FILE_TIER_LAYERS = (0, 1, 2, 3, 4, 5, 42, 43, 44, 45, 46, 47) +PINNED_SOURCE_REPOSITORY = "RadixArk/Qwen3.8-Flash-Next-NVFP4" +PINNED_SOURCE_REVISION = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" +TVM_FFI_PATCH_SHA256 = "889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec" class Qwen4ArtifactError(ValueError): @@ -51,6 +62,15 @@ def _path_from(root: Path, value: object, *, label: str) -> Path: return _resolve_z(candidate, label=label) +def _path_within_root(root: Path, value: object, *, label: str) -> Path: + path = _path_from(root, value, label=label) + try: + path.relative_to(root) + except ValueError as exc: + raise Qwen4ArtifactError(f"{label} resolves outside artifact root") from exc + return path + + def _require_mapping(value: object, *, label: str) -> Mapping[str, Any]: if not isinstance(value, Mapping): raise Qwen4ArtifactError(f"{label} must be an object") @@ -70,6 +90,48 @@ class ExpertFile: path: Path bytes: int sha256: str + source_fingerprint: str + + +@dataclass(frozen=True) +class ComponentFile: + path: Path + bytes: int + sha256: str + + +def _sha256_file(path: Path, *, chunk_bytes: int = 8 << 20) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_bytes): + digest.update(chunk) + return digest.hexdigest() + + +def _verify_component_file(entry: ComponentFile, *, label: str) -> None: + if not entry.path.is_file(): + raise Qwen4ArtifactError(f"{label} is missing: {entry.path}") + actual_bytes = entry.path.stat().st_size + if actual_bytes != entry.bytes: + raise Qwen4ArtifactError( + f"{label} length mismatch: {actual_bytes} != {entry.bytes}" + ) + actual_sha = _sha256_file(entry.path) + if actual_sha != entry.sha256: + raise Qwen4ArtifactError(f"{label} SHA-256 mismatch") + + +def _component_file(root: Path, path: Path) -> dict[str, Any]: + resolved = _resolve_z(path, label="artifact component") + try: + relative = resolved.relative_to(root) + except ValueError as exc: + raise Qwen4ArtifactError(f"artifact component resolves outside root: {resolved}") from exc + return { + "path": relative.as_posix(), + "bytes": resolved.stat().st_size, + "sha256": _sha256_file(resolved), + } @dataclass(frozen=True) @@ -83,12 +145,15 @@ class Qwen4ArtifactManifest: path: Path raw: Mapping[str, Any] active_path: Path + active_files: tuple[ComponentFile, ...] ple_manifest_path: Path ple_data_bytes: int ple_sha256: str expert_files: tuple[ExpertFile, ...] + metadata_files: tuple[ComponentFile, ...] file_tier_layers: tuple[int, ...] resident_layers: tuple[int, ...] + production_geometry: bool @property def root(self) -> Path: @@ -143,6 +208,17 @@ def file_for_layer(self, layer_id: int) -> ExpertFile: return entry raise Qwen4ArtifactError(f"manifest has no expert file for layer {layer_id}") + def verify_active(self) -> None: + for index, entry in enumerate(self.active_files): + _verify_component_file(entry, label=f"active.files[{index}]") + from freetoken.checkpoint.ftw import INDEX_NAME + + with (self.active_path / INDEX_NAME).open("r", encoding="utf-8") as handle: + index = json.load(handle) + expected = str(self.source["inventory_sha256"]).lower() + if str(index.get("source_inventory_sha256", "")).lower() != expected: + raise Qwen4ArtifactError("active FTW source inventory fingerprint mismatch") + def _validate_layers(value: object, *, label: str) -> tuple[int, ...]: if not isinstance(value, list): @@ -173,7 +249,8 @@ def _read_manifest_path(model_path: str | os.PathLike[str]) -> Path | None: def load_qwen4_artifact_manifest( - model_path: str | os.PathLike[str], *, require: bool = False + model_path: str | os.PathLike[str], *, require: bool = False, + allow_synthetic_geometry: bool = False, ) -> Qwen4ArtifactManifest | None: """Load and validate a Qwen4 modular manifest. @@ -195,28 +272,105 @@ def load_qwen4_artifact_manifest( root = _require_mapping(raw, label="Qwen4 modular manifest") if root.get("format") != FORMAT or int(root.get("version", -1)) != VERSION: raise Qwen4ArtifactError("unsupported Qwen4 modular manifest format/version") + if root.get("artifact_schema") != FORMAT: + raise Qwen4ArtifactError("unsupported Qwen4 modular artifact_schema") if root.get("text_only") is not True: raise Qwen4ArtifactError("Qwen4 modular manifest must declare text_only=true") + source = _require_mapping(root.get("source"), label="source") + if not str(source.get("repository", "")).strip() or not str(source.get("revision", "")).strip(): + raise Qwen4ArtifactError("source repository and revision are required") + source_inventory_sha256 = _require_sha( + source.get("inventory_sha256"), label="source.inventory_sha256" + ) + minimum_commit = str(root.get("minimum_freetoken_commit", "")).lower() + if len(minimum_commit) != 40 or any(char not in "0123456789abcdef" for char in minimum_commit): + raise Qwen4ArtifactError("minimum_freetoken_commit must be a 40-character Git OID") + _require_sha(root.get("tvm_ffi_patch_sha256"), label="tvm_ffi_patch_sha256") + if not allow_synthetic_geometry: + if source.get("repository") != PINNED_SOURCE_REPOSITORY: + raise Qwen4ArtifactError("source repository does not match the pinned production source") + if source.get("revision") != PINNED_SOURCE_REVISION: + raise Qwen4ArtifactError("source revision does not match the pinned production revision") + if str(root.get("tvm_ffi_patch_sha256", "")).lower() != TVM_FFI_PATCH_SHA256: + raise Qwen4ArtifactError("TVM-FFI patch does not match the frozen contract") + declared_fingerprint = _require_sha( + root.get("complete_artifact_fingerprint"), label="complete_artifact_fingerprint" + ) + unsigned = dict(root) + unsigned.pop("complete_artifact_fingerprint", None) + actual_fingerprint = hashlib.sha256( + json.dumps(unsigned, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + if actual_fingerprint != declared_fingerprint: + raise Qwen4ArtifactError("complete artifact fingerprint mismatch") + metadata = _require_mapping(root.get("metadata"), label="metadata") + if not isinstance(metadata.get("files"), list) or not metadata["files"]: + raise Qwen4ArtifactError("metadata.files must be a non-empty list") + metadata_files: list[ComponentFile] = [] + for index, item in enumerate(metadata["files"]): + entry = _require_mapping(item, label=f"metadata.files[{index}]") + component = ComponentFile( + path=_path_within_root( + path.parent, entry.get("path"), label=f"metadata.files[{index}].path" + ), + bytes=int(entry.get("bytes", -1)), + sha256=_require_sha( + entry.get("sha256"), label=f"metadata.files[{index}].sha256" + ), + ) + if component.bytes < 0: + raise Qwen4ArtifactError("metadata file bytes must be non-negative") + _verify_component_file(component, label=f"metadata.files[{index}]") + metadata_files.append(component) active = _require_mapping(root.get("active"), label="active") if active.get("format") != ACTIVE_FORMAT: raise Qwen4ArtifactError(f"unsupported active format {active.get('format')!r}") - active_path = _path_from(path.parent, active.get("path"), label="active.path") + active_path = _path_within_root(path.parent, active.get("path"), label="active.path") if "bytes" in active and int(active["bytes"]) < 0: raise Qwen4ArtifactError("active.bytes must be non-negative") - if "sha256" in active: - _require_sha(active["sha256"], label="active.sha256") + raw_active_files = active.get("files") + if not isinstance(raw_active_files, list) or not raw_active_files: + raise Qwen4ArtifactError("active.files must be a non-empty list") + active_files: list[ComponentFile] = [] + for index, item in enumerate(raw_active_files): + entry = _require_mapping(item, label=f"active.files[{index}]") + try: + size = int(entry["bytes"]) + except (KeyError, TypeError, ValueError) as exc: + raise Qwen4ArtifactError(f"active.files[{index}].bytes must be an integer") from exc + if size < 0: + raise Qwen4ArtifactError(f"active.files[{index}].bytes must be non-negative") + component = ComponentFile( + path=_path_within_root(path.parent, entry.get("path"), label=f"active.files[{index}].path"), + bytes=size, + sha256=_require_sha(entry.get("sha256"), label=f"active.files[{index}].sha256"), + ) + try: + component.path.relative_to(active_path) + except ValueError as exc: + raise Qwen4ArtifactError("active file resolves outside active.path") from exc + active_files.append(component) ple = _require_mapping(root.get("ple"), label="ple") if ple.get("format") != PLE_FORMAT: raise Qwen4ArtifactError(f"unsupported PLE format {ple.get('format')!r}") if str(ple.get("required_volume", REQUIRED_VOLUME)).upper() != REQUIRED_VOLUME: raise Qwen4ArtifactError("Qwen4 PLE sidecar must reside on Z:") - ple_manifest_path = _path_from(path.parent, ple.get("manifest"), label="ple.manifest") + ple_manifest_path = _path_within_root(path.parent, ple.get("manifest"), label="ple.manifest") ple_data_bytes = int(ple.get("data_bytes", 0)) if ple_data_bytes < 0: raise Qwen4ArtifactError("ple.data_bytes must be non-negative") ple_sha256 = _require_sha(ple.get("sha256"), label="ple.sha256") + if not ple_manifest_path.is_file(): + raise Qwen4ArtifactError(f"PLE manifest is missing: {ple_manifest_path}") + try: + with ple_manifest_path.open("r", encoding="utf-8") as handle: + ple_sidecar_manifest = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + raise Qwen4ArtifactError("cannot read Q3 PLE manifest") from exc + if str(ple_sidecar_manifest.get("source_fingerprint", "")).lower() != source_inventory_sha256: + raise Qwen4ArtifactError("Q3 PLE source fingerprint mismatch") experts = _require_mapping(root.get("experts"), label="experts") if experts.get("format") != EXPERT_FORMAT: @@ -245,26 +399,51 @@ def load_qwen4_artifact_manifest( files.append( ExpertFile( layer=layer, - path=_path_from(path.parent, entry.get("path"), label=f"experts.files[{layer}].path"), + path=_path_within_root(path.parent, entry.get("path"), label=f"experts.files[{layer}].path"), bytes=size, sha256=_require_sha(entry.get("sha256"), label=f"experts.files[{layer}].sha256"), + source_fingerprint=_require_sha( + entry.get("source_fingerprint"), + label=f"experts.files[{layer}].source_fingerprint", + ), ) ) + if any(item.source_fingerprint != source_inventory_sha256 for item in files): + raise Qwen4ArtifactError("expert sidecar source fingerprint mismatch") declared_layers = set(file_layers) | set(resident_layers) if declared_layers != seen: raise Qwen4ArtifactError( "experts.files layers must exactly match file_tier_layers + resident_layers" ) + production_geometry = not allow_synthetic_geometry + if production_geometry: + active_payload = int(active.get("payload_bytes", active.get("bytes", -1))) + if active_payload != ACTIVE_TARGET_BYTES: + raise Qwen4ArtifactError("active payload does not match frozen production bytes") + if ple_data_bytes != PLE_TARGET_BYTES: + raise Qwen4ArtifactError("Q3 PLE extent does not match frozen production bytes") + if len(files) != EXPERT_LAYERS or any(item.bytes != EXPERT_FILE_BYTES for item in files): + raise Qwen4ArtifactError("expert sidecars do not match frozen production geometry") + if file_layers != FILE_TIER_LAYERS: + raise Qwen4ArtifactError("file-tier layers do not match the frozen production policy") + expected_resident = tuple(layer for layer in range(EXPERT_LAYERS) if layer not in FILE_TIER_LAYERS) + if resident_layers != expected_resident: + raise Qwen4ArtifactError("resident layers do not match the frozen production policy") + if active_payload + ple_data_bytes + sum(item.bytes for item in files) != KNOWN_TARGET_BYTES: + raise Qwen4ArtifactError("known target component byte reconciliation failed") return Qwen4ArtifactManifest( path=path, raw=root, active_path=active_path, + active_files=tuple(active_files), ple_manifest_path=ple_manifest_path, ple_data_bytes=ple_data_bytes, ple_sha256=ple_sha256, expert_files=tuple(sorted(files, key=lambda item: item.layer)), + metadata_files=tuple(metadata_files), file_tier_layers=file_layers, resident_layers=resident_layers, + production_geometry=production_geometry, ) @@ -330,6 +509,7 @@ def build_mixed_expert_sources( file_sources[layer] = FileExpertSource( entry.path, expected_sha256=entry.sha256, + expected_source_fingerprint=entry.source_fingerprint, expected_layer_id=layer, num_experts=num_experts, verify_hash=verify_hash, @@ -337,32 +517,40 @@ def build_mixed_expert_sources( # Resident layers are bounded by one six-plane record at a time. We do not # retain a second full-layer staging tensor and close each source after fill. - for layer in sorted(manifest.resident_layers): - entry = manifest.file_for_layer(layer) - with FileExpertSource( - entry.path, - expected_sha256=entry.sha256, - expected_layer_id=layer, - num_experts=num_experts, - verify_hash=verify_hash, - ) as source: - buffers = { - name: allocator(shape=(num_experts, *shape), dtype=dtype) - for name, (shape, dtype) in FileExpertSource.plane_specs.items() - } - for expert_id in range(num_experts): - row = source.read_record(expert_id) + # If construction fails, close already-open tier sources so partial startup + # cannot retain file handles or make fixture cleanup impossible. + try: + for layer in sorted(manifest.resident_layers): + entry = manifest.file_for_layer(layer) + with FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_source_fingerprint=entry.source_fingerprint, + expected_layer_id=layer, + num_experts=num_experts, + verify_hash=verify_hash, + ) as source: + buffers = { + name: allocator(shape=(num_experts, *shape), dtype=dtype) + for name, (shape, dtype) in source.plane_specs.items() + } + for expert_id in range(num_experts): + row = source.read_record(expert_id) + for name, buffer in buffers.items(): + destination = getattr(buffer, "tensor", buffer) + destination[expert_id].copy_(row[name]) + settle = residency[layer] + for buffer in buffers.values(): + if settle == HostResidency.PINNED.value and hasattr(buffer, "pin"): + buffer.pin() + elif settle == HostResidency.LOCKED.value and hasattr(buffer, "lock"): + buffer.lock() for name, buffer in buffers.items(): - destination = getattr(buffer, "tensor", buffer) - destination[expert_id].copy_(row[name]) - settle = residency[layer] - for buffer in buffers.values(): - if settle == HostResidency.PINNED.value and hasattr(buffer, "pin"): - buffer.pin() - elif settle == HostResidency.LOCKED.value and hasattr(buffer, "lock"): - buffer.lock() - for name, buffer in buffers.items(): - sources[name][layer] = getattr(buffer, "tensor", buffer) + sources[name][layer] = getattr(buffer, "tensor", buffer) + except Exception: + for source in file_sources.values(): + source.close() + raise return sources, file_sources @@ -372,6 +560,7 @@ def configure_mixed_expert_sources( resident_sources, *, file_sources: Mapping[int, object] | None = None, + layer_residency: list[str] | None = None, ): """Wire resident bank entries and file-backed layers into an offload cache. @@ -383,8 +572,8 @@ def configure_mixed_expert_sources( if not isinstance(manifest, Qwen4ArtifactManifest): raise TypeError("manifest must be a validated Qwen4ArtifactManifest") - if cache.decode_target != "gpu": - raise ValueError("Qwen4 modular file tiers are GPU-only") + if set(manifest.file_tier_layers) & set(cache.cpu_layer_ids): + raise ValueError("Qwen4 modular file-tier layers are GPU-only") if cache.prefill_overlap: raise ValueError("Qwen4 modular file tiers require prefill_overlap=False") # The artifact's native expert geometry is 512 slots. Enforce this independently @@ -417,7 +606,7 @@ def configure_mixed_expert_sources( if values[layer] is None: raise ValueError(f"resident layer {layer} has no resident source for {name}") per_layer[name] = values - cache.set_bank_sources(per_layer) + cache.set_bank_sources(per_layer, layer_residency=layer_residency) from freetoken.moe.expert_source import FileExpertSource if file_sources is None: @@ -427,6 +616,7 @@ def configure_mixed_expert_sources( source = FileExpertSource( entry.path, expected_sha256=entry.sha256, + expected_source_fingerprint=entry.source_fingerprint, expected_layer_id=layer, num_experts=cache.num_experts, ) @@ -439,6 +629,257 @@ def configure_mixed_expert_sources( return dict(file_sources) +def build_qwen4_modular_artifact( + source_path: str | os.PathLike[str], + artifact_root: str | os.PathLike[str], + *, + source_repository: str = PINNED_SOURCE_REPOSITORY, + source_revision: str = PINNED_SOURCE_REVISION, + source_inventory_sha256: str, + minimum_freetoken_commit: str, + ple_layer_id: int = 2, + ple_split_parts: int = 128, + expert_layers: tuple[int, ...] = tuple(range(EXPERT_LAYERS)), + file_tier_layers: tuple[int, ...] = FILE_TIER_LAYERS, + expert_num_experts: int = EXPERT_NUM_EXPERTS, + expert_geometry=None, + allow_synthetic_geometry: bool = False, +) -> dict[str, Any]: + """Run the canonical C1-C5 modular conversion sequence. + + C0 (full source-file identity verification) remains a mandatory caller gate. + This function performs no network I/O and accepts only an already-local Z:-backed + source snapshot. Each component writer owns its bounded streaming/atomic contract; + the final manifest is published only after all components reopen successfully. + """ + + source = _resolve_z(source_path, label="source checkpoint") + root = _resolve_z(artifact_root, label="artifact root") + if not source.is_dir(): + raise Qwen4ArtifactError(f"source checkpoint is not a directory: {source}") + root.mkdir(parents=True, exist_ok=True) + inventory = _require_sha(source_inventory_sha256, label="source_inventory_sha256") + + from freetoken.checkpoint.convert import convert_checkpoint + from freetoken.checkpoint.q3_ple import write_q3_ple_from_safetensors + from freetoken.moe.expert_source import write_expert_sidecar_from_safetensors + + active_index = convert_checkpoint( + str(source), + str(root), + artifact_format=ARTIFACT_FORMAT, + source_inventory_sha256=inventory, + ) + write_q3_ple_from_safetensors( + source, + root / "ple-q3-000.bin", + root / "ple-q3.json", + layer_id=int(ple_layer_id), + split_parts=int(ple_split_parts), + source_fingerprint=inventory, + ) + expert_paths: dict[int, str] = {} + for layer in expert_layers: + name = f"experts-L{int(layer):02d}.nvfp4" + write_expert_sidecar_from_safetensors( + source, + root / name, + layer_id=int(layer), + source_fingerprint=inventory, + num_experts=int(expert_num_experts), + geometry=expert_geometry, + ) + expert_paths[int(layer)] = name + metadata_paths = list(active_index.get("copied_metadata", ())) + if not metadata_paths: + raise Qwen4ArtifactError("active conversion copied no target metadata") + return finalize_qwen4_modular_manifest( + root, + source_repository=source_repository, + source_revision=source_revision, + source_inventory_sha256=inventory, + minimum_freetoken_commit=minimum_freetoken_commit, + tvm_ffi_patch_sha256=TVM_FFI_PATCH_SHA256, + expert_paths=expert_paths, + file_tier_layers=file_tier_layers, + metadata_paths=metadata_paths, + expert_num_experts=expert_num_experts, + allow_synthetic_geometry=allow_synthetic_geometry, + ) + + +def finalize_qwen4_modular_manifest( + artifact_root: str | os.PathLike[str], + *, + source_repository: str, + source_revision: str, + source_inventory_sha256: str, + minimum_freetoken_commit: str, + tvm_ffi_patch_sha256: str, + active_dir: str | os.PathLike[str] = "qwen4-active-v1.ftw", + ple_manifest: str | os.PathLike[str] = "ple-q3.json", + expert_paths: Mapping[int, str | os.PathLike[str]], + file_tier_layers: list[int] | tuple[int, ...], + metadata_paths: list[str | os.PathLike[str]], + expert_num_experts: int = 512, + allow_synthetic_geometry: bool = False, +) -> dict[str, Any]: + """Validate completed components and atomically publish ``manifest.json``. + + This is the final C5 orchestration seam. It never creates weight payloads; + the C2/C3/C4 writers must already have atomically finalized their components. + Every file is length/hash inventoried here, the Q3 and FTEXPERT1 readers reopen + their formats, and only then is the complete manifest promoted. + """ + + root = _resolve_z(artifact_root, label="artifact root") + root.mkdir(parents=True, exist_ok=True) + active_root = _path_from(root, active_dir, label="active_dir") + from freetoken.checkpoint.ftw import INDEX_NAME, is_ftw_checkpoint + + if not active_root.is_dir() or not is_ftw_checkpoint(str(active_root)): + raise Qwen4ArtifactError(f"active component is not an FTW checkpoint: {active_root}") + with (active_root / INDEX_NAME).open("r", encoding="utf-8") as handle: + active_index = json.load(handle) + inventory_digest = _require_sha( + source_inventory_sha256, label="source_inventory_sha256" + ) + if str(active_index.get("source_inventory_sha256", "")).lower() != inventory_digest: + raise Qwen4ArtifactError("active FTW source inventory fingerprint mismatch") + active_payload_bytes = int(active_index.get("total_bytes", -1)) + if active_payload_bytes < 0: + raise Qwen4ArtifactError("active FTW index has no valid total_bytes") + active_files = [ + _component_file(root, item) + for item in sorted(path for path in active_root.rglob("*") if path.is_file()) + ] + + ple_path = _path_within_root(root, ple_manifest, label="ple_manifest") + from freetoken.checkpoint.q3_ple import Q3PLEReader + + with Q3PLEReader(ple_path) as ple_reader: + ple_data = ple_reader.data_path + ple_data_bytes = ple_data.stat().st_size + ple_sha256 = _sha256_file(ple_data) + ple_source_fingerprint = str(ple_reader.manifest.get("source_fingerprint", "")) + if ple_source_fingerprint.lower() != inventory_digest: + raise Qwen4ArtifactError("Q3 PLE source fingerprint mismatch") + + tiered = tuple(sorted(_validate_layers(list(file_tier_layers), label="file_tier_layers"))) + expert_items: list[dict[str, Any]] = [] + from freetoken.moe.expert_source import FileExpertSource + + for layer, value in sorted((int(layer), path) for layer, path in expert_paths.items()): + expert_path = _path_within_root(root, value, label=f"expert layer {layer}") + with FileExpertSource( + expert_path, + expected_source_fingerprint=inventory_digest, + expected_layer_id=layer, + num_experts=int(expert_num_experts), + verify_hash=True, + ) as source: + if source.layer_id != layer: + raise Qwen4ArtifactError(f"expert sidecar layer mismatch for {expert_path}") + expert_items.append( + { + "layer": layer, + **_component_file(root, expert_path), + "source_fingerprint": source.source_fingerprint, + } + ) + all_layers = tuple(item["layer"] for item in expert_items) + if all_layers != tuple(range(len(all_layers))): + raise Qwen4ArtifactError("expert sidecars must cover contiguous layers from zero") + if not set(tiered) <= set(all_layers): + raise Qwen4ArtifactError("file tier contains a layer without an expert sidecar") + resident = sorted(set(all_layers) - set(tiered)) + if not allow_synthetic_geometry: + if source_repository != PINNED_SOURCE_REPOSITORY or source_revision != PINNED_SOURCE_REVISION: + raise Qwen4ArtifactError("production artifact source pin mismatch") + commit = str(minimum_freetoken_commit).lower() + if len(commit) != 40 or any(char not in "0123456789abcdef" for char in commit): + raise Qwen4ArtifactError("production artifact requires a concrete FreeToken commit") + if str(tvm_ffi_patch_sha256).lower() != TVM_FFI_PATCH_SHA256: + raise Qwen4ArtifactError("production artifact TVM-FFI patch mismatch") + if int(expert_num_experts) != EXPERT_NUM_EXPERTS: + raise Qwen4ArtifactError("production artifact requires 512 experts per layer") + if active_payload_bytes != ACTIVE_TARGET_BYTES: + raise Qwen4ArtifactError("active FTW does not match frozen production bytes") + if ple_data_bytes != PLE_TARGET_BYTES: + raise Qwen4ArtifactError("Q3 PLE does not match frozen production extent") + if tuple(all_layers) != tuple(range(EXPERT_LAYERS)): + raise Qwen4ArtifactError("production artifact requires 48 expert sidecars") + if any(item["bytes"] != EXPERT_FILE_BYTES for item in expert_items): + raise Qwen4ArtifactError("expert sidecar does not match frozen production bytes") + if tiered != FILE_TIER_LAYERS: + raise Qwen4ArtifactError("file tier does not match the frozen production policy") + if active_payload_bytes + ple_data_bytes + sum(item["bytes"] for item in expert_items) != KNOWN_TARGET_BYTES: + raise Qwen4ArtifactError("known target component byte reconciliation failed") + + metadata_files = [ + _component_file(root, _path_within_root(root, value, label="metadata file")) + for value in metadata_paths + ] + if not any(item["path"] == "config.json" for item in metadata_files): + raise Qwen4ArtifactError("modular artifact metadata must include config.json") + config_path = root / "config.json" + with config_path.open("r", encoding="utf-8") as handle: + config = json.load(handle) + if config.get("freetoken_text_only") != TEXT_ONLY_MARKER: + raise Qwen4ArtifactError("config.json lacks the accepted text-only marker") + if config.get("freetoken_active_quant") != ACTIVE_FORMAT: + raise Qwen4ArtifactError("config.json lacks the accepted active-quant marker") + + manifest: dict[str, Any] = { + "format": FORMAT, + "version": VERSION, + "artifact_schema": FORMAT, + "text_only": True, + "source": { + "repository": str(source_repository), + "revision": str(source_revision), + "inventory_sha256": inventory_digest, + }, + "minimum_freetoken_commit": str(minimum_freetoken_commit), + "tvm_ffi_patch_sha256": _require_sha( + tvm_ffi_patch_sha256, label="tvm_ffi_patch_sha256" + ), + "active": { + "format": ACTIVE_FORMAT, + "path": active_root.relative_to(root).as_posix(), + "payload_bytes": active_payload_bytes, + "physical_file_bytes": sum(item["bytes"] for item in active_files), + "files": active_files, + }, + "ple": { + "format": PLE_FORMAT, + "manifest": ple_path.relative_to(root).as_posix(), + "data_bytes": ple_data_bytes, + "sha256": ple_sha256, + "source_fingerprint": ple_source_fingerprint, + "required_volume": REQUIRED_VOLUME, + }, + "experts": { + "format": EXPERT_FORMAT, + "files": expert_items, + "file_tier_layers": list(tiered), + "resident_layers": resident, + "required_volume": REQUIRED_VOLUME, + }, + "metadata": {"files": metadata_files}, + } + canonical = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest["complete_artifact_fingerprint"] = hashlib.sha256(canonical).hexdigest() + partial = root / f".{MANIFEST_NAME}.partial-{os.getpid()}" + with partial.open("w", encoding="utf-8", newline="\n") as handle: + json.dump(manifest, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(partial, root / MANIFEST_NAME) + return manifest + + __all__ = [ "ACTIVE_FORMAT", "ARTIFACT_FORMAT", @@ -450,8 +891,10 @@ def configure_mixed_expert_sources( "Qwen4ArtifactManifest", "TEXT_ONLY_MARKER", "VERSION", + "build_qwen4_modular_artifact", "configure_mixed_expert_sources", "build_mixed_expert_sources", + "finalize_qwen4_modular_manifest", "load_qwen4_artifact_manifest", "qwen4_text_only_marker", ] diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index ab7c0a9f1..5214619f9 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -25,7 +25,13 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: # with cache_size), so they are intentionally excluded from the per-slot growth term. # tensor[0].numel() is the per-row element count (one expert slot); see the matching # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4_paged_pool.py. - return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) + total = 0 + for per_layer in sources.values(): + representative = next((tensor for tensor in per_layer if tensor is not None), None) + if representative is None: + raise ValueError("expert bank has no resident shape for cache sizing") + total += representative[0].numel() * representative.element_size() + return total def net_cache_budget_bytes( diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index bb60c05f0..a4b5122ed 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -517,24 +517,60 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: artifact = load_qwen4_artifact_manifest(config.model_path) if artifact is not None: - if config.moe_backend != "offload": + if config.moe_backend not in ("offload",): raise ValueError( - "Qwen4 modular expert tiers are GPU-only; use --moe-backend offload" + "Qwen4 modular file tiers require --moe-backend offload; " + "CPU execution may be selected only for resident layers" ) if config.moe_prefill_overlap: object.__setattr__(config, "moe_prefill_overlap", False) # Expert sidecars are authoritative for both tiers. Resident layers are # streamed one record at a time into HostBanks; file layers remain open # FileExpertSource readers and never materialize a full layer. + from freetoken.moe.host_banks import HostResidency + from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT + + cpu_layer_ids = _resolve_cpu_layers(config, artifact.num_layers) + tiered = set(artifact.file_tier_layers) + if cpu_layer_ids & tiered: + raise ValueError( + "--moe-cpu-layers selects a file-tier layer; only resident layers " + "are CPU/hybrid eligible" + ) + resident = set(artifact.resident_layers) + pin_budget = _pin_budget_bytes() + if pin_budget is not None: + row_bytes = _BANK_BYTES_PER_EXPERT["nvfp4"]( + config.model_config.hidden_size, + config.model_config.moe_intermediate_size, + ) + max_pinned = pin_budget // (row_bytes * config.model_config.num_experts) + required_locked = max(0, len(resident - set(cpu_layer_ids)) - max_pinned) + if required_locked: + if not _cpu_moe_executor_viable(config.model_config): + raise ValueError( + "resident expert banks exceed the Windows pin budget and the " + "CPU executor is unavailable; refusing pageable GPU sources" + ) + # Deterministic outer-to-inner choice within the resident set. + candidates = sorted( + resident - set(cpu_layer_ids), + key=lambda layer: (min(layer, artifact.num_layers - 1 - layer), layer), + ) + cpu_layer_ids = frozenset(set(cpu_layer_ids) | set(candidates[:required_locked])) + residency = [HostResidency.PINNED.value] * artifact.num_layers + for layer in cpu_layer_ids: + residency[layer] = HostResidency.LOCKED.value + resident_sources, file_sources = build_mixed_expert_sources( artifact, num_experts=config.model_config.num_experts, - resident_residency=["pinned"] * artifact.num_layers, + resident_residency=residency, verify_hash=not config.use_dummy_weight, ) from freetoken.moe.expert_banks import ExpertBanks - banks = ExpertBanks("nvfp4", resident_sources) + banks = ExpertBanks("nvfp4", resident_sources, layer_residency=residency) if config.moe_cache_auto: size, pages, _overlap = self._resolve_auto_moe_cache_size(config, banks) object.__setattr__(config, "moe_cache_size", max(size, 512)) @@ -550,17 +586,23 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: prefill_overlap=False, prefill_hit_d2d=False, quant_format=banks.quant_format, - decode_target="gpu", + decode_target="cpu" if cpu_layer_ids else "gpu", hybrid_max_fetch=config.moe_hybrid_max_fetch, ) - cache.cpu_layer_ids = frozenset() + cache.cpu_layer_ids = cpu_layer_ids configure_mixed_expert_sources( - cache, artifact, banks.sources, file_sources=file_sources + cache, + artifact, + banks.sources, + file_sources=file_sources, + layer_residency=residency, ) cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) cache.collect_stats = config.moe_collect_stats layers = attach_offload_moe_cache(self.model, cache) assert len(layers) == config.model_config.num_moe_layers + if cache.decode_target == "cpu": + self._init_cpu_moe_executor(config, cache, layers) self.ctx.moe_offload_cache = cache self.moe_offload_cache = cache return cache diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 84e6bb27c..c1e507243 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -790,7 +790,11 @@ def load_q3_ple_weights(self, manifest_paths: str | dict[int, str]) -> None: def forward(self, input_ids: torch.Tensor) -> torch.Tensor: hidden = self.embed_tokens.forward(input_ids) mm_embeds = getattr(get_global_ctx().batch, "mm_embeds", None) - if mm_embeds is not None and self._image_token_id is not None: + if mm_embeds is not None and self._image_token_id is None: + raise RuntimeError( + "image inputs are not supported by this text-only Qwen4 modular artifact" + ) + if mm_embeds is not None: mask = input_ids == self._image_token_id slots = int(mask.sum().item()) if slots != mm_embeds.shape[0]: diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index c8cf2c494..52bf3ef5a 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -241,6 +241,7 @@ def load_weight( # deliberately optional so every unmarked checkpoint follows the historical path. artifact = load_qwen4_artifact_manifest(model_path) if artifact is not None: + artifact.verify_active() active_path = artifact.active_path if active_path.is_file() and active_path.name == "freetoken_weight.json": active_path = active_path.parent diff --git a/python/freetoken/moe/__init__.py b/python/freetoken/moe/__init__.py index b1edb9c1d..d18f2cf6c 100644 --- a/python/freetoken/moe/__init__.py +++ b/python/freetoken/moe/__init__.py @@ -74,6 +74,7 @@ def create_moe_backend(backend: str) -> BaseMoeBackend: "FileExpertSource", "ExpertSourceError", "write_expert_sidecar", + "write_expert_sidecar_from_safetensors", "adapt_expert_tensor_record", ] @@ -84,4 +85,5 @@ def create_moe_backend(backend: str) -> BaseMoeBackend: FileExpertSource, adapt_expert_tensor_record, write_expert_sidecar, + write_expert_sidecar_from_safetensors, ) diff --git a/python/freetoken/moe/expert_source.py b/python/freetoken/moe/expert_source.py index a7b662b55..42401cb5e 100644 --- a/python/freetoken/moe/expert_source.py +++ b/python/freetoken/moe/expert_source.py @@ -10,6 +10,7 @@ from __future__ import annotations import hashlib +import json import os import struct import threading @@ -506,6 +507,8 @@ def write_expert_sidecar( partial.unlink() payload_hash = hashlib.sha256() seen: set[int] = set() + sample_ids = tuple(sorted({0, num_experts - 1, num_experts // 4, num_experts // 2, (3 * num_experts) // 4})) + sample_raw: dict[int, bytes] = {} try: with partial.open("wb") as handle: header = bytearray(HEADER_BYTES) @@ -543,6 +546,8 @@ def write_expert_sidecar( raise TypeError(f"expert {expert_id} must be a plane mapping or raw bytes") if len(raw) != raw_record_bytes: raise ValueError(f"expert {expert_id} serialized length {len(raw)} != {raw_record_bytes}") + if expert_id in sample_ids: + sample_raw[expert_id] = raw padded = raw + b"\0" * (record_bytes - raw_record_bytes) payload_hash.update(padded) handle.write(padded) @@ -575,6 +580,21 @@ def write_expert_sidecar( pass raise whole = _sha256_path(destination) + # Reopen through the production reader and byte-compare deterministic sample + # records before reporting the sidecar complete. The bounded sample set is + # at most five native records (~13.3 MiB at real geometry). + with FileExpertSource( + destination, + expected_sha256=whole, + expected_source_fingerprint=fingerprint, + expected_layer_id=layer_id, + num_experts=num_experts, + verify_hash=True, + ) as source: + for expert_id in sample_ids: + actual = _record_from_planes(source.read_record(expert_id), specs) + if actual != sample_raw[expert_id]: + raise ExpertSourceError(f"expert {expert_id} failed writer reopen verification") return { "path": str(destination), "format": "FTEXPERT1", @@ -589,10 +609,88 @@ def write_expert_sidecar( "canonical_sha256": canonical, "whole_sha256": whole, "sha256": whole, - "sample_ids": (0, num_experts - 1), + "sample_ids": sample_ids, } +def write_expert_sidecar_from_safetensors( + model_path: str | os.PathLike[str], + path: str | os.PathLike[str], + *, + layer_id: int, + source_fingerprint: str | bytes, + num_experts: int = NUM_EXPERTS, + geometry: Mapping[str, Any] | Sequence[tuple[str, Any, Any]] | None = None, +) -> dict[str, Any]: + """Stream one official ModelOpt NVFP4 layer into ``FTEXPERT1``. + + The source index must contain exactly twelve tensors for every expert. Only + one expert's tensors and one Safetensors mapping remain live at a time; no + full layer or bank is materialized in Python memory. + """ + + folder = Path(model_path).expanduser().resolve() + if not folder.is_dir(): + raise ValueError(f"expert source must be a local checkpoint directory: {folder}") + index_path = folder / "model.safetensors.index.json" + with index_path.open("r", encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + prefix = f"model.language_model.layers.{int(layer_id)}.mlp.experts." + suffixes = tuple( + f"{projection}.{field}" + for projection in ("gate_proj", "up_proj", "down_proj") + for field in ("weight", "weight_scale", "weight_scale_2", "input_scale") + ) + expected = { + f"{prefix}{expert_id}.{suffix}" + for expert_id in range(int(num_experts)) + for suffix in suffixes + } + actual = {name for name in weight_map if name.startswith(prefix)} + missing = expected - actual + unexpected = actual - expected + if missing or unexpected: + raise ValueError( + f"expert layer {layer_id} tensor set mismatch: " + f"missing={len(missing)} unexpected={len(unexpected)}" + ) + + import safetensors + + def iter_records(): + current_file: str | None = None + current_context = None + current_handle = None + try: + for expert_id in range(int(num_experts)): + record = {} + for suffix in suffixes: + full_name = f"{prefix}{expert_id}.{suffix}" + filename = weight_map[full_name] + if filename != current_file: + if current_context is not None: + current_context.__exit__(None, None, None) + current_context = safetensors.safe_open( + folder / filename, framework="pt", device="cpu" + ) + current_handle = current_context.__enter__() + current_file = filename + record[suffix] = current_handle.get_tensor(full_name) + yield expert_id, record + finally: + if current_context is not None: + current_context.__exit__(None, None, None) + + return write_expert_sidecar( + path, + iter_records(), + layer_id=layer_id, + source_fingerprint=source_fingerprint, + num_experts=num_experts, + geometry=geometry, + ) + + class FileExpertSource: """Read fixed NVFP4 expert records from one layer sidecar. diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 35c536432..78c372e97 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -365,8 +365,15 @@ def set_file_sources(self, sources: dict[int, object]) -> None: resident host tensors and CUDA graph capture cannot contain synchronous file I/O. """ - if self.decode_target != "gpu": - raise ValueError("file-backed expert tiers are GPU-only; use resident HostBanks for CPU/hybrid") + # A mixed cache may use the CPU executor for selected resident layers, + # but a file-backed layer itself is always GPU-only. The per-layer set is + # the precise invariant; rejecting the entire cache would unnecessarily + # disable existing CPU/hybrid support for resident HostBanks. + requested = {int(layer) for layer in sources} + if requested & set(self.cpu_layer_ids): + raise ValueError( + "file-backed expert tiers are GPU-only; CPU/hybrid layers must be resident" + ) if self.prefill_overlap: raise ValueError("file-backed expert tiers require prefill_overlap=False") if not sources: diff --git a/tests/checkpoint/test_convert_metadata.py b/tests/checkpoint/test_convert_metadata.py index a58b00b7d..124d6bf93 100644 --- a/tests/checkpoint/test_convert_metadata.py +++ b/tests/checkpoint/test_convert_metadata.py @@ -5,7 +5,11 @@ import torch -from freetoken.checkpoint.convert import _copy_metadata +from freetoken.checkpoint.convert import ( + _copy_metadata, + _iter_qwen4_modular_dense_entries, + convert_checkpoint, +) from freetoken.checkpoint.ftw import FTWReader, FTWWriter, iter_ftw_weights @@ -46,6 +50,108 @@ def test_copy_metadata_keeps_only_qwen4_host_mapped_shards(tmp_path: Path) -> No ] +def test_modular_metadata_does_not_copy_source_ple_payload(tmp_path: Path) -> None: + source = tmp_path / "source" + output = tmp_path / "output" + source.mkdir() + (source / "config.json").write_text("{}", encoding="utf-8") + (source / "model-plefp8-00000.safetensors").write_bytes(b"ple") + (source / "model.safetensors.index.json").write_text( + json.dumps({ + "weight_map": { + "model.language_model.layers.0.ple.ple_embedding.ngram_embedding.shard_0.weight": "model-plefp8-00000.safetensors" + } + }), + encoding="utf-8", + ) + copied = _copy_metadata(str(source), str(output), include_host_mapped_weights=False) + assert copied == ["config.json"] + assert not (output / "model-plefp8-00000.safetensors").exists() + + +def test_modular_dense_stream_quantizes_map_and_excludes_vision() -> None: + active = torch.ones((2, 16), dtype=torch.bfloat16) + protected = torch.ones((4,), dtype=torch.bfloat16) + names = [ + name + for name, _tensor in _iter_qwen4_modular_dense_entries( + [ + ("model.layers.0.self_attn.o_proj.weight", active), + ("model.layers.0.input_layernorm.weight", protected), + ("visual.blocks.0.weight", active), + ] + ) + ] + assert names == [ + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.self_attn.o_proj.weight_scale", + "model.layers.0.self_attn.o_proj.weight_global", + "model.layers.0.input_layernorm.weight", + ] + + +def test_convert_checkpoint_builds_marked_modular_active_ftw(tmp_path: Path, monkeypatch) -> None: + source = tmp_path / "source" + output = tmp_path / "target" + source.mkdir() + (source / "config.json").write_text( + json.dumps({"architectures": ["Qwen4ExpForCausalLM"], "vision_config": {}}), + encoding="utf-8", + ) + + class FakeEngineConfig: + def __init__(self, **_kwargs): + self.model_config = type( + "ModelConfig", + (), + { + "architectures": ("Qwen4ExpForCausalLM",), + "expert_quant": "nvfp4", + "is_moe": True, + }, + )() + + active = torch.ones((2, 16), dtype=torch.bfloat16) + protected = torch.arange(4, dtype=torch.bfloat16) + + def fake_load_weight(_path, device, *, include_moe_experts): + assert device.type == "cpu" + assert include_moe_experts is False + return iter( + [ + ("model.layers.0.self_attn.o_proj.weight", active), + ("model.layers.0.input_layernorm.weight", protected), + ] + ) + + import freetoken.engine.config as engine_config + import freetoken.models.weight as weight_module + + monkeypatch.setattr(engine_config, "EngineConfig", FakeEngineConfig) + monkeypatch.setattr(weight_module, "load_weight", fake_load_weight) + inventory = "a" * 64 + index = convert_checkpoint( + str(source), + str(output), + artifact_format="qwen4_modular_v1", + source_inventory_sha256=inventory, + shard_limit=4096 * 16, + ) + + assert index["source_inventory_sha256"] == inventory + assert index["counts"] == {"weight": 4, "experts_bank": 0} + config = json.loads((output / "config.json").read_text(encoding="utf-8")) + assert config["freetoken_text_only"] == "qwen4_text_only_v1" + assert config["freetoken_active_quant"] == "nvfp4_w4a16_v1" + loaded = dict(iter_ftw_weights(str(output / "qwen4-active-v1.ftw"), workers=1)) + assert set(loaded) == { + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.self_attn.o_proj.weight_scale", + "model.layers.0.self_attn.o_proj.weight_global", + "model.layers.0.input_layernorm.weight", + } + + def test_ftw_buffered_reader_works_on_the_current_platform(tmp_path: Path) -> None: output = tmp_path / "ftw" writer = FTWWriter(str(output), shard_limit=4096) diff --git a/tests/checkpoint/test_q3_ple_writer.py b/tests/checkpoint/test_q3_ple_writer.py index b101f9ec9..f3d3c22ba 100644 --- a/tests/checkpoint/test_q3_ple_writer.py +++ b/tests/checkpoint/test_q3_ple_writer.py @@ -9,6 +9,7 @@ from pathlib import Path import pytest +import torch from freetoken.checkpoint.q3_ple import ( ALIGN, @@ -17,6 +18,7 @@ ROW_VALUES, Q3PLEReader, quantize_block, + write_q3_ple_from_safetensors, write_q3_ple_sidecar, ) @@ -72,7 +74,9 @@ def test_writer_matches_reference_vectors_and_reader(z_fixture_dir: Path) -> Non assert manifest["payload_bytes"] == 5 * ROW_BYTES assert manifest["file_bytes"] == data_path.stat().st_size assert [segment["first_row"] for segment in manifest["segments"]] == [0, 2, 4] - assert [segment["data_offset"] for segment in manifest["segments"]] == [0, ALIGN, ALIGN * 2] + assert manifest["storage_layout"] == "contiguous_rows_v1" + assert manifest["file_bytes"] == manifest["payload_bytes"] == 5 * ROW_BYTES + assert [segment["data_offset"] for segment in manifest["segments"]] == [0, 2 * ROW_BYTES, 4 * ROW_BYTES] assert all( segment["byte_length"] == (2 if segment["first_row"] < 4 else 1) * ROW_BYTES for segment in manifest["segments"] @@ -153,12 +157,54 @@ def test_writer_rejects_bad_integrity_metadata(z_fixture_dir: Path, kwargs: dict ) -def test_writer_requires_z_backing(tmp_path: Path) -> None: +def test_writer_requires_z_backing() -> None: + forbidden = Path("C:/stage7-q3-writer-must-not-create") with pytest.raises(ValueError, match="Z:"): write_q3_ple_sidecar( _rows(1), - tmp_path / "ple-q3.bin", - tmp_path / "ple-q3.json", + forbidden / "ple-q3.bin", + forbidden / "ple-q3.json", source_fingerprint="d" * 64, weight_scale=1.0, ) + + +def test_production_writer_streams_safetensor_shards_in_source_order(z_fixture_dir: Path) -> None: + from safetensors.torch import save_file + + prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding" + weight_map = {} + source_rows = [] + for part in range(2): + key = f"{prefix}.shard_{part}.weight" + filename = f"model-plefp8-{part:05d}.safetensors" + rows = ( + torch.tensor(list(_rows(2)), dtype=torch.float32) + float(part) + ).to(torch.float8_e4m3fn).contiguous() + tensors = {key: rows} + if part == 0: + tensors[prefix + ".weight_scale"] = torch.tensor(0.5, dtype=torch.bfloat16) + weight_map[prefix + ".weight_scale"] = filename + save_file(tensors, z_fixture_dir / filename) + weight_map[key] = filename + source_rows.extend(rows.float().tolist()) + (z_fixture_dir / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}), encoding="utf-8" + ) + manifest = write_q3_ple_from_safetensors( + z_fixture_dir, + z_fixture_dir / "ple-q3.bin", + z_fixture_dir / "ple-q3.json", + layer_id=2, + split_parts=2, + source_fingerprint="f" * 64, + rows_per_chunk=1, + segment_rows=3, + ) + reference = _load_authoritative_reference() + assert (z_fixture_dir / "ple-q3.bin").read_bytes() == reference.encode_table( + source_rows, refinement_passes=2, scale_dtype="bf16" + ) + assert manifest["rows"] == 4 + assert manifest["file_bytes"] == 4 * ROW_BYTES + assert manifest["weight_scale"] == 0.5 diff --git a/tests/checkpoint/test_qwen4_artifact.py b/tests/checkpoint/test_qwen4_artifact.py index 7a7778866..096d7cff9 100644 --- a/tests/checkpoint/test_qwen4_artifact.py +++ b/tests/checkpoint/test_qwen4_artifact.py @@ -1,22 +1,34 @@ from __future__ import annotations import json +import hashlib import shutil from pathlib import Path from uuid import uuid4 import pytest import torch +from types import SimpleNamespace from freetoken.checkpoint.qwen4_artifact import ( FORMAT, TEXT_ONLY_MARKER, + build_qwen4_modular_artifact, configure_mixed_expert_sources, build_mixed_expert_sources, + finalize_qwen4_modular_manifest, load_qwen4_artifact_manifest, qwen4_text_only_marker, ) -from freetoken.moe.expert_source import FileExpertSource, RAW_RECORD_BYTES +from freetoken.checkpoint.ftw import FTWWriter +from freetoken.checkpoint.q3_ple import ROW_VALUES, write_q3_ple_sidecar +from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries +from freetoken.models.weight import load_weight +from freetoken.moe.expert_source import ( + FileExpertSource, + RAW_RECORD_BYTES, + write_expert_sidecar, +) @pytest.fixture @@ -37,23 +49,50 @@ class _FakeCache: cache_size = 512 num_layers = 2 num_experts = 1 + cpu_layer_ids = frozenset() - def set_bank_sources(self, sources): + def set_bank_sources(self, sources, layer_residency=None): self.bank_sources = sources + self.layer_residency = layer_residency def set_file_sources(self, sources): self.file_sources = sources def _manifest(root: Path, sidecar: Path, digest: str) -> Path: + source_fingerprint = hashlib.sha256(b"synthetic").hexdigest() + active_root = root / "qwen4-active-v1.ftw" + active_root.mkdir(exist_ok=True) + active_index = active_root / "freetoken_weight.json" + active_index.write_text( + json.dumps({"source_inventory_sha256": source_fingerprint}), encoding="utf-8" + ) + active_digest = __import__("hashlib").sha256(active_index.read_bytes()).hexdigest() + ple_manifest = root / "ple-q3.json" + ple_manifest.write_text( + json.dumps({"source_fingerprint": source_fingerprint}), encoding="utf-8" + ) + resident_sidecar = root / "experts-L01.nvfp4" + resident_digest = FileExpertSource.create_synthetic( + resident_sidecar, num_experts=1, records=[bytes([8]) * RAW_RECORD_BYTES], layer_id=1 + ) + config_path = root / "config.json" + config_path.write_bytes(b"") + config_digest = hashlib.sha256(b"").hexdigest() data = { "format": FORMAT, "version": 1, + "artifact_schema": FORMAT, "text_only": True, - "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": "0" * 64}, + "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": source_fingerprint}, "minimum_freetoken_commit": "0" * 40, "tvm_ffi_patch_sha256": "0" * 64, - "active": {"format": "nvfp4_w4a16_v1", "path": ".", "bytes": 0, "sha256": "0" * 64}, + "active": { + "format": "nvfp4_w4a16_v1", + "path": active_root.name, + "bytes": active_index.stat().st_size, + "files": [{"path": str(active_index.relative_to(root)), "bytes": active_index.stat().st_size, "sha256": active_digest}], + }, "ple": { "format": "q3_ple_32", "manifest": "ple-q3.json", @@ -63,14 +102,19 @@ def _manifest(root: Path, sidecar: Path, digest: str) -> Path: }, "experts": { "format": "ftexpert1_nvfp4_v1", - "files": [{"layer": 0, "path": sidecar.name, "bytes": sidecar.stat().st_size, "sha256": digest}], + "files": [ + {"layer": 0, "path": sidecar.name, "bytes": sidecar.stat().st_size, "sha256": digest, "source_fingerprint": source_fingerprint}, + {"layer": 1, "path": resident_sidecar.name, "bytes": resident_sidecar.stat().st_size, "sha256": resident_digest, "source_fingerprint": source_fingerprint}, + ], "file_tier_layers": [0], "resident_layers": [1], "required_volume": "Z:", }, - "metadata": {"config_sha256": "0" * 64}, - "complete_artifact_fingerprint": "0" * 64, + "metadata": {"files": [{"path": "config.json", "bytes": 0, "sha256": config_digest}]}, } + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() path = root / "manifest.json" path.write_text(json.dumps(data), encoding="utf-8") return path @@ -86,13 +130,104 @@ class Config: qwen4_text_only_marker(Config()) +def test_active_component_hash_fails_closed(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + manifest = load_qwen4_artifact_manifest(_manifest(z_fixture_dir, sidecar, digest), require=True, allow_synthetic_geometry=True) + assert manifest is not None + manifest.verify_active() + manifest.active_files[0].path.write_bytes(b"tampered") + with pytest.raises(ValueError, match="length mismatch|SHA-256 mismatch"): + manifest.verify_active() + + +def test_manifest_unknown_schema_and_fingerprint_fail_closed(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["artifact_schema"] = "future-v99" + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="artifact_schema"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + data["artifact_schema"] = FORMAT + data["complete_artifact_fingerprint"] = "f" * 64 + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="fingerprint mismatch"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + +def test_manifest_rejects_component_source_fingerprint_drift(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["experts"]["files"][0]["source_fingerprint"] = "f" * 64 + data.pop("complete_artifact_fingerprint") + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="expert sidecar source fingerprint mismatch"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + +def test_production_loader_rejects_reduced_geometry(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["source"]["repository"] = "RadixArk/Qwen3.8-Flash-Next-NVFP4" + data["source"]["revision"] = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" + data["minimum_freetoken_commit"] = "846504bf9d81119cb72400e6c5a3cc860f2b1dd8" + data["tvm_ffi_patch_sha256"] = "889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec" + data.pop("complete_artifact_fingerprint") + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="active payload does not match"): + load_qwen4_artifact_manifest(path, require=True) + + +def test_manifest_rejects_metadata_tamper_and_out_of_root_component(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + (z_fixture_dir / "config.json").write_bytes(b"tampered") + with pytest.raises(ValueError, match="length mismatch|SHA-256 mismatch"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["experts"]["files"][0]["path"] = str(sidecar.parent.parent / sidecar.name) + data.pop("complete_artifact_fingerprint") + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="outside artifact root"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + def test_manifest_reopens_and_wires_mixed_sources(z_fixture_dir): sidecar = z_fixture_dir / "experts-L00.nvfp4" digest = FileExpertSource.create_synthetic( sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 ) manifest_path = _manifest(z_fixture_dir, sidecar, digest) - manifest = load_qwen4_artifact_manifest(manifest_path, require=True) + manifest = load_qwen4_artifact_manifest(manifest_path, require=True, allow_synthetic_geometry=True) assert manifest is not None assert manifest.file_tier_layers == (0,) assert manifest.resident_layers == (1,) @@ -108,6 +243,7 @@ def test_manifest_reopens_and_wires_mixed_sources(z_fixture_dir): def test_resident_builder_streams_one_record_without_full_layer(z_fixture_dir): + source_fingerprint = hashlib.sha256(b"synthetic").hexdigest() first_path = z_fixture_dir / "experts-L00.nvfp4" first_digest = FileExpertSource.create_synthetic( first_path, num_experts=1, records=[bytes([8]) * RAW_RECORD_BYTES], layer_id=0 @@ -119,21 +255,33 @@ def test_resident_builder_streams_one_record_without_full_layer(z_fixture_dir): manifest_data = { "format": FORMAT, "version": 1, + "artifact_schema": FORMAT, "text_only": True, - "active": {"format": "nvfp4_w4a16_v1", "path": "."}, + "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": source_fingerprint}, + "minimum_freetoken_commit": "0" * 40, + "tvm_ffi_patch_sha256": "0" * 64, + "active": {"format": "nvfp4_w4a16_v1", "path": ".", "files": [{"path": "manifest.json", "bytes": 0, "sha256": "0" * 64}]}, "ple": {"format": "q3_ple_32", "manifest": "ple-q3.json", "sha256": "0" * 64}, "experts": { "format": "ftexpert1_nvfp4_v1", - "files": [{"layer": 0, "path": first_path.name, "bytes": first_path.stat().st_size, "sha256": first_digest}, - {"layer": 1, "path": resident_path.name, "bytes": resident_path.stat().st_size, "sha256": digest}], + "files": [{"layer": 0, "path": first_path.name, "bytes": first_path.stat().st_size, "sha256": first_digest, "source_fingerprint": source_fingerprint}, + {"layer": 1, "path": resident_path.name, "bytes": resident_path.stat().st_size, "sha256": digest, "source_fingerprint": source_fingerprint}], "file_tier_layers": [], "resident_layers": [0, 1], "required_volume": "Z:", }, + "metadata": {"files": [{"path": "config.json", "bytes": 0, "sha256": hashlib.sha256(b"").hexdigest()}]}, } + manifest_data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(manifest_data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() path = z_fixture_dir / "manifest.json" + (z_fixture_dir / "config.json").write_bytes(b"") + (z_fixture_dir / "ple-q3.json").write_text( + json.dumps({"source_fingerprint": source_fingerprint}), encoding="utf-8" + ) path.write_text(json.dumps(manifest_data), encoding="utf-8") - manifest = load_qwen4_artifact_manifest(path, require=True) + manifest = load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) resident_sources, file_sources = build_mixed_expert_sources( manifest, num_experts=1, @@ -143,3 +291,200 @@ def test_resident_builder_streams_one_record_without_full_layer(z_fixture_dir): assert file_sources == {} assert resident_sources["gate_up_packed"][0].shape == (1, 1280, 1280) assert int(resident_sources["gate_up_packed"][1][0, 0, 0]) == 9 + + +def _tiny_geometry(): + return { + "gate_up_packed": ((4, 4), torch.uint8), + "gate_up_scale": ((4, 1), torch.float8_e4m3fn), + "gate_up_global": ((4,), torch.float16), + "down_packed": ((2, 16), torch.uint8), + "down_scale": ((2, 1), torch.float8_e4m3fn), + "down_global": ((2,), torch.float16), + } + + +def _tiny_planes(value: int): + return { + name: torch.full(shape, 1 if dtype == torch.float8_e4m3fn else value, dtype=dtype) + for name, (shape, dtype) in _tiny_geometry().items() + } + + +def test_end_to_end_synthetic_modular_artifact_reopens_normal_paths(z_fixture_dir, monkeypatch): + source_fingerprint = "4" * 64 + config = { + "architectures": ["Qwen4ExpForConditionalGeneration"], + "freetoken_text_only": TEXT_ONLY_MARKER, + "freetoken_active_quant": "nvfp4_w4a16_v1", + } + (z_fixture_dir / "config.json").write_text(json.dumps(config), encoding="utf-8") + (z_fixture_dir / "tokenizer_config.json").write_text("{}", encoding="utf-8") + + active = z_fixture_dir / "qwen4-active-v1.ftw" + writer = FTWWriter(str(active), shard_limit=4096 * 16) + fused = torch.arange(64, dtype=torch.float32).reshape(4, 16).to(torch.bfloat16) + protected = torch.arange(16, dtype=torch.bfloat16) + entries = iter_active_nvfp4_runtime_entries( + [ + ("model.layers.0.self_attn.qkv_proj.weight", fused), + ("model.layers.0.self_attn.index_qk_proj.weight", protected), + ] + ) + emitted_names = [] + for name, tensor in entries: + emitted_names.append(name) + writer.add_tensor(name, tensor) + writer.finalize({ + "artifact_format": "qwen4_modular_v1", + "source_inventory_sha256": source_fingerprint, + }) + assert emitted_names == [ + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight_scale", + "model.layers.0.self_attn.qkv_proj.weight_global", + "model.layers.0.self_attn.index_qk_proj.weight", + ] + + write_q3_ple_sidecar( + ([float((row + column) % 9 - 4) for column in range(ROW_VALUES)] for row in range(4)), + z_fixture_dir / "ple-q3-000.bin", + z_fixture_dir / "ple-q3.json", + source_fingerprint=source_fingerprint, + weight_scale=1.0, + segment_rows=2, + ) + experts = {} + for layer in range(2): + path = z_fixture_dir / f"experts-L{layer:02d}.nvfp4" + write_expert_sidecar( + path, + [_tiny_planes(layer + expert + 1) for expert in range(2)], + layer_id=layer, + source_fingerprint=source_fingerprint, + num_experts=2, + geometry=_tiny_geometry(), + ) + experts[layer] = path.name + + finalized = finalize_qwen4_modular_manifest( + z_fixture_dir, + source_repository="synthetic/qwen4", + source_revision="3" * 40, + source_inventory_sha256=source_fingerprint, + minimum_freetoken_commit="846504bf9d81119cb72400e6c5a3cc860f2b1dd8", + tvm_ffi_patch_sha256="889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec", + expert_paths=experts, + file_tier_layers=[0], + metadata_paths=["config.json", "tokenizer_config.json"], + expert_num_experts=2, + allow_synthetic_geometry=True, + ) + assert finalized["experts"]["resident_layers"] == [1] + manifest = load_qwen4_artifact_manifest(z_fixture_dir, require=True, allow_synthetic_geometry=True) + assert manifest is not None + manifest.verify_active() + import freetoken.checkpoint.qwen4_artifact as artifact_module + + original_loader = artifact_module.load_qwen4_artifact_manifest + monkeypatch.setattr( + artifact_module, + "load_qwen4_artifact_manifest", + lambda _path, **_kwargs: manifest, + ) + loaded = dict(load_weight(str(z_fixture_dir), torch.device("cpu"), include_moe_experts=False)) + assert set(loaded) == set(emitted_names) + resident, file_sources = build_mixed_expert_sources( + manifest, + num_experts=2, + resident_residency=["pageable", "pageable"], + allocator=lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + assert all(resident[name][0] is None for name in FileExpertSource.bank_schema) + assert all(resident[name][1] is not None for name in FileExpertSource.bank_schema) + assert set(file_sources) == {0} + file_sources[0].close() + + # Normal Qwen4 host-load dispatch selects Q3 from the manifest; no manual + # load_q3_ple_weights injection is involved. + calls = [] + fake = SimpleNamespace(load_q3_ple_weights=lambda path: calls.append(path)) + from freetoken.models.qwen4_exp.model import Qwen4ExpModel + + Qwen4ExpModel.load_host_weights(fake, str(z_fixture_dir), dummy=False) + assert calls == [str(z_fixture_dir / "ple-q3.json")] + monkeypatch.setattr(artifact_module, "load_qwen4_artifact_manifest", original_loader) + + +def test_production_orchestrator_sequences_all_modular_components(z_fixture_dir, monkeypatch): + source = z_fixture_dir / "source" + target = z_fixture_dir / "target" + source.mkdir() + inventory = "5" * 64 + + import freetoken.checkpoint.convert as convert_module + import freetoken.checkpoint.q3_ple as q3_module + import freetoken.moe.expert_source as expert_module + + def fake_convert(_source, out_dir, **kwargs): + assert kwargs["artifact_format"] == "qwen4_modular_v1" + assert kwargs["source_inventory_sha256"] == inventory + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + (out / "config.json").write_text( + json.dumps({ + "freetoken_text_only": TEXT_ONLY_MARKER, + "freetoken_active_quant": "nvfp4_w4a16_v1", + }), + encoding="utf-8", + ) + active = out / "qwen4-active-v1.ftw" + writer = FTWWriter(str(active), shard_limit=4096 * 4) + writer.add_tensor("protected.weight", torch.ones(1, dtype=torch.bfloat16)) + return writer.finalize({ + "source_inventory_sha256": inventory, + "copied_metadata": ["config.json"], + }) + + def fake_q3(_source, data_path, manifest_path, **kwargs): + assert kwargs["source_fingerprint"] == inventory + return write_q3_ple_sidecar( + ([0.0] * ROW_VALUES for _ in range(2)), + data_path, + manifest_path, + source_fingerprint=inventory, + weight_scale=1.0, + segment_rows=1, + ) + + def fake_expert(_source, path, **kwargs): + return write_expert_sidecar( + path, + [_tiny_planes(kwargs["layer_id"] + 1) for _ in range(kwargs["num_experts"])], + layer_id=kwargs["layer_id"], + source_fingerprint=kwargs["source_fingerprint"], + num_experts=kwargs["num_experts"], + geometry=kwargs["geometry"], + ) + + monkeypatch.setattr(convert_module, "convert_checkpoint", fake_convert) + monkeypatch.setattr(q3_module, "write_q3_ple_from_safetensors", fake_q3) + monkeypatch.setattr(expert_module, "write_expert_sidecar_from_safetensors", fake_expert) + manifest = build_qwen4_modular_artifact( + source, + target, + source_repository="synthetic/qwen4", + source_revision="6" * 40, + source_inventory_sha256=inventory, + minimum_freetoken_commit="7" * 40, + ple_split_parts=1, + expert_layers=(0, 1), + file_tier_layers=(0,), + expert_num_experts=2, + expert_geometry=_tiny_geometry(), + allow_synthetic_geometry=True, + ) + assert manifest["active"]["format"] == "nvfp4_w4a16_v1" + assert manifest["ple"]["source_fingerprint"] == inventory + assert [item["layer"] for item in manifest["experts"]["files"]] == [0, 1] + assert (target / "manifest.json").is_file() diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b4d..5e28d40db 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -97,6 +97,14 @@ def test_expert_bytes_per_slot_sums_row_bytes_over_banks(): assert expert_bytes_per_slot(sources) == 512 + 256 +def test_expert_bytes_per_slot_accepts_file_tier_placeholders(): + sources = { + "packed": [None, torch.empty((2, 16), dtype=torch.uint8)], + "scale": [None, torch.empty((2, 4), dtype=torch.float16)], + } + assert expert_bytes_per_slot(sources) == 16 + 8 + + def test_resolve_auto_applies_ratio_once_and_marlin_cap(): # baseline 1000, weights 100, ratio 0.9 -> budget = 900 - 100 - 0(fixed) = 800 size, pages, overlap = resolve_moe_cache_auto( diff --git a/tests/models/test_qwen4_exp_raw_config.py b/tests/models/test_qwen4_exp_raw_config.py index 966917cd0..e9d17d533 100644 --- a/tests/models/test_qwen4_exp_raw_config.py +++ b/tests/models/test_qwen4_exp_raw_config.py @@ -100,3 +100,48 @@ def test_qwen4_rejects_unknown_active_weight_marker(): assert "unsupported Qwen4 active-weight format" in str(exc) else: raise AssertionError("unknown active-weight marker was accepted") + + +def _add_minimal_vision(config: RawConfigShim) -> None: + config.vision_config = RawConfigShim( + { + "depth": 2, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 4, + "num_position_embeddings": 16, + "out_hidden_size": 2560, + "patch_size": 14, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "in_channels": 3, + "hidden_act": "gelu", + "deepstack_visual_indexes": [0], + } + ) + + +def test_qwen4_text_only_marker_disables_only_marked_artifact_vision(): + published = _raw_checkpoint_config() + _add_minimal_vision(published) + parsed = parse_config(published) + assert parsed.vision_config is not None + assert parsed.image_token_id == 248056 + + target = _raw_checkpoint_config() + _add_minimal_vision(target) + target.freetoken_text_only = "qwen4_text_only_v1" + parsed = parse_config(target) + assert parsed.vision_config is None + assert parsed.image_token_id is None + + +def test_qwen4_rejects_unknown_text_only_marker(): + config = _raw_checkpoint_config() + config.freetoken_text_only = "textish-v0" + try: + parse_config(config) + except ValueError as exc: + assert "unsupported freetoken_text_only marker" in str(exc) + else: + raise AssertionError("unknown text-only marker was accepted") diff --git a/tests/moe/test_expert_sidecar_writer.py b/tests/moe/test_expert_sidecar_writer.py index 7d9b45a32..604ab21c0 100644 --- a/tests/moe/test_expert_sidecar_writer.py +++ b/tests/moe/test_expert_sidecar_writer.py @@ -14,6 +14,7 @@ MAGIC, adapt_expert_tensor_record, write_expert_sidecar, + write_expert_sidecar_from_safetensors, ) @@ -68,7 +69,7 @@ def test_writer_reduced_geometry_reopens_and_is_deterministic(z_dir): assert result["format"] == "FTEXPERT1" assert result["raw_record_bytes"] == 66 assert result["record_bytes"] == 4096 - assert result["sample_ids"] == (0, 2) + assert result["sample_ids"] == (0, 1, 2) assert first.read_bytes() == second.read_bytes() assert result["sha256"] == hashlib.sha256(first.read_bytes()).hexdigest() assert result["sha256"] == result_copy["sha256"] @@ -93,6 +94,37 @@ def test_source_tensor_adapter_validates_twelve_names_and_expands_globals(z_dir) assert source.read_record(0)["gate_up_packed"].flatten()[0].item() == 9 +def test_production_writer_streams_indexed_safetensor_experts(z_dir): + import json + from safetensors.torch import save_file + + prefix = "model.language_model.layers.3.mlp.experts" + tensors = {} + weight_map = {} + for expert_id in range(2): + for suffix, tensor in _source_record(5 + expert_id).items(): + name = f"{prefix}.{expert_id}.{suffix}" + tensors[name] = tensor + weight_map[name] = "layer-3.safetensors" + save_file(tensors, z_dir / "layer-3.safetensors") + (z_dir / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}), encoding="utf-8" + ) + output = z_dir / "experts-L03.nvfp4" + result = write_expert_sidecar_from_safetensors( + z_dir, + output, + layer_id=3, + source_fingerprint="a" * 64, + num_experts=2, + geometry=_geometry(), + ) + assert result["sample_ids"] == (0, 1) + with FileExpertSource(output, num_experts=2, expected_layer_id=3) as source: + assert int(source.read_record(0)["gate_up_packed"][0, 0]) == 5 + assert int(source.read_record(1)["down_packed"][0, 0]) == 6 + + def test_writer_accepts_explicit_id_and_named_pairs(z_dir): named = [(name, value) for name, value in _planes(4).items()] path = z_dir / "explicit.ftex" diff --git a/tests/moe/test_file_expert_source.py b/tests/moe/test_file_expert_source.py index 226e61b6e..996003aa1 100644 --- a/tests/moe/test_file_expert_source.py +++ b/tests/moe/test_file_expert_source.py @@ -202,7 +202,7 @@ def test_file_tier_payload_hash_fails_closed(z_fixture_dir): FileExpertSource(path, num_experts=1) -def test_file_tier_rejects_cpu_hybrid_and_overlap(z_fixture_dir): +def test_file_tier_rejects_only_cpu_selected_file_layers(z_fixture_dir): path = z_fixture_dir / "experts-L00.nvfp4" digest = FileExpertSource.create_synthetic(path, num_experts=1, records=[_record(7)]) src = FileExpertSource(path, num_experts=1, expected_sha256=digest) @@ -211,11 +211,17 @@ def test_file_tier_rejects_cpu_hybrid_and_overlap(z_fixture_dir): for target in ("cpu", "hybrid"): cache = OffloadMoeCache(1, 1, 1, torch.device("cpu"), decode_target=target, quant_format="nvfp4") + cache.cpu_layer_ids = frozenset({0}) # Source registration itself fails before any source can be used; # a direct call is sufficient to prove the policy and avoids giant # synthetic resident allocations in this negative test. with pytest.raises(ValueError, match="GPU-only"): cache.set_file_sources({0: src}) + + mixed = OffloadMoeCache(2, 1, 1, torch.device("cpu"), decode_target="cpu", quant_format="nvfp4") + mixed.cpu_layer_ids = frozenset({1}) + mixed.set_file_sources({0: src}) + assert set(mixed.file_sources) == {0} finally: src.close() From deb56560bd53cf688e5ca8635a7405c3417d1015 Mon Sep 17 00:00:00 2001 From: Your Name Date: Fri, 28 Aug 2026 15:40:09 -0400 Subject: [PATCH 11/17] fix(qwen4): freeze pr257 hardware-fit builder contract --- python/freetoken/checkpoint/__init__.py | 5 +- python/freetoken/checkpoint/convert.py | 1 + python/freetoken/checkpoint/q3_ple.py | 274 ++++++++++++++++-- python/freetoken/checkpoint/qwen4_artifact.py | 9 + tests/checkpoint/test_convert_metadata.py | 1 + tests/checkpoint/test_q3_ple_writer.py | 139 +++++++++ tests/checkpoint/test_qwen4_artifact.py | 5 + 7 files changed, 413 insertions(+), 21 deletions(-) diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index f78d90ca9..60dec40be 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -15,6 +15,8 @@ from .q3_ple import ( Q3PLEReader, Q3PLESegment, + plan_q3_ple_production, + write_q3_ple_segmented_sidecar, write_q3_ple_from_safetensors, write_q3_ple_sidecar, ) @@ -23,5 +25,6 @@ "FTWReader", "FTWWriter", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", "Q3PLEReader", "Q3PLESegment", "write_q3_ple_sidecar", - "write_q3_ple_from_safetensors", + "write_q3_ple_segmented_sidecar", "write_q3_ple_from_safetensors", + "plan_q3_ple_production", ] diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 106cbdb22..dd17411dd 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -389,6 +389,7 @@ def convert_checkpoint( config_data = json.load(handle) config_data["freetoken_text_only"] = "qwen4_text_only_v1" config_data["freetoken_active_quant"] = "nvfp4_w4a16_v1" + config_data["freetoken_runtime_foundation"] = "pr257_hardware_fit_v1" tmp_config = config_path + ".tmp" with open(tmp_config, "w", encoding="utf-8") as handle: json.dump(config_data, handle, indent=2, sort_keys=True) diff --git a/python/freetoken/checkpoint/q3_ple.py b/python/freetoken/checkpoint/q3_ple.py index cd62cc5b2..84eec13b5 100644 --- a/python/freetoken/checkpoint/q3_ple.py +++ b/python/freetoken/checkpoint/q3_ple.py @@ -35,6 +35,28 @@ REFINEMENT_PASSES = 2 DEFAULT_SEGMENT_ROWS = 128 +# Production Qwen4 PLE geometry. ``segment_count`` is a logical source +# tensor count; it is deliberately independent from the bounded row chunk used +# while reading a Safetensors tensor. +PRODUCTION_SEGMENT_COUNT = 128 +PRODUCTION_ROWS_PER_SEGMENT = 2_500_012 +PRODUCTION_TOTAL_ROWS = 320_001_536 +PRODUCTION_SEGMENT_BYTES = 175_000_840 +PRODUCTION_TOTAL_BYTES = 22_400_107_520 + + +def plan_q3_ple_production() -> dict[str, int]: + """Return the frozen real-table plan without allocating any payload.""" + + return { + "segment_count": PRODUCTION_SEGMENT_COUNT, + "rows_per_segment": PRODUCTION_ROWS_PER_SEGMENT, + "total_rows": PRODUCTION_TOTAL_ROWS, + "row_bytes": ROW_BYTES, + "segment_bytes": PRODUCTION_SEGMENT_BYTES, + "total_bytes": PRODUCTION_TOTAL_BYTES, + } + def _z_path(path: str | os.PathLike[str]) -> Path: """Resolve *path* and fail closed unless its physical drive is ``Z:``.""" @@ -266,6 +288,8 @@ def __init__(self, manifest_path: str | os.PathLike[str], *, data_path: str | os if not isinstance(raw_segments, list) or not raw_segments: raise ValueError("Q3_PLE_32 segment directory is empty") self.segments: tuple[Q3PLESegment, ...] = tuple(self._parse_segment(item) for item in raw_segments) + if "segment_count" in manifest and int(manifest["segment_count"]) != len(self.segments): + raise ValueError("Q3_PLE_32 segment_count does not match the segment directory") self._validate_segments() stat = self.data_path.stat() expected_file_bytes = int(manifest.get("file_bytes", stat.st_size)) @@ -630,6 +654,148 @@ def write_q3_ple_sidecar( return manifest +def write_q3_ple_segmented_sidecar( + segments: Iterable[Iterable[Sequence[float] | torch.Tensor]], + data_path: str | os.PathLike[str], + manifest_path: str | os.PathLike[str], + *, + source_fingerprint: str, + weight_scale: float, + segment_count: int, + rows_per_segment: int | None = None, +) -> dict: + """Write Q3 data with explicit logical source-segment boundaries. + + Each item in ``segments`` represents one source tensor (for production, + ``shard_0`` through ``shard_127``). Segment identity is therefore stable + regardless of the internal row-chunk size used by the caller. The legacy + :func:`write_q3_ple_sidecar` API remains row-count based for compatibility + with small historical fixtures; production Safetensors conversion uses + this explicit API instead. + """ + + data_final = _z_output_path(data_path) + manifest_final = _z_output_path(manifest_path) + if data_final == manifest_final: + raise ValueError("Q3_PLE_32 data_path and manifest_path must differ") + source_digest = _validate_source_fingerprint(source_fingerprint) + if isinstance(segment_count, bool): + raise ValueError("segment_count must be a positive integer") + try: + expected_segments = operator.index(segment_count) + except TypeError as exc: + raise ValueError("segment_count must be a positive integer") from exc + if expected_segments <= 0: + raise ValueError("segment_count must be a positive integer") + expected_rows = None if rows_per_segment is None else operator.index(rows_per_segment) + if expected_rows is not None and expected_rows <= 0: + raise ValueError("rows_per_segment must be a positive integer") + try: + global_scale = float(weight_scale) + except (TypeError, ValueError) as exc: + raise ValueError("weight_scale must be finite") from exc + if not math.isfinite(global_scale): + raise ValueError("weight_scale must be finite") + + token = uuid.uuid4().hex + data_partial = _partial_path(data_final, token) + manifest_partial = _partial_path(manifest_final, token) + segments_manifest: list[dict[str, int | str]] = [] + whole_digest = hashlib.sha256() + payload_digest = hashlib.sha256() + rows_written = 0 + file_offset = 0 + + try: + with data_partial.open("wb") as output: + for segment_index, source_segment in enumerate(segments): + if segment_index >= expected_segments: + raise ValueError( + f"Q3_PLE_32 expected {expected_segments} logical segments, got more" + ) + first_row = rows_written + segment_offset = file_offset + segment_digest = hashlib.sha256() + segment_rows = 0 + for source_row in source_segment: + row_values = _materialize_row(source_row) + encoded_row = quantize_row(row_values, refinement_passes=REFINEMENT_PASSES) + if len(encoded_row) != ROW_BYTES: + raise AssertionError(f"Q3_PLE_32 row has wrong size: {len(encoded_row)}") + output.write(encoded_row) + whole_digest.update(encoded_row) + payload_digest.update(encoded_row) + segment_digest.update(encoded_row) + file_offset += len(encoded_row) + rows_written += 1 + segment_rows += 1 + if segment_rows <= 0: + raise ValueError(f"Q3_PLE_32 logical segment {segment_index} is empty") + if expected_rows is not None and segment_rows != expected_rows: + raise ValueError( + f"Q3_PLE_32 logical segment {segment_index} has {segment_rows} rows, " + f"expected {expected_rows}" + ) + segments_manifest.append( + { + "first_row": first_row, + "end_row": rows_written, + "data_offset": segment_offset, + "byte_length": segment_rows * ROW_BYTES, + "sha256": segment_digest.hexdigest(), + } + ) + if len(segments_manifest) != expected_segments: + raise ValueError( + f"Q3_PLE_32 expected {expected_segments} logical segments, " + f"got {len(segments_manifest)}" + ) + _fsync(output) + except Exception: + data_partial.unlink(missing_ok=True) + manifest_partial.unlink(missing_ok=True) + raise + + file_bytes = data_partial.stat().st_size + if file_bytes != file_offset: + data_partial.unlink(missing_ok=True) + raise OSError(f"Q3_PLE_32 partial length mismatch: {file_bytes} != {file_offset}") + _validate_segment_directory(segments_manifest, rows_written, file_bytes) + manifest = { + "format": FORMAT, + "version": VERSION, + "endianness": "little", + "block_values": BLOCK_VALUES, + "block_bytes": BLOCK_BYTES, + "row_values": ROW_VALUES, + "row_bytes": ROW_BYTES, + "rows": rows_written, + "payload_bytes": rows_written * ROW_BYTES, + "file_bytes": file_bytes, + "storage_layout": "contiguous_rows_v1", + "data_file": os.path.relpath(data_final, manifest_final.parent), + "weight_scale": global_scale, + "source_fingerprint": source_digest, + "sha256": whole_digest.hexdigest(), + "payload_sha256": payload_digest.hexdigest(), + "segments": segments_manifest, + "segment_count": expected_segments, + "segment_identity": "source_tensor_numeric_suffix_v1", + } + try: + with manifest_partial.open("w", encoding="utf-8", newline="\n") as manifest_handle: + json.dump(manifest, manifest_handle, ensure_ascii=False, indent=2, sort_keys=True) + manifest_handle.write("\n") + _fsync(manifest_handle) + os.replace(data_partial, data_final) + os.replace(manifest_partial, manifest_final) + except Exception: + data_partial.unlink(missing_ok=True) + manifest_partial.unlink(missing_ok=True) + raise + return manifest + + def write_q3_ple_from_safetensors( model_path: str | os.PathLike[str], data_path: str | os.PathLike[str], @@ -639,7 +805,10 @@ def write_q3_ple_from_safetensors( split_parts: int, source_fingerprint: str, rows_per_chunk: int = 8192, - segment_rows: int = DEFAULT_SEGMENT_ROWS, + segment_rows: int | None = None, + processing_chunk_rows: int | None = None, + segment_count: int | None = None, + rows_per_segment: int | None = None, ) -> dict: """Stream the official FP8 PLE shards into the native Q3 sidecar. @@ -652,8 +821,24 @@ def write_q3_ple_from_safetensors( folder = Path(model_path).expanduser().resolve() if not folder.is_dir(): raise ValueError(f"Q3 PLE source must be a local checkpoint directory: {folder}") + if processing_chunk_rows is not None: + if rows_per_chunk != 8192: + raise ValueError("specify only one of rows_per_chunk or processing_chunk_rows") + rows_per_chunk = processing_chunk_rows if rows_per_chunk <= 0 or split_parts <= 0: - raise ValueError("rows_per_chunk and split_parts must be positive") + raise ValueError("processing chunk rows and split_parts must be positive") + if segment_rows is not None: + # ``segment_rows`` was the legacy flat-stream API. Keep it for small + # fixtures, but production conversion must use explicit source tensor + # segments so a value of 128 can never mean 128 rows per segment. + if segment_count is not None or rows_per_segment is not None: + raise ValueError("segment_rows cannot be combined with explicit segment geometry") + if int(split_parts) == PRODUCTION_SEGMENT_COUNT and int(segment_rows) == DEFAULT_SEGMENT_ROWS: + # Historical callers passed ``segment_rows=128`` intending the + # production 128 logical tensors. Treat that exact combination as + # the explicit segmented contract; it must never create 128-row + # chunks across the complete flattened table. + segment_rows = None index_path = folder / "model.safetensors.index.json" with index_path.open("r", encoding="utf-8") as handle: weight_map = json.load(handle)["weight_map"] @@ -662,6 +847,26 @@ def write_q3_ple_from_safetensors( "ngram_embedding" ) shard_keys = [f"{prefix}.shard_{part}.weight" for part in range(int(split_parts))] + shard_prefix = prefix + ".shard_" + indexed_keys: dict[int, str] = {} + for key in weight_map: + if not isinstance(key, str) or not key.startswith(shard_prefix): + continue + suffix = key[len(shard_prefix) :] + if not suffix.endswith(".weight"): + raise ValueError(f"malformed PLE source tensor suffix: {key}") + index_text = suffix[: -len(".weight")] + if not index_text.isdigit(): + raise ValueError(f"malformed PLE source tensor suffix: {key}") + index = int(index_text) + if index in indexed_keys: + raise ValueError(f"duplicate PLE source tensor index: {index}") + if not 0 <= index < int(split_parts): + raise ValueError(f"PLE source tensor index outside 0..{int(split_parts) - 1}: {index}") + indexed_keys[index] = key + if set(indexed_keys) != set(range(int(split_parts))): + missing_indices = sorted(set(range(int(split_parts))) - set(indexed_keys)) + raise ValueError(f"missing PLE source tensor indices: {missing_indices}") missing = [key for key in shard_keys if key not in weight_map] scale_key = prefix + ".weight_scale" if missing or scale_key not in weight_map: @@ -674,28 +879,50 @@ def write_q3_ple_from_safetensors( scale = handle.get_tensor(scale_key).reshape(()) weight_scale = float(scale.float().item()) - def iter_rows(): - for key in shard_keys: - source_file = folder / weight_map[key] - with safetensors.safe_open(source_file, framework="pt", device="cpu") as handle: - sliced = handle.get_slice(key) - shape = tuple(int(value) for value in sliced.get_shape()) - if len(shape) != 2 or shape[1] != ROW_VALUES: - raise ValueError(f"unexpected PLE source shape for {key}: {shape}") - for start in range(0, shape[0], int(rows_per_chunk)): - chunk = sliced[start : min(start + int(rows_per_chunk), shape[0])] - if chunk.dtype != torch.float8_e4m3fn: - raise ValueError(f"unexpected PLE source dtype for {key}: {chunk.dtype}") - for row in chunk.float(): - yield row - - return write_q3_ple_sidecar( - iter_rows(), + def iter_segment_rows(key: str): + source_file = folder / weight_map[key] + with safetensors.safe_open(source_file, framework="pt", device="cpu") as handle: + sliced = handle.get_slice(key) + shape = tuple(int(value) for value in sliced.get_shape()) + if len(shape) != 2 or shape[1] != ROW_VALUES: + raise ValueError(f"unexpected PLE source shape for {key}: {shape}") + if rows_per_segment is not None and shape[0] != int(rows_per_segment): + raise ValueError( + f"unexpected PLE source row count for {key}: {shape[0]} != {rows_per_segment}" + ) + for start in range(0, shape[0], int(rows_per_chunk)): + chunk = sliced[start : min(start + int(rows_per_chunk), shape[0])] + if chunk.dtype != torch.float8_e4m3fn: + raise ValueError(f"unexpected PLE source dtype for {key}: {chunk.dtype}") + for row in chunk.float(): + yield row + + # Explicit segmented mode is the production contract. Legacy callers can + # request flat row segmentation by passing ``segment_rows`` explicitly. + if segment_rows is not None: + def iter_rows(): + for key in shard_keys: + yield from iter_segment_rows(key) + + return write_q3_ple_sidecar( + iter_rows(), data_path, manifest_path, + source_fingerprint=source_fingerprint, + weight_scale=weight_scale, + segment_rows=segment_rows, + ) + + logical_count = int(segment_count if segment_count is not None else split_parts) + if logical_count != int(split_parts): + raise ValueError("segment_count must equal split_parts for PLE Safetensors conversion") + ordered_segments = (iter_segment_rows(indexed_keys[index]) for index in range(logical_count)) + return write_q3_ple_segmented_sidecar( + ordered_segments, data_path, manifest_path, source_fingerprint=source_fingerprint, weight_scale=weight_scale, - segment_rows=segment_rows, + segment_count=logical_count, + rows_per_segment=rows_per_segment, ) @@ -711,8 +938,15 @@ def iter_rows(): "ROW_BYTES", "ROW_VALUES", "VERSION", + "PRODUCTION_SEGMENT_COUNT", + "PRODUCTION_ROWS_PER_SEGMENT", + "PRODUCTION_TOTAL_ROWS", + "PRODUCTION_SEGMENT_BYTES", + "PRODUCTION_TOTAL_BYTES", + "plan_q3_ple_production", "quantize_block", "quantize_row", "write_q3_ple_sidecar", + "write_q3_ple_segmented_sidecar", "write_q3_ple_from_safetensors", ] diff --git a/python/freetoken/checkpoint/qwen4_artifact.py b/python/freetoken/checkpoint/qwen4_artifact.py index 4e2096286..cd84df16e 100644 --- a/python/freetoken/checkpoint/qwen4_artifact.py +++ b/python/freetoken/checkpoint/qwen4_artifact.py @@ -39,6 +39,7 @@ PINNED_SOURCE_REPOSITORY = "RadixArk/Qwen3.8-Flash-Next-NVFP4" PINNED_SOURCE_REVISION = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" TVM_FFI_PATCH_SHA256 = "889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec" +RUNTIME_FOUNDATION_MARKER = "pr257_hardware_fit_v1" class Qwen4ArtifactError(ValueError): @@ -276,6 +277,8 @@ def load_qwen4_artifact_manifest( raise Qwen4ArtifactError("unsupported Qwen4 modular artifact_schema") if root.get("text_only") is not True: raise Qwen4ArtifactError("Qwen4 modular manifest must declare text_only=true") + if root.get("runtime_foundation") != RUNTIME_FOUNDATION_MARKER: + raise Qwen4ArtifactError("Qwen4 modular manifest lacks the accepted runtime foundation") source = _require_mapping(root.get("source"), label="source") if not str(source.get("repository", "")).strip() or not str(source.get("revision", "")).strip(): raise Qwen4ArtifactError("source repository and revision are required") @@ -677,6 +680,7 @@ def build_qwen4_modular_artifact( layer_id=int(ple_layer_id), split_parts=int(ple_split_parts), source_fingerprint=inventory, + rows_per_segment=(None if allow_synthetic_geometry else 2_500_012), ) expert_paths: dict[int, str] = {} for layer in expert_layers: @@ -829,12 +833,16 @@ def finalize_qwen4_modular_manifest( raise Qwen4ArtifactError("config.json lacks the accepted text-only marker") if config.get("freetoken_active_quant") != ACTIVE_FORMAT: raise Qwen4ArtifactError("config.json lacks the accepted active-quant marker") + config_runtime_foundation = config.get("freetoken_runtime_foundation") + if config_runtime_foundation != RUNTIME_FOUNDATION_MARKER: + raise Qwen4ArtifactError("config.json lacks the accepted runtime-foundation marker") manifest: dict[str, Any] = { "format": FORMAT, "version": VERSION, "artifact_schema": FORMAT, "text_only": True, + "runtime_foundation": RUNTIME_FOUNDATION_MARKER, "source": { "repository": str(source_repository), "revision": str(source_revision), @@ -890,6 +898,7 @@ def finalize_qwen4_modular_manifest( "Qwen4ArtifactError", "Qwen4ArtifactManifest", "TEXT_ONLY_MARKER", + "RUNTIME_FOUNDATION_MARKER", "VERSION", "build_qwen4_modular_artifact", "configure_mixed_expert_sources", diff --git a/tests/checkpoint/test_convert_metadata.py b/tests/checkpoint/test_convert_metadata.py index 124d6bf93..1bf95ecc4 100644 --- a/tests/checkpoint/test_convert_metadata.py +++ b/tests/checkpoint/test_convert_metadata.py @@ -143,6 +143,7 @@ def fake_load_weight(_path, device, *, include_moe_experts): config = json.loads((output / "config.json").read_text(encoding="utf-8")) assert config["freetoken_text_only"] == "qwen4_text_only_v1" assert config["freetoken_active_quant"] == "nvfp4_w4a16_v1" + assert config["freetoken_runtime_foundation"] == "pr257_hardware_fit_v1" loaded = dict(iter_ftw_weights(str(output / "qwen4-active-v1.ftw"), workers=1)) assert set(loaded) == { "model.layers.0.self_attn.o_proj.weight", diff --git a/tests/checkpoint/test_q3_ple_writer.py b/tests/checkpoint/test_q3_ple_writer.py index f3d3c22ba..2b9835913 100644 --- a/tests/checkpoint/test_q3_ple_writer.py +++ b/tests/checkpoint/test_q3_ple_writer.py @@ -17,9 +17,16 @@ ROW_BYTES, ROW_VALUES, Q3PLEReader, + PRODUCTION_SEGMENT_BYTES, + PRODUCTION_SEGMENT_COUNT, + PRODUCTION_ROWS_PER_SEGMENT, + PRODUCTION_TOTAL_BYTES, + PRODUCTION_TOTAL_ROWS, quantize_block, + plan_q3_ple_production, write_q3_ple_from_safetensors, write_q3_ple_sidecar, + write_q3_ple_segmented_sidecar, ) @@ -208,3 +215,135 @@ def test_production_writer_streams_safetensor_shards_in_source_order(z_fixture_d assert manifest["rows"] == 4 assert manifest["file_bytes"] == 4 * ROW_BYTES assert manifest["weight_scale"] == 0.5 + + +def _write_128_segment_source(root: Path, *, malformed: str | None = None) -> str: + """Create a tiny source inventory with one deterministic row per shard.""" + from safetensors.torch import save_file + + prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding" + tensors = {} + weight_map = {} + source_file = root / "model-plefp8-00000.safetensors" + for index in range(PRODUCTION_SEGMENT_COUNT): + suffix = f"shard_{index}.weight" + if malformed == "duplicate" and index == 18: + # A JSON object cannot contain a literal duplicate key. A + # different spelling of the same numeric suffix exercises the + # production duplicate-index guard without relying on parser + # behavior for duplicate object members. + suffix = "shard_017.weight" + if malformed == "outside" and index == 18: + suffix = "shard_128.weight" + if malformed == "malformed" and index == 18: + suffix = "shard_bad.weight" + key = f"{prefix}.{suffix}" + values = torch.full((1, ROW_VALUES), float(index + 1), dtype=torch.float32) + tensors[key] = values.to(torch.float8_e4m3fn) + weight_map[key] = source_file.name + if malformed == "missing": + weight_map.pop(f"{prefix}.shard_63.weight") + tensors.pop(f"{prefix}.shard_63.weight") + if malformed == "shape": + tensors[f"{prefix}.shard_7.weight"] = torch.zeros((2, ROW_VALUES), dtype=torch.float8_e4m3fn) + tensors[f"{prefix}.weight_scale"] = torch.tensor(0.5, dtype=torch.bfloat16) + weight_map[f"{prefix}.weight_scale"] = source_file.name + save_file(tensors, source_file) + (root / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}), encoding="utf-8" + ) + return prefix + + +def test_q3_segmented_writer_has_stable_128_source_segments_and_boundaries(z_fixture_dir: Path) -> None: + rows = [ + [[float(index + 1)] * ROW_VALUES] + for index in range(PRODUCTION_SEGMENT_COUNT) + ] + first = write_q3_ple_segmented_sidecar( + (iter(segment) for segment in rows), + z_fixture_dir / "a.bin", + z_fixture_dir / "a.json", + source_fingerprint="1" * 64, + weight_scale=0.5, + segment_count=PRODUCTION_SEGMENT_COUNT, + ) + second = write_q3_ple_segmented_sidecar( + (iter(segment) for segment in rows), + z_fixture_dir / "b.bin", + z_fixture_dir / "b.json", + source_fingerprint="1" * 64, + weight_scale=0.5, + segment_count=PRODUCTION_SEGMENT_COUNT, + ) + assert len(first["segments"]) == PRODUCTION_SEGMENT_COUNT + assert [item["first_row"] for item in first["segments"]] == list(range(128)) + assert first["segments"][0]["data_offset"] == 0 + assert first["segments"][1]["data_offset"] == ROW_BYTES + assert first["segments"][-1]["first_row"] == 127 + assert first["segments"][-1]["end_row"] == 128 + # Data and semantic segment directory are byte-identical across runs. + assert (z_fixture_dir / "a.bin").read_bytes() == (z_fixture_dir / "b.bin").read_bytes() + assert {k: v for k, v in first.items() if k != "data_file"} == { + k: v for k, v in second.items() if k != "data_file" + } + + +def test_q3_safetensor_chunking_cannot_change_logical_segment_directory(z_fixture_dir: Path) -> None: + source = z_fixture_dir / "source" + source.mkdir() + _write_128_segment_source(source) + left = write_q3_ple_from_safetensors( + source, z_fixture_dir / "left.bin", z_fixture_dir / "left.json", + layer_id=2, split_parts=128, source_fingerprint="2" * 64, + processing_chunk_rows=1, + ) + right = write_q3_ple_from_safetensors( + source, z_fixture_dir / "right.bin", z_fixture_dir / "right.json", + layer_id=2, split_parts=128, source_fingerprint="2" * 64, + processing_chunk_rows=2, + ) + assert (z_fixture_dir / "left.bin").read_bytes() == (z_fixture_dir / "right.bin").read_bytes() + assert left["segment_count"] == right["segment_count"] == 128 + assert left["segments"] == right["segments"] + + +@pytest.mark.parametrize("bad", ["missing", "duplicate", "outside", "malformed", "shape"]) +def test_q3_safetensor_source_rejects_bad_logical_segments(z_fixture_dir: Path, bad: str) -> None: + source = z_fixture_dir / bad + source.mkdir() + _write_128_segment_source(source, malformed=bad) + with pytest.raises(ValueError): + write_q3_ple_from_safetensors( + source, z_fixture_dir / f"{bad}.bin", z_fixture_dir / f"{bad}.json", + layer_id=2, split_parts=128, source_fingerprint="3" * 64, + rows_per_segment=1, + ) + + +@pytest.mark.parametrize("segment_count", [127, 129]) +def test_q3_segmented_writer_rejects_wrong_segment_count(z_fixture_dir: Path, segment_count: int) -> None: + segments = ([float(index)] * ROW_VALUES for index in range(128)) + with pytest.raises(ValueError): + write_q3_ple_segmented_sidecar( + ((row,) for row in segments), + z_fixture_dir / f"{segment_count}.bin", + z_fixture_dir / f"{segment_count}.json", + source_fingerprint="4" * 64, + weight_scale=1.0, + segment_count=segment_count, + ) + + +def test_q3_production_planner_is_exact_without_allocating_payload() -> None: + plan = plan_q3_ple_production() + assert plan == { + "segment_count": 128, + "rows_per_segment": 2_500_012, + "total_rows": 320_001_536, + "row_bytes": ROW_BYTES, + "segment_bytes": 175_000_840, + "total_bytes": 22_400_107_520, + } + assert PRODUCTION_SEGMENT_COUNT * PRODUCTION_ROWS_PER_SEGMENT == PRODUCTION_TOTAL_ROWS + assert PRODUCTION_SEGMENT_COUNT * PRODUCTION_SEGMENT_BYTES == PRODUCTION_TOTAL_BYTES diff --git a/tests/checkpoint/test_qwen4_artifact.py b/tests/checkpoint/test_qwen4_artifact.py index 096d7cff9..9f67db5ae 100644 --- a/tests/checkpoint/test_qwen4_artifact.py +++ b/tests/checkpoint/test_qwen4_artifact.py @@ -12,6 +12,7 @@ from freetoken.checkpoint.qwen4_artifact import ( FORMAT, + RUNTIME_FOUNDATION_MARKER, TEXT_ONLY_MARKER, build_qwen4_modular_artifact, configure_mixed_expert_sources, @@ -84,6 +85,7 @@ def _manifest(root: Path, sidecar: Path, digest: str) -> Path: "version": 1, "artifact_schema": FORMAT, "text_only": True, + "runtime_foundation": RUNTIME_FOUNDATION_MARKER, "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": source_fingerprint}, "minimum_freetoken_commit": "0" * 40, "tvm_ffi_patch_sha256": "0" * 64, @@ -257,6 +259,7 @@ def test_resident_builder_streams_one_record_without_full_layer(z_fixture_dir): "version": 1, "artifact_schema": FORMAT, "text_only": True, + "runtime_foundation": RUNTIME_FOUNDATION_MARKER, "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": source_fingerprint}, "minimum_freetoken_commit": "0" * 40, "tvm_ffi_patch_sha256": "0" * 64, @@ -317,6 +320,7 @@ def test_end_to_end_synthetic_modular_artifact_reopens_normal_paths(z_fixture_di "architectures": ["Qwen4ExpForConditionalGeneration"], "freetoken_text_only": TEXT_ONLY_MARKER, "freetoken_active_quant": "nvfp4_w4a16_v1", + "freetoken_runtime_foundation": "pr257_hardware_fit_v1", } (z_fixture_dir / "config.json").write_text(json.dumps(config), encoding="utf-8") (z_fixture_dir / "tokenizer_config.json").write_text("{}", encoding="utf-8") @@ -435,6 +439,7 @@ def fake_convert(_source, out_dir, **kwargs): json.dumps({ "freetoken_text_only": TEXT_ONLY_MARKER, "freetoken_active_quant": "nvfp4_w4a16_v1", + "freetoken_runtime_foundation": "pr257_hardware_fit_v1", }), encoding="utf-8", ) From 66c84d51b07aad40aa8df1a705c33e92da299704 Mon Sep 17 00:00:00 2001 From: Your Name Date: Fri, 28 Aug 2026 18:16:39 -0400 Subject: [PATCH 12/17] feat(qwen4): add restartable step9b executor --- docs/FREETOKEN_STAGE7F_STEP9B_EXECUTOR.md | 108 ++ docs/plans/FREETOKEN-QWEN4-001-STAGE7F.md | 22 + .../freetoken/checkpoint/step9b_executor.py | 1278 +++++++++++++++++ tests/checkpoint/test_step9b_executor.py | 98 ++ .../test_step9b_executor_contract.py | 761 ++++++++++ 5 files changed, 2267 insertions(+) create mode 100644 docs/FREETOKEN_STAGE7F_STEP9B_EXECUTOR.md create mode 100644 docs/plans/FREETOKEN-QWEN4-001-STAGE7F.md create mode 100644 python/freetoken/checkpoint/step9b_executor.py create mode 100644 tests/checkpoint/test_step9b_executor.py create mode 100644 tests/checkpoint/test_step9b_executor_contract.py diff --git a/docs/FREETOKEN_STAGE7F_STEP9B_EXECUTOR.md b/docs/FREETOKEN_STAGE7F_STEP9B_EXECUTOR.md new file mode 100644 index 000000000..80f750fa2 --- /dev/null +++ b/docs/FREETOKEN_STAGE7F_STEP9B_EXECUTOR.md @@ -0,0 +1,108 @@ +# FreeToken Qwen4 Stage 7F Step 9B executor + +Gate: `FREETOKEN-QWEN4-001 / STAGE7F-STEP9B-EXECUTOR-IMPLEMENTATION` + +This branch adds the missing execution controller around the accepted Stage 7E +component writers. It does not change Q3, FTEXPERT1, active FTW v1, or the +pinned PR257 runtime. Stage 7F itself is a zero-real-payload gate. + +## Safety model + +`Step9BExecutor` is dry-run by default. A network response body is unreachable +unless both execution mode and explicit body permission are enabled. The CLI +therefore requires `--execute --allow-network-body`; `--execute` alone fails +before acquisition. The frozen manifest is the only source-file authority. + +The downloader: + +- resolves the immutable Hugging Face commit and checks file size, ETag, LFS + OID, Xet identity, and bounded Safetensors-header identity before transfer; +- streams directly to `.partial` on Z:, with no Hugging Face cache; +- binds every partial to a durable identity sidecar and resumes only with an + exact `Range` plus `If-Range` contract; +- rejects an ignored range, malformed `Content-Range`, missing or changed body + ETag, oversized or undersized body, hash mismatch, and header mismatch; +- persists actual response-body bytes before accepting each chunk, including + bytes received in a failed or over-cap attempt; +- enforces no more than two active response bodies and provides cancellation; +- promotes a source atomically only after all identities validate, then writes + an atomic source receipt. + +The controller rejects source payload without matching durable transfer-budget +provenance. This prevents an orphan partial or final source from bypassing the +upstream byte cap after a restart. + +## Stage controller + +The controller exposes explicit boundaries for: + +1. B1: nine metadata/config/tokenizer files and a bound receipt. +2. B2: ten PLE source files, the accepted 128-segment Q3 writer, reader reopen, + exact 22,400,107,520-byte validation, precommit, and final receipt. +3. B3: 48 ordered expert transactions, four source files per layer, accepted + FTEXPERT1 writer, exact 1,419,776,000-byte extent, FileExpertSource reopen, + and an independent layer receipt. +4. B4: four BF16 source files, accepted active conversion, exact + 4,804,403,200-byte FTW v1 contract, and receipt. +5. B5: final modular-manifest validation and exact 95,353,758,720-byte known + component reconciliation. +6. C6: a process-isolated static reopen using only the pinned PR257 worktree. + +Existing valid targets are never trusted by name or receipt alone. The +controller rehashes and reopens them, then can recover a missing final receipt. +Accepted component writers own atomic target promotion; the controller adds a +validated precommit document followed by an atomic final receipt. Incomplete +component partials are never promoted by the controller. + +Every component receipt binds the builder commit, runtime commit where +relevant, source revision, source-inventory fingerprint, receipt-backed source +file hashes, target length/hash, format, and validation results. + +## Capacity and environment gates + +The disk gate is restart-aware. It computes: + +`remaining verified source + remaining target + common conversion allowance + 64 GiB reserve` + +rather than requiring the original peak-free threshold after every completed +file. With an empty source and target, the formula exactly reconciles to +309,257,827,893 bytes. Physical host availability is measured independently +and must remain at least 6,442,450,944 bytes; pagefile use is recorded but never +counted as model capacity. + +The executor verifies the isolated `_pinned_tensor` extension hash and, in real +execution preflight, runs a bounded HostBank pin/device-alias probe. Triton, +TVM-FFI, Torch, compiler, and temporary caches are forced to the supplied Z: +toolchain root. + +Source retirement is explicitly rejected by this controller. The real handoff +uses `source_retirement_authorized=false` and retains every verified source. + +## C6 isolation + +C6 launches a fresh child process with `PYTHONPATH` beginning at exactly the +pinned PR257 runtime source. It performs static-only checks of: + +- hardware-fit markers and active FTW verification; +- Q3PLEFileTable construction; +- all 48 expert metadata entries and the 12/36 placement policy; +- QD4 FileExpertSource construction for each file-tier layer; +- graph-disabled and prefill-overlap-disabled policy, including forced-graph + rejection; +- a tiny resident HostBank pin/device-pointer probe. + +C6 does not instantiate the complete model, resident expert banks, production +GPU cache, KV cache, a model layer, a forward pass, generation, or a server. + +## Validation + +The Stage 7F suite uses a local deterministic HTTP server for clean transfer, +interruption, exact resume, ETag drift, ignored and malformed ranges, +over/undersized bodies, hashes, Safetensors headers, cancellation, concurrency, +failure isolation, receipt recovery, and byte-cap behavior. A hard kill-switch +test proves dry run cannot issue a real Hugging Face body GET. The full frozen +manifest test plans 9 metadata files, 206 weight files, 48 expert boundaries, +135,195,303,851 weight bytes, and the complete B1-B5/C6 order. + +The production command is intentionally emitted only in the separately +regenerated Step 9 handoff. Stage 7F does not execute it. diff --git a/docs/plans/FREETOKEN-QWEN4-001-STAGE7F.md b/docs/plans/FREETOKEN-QWEN4-001-STAGE7F.md new file mode 100644 index 000000000..e023131c8 --- /dev/null +++ b/docs/plans/FREETOKEN-QWEN4-001-STAGE7F.md @@ -0,0 +1,22 @@ +# FREETOKEN-QWEN4-001 Stage 7F execution plan + +Task: `STAGE7F-STEP9B-EXECUTOR-IMPLEMENTATION` + +Definition of done: a dry-run-default, restartable Step 9B controller coordinates immutable source acquisition, B1-B5 conversion, receipts, stop gates, retained-source policy, and isolated C6 validation without changing frozen component formats or the pinned PR257 runtime. Synthetic transport/controller tests and the real 206-row manifest dry run pass with zero upstream model payload bytes, then one local commit and a non-executed Step 9 handoff are produced. + +Dependencies: accepted builder `b64a342ea8e5ccac39a7619747b4a7b3e37466f3`; runtime `0307a6114c57b0efc61bc17688f3288fe0bf1dc7`; source manifest revision `7b719225242aacd3dbd3f9407468c2ee9a9d2594`. + +Validation: focused downloader, receipt/recovery, staged-controller, C6 subprocess, and accepted component-writer tests; full real-manifest dry run; `python -m compileall -q python/freetoken`; `git diff --check`; commit parent and clean-worktree checks. + +| Step | Status | Work | +|---|---|---| +| F0 | DONE | Verify authorities, PR257 head, zero-payload state, and create isolated branch. | +| F1 | DONE | Map accepted writer/manifest/runtime seams and freeze executor contract. | +| F2 | DONE | Implement immutable transport, downloader, partial identity, byte budget, receipts, and logging. | +| F3 | DONE | Implement B1-B5 state controller, disk/RAM gates, retained-source policy, and CLI. | +| F4 | DONE | Implement isolated C6 static-validation subprocess controller and environment check. | +| F5 | DONE | Add synthetic HTTP, range/cap/cancellation, receipt crash-recovery, and controller tests. | +| F6 | DONE | Run synthetic end-to-end and real 206-row no-body dry run; capture required evidence. | +| F7 | DONE | Run accepted regressions, compileall, diff checks, and adversarial final review. | +| F8 | DONE | Commit the executor branch and regenerate the non-executed Step 9 handoff. | +| F9 | DONE | Verify clean worktrees, zero payload, evidence completeness, and final gate decision. | diff --git a/python/freetoken/checkpoint/step9b_executor.py b/python/freetoken/checkpoint/step9b_executor.py new file mode 100644 index 000000000..50b2420c3 --- /dev/null +++ b/python/freetoken/checkpoint/step9b_executor.py @@ -0,0 +1,1278 @@ +"""Contract-complete, restartable Step 9B acquisition controller. + +This module deliberately owns orchestration and transport only. The accepted +Q3, FTEXPERT1 and FTW writers remain the byte-level authorities. In particular, +the default mode is a manifest-only dry run: a body GET is impossible unless +both ``execute`` and ``allow_network_body`` are explicitly enabled. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import struct +import subprocess +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, BinaryIO, Iterable, Mapping, Protocol + + +PINNED_REPOSITORY = "RadixArk/Qwen3.8-Flash-Next-NVFP4" +PINNED_REVISION = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" +MAX_TRANSFER_BYTES = 135_252_480_565 +MIN_DISK_FREE_BYTES = 309_257_827_893 +MIN_DISK_RESERVE_BYTES = 68_719_476_736 +COMMON_PEAK_ALLOWANCE_BYTES = 9_989_288_586 +MIN_HOST_FREE_BYTES = 6_442_450_944 +KNOWN_TARGET_BYTES = 95_353_758_720 +Q3_BYTES = 22_400_107_520 +ACTIVE_BYTES = 4_804_403_200 +EXPERT_BYTES = 1_419_776_000 +MAX_DOWNLOADS = 2 +MAX_SAFETENSORS_HEADER_BYTES = 256 << 20 +ACCEPTED_SOURCE_INVENTORY = "8572d200e31b344faff0fda f0dc72aa4726c1f062443d4109531b62ca63f66eb".replace(" ", "") + + +class ExecutorError(RuntimeError): + """A stop-gate or validation failure; callers must preserve evidence.""" + + +class BodyTransferDisabled(ExecutorError): + """Raised before a network body request when explicit authorization is absent.""" + + +class ResumeRejected(ExecutorError): + """Raised when the server cannot prove an identity-safe range response.""" + + +def _sha256(path: Path, chunk: int = 8 << 20) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while data := handle.read(chunk): + digest.update(data) + return digest.hexdigest() + + +def _tree_sha256(root: Path) -> tuple[int, str]: + digest = hashlib.sha256() + total = 0 + for path in sorted(item for item in root.rglob("*") if item.is_file()): + relative = path.relative_to(root).as_posix().encode("utf-8") + file_digest = bytes.fromhex(_sha256(path)) + size = path.stat().st_size + digest.update(len(relative).to_bytes(4, "little")) + digest.update(relative) + digest.update(size.to_bytes(8, "little")) + digest.update(file_digest) + total += size + return total, digest.hexdigest() + + +def _z_path(path: str | os.PathLike[str], *, must_exist: bool = False) -> Path: + candidate = Path(path).expanduser() + resolved = candidate.resolve(strict=must_exist) + drive = (resolved.drive or os.path.splitdrive(str(resolved))[0]).upper() + # Tests may use a POSIX-mounted /z volume; production Windows requires Z:. + if drive != "Z:" and not str(resolved).lower().startswith("/z/"): + raise ExecutorError(f"{path} must resolve physically to Z:, got {resolved}") + return resolved + + +def _atomic_json(path: Path, value: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + partial = path.with_name(f".{path.name}.partial-{os.getpid()}-{threading.get_ident()}") + with partial.open("w", encoding="utf-8", newline="\n") as handle: + json.dump(value, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + try: + os.fsync(handle.fileno()) + except OSError: + pass + os.replace(partial, path) + + +def _publish_component_receipt(receipt_path: Path, value: Mapping[str, Any]) -> dict[str, Any]: + final = dict(value) + final["completion"] = "COMPONENT_COMPLETE" + precommit = receipt_path.with_suffix(receipt_path.suffix + ".precommit") + _atomic_json(precommit, {**final, "completion": "VALIDATED_PRECOMMIT"}) + _atomic_json(receipt_path, final) + return final + + +def _receipt_matches(path: Path, expected: Mapping[str, Any]) -> bool: + if not path.is_file(): + return False + try: + with path.open("r", encoding="utf-8") as handle: + value = json.load(handle) + except (OSError, ValueError, TypeError, json.JSONDecodeError): + return False + return value.get("completion") == "COMPONENT_COMPLETE" and all(value.get(key) == item for key, item in expected.items()) + + +@dataclass(frozen=True) +class SourceEntry: + """Normalized source row from ``source_weight_shards`` or metadata.""" + + filename: str + byte_length: int + source_class: str + acquisition_order: int + repository: str + revision: str + accepted_etag: str | None = None + lfs_oid_sha256: str | None = None + accepted_header_length: int | None = None + accepted_header_sha256: str | None = None + layer_id: int | None = None + tensor_payload_bytes: int | None = None + git_blob_id: str | None = None + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any], *, order: int | None = None, metadata: bool = False) -> "SourceEntry": + filename = str(raw["filename"]) + if not filename or Path(filename).name != filename or filename in {".", ".."}: + raise ExecutorError(f"unsafe manifest filename: {filename!r}") + repository = str(raw.get("repository", PINNED_REPOSITORY)) + revision = str(raw.get("revision", PINNED_REVISION)) + if repository != PINNED_REPOSITORY or revision != PINNED_REVISION: + raise ExecutorError(f"source identity mismatch for {filename}") + return cls( + filename=filename, + byte_length=int(raw["byte_length"]), + source_class=("METADATA" if metadata else str(raw["source_class"])), + acquisition_order=int(raw.get("acquisition_order", order or 0)), + repository=repository, + revision=revision, + accepted_etag=(raw.get("accepted_etag") or (raw.get("git_blob_id") if metadata else None)), + lfs_oid_sha256=raw.get("lfs_oid_sha256"), + accepted_header_length=raw.get("accepted_header_length"), + accepted_header_sha256=raw.get("accepted_header_sha256"), + layer_id=(None if raw.get("layer_id") is None else int(raw["layer_id"])), + tensor_payload_bytes=(None if raw.get("tensor_payload_bytes") is None else int(raw["tensor_payload_bytes"])), + git_blob_id=(None if raw.get("git_blob_id") is None else str(raw["git_blob_id"])), + ) + + +@dataclass(frozen=True) +class AcquisitionManifest: + repository: str + revision: str + entries: tuple[SourceEntry, ...] + metadata: tuple[SourceEntry, ...] + source_inventory_fingerprint: str + expected_weight_bytes: int + transfer_cap: int = MAX_TRANSFER_BYTES + + @classmethod + def load(cls, path: str | os.PathLike[str]) -> "AcquisitionManifest": + with Path(path).open("r", encoding="utf-8") as handle: + raw = json.load(handle) + if raw.get("schema") != "freetoken-step9-acquisition-v1": + raise ExecutorError("unsupported acquisition manifest schema") + repo, revision = str(raw.get("repository")), str(raw.get("revision")) + if repo != PINNED_REPOSITORY or revision != PINNED_REVISION: + raise ExecutorError("manifest source pin does not match the frozen revision") + rows = tuple(SourceEntry.from_mapping(row) for row in raw.get("source_weight_shards", ())) + metadata = tuple(SourceEntry.from_mapping(row, order=i + 1, metadata=True) for i, row in enumerate(raw.get("required_small_metadata", ()))) + if len(rows) != 206 or len(metadata) != 9: + raise ExecutorError(f"manifest requires 206 weights and 9 metadata files, got {len(rows)} / {len(metadata)}") + orders = [row.acquisition_order for row in rows] + if orders != list(range(1, len(rows) + 1)): + raise ExecutorError("weight acquisition order is not contiguous") + expected = int(raw["reconciliation"]["expected_source_file_bytes"]) + if sum(row.byte_length for row in rows) != expected: + raise ExecutorError("source weight byte reconciliation failed") + class_contract = { + "BF16": (4, 16_007_756_462), + "PLE": (10, 51_200_267_901), + "EXPERT": (192, 67_987_279_488), + } + for source_class, (count, byte_length) in class_contract.items(): + selected = tuple(row for row in rows if row.source_class.upper() == source_class) + if len(selected) != count or sum(row.byte_length for row in selected) != byte_length: + raise ExecutorError(f"{source_class} source inventory mismatch") + for layer in range(48): + selected = tuple(row for row in rows if row.source_class.upper() == "EXPERT" and row.layer_id == layer) + if len(selected) != 4: + raise ExecutorError(f"expert layer {layer} must have exactly four source files") + if sum(row.byte_length for row in metadata) != 57_176_714: + raise ExecutorError("metadata inventory byte reconciliation failed") + inventory = str(raw.get("source_inventory_sha256") or ACCEPTED_SOURCE_INVENTORY).lower() + if inventory != ACCEPTED_SOURCE_INVENTORY: + raise ExecutorError("source tensor inventory fingerprint mismatch") + return cls(repo, revision, rows, metadata, inventory, expected, MAX_TRANSFER_BYTES) + + @property + def all_entries(self) -> tuple[SourceEntry, ...]: + return self.metadata + self.entries + + def rows_for_stage(self, stage: str) -> tuple[SourceEntry, ...]: + key = stage.upper() + if key == "B1": + return self.metadata + if key == "B2": + return tuple(row for row in self.entries if row.source_class.upper() == "PLE") + if key == "B4": + return tuple(row for row in self.entries if row.source_class.upper() == "BF16") + return self.entries + + +@dataclass +class TransferBudget: + cap: int + transferred: int = 0 + state_path: Path | None = None + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + _inflight_reserved: int = field(default=0, repr=False) + + def __post_init__(self) -> None: + if self.state_path and self.state_path.is_file(): + try: + with self.state_path.open("r", encoding="utf-8") as handle: + prior = json.load(handle) + if int(prior.get("cap", self.cap)) != self.cap: + raise ExecutorError("transfer budget cap changed across restart") + self.transferred = int(prior.get("transferred", 0)) + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ExecutorError("invalid persisted transfer budget") from exc + + def _persist(self) -> None: + if self.state_path: + _atomic_json(self.state_path, {"cap": self.cap, "transferred": self.transferred}) + + def reserve(self, amount: int) -> None: + amount = int(amount) + if amount < 0: + raise ValueError("transfer amount cannot be negative") + with self._lock: + self.transferred += amount + self._persist() + if self.transferred > self.cap: + raise ExecutorError(f"transfer cap exceeded: {self.transferred} > {self.cap}") + + def admit(self, maximum_response_bytes: int) -> None: + """Atomically reserve room before a body request is opened.""" + amount = int(maximum_response_bytes) + if amount < 0: + raise ValueError("admission amount cannot be negative") + with self._lock: + if self.transferred + self._inflight_reserved + amount > self.cap: + raise ExecutorError("transfer budget cannot admit response body") + self._inflight_reserved += amount + + def record_received(self, amount: int) -> None: + """Persist actual bytes received and consume their inflight reservation.""" + amount = int(amount) + if amount < 0: + raise ValueError("received amount cannot be negative") + with self._lock: + admitted = min(amount, self._inflight_reserved) + self._inflight_reserved -= admitted + self.transferred += amount + self._persist() + if amount > admitted: + raise ExecutorError("received bytes exceed admitted response budget") + if self.transferred > self.cap: + raise ExecutorError(f"transfer cap exceeded: {self.transferred} > {self.cap}") + + def release_admission(self, unused_bytes: int) -> None: + amount = int(unused_bytes) + if amount < 0: + raise ValueError("unused admission cannot be negative") + with self._lock: + if amount > self._inflight_reserved: + raise ExecutorError("released admission exceeds inflight reservation") + self._inflight_reserved -= amount + + @property + def remaining(self) -> int: + with self._lock: + return self.cap - self.transferred + + +class TransportResponse(Protocol): + status: int + headers: Mapping[str, str] + + def iter_bytes(self, chunk_bytes: int = 8 << 20) -> Iterable[bytes]: ... + def close(self) -> None: ... + + +class Transport(Protocol): + def head(self, url: str, *, headers: Mapping[str, str] | None = None) -> TransportResponse: ... + def get(self, url: str, *, headers: Mapping[str, str] | None = None, allow_body: bool = False) -> TransportResponse: ... + + +class _UrllibResponse: + def __init__(self, response: Any): + self._response = response + self.status = int(response.status) + self.headers = {str(k): str(v) for k, v in response.headers.items()} + + def iter_bytes(self, chunk_bytes: int = 8 << 20) -> Iterable[bytes]: + while data := self._response.read(chunk_bytes): + yield data + + def close(self) -> None: + self._response.close() + + +class _SafeRedirectHandler(urllib.request.HTTPRedirectHandler): + _allowed_hosts = ("huggingface.co", "hf.co") + + def redirect_request(self, req, fp, code, msg, headers, newurl): + host = (urllib.parse.urlparse(newurl).hostname or "").lower() + if not any(host == suffix or host.endswith("." + suffix) for suffix in self._allowed_hosts): + raise ExecutorError(f"refusing model redirect to unapproved host: {host or ''}") + return super().redirect_request(req, fp, code, msg, headers, newurl) + + +class UrllibTransport: + """Small dependency-free HTTP transport; body permission is explicit.""" + + def __init__(self) -> None: + self._opener = urllib.request.build_opener(_SafeRedirectHandler()) + + def _request(self, method: str, url: str, headers: Mapping[str, str] | None) -> _UrllibResponse: + request = urllib.request.Request(url, method=method, headers=dict(headers or {})) + try: + return _UrllibResponse(self._opener.open(request, timeout=60)) + except urllib.error.HTTPError as exc: + return _UrllibResponse(exc) + + def head(self, url: str, *, headers: Mapping[str, str] | None = None) -> _UrllibResponse: + return self._request("HEAD", url, headers) + + def get(self, url: str, *, headers: Mapping[str, str] | None = None, allow_body: bool = False) -> _UrllibResponse: + if not allow_body: + raise BodyTransferDisabled("GET body blocked; require --execute and --allow-network-body") + return self._request("GET", url, headers) + + +class JsonlLogger: + def __init__(self, path: str | os.PathLike[str]): + self.path = _z_path(path) + self._lock = threading.Lock() + + def event(self, event: str, **fields: Any) -> None: + # Never persist credentials, signed URLs, or cookies. + safe = {k: ("" if any(x in k.lower() for x in ("token", "cookie", "authorization", "url")) else v) for k, v in fields.items()} + safe.update(event=event, timestamp=time.time()) + self.path.parent.mkdir(parents=True, exist_ok=True) + with self._lock, self.path.open("a", encoding="utf-8", newline="\n") as handle: + handle.write(json.dumps(safe, sort_keys=True) + "\n") + + +class Downloader: + def __init__(self, source_root: str | os.PathLike[str], manifest: AcquisitionManifest, *, receipt_root: str | os.PathLike[str] | None = None, transport: Transport | None = None, execute: bool = False, allow_network_body: bool = False, max_concurrent: int = MAX_DOWNLOADS, budget: TransferBudget | None = None, logger: JsonlLogger | None = None): + self.root = _z_path(source_root) + self.receipt_root = _z_path(receipt_root or (self.root / ".step9b-receipts")) + self.manifest = manifest + self.transport = transport or UrllibTransport() + self.execute = bool(execute) + self.allow_network_body = bool(allow_network_body) + if max_concurrent < 1 or max_concurrent > MAX_DOWNLOADS: + raise ValueError("max_concurrent must be between 1 and 2") + self.max_concurrent = max_concurrent + self.budget = budget or TransferBudget(manifest.transfer_cap) + self.logger = logger + self.active = 0 + self.max_active = 0 + self._active_lock = threading.Lock() + self._semaphore = threading.BoundedSemaphore(max_concurrent) + self._cancel = threading.Event() + self._file_locks: dict[str, threading.Lock] = {} + self._file_locks_guard = threading.Lock() + + def _url(self, row: SourceEntry) -> str: + return f"https://huggingface.co/{row.repository}/resolve/{row.revision}/{row.filename}" + + def cancel(self) -> None: + self._cancel.set() + + def _identity(self, row: SourceEntry, length: int, remote: Mapping[str, Any] | None = None) -> dict[str, Any]: + remote = dict(remote or {}) + return {"repository": row.repository, "revision": row.revision, "resolved_commit": remote.get("commit") or remote.get("resolved_commit") or row.revision, "filename": row.filename, "expected_length": row.byte_length, "expected_etag": row.accepted_etag, "expected_lfs_oid": row.lfs_oid_sha256, "expected_header_length": row.accepted_header_length, "expected_header_sha256": row.accepted_header_sha256, "observed_etag": remote.get("etag"), "observed_xet_file_hash": remote.get("xet_file_hash"), "partial_length": length, "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, "acquisition_order": row.acquisition_order} + + def validate_metadata(self, row: SourceEntry, response: TransportResponse) -> dict[str, Any]: + headers = {k.lower(): v for k, v in response.headers.items()} + observed_length = int(headers.get("content-length", "-1")) + observed_etag = headers.get("etag") + if observed_length != row.byte_length: + raise ExecutorError(f"{row.filename}: content length mismatch") + # Hugging Face uses two identities for Xet/LFS files. The manifest's + # accepted_etag is the Xet file hash (often quoted), while lfs_oid is + # surfaced as HfFileMetadata.etag and may be returned as HTTP ETag by + # the CDN. A transport may expose either, so accept only either exact + # frozen identity and never a merely non-empty header. + accepted = {str(value).strip('"') for value in (row.accepted_etag, row.lfs_oid_sha256) if value} + if accepted and not observed_etag: + raise ExecutorError(f"{row.filename}: metadata response omitted required ETag") + if observed_etag and accepted and observed_etag.strip('"') not in accepted: + raise ExecutorError(f"{row.filename}: ETag/Xet identity mismatch") + return {"resolved_commit": row.revision, "length": observed_length, "etag": observed_etag} + + def resolve_hf_metadata(self, row: SourceEntry) -> dict[str, Any]: + """Resolve immutable HF metadata without requesting a response body. + + ``huggingface_hub`` is imported lazily so synthetic transports and the + dry-run controller remain usable in minimal environments. For Xet + entries, ``xet_file_data.file_hash`` is checked against the manifest's + accepted_etag and ``etag`` against the LFS OID. Metadata files use the + Git blob etag instead. + """ + try: + from huggingface_hub import get_hf_file_metadata + except ImportError as exc: # pragma: no cover - environment-specific + raise ExecutorError("huggingface_hub is required for HF metadata identity") from exc + url = self._url(row) + try: + meta = get_hf_file_metadata(url=url, token=None) + except TypeError: + # Older hub releases use ``filename``/``repo_id`` but retain the + # immutable resolve URL contract; only this metadata-only fallback + # is allowed. + meta = get_hf_file_metadata(url) + commit = str(getattr(meta, "commit_hash", "") or "") + size = int(getattr(meta, "size", -1) or -1) + etag = str(getattr(meta, "etag", "") or "").strip('"') + if commit != row.revision: + raise ExecutorError(f"{row.filename}: resolved commit mismatch") + if size != row.byte_length: + raise ExecutorError(f"{row.filename}: metadata length mismatch") + xet = getattr(meta, "xet_file_data", None) + xet_hash = str(getattr(xet, "file_hash", "") or "").strip('"') + if row.lfs_oid_sha256: + if etag != row.lfs_oid_sha256.lower(): + raise ExecutorError(f"{row.filename}: metadata LFS OID mismatch") + if row.accepted_etag and xet_hash != row.accepted_etag.strip('"').lower(): + raise ExecutorError(f"{row.filename}: metadata Xet hash mismatch") + elif row.accepted_etag: + # Git-backed metadata has no Xet data; accepted_etag is the blob id. + if etag != row.accepted_etag.strip('"').lower(): + raise ExecutorError(f"{row.filename}: metadata Git identity mismatch") + return {"url": url, "commit": commit, "size": size, "etag": etag, "xet_file_hash": xet_hash or None, "body_bytes": 0} + + def _validate_partial_identity(self, row: SourceEntry, meta: Path, length: int, remote: Mapping[str, Any]) -> None: + if not meta.is_file(): + raise ResumeRejected(f"partial identity sidecar missing for {row.filename}") + with meta.open("r", encoding="utf-8") as handle: + identity = json.load(handle) + expected = self._identity(row, length, remote) + for key in ("repository", "revision", "resolved_commit", "filename", "expected_length", "expected_etag", "expected_lfs_oid", "observed_etag", "observed_xet_file_hash", "source_inventory_fingerprint", "acquisition_order"): + if identity.get(key) != expected.get(key): + raise ResumeRejected(f"partial identity mismatch: {key}") + if int(identity.get("partial_length", -1)) != length: + raise ResumeRejected("partial length identity mismatch") + + def _validate_safetensors_header(self, row: SourceEntry, path: Path) -> None: + if row.accepted_header_length is None or row.accepted_header_sha256 is None: + return + with path.open("rb") as handle: + prefix = handle.read(8) + if len(prefix) != 8: + raise ExecutorError(f"{row.filename}: missing Safetensors framing") + (header_length,) = struct.unpack(" MAX_SAFETENSORS_HEADER_BYTES or header_length > row.byte_length - 8: + raise ExecutorError(f"{row.filename}: unsafe Safetensors header length") + header = handle.read(header_length) + if hashlib.sha256(header).hexdigest() != row.accepted_header_sha256: + raise ExecutorError(f"{row.filename}: Safetensors header hash mismatch") + + def validate_existing(self, row: SourceEntry, final: Path) -> dict[str, Any]: + if not final.is_file() or final.stat().st_size != row.byte_length: + raise ExecutorError(f"{row.filename}: final source is missing or wrong length") + digest = _sha256(final) + if row.lfs_oid_sha256 and digest.lower() != row.lfs_oid_sha256.lower(): + raise ExecutorError(f"{row.filename}: source SHA/LFS OID mismatch") + if row.git_blob_id and not row.lfs_oid_sha256: + git = hashlib.sha1(f"blob {row.byte_length}\0".encode()) + with final.open("rb") as source: + while data := source.read(8 << 20): + git.update(data) + git_digest = git.hexdigest() + if git_digest.lower() != row.git_blob_id.lower(): + raise ExecutorError(f"{row.filename}: Git blob identity mismatch") + self._validate_safetensors_header(row, final) + return {"state": "VALID", "bytes": row.byte_length, "sha256": digest} + + def _validate_promote_partial( + self, + row: SourceEntry, + partial: Path, + identity: Path, + final: Path, + receipt: Path, + remote_meta: Mapping[str, Any], + *, + resumed_from: int, + body_bytes_this_run: int, + recovered_complete_partial: bool = False, + ) -> dict[str, Any]: + if partial.stat().st_size != row.byte_length: + raise ExecutorError("partial is not complete enough to promote") + self._validate_safetensors_header(row, partial) + digest = _sha256(partial) + if row.lfs_oid_sha256 and digest.lower() != row.lfs_oid_sha256.lower(): + raise ExecutorError("source SHA/LFS OID mismatch") + if row.git_blob_id and not row.lfs_oid_sha256: + git = hashlib.sha1() + git.update(f"blob {row.byte_length}\0".encode()) + with partial.open("rb") as source: + while data := source.read(8 << 20): + git.update(data) + if git.hexdigest().lower() != row.git_blob_id.lower(): + raise ExecutorError("Git blob identity mismatch") + os.replace(partial, final) + identity.unlink(missing_ok=True) + result = { + "state": "SOURCE_COMPLETE", + "final_path": str(final), + "expected_bytes": row.byte_length, + "bytes": row.byte_length, + "sha256": digest, + "resumed_from": resumed_from, + "body_bytes": row.byte_length, + "body_bytes_this_run": body_bytes_this_run, + "resolved_commit": row.revision, + "expected_etag": row.accepted_etag, + "observed_etag": remote_meta.get("etag"), + "observed_xet_file_hash": remote_meta.get("xet_file_hash"), + "expected_lfs_oid": row.lfs_oid_sha256, + "observed_lfs_oid": row.lfs_oid_sha256, + "observed_header_length": row.accepted_header_length, + "observed_header_sha256": row.accepted_header_sha256, + "recovered_complete_partial": recovered_complete_partial, + } + _atomic_json(receipt, {"entry": asdict(row), **result, "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, "completion": "SOURCE_COMPLETE"}) + return result + + def acquire(self, row: SourceEntry) -> dict[str, Any]: + with self._file_locks_guard: + lock = self._file_locks.setdefault(row.filename, threading.Lock()) + with lock: + return self._acquire_locked(row) + + def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: + self.root.mkdir(parents=True, exist_ok=True) + final = self.root / row.filename + partial = self.root / f"{row.filename}.partial" + identity = partial.with_name(partial.name + ".meta.json") + receipt = self.receipt_root / f"{row.acquisition_order:03d}-{row.filename}.receipt.json" + if final.exists(): + result = self.validate_existing(row, final) + receipt_binding = { + "entry": asdict(row), + "final_path": str(final), + "expected_bytes": row.byte_length, + "bytes": row.byte_length, + "resolved_commit": row.revision, + "expected_etag": row.accepted_etag, + "expected_lfs_oid": row.lfs_oid_sha256, + "observed_lfs_oid": row.lfs_oid_sha256, + "observed_header_length": row.accepted_header_length, + "observed_header_sha256": row.accepted_header_sha256, + "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, + "sha256": result["sha256"], + } + if receipt.exists(): + try: + with receipt.open("r", encoding="utf-8") as handle: + prior = json.load(handle) + if prior.get("completion") != "SOURCE_COMPLETE" or any(prior.get(key) != value for key, value in receipt_binding.items()): + raise ExecutorError(f"{row.filename}: source receipt binding mismatch") + accepted_observed = {str(value).strip('"').lower() for value in (row.accepted_etag, row.lfs_oid_sha256) if value} + if str(prior.get("observed_etag", "")).strip('"').lower() not in accepted_observed: + raise ExecutorError(f"{row.filename}: source receipt ETag binding mismatch") + if row.lfs_oid_sha256 and str(prior.get("observed_xet_file_hash", "")).strip('"').lower() != str(row.accepted_etag).strip('"').lower(): + raise ExecutorError(f"{row.filename}: source receipt Xet binding mismatch") + resumed = prior.get("resumed_from") + if resumed is not None and not 0 <= int(resumed) <= row.byte_length: + raise ExecutorError(f"{row.filename}: invalid source receipt resume offset") + if int(prior.get("body_bytes", -1)) < 0: + raise ExecutorError(f"{row.filename}: invalid source receipt body byte count") + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ExecutorError(f"{row.filename}: invalid source receipt") from exc + else: + _atomic_json(receipt, {**receipt_binding, "observed_etag": row.lfs_oid_sha256 or row.accepted_etag, "observed_xet_file_hash": row.accepted_etag, "resumed_from": None, "body_bytes": 0, "completion": "SOURCE_COMPLETE", "recovered_after_promotion": True}) + return {"filename": row.filename, **result, "state": "SKIP_VALID_FINAL"} + if self.execute and not self.allow_network_body: + raise BodyTransferDisabled("execution mode requires explicit network-body authorization") + if not self.execute: + if partial.exists() or identity.exists(): + raise BodyTransferDisabled("dry run cannot inspect/resume body partials") + return {"filename": row.filename, "state": "PLANNED", "bytes": row.byte_length} + if self._cancel.is_set(): + raise ExecutorError("source acquisition cancelled") + if isinstance(self.transport, UrllibTransport): + remote_meta = self.resolve_hf_metadata(row) + else: + head = self.transport.head(self._url(row), headers={}) + try: + remote_meta = self.validate_metadata(row, head) + finally: + head.close() + current = partial.stat().st_size if partial.exists() else 0 + if partial.exists() != identity.exists(): + raise ResumeRejected(f"partial and identity sidecar must exist together for {row.filename}") + if current or partial.exists(): + self._validate_partial_identity(row, identity, current, remote_meta) + if current == row.byte_length: + result = self._validate_promote_partial( + row, + partial, + identity, + final, + receipt, + remote_meta, + resumed_from=current, + body_bytes_this_run=0, + recovered_complete_partial=True, + ) + return {"filename": row.filename, **result} + headers: dict[str, str] = {} + if current: + validated_etag = str(remote_meta.get("etag") or "") + if not validated_etag: + raise ResumeRejected("validated remote ETag unavailable for If-Range") + headers = {"Range": f"bytes={current}-", "If-Range": validated_etag} + self._semaphore.acquire() + with self._active_lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + response = None + received = 0 + admitted = 0 + try: + # Reserve the advertised maximum before opening the body; this + # prevents a response from starting after the transfer cap is spent. + remaining = row.byte_length - current + if remaining < 0 or remaining > self.budget.remaining: + raise ExecutorError(f"transfer budget cannot admit {row.filename}") + self.budget.admit(remaining) + admitted = remaining + if self._cancel.is_set(): + raise ExecutorError("source acquisition cancelled") + response = self.transport.get(self._url(row), headers=headers, allow_body=True) + rh = {k.lower(): v for k, v in response.headers.items()} + if current: + if response.status != 206: + raise ResumeRejected("server ignored Range; refusing to append a 200 response") + content_range = rh.get("content-range", "") + expected_prefix = f"bytes {current}-" + try: + span, total = content_range.removeprefix("bytes ").split("/") + start, end = (int(value) for value in span.split("-")) + except (ValueError, TypeError): + raise ResumeRejected("malformed Content-Range") + content_length = int(rh.get("content-length", "-1")) + if (not content_range.startswith(expected_prefix) or total != str(row.byte_length) or start != current or end != row.byte_length - 1 or content_length != remaining): + raise ResumeRejected("malformed Content-Range") + elif response.status != 200: + raise ExecutorError(f"unexpected initial body status {response.status}") + elif int(rh.get("content-length", "-1")) != row.byte_length: + raise ExecutorError("initial Content-Length mismatch") + etag = rh.get("etag") + accepted = {str(value).strip('"').lower() for value in (row.accepted_etag, row.lfs_oid_sha256) if value} + if accepted and not etag: + raise ExecutorError("body response omitted required ETag identity") + if etag and accepted and etag.strip('"').lower() not in accepted: + raise ExecutorError("body ETag/Xet identity changed") + mode = "ab" if current else "wb" + with partial.open(mode) as target: + if not current: + _atomic_json(identity, self._identity(row, 0, remote_meta)) + for chunk in response.iter_bytes(): + if not chunk: + continue + self.budget.record_received(len(chunk)) + admitted -= len(chunk) + if self._cancel.is_set(): + raise ExecutorError("source acquisition cancelled") + if received + len(chunk) > remaining or current + received + len(chunk) > row.byte_length: + raise ExecutorError("response body exceeds expected length") + target.write(chunk) + received += len(chunk) + _atomic_json(identity, self._identity(row, current + received, remote_meta)) + target.flush() + os.fsync(target.fileno()) + if current + received != row.byte_length: + raise ExecutorError("undersized final body") + result = self._validate_promote_partial( + row, + partial, + identity, + final, + receipt, + remote_meta, + resumed_from=current, + body_bytes_this_run=received, + ) + return {"filename": row.filename, **result} + finally: + if response is not None: + response.close() + if admitted: + self.budget.release_admission(admitted) + with self._active_lock: + self.active -= 1 + self._semaphore.release() + + +class Step9BExecutor: + """Explicit B1-B5 state machine; conversion methods delegate to accepted writers.""" + + def __init__(self, manifest_path: str | os.PathLike[str], source_root: str | os.PathLike[str], target_root: str | os.PathLike[str], scratch_root: str | os.PathLike[str], logs_root: str | os.PathLike[str], *, builder_commit: str, runtime_worktree: str | os.PathLike[str], runtime_commit: str, source_inventory_fingerprint: str | None = None, source_revision: str = PINNED_REVISION, transfer_cap: int = MAX_TRANSFER_BYTES, toolchain_root: str | os.PathLike[str] | None = None, execute: bool = False, allow_network_body: bool = False, source_retirement_authorized: bool = False, min_disk_free: int = MIN_DISK_RESERVE_BYTES, min_host_free: int = MIN_HOST_FREE_BYTES, max_concurrent_downloads: int = MAX_DOWNLOADS, transport: Transport | None = None): + self.manifest_path = Path(manifest_path) + self.manifest = AcquisitionManifest.load(manifest_path) + if str(source_revision) != self.manifest.revision: + raise ExecutorError("source revision does not match acquisition manifest") + if int(transfer_cap) != self.manifest.transfer_cap: + raise ExecutorError("transfer cap does not match frozen authorization") + self.source_root = _z_path(source_root) + self.target_root = _z_path(target_root) + self.scratch_root = _z_path(scratch_root) + self.logs_root = _z_path(logs_root) + self.toolchain_root = _z_path(toolchain_root or (self.logs_root.resolve().parents[2] / "artifacts" / "toolchain" / "step9b-pr257")) + self.builder_commit = str(builder_commit).lower() + self.source_inventory_fingerprint = str(source_inventory_fingerprint or self.manifest.source_inventory_fingerprint).lower() + if len(self.source_inventory_fingerprint) != 64 or any(char not in "0123456789abcdef" for char in self.source_inventory_fingerprint): + raise ExecutorError("source inventory fingerprint must be a SHA-256 digest") + if self.source_inventory_fingerprint != self.manifest.source_inventory_fingerprint: + raise ExecutorError("source inventory fingerprint does not match frozen manifest") + self.runtime_worktree = _z_path(runtime_worktree, must_exist=True) + self.runtime_commit = str(runtime_commit).lower() + self.execute = bool(execute) + self.allow_network_body = bool(allow_network_body) + self.source_retirement_authorized = bool(source_retirement_authorized) + self.min_disk_free = int(min_disk_free) + self.min_host_free = int(min_host_free) + self.logger = JsonlLogger(self.logs_root / "events.jsonl") + self.budget = TransferBudget(self.manifest.transfer_cap, state_path=self.scratch_root / "transfer-budget.json") + self.downloader = Downloader(self.source_root, self.manifest, receipt_root=self.scratch_root / "receipts" / "sources", transport=transport, execute=execute, allow_network_body=allow_network_body, max_concurrent=max_concurrent_downloads, budget=self.budget, logger=self.logger) + self.state_path = self.scratch_root / "executor-state.json" + self.state: dict[str, Any] = {"mode": "EXECUTE" if execute else "DRY_RUN", "source_retirement_authorized": self.source_retirement_authorized, "stages": {}} + + def preflight(self) -> dict[str, Any]: + if self.execute and not self.allow_network_body: + raise BodyTransferDisabled("--execute requires --allow-network-body") + if self.source_retirement_authorized: + # Real Step 9 handoff intentionally sets this false. A caller may + # test true only with an explicit, separately reviewed controller. + raise ExecutorError("source retirement is disabled for this executor invocation") + self._validate_source_workspace() + source_payload_exists = self.source_root.exists() and any( + path.is_file() and (path.name in {row.filename for row in self.manifest.all_entries} or path.name.endswith(".partial")) + for path in self.source_root.iterdir() + ) + if source_payload_exists and not (self.scratch_root / "transfer-budget.json").is_file(): + raise ExecutorError("source payload exists without durable transfer-budget provenance") + usage = shutil.disk_usage(self.source_root.anchor or self.source_root) + required_free = self._projected_required_free() + if usage.free < required_free: + raise ExecutorError(f"Z: free space below projected retained-source gate: {usage.free} < {required_free}") + # Builder/runtime authority checks happen before any irreversible work. + builder_repo = Path(__file__).resolve().parents[3] + try: + builder_head = subprocess.check_output(["git", "-C", str(builder_repo), "rev-parse", "HEAD"], text=True, stderr=subprocess.DEVNULL).strip().lower() + dirty = bool(subprocess.check_output(["git", "-C", str(builder_repo), "status", "--porcelain"], text=True, stderr=subprocess.DEVNULL).strip()) + except (OSError, subprocess.CalledProcessError) as exc: + raise ExecutorError("cannot inspect builder Git authority") from exc + if builder_head != self.builder_commit or dirty: + raise ExecutorError(f"builder authority mismatch/dirty: {builder_head} dirty={dirty}") + try: + runtime_head = subprocess.check_output(["git", "-C", str(self.runtime_worktree), "rev-parse", "HEAD"], text=True, stderr=subprocess.DEVNULL).strip().lower() + runtime_dirty = bool(subprocess.check_output(["git", "-C", str(self.runtime_worktree), "status", "--porcelain"], text=True, stderr=subprocess.DEVNULL).strip()) + except (OSError, subprocess.CalledProcessError) as exc: + raise ExecutorError("cannot inspect runtime Git authority") from exc + if runtime_head != self.runtime_commit or runtime_dirty: + raise ExecutorError(f"runtime authority mismatch/dirty: {runtime_head} dirty={runtime_dirty}") + # Keep every JIT/toolchain cache on Z for this execution. + cache_root = self.toolchain_root + cache_root.mkdir(parents=True, exist_ok=True) + for name in ("TRITON_CACHE_DIR", "TRITON_HOME", "TRITON_DUMP_DIR", "TRITON_OVERRIDE_DIR", "FREETOKEN_KERNEL_CACHE_DIR", "TVM_FFI_CACHE_DIR", "XDG_CACHE_HOME", "TORCH_EXTENSIONS_DIR", "TORCHINDUCTOR_CACHE_DIR", "TEMP", "TMP", "TMPDIR"): + os.environ[name] = str(cache_root) + extension = cache_root / "build-lib" / "pinned-only" / "freetoken" / "kernel" / "_pinned_tensor.cp312-win_amd64.pyd" + extension_sha = "58c2ea0b7f74e457a48707eeb88012dfd28c943786e43cfa2896970e01f26b14" + extension_ok = extension.is_file() and _sha256(extension).lower() == extension_sha + if self.execute and not extension_ok: + raise ExecutorError("resident HostBank pinning extension missing or hash mismatch") + pin_probe = {"executed": False, "ok": extension_ok} + if self.execute: + probe_env = os.environ.copy() + probe_env["PYTHONPATH"] = str(self.runtime_worktree / "python") + probe_env["FREETOKEN_PINNED_EXTENSION_DIR"] = str(extension.parent) + probe = subprocess.run( + [sys.executable, "-c", "import os,torch,freetoken.kernel; freetoken.kernel.__path__.append(os.environ['FREETOKEN_PINNED_EXTENSION_DIR']); from freetoken.moe.host_banks import HostBank,HostResidency; from freetoken.kernel.pinned import device_ptr; b=HostBank((64,32),torch.bfloat16,backing='cuda'); assert b.residency is HostResidency.PINNED and b.addr%4096==0 and device_ptr(b.tensor)>0; del b; torch.cuda.synchronize()"], + cwd=self.runtime_worktree, + env=probe_env, + text=True, + capture_output=True, + timeout=120, + ) + if probe.returncode != 0: + raise ExecutorError(f"resident HostBank pinning probe failed: {probe.stderr[-1000:]}") + pin_probe = {"executed": True, "ok": True} + host = self._host_gate() + host_available, pagefile_used = host["available_bytes"], host["pagefile_used_bytes"] + if pagefile_used > 0 and self.execute: + # Swap use is evidence, not capacity; continue only if reserve is + # still measured physically, but record the condition for audit. + self.logger.event("pagefile_observed", bytes=pagefile_used) + self.state["preflight"] = {"free_bytes": usage.free, "payload_bytes": 0, "manifest": str(self.manifest_path)} + self.state["preflight"].update({"builder_head": builder_head, "runtime_head": runtime_head, "host_available": host_available, "pagefile_used": pagefile_used, "pinned_extension": str(extension), "pinned_extension_ok": extension_ok, "pin_probe": pin_probe, "cache_root": str(cache_root), "projected_required_free": required_free}) + _atomic_json(self.state_path, self.state) + self.logger.event("preflight", free_bytes=usage.free, transfer_cap=self.manifest.transfer_cap) + return self.state["preflight"] + + def _download_stage(self, stage: str, rows: Iterable[SourceEntry]) -> list[dict[str, Any]]: + results = [] + for row in rows: + self._disk_gate() + self.logger.event("file_acquisition_start", filename=row.filename, stage=stage, planned_bytes=row.byte_length) + result = self.downloader.acquire(row) + results.append(result) + self.logger.event("file_acquisition_end", filename=row.filename, stage=stage, state=result["state"], body_bytes=result.get("body_bytes", 0)) + self.state["stages"].setdefault(stage, {})["sources"] = results + _atomic_json(self.state_path, self.state) + return results + + def _remaining_weight_bytes(self) -> int: + remaining = 0 + for row in self.manifest.entries: + final = self.source_root / row.filename + if final.is_file() and final.stat().st_size == row.byte_length: + continue + partial = self.source_root / f"{row.filename}.partial" + present = partial.stat().st_size if partial.is_file() else 0 + remaining += max(0, row.byte_length - present) + return remaining + + def _validate_source_workspace(self) -> None: + if not self.source_root.exists(): + return + allowed = {row.filename for row in self.manifest.all_entries} + for path in self.source_root.iterdir(): + if path.is_dir() and path.name == ".step9b-receipts": + continue + if not path.is_file(): + raise ExecutorError(f"unexpected source workspace entry: {path.name}") + name = path.name + if name in allowed: + continue + if name.endswith(".partial") and name[:-8] in allowed: + continue + if name.endswith(".partial.meta.json") and name[:-18] in allowed: + continue + raise ExecutorError(f"unexpected source workspace file: {name}") + + def _remaining_target_bytes(self) -> int: + remaining = 0 + q3 = self.target_root / "ple-q3-000.bin" + if not q3.is_file() or q3.stat().st_size != Q3_BYTES: + remaining += Q3_BYTES + for layer in range(48): + sidecar = self.target_root / f"experts-L{layer:02d}.nvfp4" + if not sidecar.is_file() or sidecar.stat().st_size != EXPERT_BYTES: + remaining += EXPERT_BYTES + active_index = self.target_root / "qwen4-active-v1.ftw" + if not active_index.is_dir(): + remaining += ACTIVE_BYTES + return remaining + + def _projected_required_free(self) -> int: + return self._remaining_weight_bytes() + self._remaining_target_bytes() + COMMON_PEAK_ALLOWANCE_BYTES + self.min_disk_free + + def _disk_gate(self) -> int: + usage = shutil.disk_usage(self.source_root.anchor or self.source_root) + required = self._projected_required_free() + if usage.free < required: + raise ExecutorError(f"Z: projected reserve would be threatened: free={usage.free} required={required}") + self.logger.event("disk_gate", free_bytes=usage.free, projected_required_bytes=required, reserve_bytes=self.min_disk_free) + return usage.free + + def _host_gate(self) -> dict[str, int]: + try: + import psutil + except ImportError: + if os.name != "nt": + raise ExecutorError("physical host reserve probe unavailable without psutil") + command = ( + "$m=Get-CimInstance Win32_OperatingSystem;" + "$pf=(Get-CimInstance Win32_PageFileUsage|Measure-Object CurrentUsage -Sum).Sum;" + f"$p=Get-Process -Id {os.getpid()};" + "[pscustomobject]@{total_bytes=[int64]$m.TotalVisibleMemorySize*1024;" + "available_bytes=[int64]$m.FreePhysicalMemory*1024;" + "process_rss_bytes=[int64]$p.WorkingSet64;" + "pagefile_used_bytes=[int64]$pf*1MB}|ConvertTo-Json -Compress" + ) + try: + completed = subprocess.run( + ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", command], + text=True, + capture_output=True, + timeout=30, + check=True, + ) + result = {key: int(value) for key, value in json.loads(completed.stdout).items()} + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ExecutorError("Windows physical host reserve probe failed") from exc + available = result["available_bytes"] + else: + process = psutil.Process() + available = int(psutil.virtual_memory().available) + result = {"total_bytes": int(psutil.virtual_memory().total), "available_bytes": available, "process_rss_bytes": int(process.memory_info().rss), "pagefile_used_bytes": int(psutil.swap_memory().used)} + if available < self.min_host_free: + raise ExecutorError(f"physical host reserve below gate: {available} < {self.min_host_free}") + self.logger.event("host_gate", **result, minimum_available_bytes=self.min_host_free) + return result + + def acquire_metadata(self) -> list[dict[str, Any]]: + rows = self._download_stage("B1", self.manifest.metadata) + expected = sum(row.byte_length for row in self.manifest.metadata) + if expected != 57_176_714: + raise ExecutorError("metadata byte cap mismatch") + if self.execute: + receipt_path = self.scratch_root / "receipts" / "B1-metadata.json" + source_hashes = {row.filename: self.downloader.validate_existing(row, self.source_root / row.filename)["sha256"] for row in self.manifest.metadata} + binding = {"stage": "B1", "files": len(rows), "bytes": expected, "source_inventory_fingerprint": self.source_inventory_fingerprint, "source_revision": self.manifest.revision, "builder_commit": self.builder_commit, "source_hashes": source_hashes} + if not _receipt_matches(receipt_path, binding): + _publish_component_receipt(receipt_path, {**binding, "validation": {"all_source_receipts": True, "metadata_total": expected}}) + return rows + + def acquire_ple(self) -> list[dict[str, Any]]: + rows = self.manifest.rows_for_stage("B2") + if len(rows) != 10 or sum(row.byte_length for row in rows) != 51_200_267_901: + raise ExecutorError("PLE source inventory mismatch") + return self._download_stage("B2", rows) + + def acquire_expert_layer(self, layer: int) -> list[dict[str, Any]]: + rows = tuple(row for row in self.manifest.entries if row.source_class.upper() == "EXPERT" and row.layer_id == layer) + if len(rows) != 4: + raise ExecutorError(f"layer {layer} requires exactly four source rows") + return self._download_stage(f"B3-L{layer:02d}", rows) + + def acquire_active(self) -> list[dict[str, Any]]: + rows = self.manifest.rows_for_stage("B4") + if len(rows) != 4 or sum(row.byte_length for row in rows) != 16_007_756_462: + raise ExecutorError("BF16 source inventory mismatch") + return self._download_stage("B4", rows) + + def _source_bindings(self, rows: Iterable[SourceEntry]) -> list[dict[str, Any]]: + """Return receipt-backed source hashes for a component transaction.""" + bindings: list[dict[str, Any]] = [] + for row in rows: + current = self.downloader.validate_existing(row, self.source_root / row.filename) + receipt = self.scratch_root / "receipts" / "sources" / f"{row.acquisition_order:03d}-{row.filename}.receipt.json" + try: + with receipt.open("r", encoding="utf-8") as handle: + value = json.load(handle) + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ExecutorError(f"{row.filename}: source receipt unavailable for component binding") from exc + if (value.get("completion") != "SOURCE_COMPLETE" + or value.get("entry") != asdict(row) + or value.get("source_inventory_fingerprint") != self.source_inventory_fingerprint + or value.get("resolved_commit") != self.manifest.revision + or int(value.get("bytes", -1)) != row.byte_length + or str(value.get("sha256", "")).lower() != str(current["sha256"]).lower()): + raise ExecutorError(f"{row.filename}: source receipt cannot bind component input") + bindings.append({"filename": row.filename, "bytes": row.byte_length, "sha256": str(current["sha256"]).lower()}) + return bindings + + def dry_run_plan(self) -> dict[str, Any]: + staged_rows = [*self.manifest.metadata, *self.manifest.rows_for_stage("B2")] + for layer in range(48): + staged_rows.extend(row for row in self.manifest.entries if row.source_class.upper() == "EXPERT" and row.layer_id == layer) + staged_rows.extend(self.manifest.rows_for_stage("B4")) + source_receipt_root = self.scratch_root / "receipts" / "sources" + component_receipts = [ + self.scratch_root / "receipts" / "B1-metadata.json", + self.scratch_root / "receipts" / "B2-q3.json", + *(self.scratch_root / "receipts" / f"B3-L{layer:02d}.json" for layer in range(48)), + self.scratch_root / "receipts" / "B4-active.json", + self.scratch_root / "receipts" / "B5-manifest.json", + self.scratch_root / "receipts" / "C6-static.json", + ] + return { + "mode": "DRY_RUN", + "body_requests": 0, + "real_model_payload_bytes": 0, + "metadata_files": len(self.manifest.metadata), + "weight_files": len(self.manifest.entries), + "weight_bytes": self.manifest.expected_weight_bytes, + "metadata_bytes": sum(row.byte_length for row in self.manifest.metadata), + "transfer_cap": self.manifest.transfer_cap, + "known_target_bytes": KNOWN_TARGET_BYTES, + "projected_required_free": self._projected_required_free(), + "disk_reserve_bytes": self.min_disk_free, + "host_reserve_bytes": self.min_host_free, + "max_concurrent_downloads": self.downloader.max_concurrent, + "stage_order": ["B1", "B2", "B3-L00..L47", "B4", "B5", "C6"], + "expert_layers": list(range(48)), + "planned_files": [row.filename for row in staged_rows], + "source_receipts": [str(source_receipt_root / f"{row.acquisition_order:03d}-{row.filename}.receipt.json") for row in staged_rows], + "component_receipts": [str(path) for path in component_receipts], + } + + def convert_and_validate_q3(self) -> dict[str, Any]: + from freetoken.checkpoint.q3_ple import plan_q3_ple_production + plan = plan_q3_ple_production() + if plan["segment_count"] != 128 or plan["total_bytes"] != Q3_BYTES: + raise ExecutorError("Q3 production plan mismatch") + if not self.execute: + return {"format": "q3_ple_32", **plan, "state": "PLANNED"} + self._disk_gate() + self._host_gate() + self.target_root.mkdir(parents=True, exist_ok=True) + from freetoken.checkpoint.q3_ple import Q3PLEReader, write_q3_ple_from_safetensors + data_path = self.target_root / "ple-q3-000.bin" + manifest_path = self.target_root / "ple-q3.json" + receipt_path = self.scratch_root / "receipts" / "B2-q3.json" + result: dict[str, Any] = {} + if data_path.exists() != manifest_path.exists(): + raise ExecutorError("incomplete Q3 final target pair") + if not data_path.exists(): + result = write_q3_ple_from_safetensors(self.source_root, data_path, manifest_path, layer_id=2, split_parts=128, source_fingerprint=self.source_inventory_fingerprint, rows_per_segment=2_500_012, processing_chunk_rows=8192) + if data_path.stat().st_size != Q3_BYTES: + raise ExecutorError("Q3 target extent mismatch") + with Q3PLEReader(manifest_path) as reader: + if int(reader.manifest.get("segment_count", -1)) != 128: + raise ExecutorError("Q3 manifest must contain exactly 128 logical segments") + reader.gather([0, 2_500_011, 2_500_012, 160_000_768, 317_501_524, 320_001_535]) + source_inputs = self._source_bindings(self.manifest.rows_for_stage("B2")) + expected = {"stage": "B2", "format": "q3_ple_32", "target": str(data_path), "target_bytes": Q3_BYTES, "target_sha256": _sha256(data_path), "manifest_sha256": _sha256(manifest_path), "source_inventory_fingerprint": self.source_inventory_fingerprint, "source_revision": self.manifest.revision, "builder_commit": self.builder_commit, "source_inputs": source_inputs} + recovered = not bool(result) + if not _receipt_matches(receipt_path, expected): + _publish_component_receipt(receipt_path, {**expected, "validation": {"reopen": True, "segment_count": 128, "sample_rows": [0, 2_500_011, 2_500_012, 160_000_768, 317_501_524, 320_001_535]}, "recovered_after_promotion": recovered}) + return {**expected, "state": "COMPLETE", "recovered_after_promotion": recovered} + + def convert_and_validate_expert(self, layer: int) -> dict[str, Any]: + if not 0 <= int(layer) < 48: + raise ValueError("expert layer outside 0..47") + if not self.execute: + return {"layer": int(layer), "format": "ftexpert1_nvfp4_v1", "target_bytes": EXPERT_BYTES, "state": "PLANNED"} + self._disk_gate() + self._host_gate() + self.target_root.mkdir(parents=True, exist_ok=True) + from freetoken.moe.expert_source import FileExpertSource, write_expert_sidecar_from_safetensors + path = self.target_root / f"experts-L{int(layer):02d}.nvfp4" + receipt_path = self.scratch_root / "receipts" / f"B3-L{int(layer):02d}.json" + created = False + if not path.exists(): + write_expert_sidecar_from_safetensors(self.source_root, path, layer_id=int(layer), source_fingerprint=self.source_inventory_fingerprint) + created = True + if path.stat().st_size != EXPERT_BYTES: + raise ExecutorError(f"expert layer {layer} target extent mismatch") + with FileExpertSource(path, expected_source_fingerprint=self.manifest.source_inventory_fingerprint, expected_layer_id=int(layer), verify_hash=True) as source: + source.read_records([0, 511]) + layer_rows = tuple(row for row in self.manifest.entries if row.source_class.upper() == "EXPERT" and row.layer_id == int(layer)) + expected = {"stage": f"B3-L{int(layer):02d}", "layer": int(layer), "format": "ftexpert1_nvfp4_v1", "target": str(path), "target_bytes": EXPERT_BYTES, "target_sha256": _sha256(path), "source_inventory_fingerprint": self.source_inventory_fingerprint, "source_revision": self.manifest.revision, "builder_commit": self.builder_commit, "source_inputs": self._source_bindings(layer_rows)} + if not _receipt_matches(receipt_path, expected): + _publish_component_receipt(receipt_path, {**expected, "validation": {"reopen": True, "sample_experts": [0, 511]}, "recovered_after_promotion": not created}) + return {**expected, "state": "COMPLETE", "recovered_after_promotion": not created} + + def convert_and_validate_active(self) -> dict[str, Any]: + if not self.execute: + return {"format": "nvfp4_w4a16_v1", "target_bytes": ACTIVE_BYTES, "state": "PLANNED"} + self._disk_gate() + self._host_gate() + self.target_root.mkdir(parents=True, exist_ok=True) + from freetoken.checkpoint.convert import convert_checkpoint + active = self.target_root / "qwen4-active-v1.ftw" + receipt_path = self.scratch_root / "receipts" / "B4-active.json" + result: dict[str, Any] = {} + if not active.exists(): + result = convert_checkpoint(str(self.source_root), str(self.target_root), artifact_format="qwen4_modular_v1", source_inventory_sha256=self.source_inventory_fingerprint) + from freetoken.checkpoint.ftw import INDEX_NAME + index = active / INDEX_NAME + if not active.is_dir() or not index.is_file(): + raise ExecutorError("active FTW was not created") + with index.open("r", encoding="utf-8") as handle: + index_data = json.load(handle) + if int(index_data.get("total_bytes", -1)) != ACTIVE_BYTES: + raise ExecutorError("active FTW extent mismatch") + tree_bytes, tree_sha = _tree_sha256(active) + expected = {"stage": "B4", "format": "nvfp4_w4a16_v1", "target": str(active), "target_bytes": ACTIVE_BYTES, "target_tree_bytes": tree_bytes, "target_sha256": tree_sha, "source_inventory_fingerprint": self.source_inventory_fingerprint, "source_revision": self.manifest.revision, "builder_commit": self.builder_commit, "runtime_commit": self.runtime_commit, "source_inputs": self._source_bindings(self.manifest.rows_for_stage("B4"))} + if not _receipt_matches(receipt_path, expected): + _publish_component_receipt(receipt_path, {**expected, "copied_metadata": list(result.get("copied_metadata", ())), "validation": {"ftw_index": True}, "recovered_after_promotion": not bool(result)}) + return {**expected, "state": "COMPLETE", "recovered_after_promotion": not bool(result)} + + def finalize_artifact(self) -> dict[str, Any]: + if not self.execute: + return {"known_target_bytes": KNOWN_TARGET_BYTES, "reconciliation_error": 0, "state": "PLANNED"} + self._disk_gate() + self._host_gate() + self.target_root.mkdir(parents=True, exist_ok=True) + from freetoken.checkpoint.qwen4_artifact import finalize_qwen4_modular_manifest + manifest_path = self.target_root / "manifest.json" + receipt_path = self.scratch_root / "receipts" / "B5-manifest.json" + expert_paths = {layer: f"experts-L{layer:02d}.nvfp4" for layer in range(48)} + metadata = [] + active_receipt = self.scratch_root / "receipts" / "B4-active.json" + if active_receipt.is_file(): + with active_receipt.open("r", encoding="utf-8") as handle: + copied = (json.load(handle).get("result") or {}).get("copied_metadata", ()) + metadata = [str(value) for value in copied if value and (self.target_root / str(value)).is_file()] + if not metadata: + metadata = [str(path.relative_to(self.target_root)) for path in sorted(self.target_root.iterdir()) if path.is_file() and path.name not in {"manifest.json", "ple-q3.json", "ple-q3-000.bin", "expert-placement.json"}] + created = False + if not manifest_path.exists(): + finalize_qwen4_modular_manifest(self.target_root, source_repository=self.manifest.repository, source_revision=self.manifest.revision, source_inventory_sha256=self.source_inventory_fingerprint, minimum_freetoken_commit=self.builder_commit, tvm_ffi_patch_sha256="889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec", expert_paths=expert_paths, file_tier_layers=(0, 1, 2, 3, 4, 5, 42, 43, 44, 45, 46, 47), metadata_paths=metadata) + created = True + from freetoken.checkpoint.qwen4_artifact import load_qwen4_artifact_manifest + artifact = load_qwen4_artifact_manifest(self.target_root, require=True) + artifact.verify_active() + expected = {"stage": "B5", "known_target_bytes": KNOWN_TARGET_BYTES, "reconciliation_error": 0, "builder_commit": self.builder_commit, "runtime_commit": self.runtime_commit, "manifest_sha256": _sha256(manifest_path), "manifest_fingerprint": artifact.raw.get("complete_artifact_fingerprint"), "source_inventory_fingerprint": self.source_inventory_fingerprint} + if not _receipt_matches(receipt_path, expected): + _publish_component_receipt(receipt_path, {**expected, "source_inputs": ["B2", *[f"B3-L{layer:02d}" for layer in range(48)], "B4"], "validation": {"manifest_reopen": True, "active": True, "experts": 48, "q3": True}, "recovered_after_promotion": not created}) + return {**expected, "state": "COMPLETE", "recovered_after_promotion": not created} + + def run_c6_static_reopen(self) -> dict[str, Any]: + if not self.runtime_worktree.is_dir(): + raise ExecutorError("runtime worktree missing") + if not self.execute: + return {"state": "PLANNED", "runtime_worktree": str(self.runtime_worktree), "runtime_commit": self.runtime_commit, "inference": False} + self._disk_gate() + self._host_gate() + runtime_python = self.runtime_worktree / "python" + extension_root = self.toolchain_root / "build-lib" / "pinned-only" + script = r""" +import json, os +from pathlib import Path +from types import SimpleNamespace +import torch, freetoken.kernel +freetoken.kernel.__path__.append(os.environ['FREETOKEN_PINNED_EXTENSION_DIR']) +from freetoken.checkpoint.qwen4_artifact import load_qwen4_artifact_manifest, load_qwen4_expert_placement_policy +from freetoken.models.qwen4_exp.ple import Q3PLEFileTable +from freetoken.moe.expert_source import FileExpertSource +from freetoken.moe.host_banks import HostBank, HostResidency +from freetoken.kernel.pinned import device_ptr +from freetoken.engine.engine import _apply_pr257_hardware_fit_policy +root = Path(os.environ['FREETOKEN_C6_ARTIFACT']) +manifest = load_qwen4_artifact_manifest(root, require=True) +if not manifest.is_pr257_hardware_fit or not manifest.production_geometry: + raise RuntimeError('marked runtime foundation not recognized') +if manifest.raw.get('runtime_foundation') != 'pr257_hardware_fit_v1' or not manifest.text_only or manifest.active_format != 'nvfp4_w4a16_v1': + raise RuntimeError('hardware-fit manifest markers mismatch') +manifest.verify_active() +policy = load_qwen4_expert_placement_policy(manifest) +expected_tier = (0,1,2,3,4,5,42,43,44,45,46,47) +if policy.file_tier_layers != expected_tier or len(policy.resident_layers) != 36 or policy.file_expert_queue_depth != 4: + raise RuntimeError('placement policy mismatch') +table = Q3PLEFileTable(str(manifest.ple_manifest_path), expected_sha256=manifest.ple_sha256, expected_source_fingerprint=manifest.source['inventory_sha256']) +table.close() +for layer in policy.file_tier_layers: + entry = manifest.file_for_layer(layer) + source = FileExpertSource(entry.path, expected_source_fingerprint=manifest.source['inventory_sha256'], expected_layer_id=layer, max_queue_depth=4, verify_hash=False) + if source.requested_queue_depth != 4: + raise RuntimeError('file source queue depth mismatch') + source.close() +config = json.loads((root / 'config.json').read_text(encoding='utf-8')) +if config.get('freetoken_runtime_foundation') != 'pr257_hardware_fit_v1' or config.get('freetoken_text_only') != 'qwen4_text_only_v1' or config.get('freetoken_active_quant') != 'nvfp4_w4a16_v1': + raise RuntimeError('config markers mismatch') +cfg = SimpleNamespace(model_config=SimpleNamespace(freetoken_runtime_foundation='pr257_hardware_fit_v1'), moe_prefill_overlap=True, cuda_graph_bs=None, cuda_graph_max_bs=None) +if not _apply_pr257_hardware_fit_policy(cfg, graph_requested=False) or cfg.cuda_graph_bs != [] or cfg.cuda_graph_max_bs != 0 or cfg.moe_prefill_overlap is not False: + raise RuntimeError('eager graph policy mismatch') +try: + _apply_pr257_hardware_fit_policy(cfg, graph_requested=True) +except ValueError: + pass +else: + raise RuntimeError('forced graph policy did not fail closed') +bank = HostBank((64, 32), torch.bfloat16, backing='cuda') +if bank.residency is not HostResidency.PINNED or bank.addr % 4096 or device_ptr(bank.tensor) <= 0: + raise RuntimeError('resident pin capability unavailable') +del bank +if torch.cuda.is_available(): + torch.cuda.synchronize() +print(json.dumps({'status':'C6_STATIC_OK','experts':len(manifest.expert_files),'file_tier':len(policy.file_tier_layers),'resident':len(policy.resident_layers),'queue_depth':policy.file_expert_queue_depth,'graphs':cfg.cuda_graph_bs,'prefill_overlap':cfg.moe_prefill_overlap,'inference':False}, sort_keys=True)) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join((str(runtime_python), str(extension_root))) + env["FREETOKEN_C6_ARTIFACT"] = str(self.target_root) + env["FREETOKEN_C6_NO_INFERENCE"] = "1" + env["FREETOKEN_PINNED_EXTENSION_DIR"] = str(extension_root / "freetoken" / "kernel") + proc = subprocess.run([sys.executable, "-c", script], env=env, cwd=str(self.runtime_worktree), text=True, capture_output=True, timeout=600) + if proc.returncode != 0 or "C6_STATIC_OK" not in proc.stdout: + raise ExecutorError(f"isolated C6 static reopen failed: {proc.stderr[-1000:]}") + receipt = {"state": "COMPLETE", "runtime_worktree": str(self.runtime_worktree), "runtime_commit": self.runtime_commit, "artifact": str(self.target_root), "inference": False, "stdout": proc.stdout.strip()} + _atomic_json(self.scratch_root / "receipts" / "C6-static.json", receipt) + return receipt + + def closeout(self) -> dict[str, Any]: + result = {"real_model_payload_bytes": 0 if not self.execute else self.budget.transferred, "source_retirement_authorized": False, "transfer_bytes": self.budget.transferred, "state": "CLOSED"} + self.state["closeout"] = result + _atomic_json(self.state_path, self.state) + return result + + def run(self) -> dict[str, Any]: + """Run explicit stage boundaries. ``DRY_RUN`` never opens a body GET.""" + if not self.execute: + self.preflight() + return self.dry_run_plan() + self.preflight() + self.acquire_metadata() + self.acquire_ple() + self.convert_and_validate_q3() + for layer in range(48): + self.acquire_expert_layer(layer) + self.convert_and_validate_expert(layer) + self.acquire_active() + self.convert_and_validate_active() + result = self.finalize_artifact() + result["c6"] = self.run_c6_static_reopen() + return self.closeout() | result + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Safe Step 9B staged acquisition/conversion controller") + parser.add_argument("--manifest", required=True) + parser.add_argument("--source-root", required=True) + parser.add_argument("--target-root", required=True) + parser.add_argument("--scratch-root", required=True) + parser.add_argument("--logs-root", required=True) + parser.add_argument("--builder-commit", required=True) + parser.add_argument("--runtime-worktree", required=True) + parser.add_argument("--runtime-commit", required=True) + parser.add_argument("--source-revision", required=True) + parser.add_argument("--source-inventory-fingerprint", required=True) + parser.add_argument("--transfer-cap", required=True, type=int) + parser.add_argument("--toolchain-root", required=True) + parser.add_argument("--dry-run", action="store_true", help="plan only (default)") + parser.add_argument("--execute", action="store_true") + parser.add_argument("--allow-network-body", action="store_true") + parser.add_argument("--max-concurrent-downloads", type=int, default=2) + parser.add_argument("--min-disk-reserve", type=int, required=True) + parser.add_argument("--min-host-reserve", type=int, required=True) + parser.add_argument("--source-retirement-policy", choices=("false", "true"), default="false") + return parser + + +def main(argv: list[str] | None = None) -> int: + ns = _parser().parse_args(argv) + execute = bool(ns.execute and not ns.dry_run) + executor = Step9BExecutor(ns.manifest, ns.source_root, ns.target_root, ns.scratch_root, ns.logs_root, builder_commit=ns.builder_commit, runtime_worktree=ns.runtime_worktree, runtime_commit=ns.runtime_commit, source_revision=ns.source_revision, source_inventory_fingerprint=ns.source_inventory_fingerprint, transfer_cap=ns.transfer_cap, toolchain_root=ns.toolchain_root, execute=execute, allow_network_body=ns.allow_network_body, source_retirement_authorized=(ns.source_retirement_policy == "true"), min_disk_free=ns.min_disk_reserve, min_host_free=ns.min_host_reserve, max_concurrent_downloads=ns.max_concurrent_downloads) + result = executor.run() + print(json.dumps(result, indent=2, sort_keys=True)) + return 0 + + +__all__ = ["AcquisitionManifest", "BodyTransferDisabled", "Downloader", "ExecutorError", "JsonlLogger", "MAX_TRANSFER_BYTES", "SourceEntry", "Step9BExecutor", "TransferBudget", "UrllibTransport", "main"] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/checkpoint/test_step9b_executor.py b/tests/checkpoint/test_step9b_executor.py new file mode 100644 index 000000000..68ff59893 --- /dev/null +++ b/tests/checkpoint/test_step9b_executor.py @@ -0,0 +1,98 @@ +"""Bounded, payload-free tests for the Step 9B controller.""" +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import uuid + +import pytest + +from freetoken.checkpoint.step9b_executor import ( + AcquisitionManifest, + BodyTransferDisabled, + Downloader, + SourceEntry, + TransferBudget, +) + + +class Response: + def __init__(self, status, body=b"", *, headers=None): + self.status = status + self._body = body + self.headers = headers or {} + + def iter_bytes(self, chunk_bytes=8 << 20): + for offset in range(0, len(self._body), 3): + yield self._body[offset : offset + 3] + + def close(self): + pass + + +class Transport: + def __init__(self, body, etag='"etag"'): + self.body = body + self.etag = etag + self.get_headers = [] + + def head(self, url, *, headers=None): + return Response(200, headers={"Content-Length": str(len(self.body)), "ETag": self.etag}) + + def get(self, url, *, headers=None, allow_body=False): + assert allow_body + self.get_headers.append(dict(headers or {})) + start = int((headers or {}).get("Range", "bytes=0-").split("=")[1].split("-")[0]) + if start: + return Response(206, self.body[start:], headers={"Content-Length": str(len(self.body) - start), "Content-Range": f"bytes {start}-{len(self.body)-1}/{len(self.body)}", "ETag": self.etag}) + return Response(200, self.body, headers={"Content-Length": str(len(self.body)), "ETag": self.etag}) + + +def entry(body: bytes) -> SourceEntry: + return SourceEntry("fixture.bin", len(body), "METADATA", 1, "RadixArk/Qwen3.8-Flash-Next-NVFP4", "7b719225242aacd3dbd3f9407468c2ee9a9d2594", accepted_etag='"etag"', lfs_oid_sha256=hashlib.sha256(body).hexdigest()) + + +def manifest_for(row: SourceEntry) -> AcquisitionManifest: + return AcquisitionManifest("RadixArk/Qwen3.8-Flash-Next-NVFP4", "7b719225242aacd3dbd3f9407468c2ee9a9d2594", (row,) * 206, (row,) * 9, "8572d200e31b344faff0fda f0dc72aa4726c1f062443d4109531b62ca63f66eb".replace(" ", ""), row.byte_length * 206, 10_000) + + +def z_test_root(label: str) -> Path: + root = Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7f-test-fixtures") / f"{label}-{uuid.uuid4().hex}" + root.mkdir(parents=True) + return root + + +def test_dry_run_cannot_request_body(tmp_path): + body = b"hello world" + row = entry(body) + root = z_test_root("dry-run") + downloader = Downloader(root, manifest_for(row), transport=Transport(body), execute=False) + assert downloader.acquire(row)["state"] == "PLANNED" + with pytest.raises(BodyTransferDisabled): + Downloader(root, manifest_for(row), transport=Transport(body), execute=True, allow_network_body=False).acquire(row) + + +def test_download_resume_and_cap(tmp_path): + root = z_test_root("resume") + body = b"0123456789abcdef" + row = entry(body) + manifest = manifest_for(row) + transport = Transport(body) + downloader = Downloader(root, manifest, transport=transport, execute=True, allow_network_body=True, budget=TransferBudget(10_000)) + partial = root / "fixture.bin.partial" + partial.write_bytes(body[:5]) + from freetoken.checkpoint.step9b_executor import _atomic_json + _atomic_json(partial.with_name(partial.name + ".meta.json"), downloader._identity(row, 5, {"resolved_commit": row.revision, "etag": '"etag"'})) + result = downloader.acquire(row) + assert result["state"] == "SOURCE_COMPLETE" + assert (root / "fixture.bin").read_bytes() == body + assert transport.get_headers[-1]["Range"] == "bytes=5-" + assert downloader.budget.transferred == len(body) - 5 + + +def test_transfer_budget_rejects_extra_byte(tmp_path): + budget = TransferBudget(3) + budget.reserve(3) + with pytest.raises(Exception): + budget.reserve(1) diff --git a/tests/checkpoint/test_step9b_executor_contract.py b/tests/checkpoint/test_step9b_executor_contract.py new file mode 100644 index 000000000..682b4bfca --- /dev/null +++ b/tests/checkpoint/test_step9b_executor_contract.py @@ -0,0 +1,761 @@ +from __future__ import annotations + +import hashlib +import http.server +import json +import socket +import struct +import threading +import time +import uuid +from pathlib import Path + +import pytest + +import freetoken.checkpoint.step9b_executor as executor_module +from freetoken.checkpoint.step9b_executor import ( + ACCEPTED_SOURCE_INVENTORY, + AcquisitionManifest, + BodyTransferDisabled, + Downloader, + ExecutorError, + MIN_DISK_RESERVE_BYTES, + ResumeRejected, + SourceEntry, + Step9BExecutor, + TransferBudget, + UrllibTransport, + _atomic_json, + _publish_component_receipt, + _receipt_matches, +) + + +PIN = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" +REPO = "RadixArk/Qwen3.8-Flash-Next-NVFP4" +WORKSPACE = Path("Z:/Qwen38-FlashNext-Cluster") + + +def zroot(label: str) -> Path: + path = WORKSPACE / "artifacts" / "stage7f-test-fixtures" / f"{label}-{uuid.uuid4().hex}" + path.mkdir(parents=True) + return path + + +def source_entry(name: str, body: bytes, *, etag: str = "fixture-etag", safetensors: bool = False) -> SourceEntry: + header_length = header_sha = None + if safetensors: + (header_length,) = struct.unpack(" AcquisitionManifest: + return AcquisitionManifest(REPO, PIN, (row,), (), ACCEPTED_SOURCE_INVENTORY, row.byte_length, cap) + + +class FixtureServer: + def __init__(self, body: bytes, *, etag: str = "fixture-etag", mode: str = "normal", interrupt_at: int = 0): + self.body = body + self.etag = etag + self.mode = mode + self.interrupt_at = interrupt_at + self.requests: list[tuple[str, str | None, str | None]] = [] + owner = self + + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, *args): + return + + def do_HEAD(self): + owner.requests.append(("HEAD", self.headers.get("Range"), self.headers.get("If-Range"))) + self.send_response(200) + self.send_header("Content-Length", str(len(owner.body))) + if owner.mode != "missing_head_etag": + self.send_header("ETag", f'"{owner.etag}"') + self.end_headers() + + def do_GET(self): + range_header = self.headers.get("Range") + owner.requests.append(("GET", range_header, self.headers.get("If-Range"))) + start = int(range_header.removeprefix("bytes=").split("-")[0]) if range_header else 0 + if range_header and owner.mode == "ignore_range": + start = 0 + status = 200 + else: + status = 206 if range_header else 200 + payload = owner.body[start:] + if owner.mode == "oversized": + payload += b"!" + if owner.mode == "undersized": + payload = payload[:-1] + self.send_response(status) + if owner.mode != "missing_etag": + self.send_header("ETag", f'"{owner.etag}"') + self.send_header("Content-Length", str(len(payload))) + if status == 206: + end = len(owner.body) - 1 + if owner.mode == "malformed_range": + end -= 1 + self.send_header("Content-Range", f"bytes {start}-{end}/{len(owner.body)}") + self.end_headers() + if owner.mode == "interrupt": + self.wfile.write(payload[: owner.interrupt_at]) + self.wfile.flush() + self.connection.shutdown(socket.SHUT_RDWR) + self.connection.close() + return + self.wfile.write(payload) + + self.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + + @property + def url(self) -> str: + host, port = self.server.server_address + return f"http://{host}:{port}/fixture" + + def __enter__(self): + self.thread.start() + return self + + def __exit__(self, *args): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + + +class LocalDownloader(Downloader): + def __init__(self, *args, server: FixtureServer, **kwargs): + self.server = server + super().__init__(*args, transport=UrllibTransport(), **kwargs) + + def _url(self, row: SourceEntry) -> str: + return self.server.url + + def resolve_hf_metadata(self, row: SourceEntry) -> dict: + return {"commit": PIN, "size": row.byte_length, "etag": row.lfs_oid_sha256, "xet_file_hash": row.accepted_etag.strip('"'), "body_bytes": 0} + + +def local_downloader(server: FixtureServer, row: SourceEntry, root: Path, *, cap: int = 1_000_000) -> LocalDownloader: + return LocalDownloader( + root, + small_manifest(row, cap=cap), + server=server, + execute=True, + allow_network_body=True, + budget=TransferBudget(cap, state_path=root / "budget.json"), + ) + + +def test_local_http_clean_download_promotes_and_receipts(): + body = b"contract-complete-body" + row = source_entry("fixture.bin", body) + root = zroot("http-clean") + with FixtureServer(body) as server: + result = local_downloader(server, row, root).acquire(row) + assert result["state"] == "SOURCE_COMPLETE" + assert (root / row.filename).read_bytes() == body + assert not (root / f"{row.filename}.partial").exists() + receipts = list((root / ".step9b-receipts").glob("*.receipt.json")) + assert len(receipts) == 1 + receipt = json.loads(receipts[0].read_text()) + assert receipt["completion"] == "SOURCE_COMPLETE" + assert receipt["resolved_commit"] == PIN + assert receipt["source_inventory_fingerprint"] == ACCEPTED_SOURCE_INVENTORY + + +def test_body_without_required_etag_fails_before_promotion(): + body = b"missing-etag" + row = source_entry("missing-etag.bin", body) + root = zroot("missing-etag") + with FixtureServer(body, mode="missing_etag") as server: + with pytest.raises(ExecutorError, match="omitted required ETag"): + local_downloader(server, row, root).acquire(row) + assert not (root / row.filename).exists() + + +def test_head_without_required_etag_fails_before_body(): + body = b"missing-head-etag" + row = source_entry("missing-head-etag.bin", body) + root = zroot("missing-head-etag") + with FixtureServer(body, mode="missing_head_etag") as server: + downloader = local_downloader(server, row, root) + downloader.resolve_hf_metadata = lambda _row: downloader.validate_metadata(_row, downloader.transport.head(server.url)) + with pytest.raises(ExecutorError, match="omitted required ETag"): + downloader.acquire(row) + assert not any(request[0] == "GET" for request in server.requests) + + +def test_transfer_budget_exact_cap_passes_and_next_received_byte_is_counted(): + body = b"exact-cap" + row = source_entry("exact-cap.bin", body) + root = zroot("exact-cap") + with FixtureServer(body) as server: + result = local_downloader(server, row, root, cap=len(body)).acquire(row) + assert result["state"] == "SOURCE_COMPLETE" + persisted = TransferBudget(len(body), state_path=root / "budget.json") + assert persisted.transferred == len(body) + with pytest.raises(ExecutorError, match="transfer cap exceeded"): + persisted.reserve(1) + assert TransferBudget(len(body), state_path=root / "budget.json").transferred == len(body) + 1 + + +def test_interrupted_download_persists_partial_then_exact_resume(): + body = b"0123456789abcdef" + row = source_entry("resume.bin", body) + root = zroot("http-resume") + with FixtureServer(body, mode="interrupt", interrupt_at=5) as server: + downloader = local_downloader(server, row, root) + with pytest.raises(ExecutorError): + downloader.acquire(row) + partial = root / "resume.bin.partial" + meta = root / "resume.bin.partial.meta.json" + assert partial.read_bytes() == body[:5] + assert json.loads(meta.read_text())["partial_length"] == 5 + with FixtureServer(body) as server: + result = local_downloader(server, row, root).acquire(row) + get = [request for request in server.requests if request[0] == "GET"][-1] + assert get[1] == "bytes=5-" + assert get[2] == row.lfs_oid_sha256 + assert result["resumed_from"] == 5 + assert (root / "resume.bin").read_bytes() == body + + +def test_complete_partial_is_revalidated_and_promoted_without_another_body(): + body = b"complete-partial" + row = source_entry("complete-partial.bin", body) + root = zroot("complete-partial") + partial = root / f"{row.filename}.partial" + partial.write_bytes(body) + with FixtureServer(body) as server: + downloader = local_downloader(server, row, root) + remote = downloader.resolve_hf_metadata(row) + _atomic_json(partial.with_name(partial.name + ".meta.json"), downloader._identity(row, len(body), remote)) + result = downloader.acquire(row) + assert result["recovered_complete_partial"] is True + assert result["body_bytes_this_run"] == 0 + assert (root / row.filename).read_bytes() == body + assert not any(request[0] == "GET" for request in server.requests) + + +def test_crash_after_source_promotion_before_receipt_recovers_without_body(monkeypatch): + body = b"promotion-crash" + row = source_entry("promotion-crash.bin", body) + root = zroot("promotion-crash") + partial = root / f"{row.filename}.partial" + partial.write_bytes(body) + with FixtureServer(body) as server: + downloader = local_downloader(server, row, root) + remote = downloader.resolve_hf_metadata(row) + identity = partial.with_name(partial.name + ".meta.json") + _atomic_json(identity, downloader._identity(row, len(body), remote)) + real_atomic = executor_module._atomic_json + + def fail_receipt(path, value): + if path.name.endswith(".receipt.json"): + raise OSError("injected receipt publication crash") + return real_atomic(path, value) + + monkeypatch.setattr(executor_module, "_atomic_json", fail_receipt) + with pytest.raises(OSError, match="publication crash"): + downloader.acquire(row) + assert (root / row.filename).read_bytes() == body + assert not list((root / ".step9b-receipts").glob("*.receipt.json")) + monkeypatch.setattr(executor_module, "_atomic_json", real_atomic) + recovered = local_downloader(server, row, root).acquire(row) + assert recovered["state"] == "SKIP_VALID_FINAL" + assert not any(request[0] == "GET" for request in server.requests) + + +@pytest.mark.parametrize("mode,error", [("ignore_range", ResumeRejected), ("malformed_range", ResumeRejected)]) +def test_resume_rejects_invalid_range_semantics(mode, error): + body = b"abcdefghijk" + row = source_entry("range.bin", body) + root = zroot(mode) + partial = root / "range.bin.partial" + partial.write_bytes(body[:3]) + with FixtureServer(body, mode=mode) as server: + downloader = local_downloader(server, row, root) + _atomic_json(partial.with_name(partial.name + ".meta.json"), downloader._identity(row, 3, downloader.resolve_hf_metadata(row))) + with pytest.raises(error): + downloader.acquire(row) + assert partial.read_bytes() == body[:3] + + +def test_etag_drift_rejects_partial_before_body(): + body = b"etag-drift" + row = source_entry("etag.bin", body, etag="old") + root = zroot("etag") + partial = root / "etag.bin.partial" + partial.write_bytes(body[:2]) + with FixtureServer(body, etag="new") as server: + downloader = local_downloader(server, row, root) + old = {"commit": PIN, "etag": '"old"', "xet_file_hash": "old"} + _atomic_json(partial.with_name(partial.name + ".meta.json"), downloader._identity(row, 2, old)) + downloader.resolve_hf_metadata = lambda _row: {"commit": PIN, "etag": '"new"', "xet_file_hash": "new"} + with pytest.raises(ResumeRejected): + downloader.acquire(row) + assert not any(request[0] == "GET" for request in server.requests) + + +@pytest.mark.parametrize("mode", ["oversized", "undersized"]) +def test_wrong_body_length_never_promotes(mode): + body = b"length-contract" + row = source_entry("length.bin", body) + root = zroot(mode) + with FixtureServer(body, mode=mode) as server: + with pytest.raises(ExecutorError): + local_downloader(server, row, root).acquire(row) + assert not (root / row.filename).exists() + + +def test_sha_and_safetensors_header_mismatch_never_promote(): + header = b'{"tensor":{"dtype":"U8","shape":[1],"data_offsets":[0,1]}}' + body = struct.pack(" Date: Fri, 28 Aug 2026 21:54:47 -0400 Subject: [PATCH 13/17] fix(qwen4): separate step9b storage identities --- docs/plans/FREETOKEN-QWEN4-001-STAGE7G.md | 35 ++ .../freetoken/checkpoint/step9b_executor.py | 489 ++++++++++++++++-- tests/checkpoint/test_step9b_identity_v2.py | 310 +++++++++++ 3 files changed, 791 insertions(+), 43 deletions(-) create mode 100644 docs/plans/FREETOKEN-QWEN4-001-STAGE7G.md create mode 100644 tests/checkpoint/test_step9b_identity_v2.py diff --git a/docs/plans/FREETOKEN-QWEN4-001-STAGE7G.md b/docs/plans/FREETOKEN-QWEN4-001-STAGE7G.md new file mode 100644 index 000000000..972f17f3f --- /dev/null +++ b/docs/plans/FREETOKEN-QWEN4-001-STAGE7G.md @@ -0,0 +1,35 @@ +# FREETOKEN-QWEN4-001 / Stage 7G Storage Identity Remediation + +## Definition of done + +Separate Git blob, LFS OID/metadata ETag, Xet hash, and transport-body ETag throughout the Step 9B manifest, downloader, partial/resume state, and receipts; preserve the lifetime transfer ledger at 14,137 bytes; migrate the three existing source receipts without network body transfer; pass a body-disabled real restart rehearsal and the complete 215-file dry run; commit one clean remediation commit without changing any binary artifact/runtime contract. + +## Dependencies and hard boundaries + +- Parent executor commit: `6b62c5057ae13deb76bf92a80e8f925e90d55929`. +- Runtime remains byte-identical at `0307a6114c57b0efc61bc17688f3288fe0bf1dc7`. +- Source inventory fingerprint remains `8572d200e31b344faff0fdaf0dc72aa4726c1f062443d4109531b62ca63f66eb`. +- Lifetime transfer ledger remains `14,137 / 135,252,480,565` bytes. +- HEAD/API metadata only; any new source response-body byte is a stop condition. +- Q3, FTEXPERT1, active FTW, PR257 runtime, placement, conversion arithmetic, and target byte totals are out of scope. + +## Execution plan + +| Step | Status | Validation | +|---|---|---| +| Preserve and hash v1 manifest, receipts, transfer ledger, source files, and accepted worktrees | DONE | Exact SHA/length/state inventory | +| Freeze all nine metadata identities and representative BF16/PLE/EXPERT identities using metadata-only queries | DONE | Commit/size/Git/LFS/Xet semantic checks; zero body bytes | +| Implement manifest v2 identity fields and generation/migration | DONE | 9 metadata + 206 weights; totals/fingerprint unchanged | +| Implement executor v2 metadata, body ETag, partial, resume, and receipt semantics | DONE | Focused Git/LFS/Xet and transport tests | +| Revalidate three existing files and migrate receipts with predecessor hashes | DONE | Local full hashes passed; zero new body bytes; ledger unchanged | +| Run body-disabled real restart rehearsal and full 215-file dry run | DONE | Reached tokenizer body boundary; no Xet mismatch; no body GET | +| Run regressions, compileall, diff-check, source-delta audit, and adversarial review | DONE | 168 non-overlapping tests passed; compileall/diff-check passed; independent findings remediated | +| Create evidence, regenerate handoff, commit, and verify clean authorities | DONE | One commit with exact parent; runtime/history unchanged | + +## Validation commands + +- Focused and complete pytest through the accepted Stage 7F Python environment. +- `python -m compileall -q python/freetoken` +- `git diff --check` +- Stage 7G executor `--execute` without `--allow-network-body` against the actual restart state. +- Stage 7G executor `--dry-run` against all 215 manifest rows. diff --git a/python/freetoken/checkpoint/step9b_executor.py b/python/freetoken/checkpoint/step9b_executor.py index 50b2420c3..36f71a8fe 100644 --- a/python/freetoken/checkpoint/step9b_executor.py +++ b/python/freetoken/checkpoint/step9b_executor.py @@ -40,6 +40,10 @@ MAX_DOWNLOADS = 2 MAX_SAFETENSORS_HEADER_BYTES = 256 << 20 ACCEPTED_SOURCE_INVENTORY = "8572d200e31b344faff0fda f0dc72aa4726c1f062443d4109531b62ca63f66eb".replace(" ", "") +ACQUISITION_MANIFEST_V1 = "freetoken-step9-acquisition-v1" +ACQUISITION_MANIFEST_V2 = "freetoken-step9-acquisition-v2" +SOURCE_IDENTITY_VERSION = 2 +SOURCE_RECEIPT_VERSION = 2 class ExecutorError(RuntimeError): @@ -121,6 +125,31 @@ def _receipt_matches(path: Path, expected: Mapping[str, Any]) -> bool: return value.get("completion") == "COMPONENT_COMPLETE" and all(value.get(key) == item for key, item in expected.items()) +def _clean_identity(value: Any) -> str | None: + """Return a normalized hex identity, retaining non-hex v1 fixtures.""" + if value is None: + return None + text = str(value).strip().strip('"') + return text.lower() or None + + +def _identity_values(*values: Any) -> tuple[str, ...]: + """Deduplicate identity values while preserving their declared order.""" + result: list[str] = [] + for value in values: + text = _clean_identity(value) + if text and text not in result: + result.append(text) + return tuple(result) + + +def _validate_hex_identity(name: str, value: str | None, length: int) -> None: + if value is None: + return + if len(value) != length or any(char not in "0123456789abcdef" for char in value.lower()): + raise ExecutorError(f"{name} must be a {length}-character hexadecimal digest") + + @dataclass(frozen=True) class SourceEntry: """Normalized source row from ``source_weight_shards`` or metadata.""" @@ -138,9 +167,60 @@ class SourceEntry: layer_id: int | None = None tensor_payload_bytes: int | None = None git_blob_id: str | None = None + xet_file_hash: str | None = None + allowed_body_etags: tuple[str, ...] = () + identity_version: int = 1 + + def __post_init__(self) -> None: + # v1 fixtures intentionally use short synthetic ETags. Strict length + # and source-kind checks are applied only to normalized v2 rows. + if int(self.identity_version) >= SOURCE_IDENTITY_VERSION: + _validate_hex_identity("git_blob_id", self.git_blob_id, 40) + _validate_hex_identity("lfs_oid_sha256", self.lfs_oid_sha256, 64) + _validate_hex_identity("xet_file_hash", self.xet_file_hash, 64) + if not self.git_blob_id: + raise ExecutorError(f"{self.filename}: v2 row requires git_blob_id provenance") + if self.source_class.upper() != "METADATA" and (not self.lfs_oid_sha256 or not self.xet_file_hash): + raise ExecutorError(f"{self.filename}: v2 weight row requires LFS OID and Xet hash") + if self.lfs_oid_sha256 is not None and self.xet_file_hash is None: + raise ExecutorError(f"{self.filename}: v2 LFS row requires xet_file_hash") + if self.lfs_oid_sha256 is None and self.xet_file_hash is not None: + raise ExecutorError(f"{self.filename}: Xet hash requires an LFS OID") + if not self.allowed_body_etags: + raise ExecutorError(f"{self.filename}: v2 row requires allowed_body_etags") + for etag in self.allowed_body_etags: + if not str(etag).strip(): + raise ExecutorError(f"{self.filename}: body ETag cannot be empty") + expected_body = ( + {self.git_blob_id.lower()} + if self.lfs_oid_sha256 is None and self.git_blob_id + else {self.lfs_oid_sha256.lower(), self.xet_file_hash.lower()} + ) + if {str(etag).strip('"').lower() for etag in self.allowed_body_etags} != expected_body: + raise ExecutorError(f"{self.filename}: allowed_body_etags must match its Git or LFS/Xet identities") + + @property + def body_etags(self) -> tuple[str, ...]: + """Canonical body ETag allow-list (v1 alias retained for callers).""" + return tuple(self.allowed_body_etags) + + @property + def metadata_etag(self) -> str | None: + """The immutable metadata ETag, never the transport-body ETag.""" + return self.lfs_oid_sha256 or self.git_blob_id + + @property + def semantic_kind(self) -> str: + if self.lfs_oid_sha256 and self.xet_file_hash: + return "LFS_XET" + if self.lfs_oid_sha256: + return "LFS" + if self.git_blob_id: + return "GIT" + return "LEGACY" @classmethod - def from_mapping(cls, raw: Mapping[str, Any], *, order: int | None = None, metadata: bool = False) -> "SourceEntry": + def from_mapping(cls, raw: Mapping[str, Any], *, order: int | None = None, metadata: bool = False, schema_version: int = 1, legacy_migration: bool = False) -> "SourceEntry": filename = str(raw["filename"]) if not filename or Path(filename).name != filename or filename in {".", ".."}: raise ExecutorError(f"unsafe manifest filename: {filename!r}") @@ -148,6 +228,27 @@ def from_mapping(cls, raw: Mapping[str, Any], *, order: int | None = None, metad revision = str(raw.get("revision", PINNED_REVISION)) if repository != PINNED_REPOSITORY or revision != PINNED_REVISION: raise ExecutorError(f"source identity mismatch for {filename}") + raw_git = _clean_identity(raw.get("git_blob_id")) + raw_lfs = _clean_identity(raw.get("lfs_oid_sha256")) + raw_xet = _clean_identity(raw.get("xet_file_hash")) + legacy_etag = raw.get("accepted_etag") + # v1 used accepted_etag for Xet (and occasionally for Git). It is + # migrated into an explicit xet_file_hash/body allow-list but never + # consulted by v2 execution paths. + if raw_xet is None and raw_lfs and legacy_etag and "xet_file_hash" not in raw and (int(schema_version) < SOURCE_IDENTITY_VERSION or legacy_migration): + # Migration of a v1 row: accepted_etag was the Xet hash. Once a + # v2 row carries an explicit xet_file_hash (including null), the + # legacy field is intentionally ignored. + raw_xet = _clean_identity(legacy_etag) + body_values = raw.get("allowed_body_etags", raw.get("body_etags", raw.get("accepted_body_etags"))) + if body_values is None: + if int(schema_version) >= SOURCE_IDENTITY_VERSION: + body_values = _identity_values(raw_lfs, raw_xet) if raw_lfs else _identity_values(raw_git) + else: + body_values = _identity_values(legacy_etag, raw_lfs, raw_git) + elif isinstance(body_values, str): + body_values = (body_values,) + body_etags = _identity_values(*tuple(body_values)) return cls( filename=filename, byte_length=int(raw["byte_length"]), @@ -156,12 +257,15 @@ def from_mapping(cls, raw: Mapping[str, Any], *, order: int | None = None, metad repository=repository, revision=revision, accepted_etag=(raw.get("accepted_etag") or (raw.get("git_blob_id") if metadata else None)), - lfs_oid_sha256=raw.get("lfs_oid_sha256"), + lfs_oid_sha256=raw_lfs, accepted_header_length=raw.get("accepted_header_length"), accepted_header_sha256=raw.get("accepted_header_sha256"), layer_id=(None if raw.get("layer_id") is None else int(raw["layer_id"])), tensor_payload_bytes=(None if raw.get("tensor_payload_bytes") is None else int(raw["tensor_payload_bytes"])), - git_blob_id=(None if raw.get("git_blob_id") is None else str(raw["git_blob_id"])), + git_blob_id=raw_git, + xet_file_hash=raw_xet, + allowed_body_etags=body_etags, + identity_version=int(schema_version), ) @@ -174,18 +278,25 @@ class AcquisitionManifest: source_inventory_fingerprint: str expected_weight_bytes: int transfer_cap: int = MAX_TRANSFER_BYTES + schema_version: int = 1 @classmethod - def load(cls, path: str | os.PathLike[str]) -> "AcquisitionManifest": + def load(cls, path: str | os.PathLike[str], *, require_v2: bool = False) -> "AcquisitionManifest": with Path(path).open("r", encoding="utf-8") as handle: raw = json.load(handle) - if raw.get("schema") != "freetoken-step9-acquisition-v1": + schema = str(raw.get("schema") or "") + if schema not in {ACQUISITION_MANIFEST_V1, ACQUISITION_MANIFEST_V2}: raise ExecutorError("unsupported acquisition manifest schema") + schema_version = 2 if schema == ACQUISITION_MANIFEST_V2 else 1 + if schema_version >= SOURCE_IDENTITY_VERSION and raw.get("identity_schema") != "source_identity_v2": + raise ExecutorError("acquisition manifest v2 requires identity_schema=source_identity_v2") + if require_v2 and schema_version < SOURCE_IDENTITY_VERSION: + raise ExecutorError("canonical Step 9B execution requires acquisition manifest v2") repo, revision = str(raw.get("repository")), str(raw.get("revision")) if repo != PINNED_REPOSITORY or revision != PINNED_REVISION: raise ExecutorError("manifest source pin does not match the frozen revision") - rows = tuple(SourceEntry.from_mapping(row) for row in raw.get("source_weight_shards", ())) - metadata = tuple(SourceEntry.from_mapping(row, order=i + 1, metadata=True) for i, row in enumerate(raw.get("required_small_metadata", ()))) + rows = tuple(SourceEntry.from_mapping(row, schema_version=schema_version) for row in raw.get("source_weight_shards", ())) + metadata = tuple(SourceEntry.from_mapping(row, order=i + 1, metadata=True, schema_version=schema_version) for i, row in enumerate(raw.get("required_small_metadata", ()))) if len(rows) != 206 or len(metadata) != 9: raise ExecutorError(f"manifest requires 206 weights and 9 metadata files, got {len(rows)} / {len(metadata)}") orders = [row.acquisition_order for row in rows] @@ -212,7 +323,7 @@ def load(cls, path: str | os.PathLike[str]) -> "AcquisitionManifest": inventory = str(raw.get("source_inventory_sha256") or ACCEPTED_SOURCE_INVENTORY).lower() if inventory != ACCEPTED_SOURCE_INVENTORY: raise ExecutorError("source tensor inventory fingerprint mismatch") - return cls(repo, revision, rows, metadata, inventory, expected, MAX_TRANSFER_BYTES) + return cls(repo, revision, rows, metadata, inventory, expected, MAX_TRANSFER_BYTES, schema_version) @property def all_entries(self) -> tuple[SourceEntry, ...]: @@ -229,6 +340,155 @@ def rows_for_stage(self, stage: str) -> tuple[SourceEntry, ...]: return self.entries +def normalize_source_entry_mapping(raw: Mapping[str, Any], *, metadata: bool = False, order: int | None = None, schema_version: int = SOURCE_IDENTITY_VERSION, identity_overrides: Mapping[str, Mapping[str, Any]] | None = None, legacy_migration: bool = False) -> dict[str, Any]: + """Normalize one v1/v2 manifest row into explicit source identities. + + ``accepted_etag`` is retained solely as a legacy audit field. v2 callers + must use ``git_blob_id``, ``lfs_oid_sha256``, ``xet_file_hash`` and the + explicit ``allowed_body_etags`` list. + """ + normalized_raw = dict(raw) + override = dict((identity_overrides or {}).get(str(raw.get("filename")), {})) + for key in ("git_blob_id", "lfs_oid_sha256", "xet_file_hash", "allowed_body_etags"): + if key in override: + normalized_raw[key] = override[key] + entry = SourceEntry.from_mapping(normalized_raw, metadata=metadata, order=order, schema_version=schema_version, legacy_migration=legacy_migration) + value = dict(normalized_raw) + value["filename"] = entry.filename + value["byte_length"] = entry.byte_length + value["source_class"] = entry.source_class + value["acquisition_order"] = entry.acquisition_order + value["repository"] = entry.repository + value["revision"] = entry.revision + value["git_blob_id"] = entry.git_blob_id + value["lfs_oid_sha256"] = entry.lfs_oid_sha256 + value["xet_file_hash"] = entry.xet_file_hash + value["allowed_body_etags"] = list(entry.allowed_body_etags) + # Preserve accepted_etag for v1 audit/replay only; no v2 execution path + # reads it as an identity. + return value + + +def generate_acquisition_manifest_v2(raw: Mapping[str, Any], *, identity_overrides: Mapping[str, Mapping[str, Any]] | None = None) -> dict[str, Any]: + """Generate a v2 manifest from a v1-shaped mapping without payload I/O.""" + if str(raw.get("schema") or "") not in {ACQUISITION_MANIFEST_V1, ACQUISITION_MANIFEST_V2}: + raise ExecutorError("unsupported acquisition manifest schema") + result = dict(raw) + legacy_migration = str(raw.get("schema") or "") == ACQUISITION_MANIFEST_V1 + result["schema"] = ACQUISITION_MANIFEST_V2 + result["identity_schema"] = "source_identity_v2" + result["required_small_metadata"] = [ + normalize_source_entry_mapping(row, metadata=True, order=index + 1, schema_version=SOURCE_IDENTITY_VERSION, identity_overrides=identity_overrides, legacy_migration=legacy_migration) + for index, row in enumerate(raw.get("required_small_metadata", ())) + ] + result["source_weight_shards"] = [ + normalize_source_entry_mapping(row, schema_version=SOURCE_IDENTITY_VERSION, identity_overrides=identity_overrides, legacy_migration=legacy_migration) + for row in raw.get("source_weight_shards", ()) + ] + return result + + +def migrate_acquisition_manifest_v1_to_v2(path: str | os.PathLike[str], output_path: str | os.PathLike[str] | None = None, *, identity_overrides: Mapping[str, Mapping[str, Any]] | None = None) -> Path: + """Write a durable v2 manifest beside the v1 source (metadata-only).""" + source = Path(path) + destination = Path(output_path) if output_path is not None else source.with_name(f"{source.stem}.v2{source.suffix}") + with source.open("r", encoding="utf-8") as handle: + raw = json.load(handle) + converted = generate_acquisition_manifest_v2(raw, identity_overrides=identity_overrides) + _atomic_json(destination, converted) + return destination + + +def _bytes_sha256(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def migrate_source_receipt_v1_to_v2( + receipt_path: str | os.PathLike[str], + *, + row: SourceEntry, + source_inventory_fingerprint: str, + predecessor_path: str | os.PathLike[str] | None = None, + observed_metadata_etag: str | None = None, + observed_xet_file_hash: str | None = None, + observed_body_etag: str | None = None, + body_bytes: int = 0, +) -> dict[str, Any]: + """Migrate one completed v1 source receipt without transferring bytes. + + The original JSON bytes are copied to an immutable ``.v1`` sibling and + hashed into the v2 receipt. Callers must validate the final source and + immutable metadata before invoking this helper. + """ + path = Path(receipt_path) + predecessor = Path(predecessor_path) if predecessor_path is not None else path.with_suffix(path.suffix + ".v1") + if not path.is_file(): + raise ExecutorError(f"source receipt missing: {path}") + original = path.read_bytes() + predecessor_sha = _bytes_sha256(original) + try: + prior = json.loads(original.decode("utf-8")) + except (UnicodeDecodeError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ExecutorError("invalid v1 source receipt") from exc + if prior.get("completion") != "SOURCE_COMPLETE": + raise ExecutorError("only completed v1 source receipts can be migrated") + if prior.get("receipt_version") == SOURCE_RECEIPT_VERSION: + predecessor_hash = _bytes_sha256(predecessor.read_bytes()) if predecessor.exists() else None + if prior.get("predecessor_receipt_sha256") and predecessor_hash != prior["predecessor_receipt_sha256"]: + raise ExecutorError("v2 receipt predecessor hash changed") + return prior + if predecessor.exists(): + if predecessor.read_bytes() != original: + raise ExecutorError("existing predecessor receipt differs from current v1 bytes") + else: + predecessor.parent.mkdir(parents=True, exist_ok=True) + temporary = predecessor.with_name(f".{predecessor.name}.partial-{os.getpid()}-{threading.get_ident()}") + temporary.write_bytes(original) + try: + with temporary.open("rb") as handle: + os.fsync(handle.fileno()) + except OSError: + pass + os.replace(temporary, predecessor) + metadata_etag = observed_metadata_etag or prior.get("observed_metadata_etag") or prior.get("observed_etag") or row.metadata_etag + xet = observed_xet_file_hash or prior.get("observed_xet_file_hash") + if row.xet_file_hash: + if not xet: + raise ExecutorError("LFS/Xet v1 receipt migration requires explicitly validated Xet metadata") + if _clean_identity(xet) != _clean_identity(row.xet_file_hash): + raise ExecutorError("v1 receipt Xet identity does not match manifest") + body_value = observed_body_etag or prior.get("observed_body_etag") or prior.get("observed_etag") or (row.allowed_body_etags[0] if row.allowed_body_etags else None) + body = str(body_value) if body_value is not None else None + receipt_entry = asdict(row) + receipt_entry["allowed_body_etags"] = list(row.allowed_body_etags) + migrated = { + **prior, + "receipt_version": SOURCE_RECEIPT_VERSION, + "identity_version": SOURCE_IDENTITY_VERSION, + "entry": receipt_entry, + "source_inventory_fingerprint": source_inventory_fingerprint, + "resolved_commit": row.revision, + "expected_git_blob_id": row.git_blob_id, + "expected_lfs_oid_sha256": row.lfs_oid_sha256, + "expected_xet_file_hash": row.xet_file_hash, + "allowed_body_etags": list(row.allowed_body_etags), + "observed_metadata_etag": metadata_etag, + "observed_xet_file_hash": xet, + "observed_body_etag": body, + "body_bytes": int(body_bytes), + "original_body_bytes": int(prior.get("body_bytes", 0) or 0), + "original_completion": prior.get("completion"), + "new_body_bytes": 0, + "lifetime_transfer_accounting": "unchanged", + "migration_reason": "storage_identity_v2", + "predecessor_receipt_sha256": predecessor_sha, + "predecessor_receipt_path": str(predecessor), + "completion": "SOURCE_COMPLETE", + } + _atomic_json(path, migrated) + return migrated + + @dataclass class TransferBudget: cap: int @@ -402,14 +662,49 @@ def _url(self, row: SourceEntry) -> str: def cancel(self) -> None: self._cancel.set() - def _identity(self, row: SourceEntry, length: int, remote: Mapping[str, Any] | None = None) -> dict[str, Any]: + def _identity(self, row: SourceEntry, length: int, remote: Mapping[str, Any] | None = None, *, observed_body_etag: str | None = None) -> dict[str, Any]: remote = dict(remote or {}) - return {"repository": row.repository, "revision": row.revision, "resolved_commit": remote.get("commit") or remote.get("resolved_commit") or row.revision, "filename": row.filename, "expected_length": row.byte_length, "expected_etag": row.accepted_etag, "expected_lfs_oid": row.lfs_oid_sha256, "expected_header_length": row.accepted_header_length, "expected_header_sha256": row.accepted_header_sha256, "observed_etag": remote.get("etag"), "observed_xet_file_hash": remote.get("xet_file_hash"), "partial_length": length, "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, "acquisition_order": row.acquisition_order} + metadata_etag = remote.get("metadata_etag") or remote.get("etag") + body_etag = observed_body_etag if observed_body_etag is not None else remote.get("observed_body_etag") + value = { + "identity_version": SOURCE_IDENTITY_VERSION if self.manifest.schema_version >= SOURCE_IDENTITY_VERSION else 1, + "repository": row.repository, + "revision": row.revision, + "resolved_commit": remote.get("commit") or remote.get("resolved_commit") or row.revision, + "filename": row.filename, + "expected_length": row.byte_length, + "expected_git_blob_id": row.git_blob_id, + "expected_lfs_oid_sha256": row.lfs_oid_sha256, + "expected_xet_file_hash": row.xet_file_hash, + "allowed_body_etags": list(row.allowed_body_etags), + # Legacy aliases remain readable during v1->v2 migration. + "expected_etag": row.accepted_etag, + "expected_lfs_oid": row.lfs_oid_sha256, + "expected_header_length": row.accepted_header_length, + "expected_header_sha256": row.accepted_header_sha256, + "observed_metadata_etag": metadata_etag, + "observed_etag": metadata_etag, + "observed_xet_file_hash": remote.get("xet_file_hash"), + "observed_body_etag": body_etag, + "partial_length": length, + "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, + "acquisition_order": row.acquisition_order, + } + return value def validate_metadata(self, row: SourceEntry, response: TransportResponse) -> dict[str, Any]: headers = {k.lower(): v for k, v in response.headers.items()} observed_length = int(headers.get("content-length", "-1")) observed_etag = headers.get("etag") + observed_xet = headers.get("x-xet-hash") or headers.get("xet-file-hash") or headers.get("xet_file_hash") + observed_commit = headers.get("x-repo-commit") or headers.get("x-linked-commit") + if row.identity_version >= SOURCE_IDENTITY_VERSION: + if not observed_commit: + raise ExecutorError(f"{row.filename}: metadata response omitted immutable commit identity") + if observed_commit.strip() != row.revision: + raise ExecutorError(f"{row.filename}: metadata commit identity mismatch") + elif observed_commit and observed_commit.strip() != row.revision: + raise ExecutorError(f"{row.filename}: metadata commit identity mismatch") if observed_length != row.byte_length: raise ExecutorError(f"{row.filename}: content length mismatch") # Hugging Face uses two identities for Xet/LFS files. The manifest's @@ -417,6 +712,19 @@ def validate_metadata(self, row: SourceEntry, response: TransportResponse) -> di # surfaced as HfFileMetadata.etag and may be returned as HTTP ETag by # the CDN. A transport may expose either, so accept only either exact # frozen identity and never a merely non-empty header. + if row.identity_version >= SOURCE_IDENTITY_VERSION: + expected_metadata_etag = row.metadata_etag + if expected_metadata_etag and not observed_etag: + raise ExecutorError(f"{row.filename}: metadata response omitted required metadata ETag") + if expected_metadata_etag and observed_etag and observed_etag.strip('"').lower() != expected_metadata_etag.strip('"').lower(): + raise ExecutorError(f"{row.filename}: metadata ETag identity mismatch") + if row.xet_file_hash and observed_xet and observed_xet.strip('"').lower() != row.xet_file_hash.lower(): + raise ExecutorError(f"{row.filename}: metadata Xet identity mismatch") + if row.xet_file_hash and not observed_xet: + raise ExecutorError(f"{row.filename}: metadata response omitted required Xet identity") + if not row.xet_file_hash and observed_xet: + raise ExecutorError(f"{row.filename}: Git-backed metadata unexpectedly exposed Xet identity") + return {"resolved_commit": row.revision, "length": observed_length, "etag": observed_etag, "metadata_etag": observed_etag, "xet_file_hash": observed_xet, "body_bytes": 0} accepted = {str(value).strip('"') for value in (row.accepted_etag, row.lfs_oid_sha256) if value} if accepted and not observed_etag: raise ExecutorError(f"{row.filename}: metadata response omitted required ETag") @@ -454,7 +762,15 @@ def resolve_hf_metadata(self, row: SourceEntry) -> dict[str, Any]: raise ExecutorError(f"{row.filename}: metadata length mismatch") xet = getattr(meta, "xet_file_data", None) xet_hash = str(getattr(xet, "file_hash", "") or "").strip('"') - if row.lfs_oid_sha256: + if row.identity_version >= SOURCE_IDENTITY_VERSION: + expected_metadata_etag = (row.lfs_oid_sha256 or row.git_blob_id or "").lower() + if expected_metadata_etag and etag != expected_metadata_etag: + raise ExecutorError(f"{row.filename}: metadata {'LFS OID' if row.lfs_oid_sha256 else 'Git blob'} identity mismatch") + if row.xet_file_hash and xet_hash != row.xet_file_hash.lower(): + raise ExecutorError(f"{row.filename}: metadata Xet hash mismatch") + if not row.lfs_oid_sha256 and xet_hash: + raise ExecutorError(f"{row.filename}: Git metadata unexpectedly carries Xet identity") + elif row.lfs_oid_sha256: if etag != row.lfs_oid_sha256.lower(): raise ExecutorError(f"{row.filename}: metadata LFS OID mismatch") if row.accepted_etag and xet_hash != row.accepted_etag.strip('"').lower(): @@ -463,7 +779,7 @@ def resolve_hf_metadata(self, row: SourceEntry) -> dict[str, Any]: # Git-backed metadata has no Xet data; accepted_etag is the blob id. if etag != row.accepted_etag.strip('"').lower(): raise ExecutorError(f"{row.filename}: metadata Git identity mismatch") - return {"url": url, "commit": commit, "size": size, "etag": etag, "xet_file_hash": xet_hash or None, "body_bytes": 0} + return {"url": url, "commit": commit, "size": size, "etag": etag, "metadata_etag": etag, "xet_file_hash": xet_hash or None, "body_bytes": 0} def _validate_partial_identity(self, row: SourceEntry, meta: Path, length: int, remote: Mapping[str, Any]) -> None: if not meta.is_file(): @@ -471,11 +787,21 @@ def _validate_partial_identity(self, row: SourceEntry, meta: Path, length: int, with meta.open("r", encoding="utf-8") as handle: identity = json.load(handle) expected = self._identity(row, length, remote) - for key in ("repository", "revision", "resolved_commit", "filename", "expected_length", "expected_etag", "expected_lfs_oid", "observed_etag", "observed_xet_file_hash", "source_inventory_fingerprint", "acquisition_order"): + keys = ("repository", "revision", "resolved_commit", "filename", "expected_length", "source_inventory_fingerprint", "acquisition_order") + if row.identity_version >= SOURCE_IDENTITY_VERSION: + keys += ("identity_version", "expected_git_blob_id", "expected_lfs_oid_sha256", "expected_xet_file_hash", "allowed_body_etags", "observed_metadata_etag", "observed_xet_file_hash") + else: + keys += ("expected_etag", "expected_lfs_oid", "observed_etag", "observed_xet_file_hash") + for key in keys: if identity.get(key) != expected.get(key): raise ResumeRejected(f"partial identity mismatch: {key}") if int(identity.get("partial_length", -1)) != length: raise ResumeRejected("partial length identity mismatch") + if row.identity_version >= SOURCE_IDENTITY_VERSION: + observed_body = _clean_identity(identity.get("observed_body_etag")) + allowed = {_clean_identity(value) for value in row.allowed_body_etags} + if length > 0 and observed_body not in allowed: + raise ResumeRejected("partial identity missing allowed observed_body_etag") def _validate_safetensors_header(self, row: SourceEntry, path: Path) -> None: if row.accepted_header_length is None or row.accepted_header_sha256 is None: @@ -521,6 +847,7 @@ def _validate_promote_partial( *, resumed_from: int, body_bytes_this_run: int, + observed_body_etag: str | None = None, recovered_complete_partial: bool = False, ) -> dict[str, Any]: if partial.stat().st_size != row.byte_length: @@ -550,15 +877,21 @@ def _validate_promote_partial( "body_bytes_this_run": body_bytes_this_run, "resolved_commit": row.revision, "expected_etag": row.accepted_etag, + "expected_git_blob_id": row.git_blob_id, + "expected_lfs_oid_sha256": row.lfs_oid_sha256, + "expected_xet_file_hash": row.xet_file_hash, + "allowed_body_etags": list(row.allowed_body_etags), "observed_etag": remote_meta.get("etag"), + "observed_metadata_etag": remote_meta.get("metadata_etag") or remote_meta.get("etag"), "observed_xet_file_hash": remote_meta.get("xet_file_hash"), + "observed_body_etag": observed_body_etag, "expected_lfs_oid": row.lfs_oid_sha256, "observed_lfs_oid": row.lfs_oid_sha256, "observed_header_length": row.accepted_header_length, "observed_header_sha256": row.accepted_header_sha256, "recovered_complete_partial": recovered_complete_partial, } - _atomic_json(receipt, {"entry": asdict(row), **result, "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, "completion": "SOURCE_COMPLETE"}) + _atomic_json(receipt, {"entry": asdict(row), "receipt_version": SOURCE_RECEIPT_VERSION, "identity_version": row.identity_version, **result, "source_inventory_fingerprint": self.manifest.source_inventory_fingerprint, "completion": "SOURCE_COMPLETE"}) return result def acquire(self, row: SourceEntry) -> dict[str, Any]: @@ -575,13 +908,22 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: receipt = self.receipt_root / f"{row.acquisition_order:03d}-{row.filename}.receipt.json" if final.exists(): result = self.validate_existing(row, final) + receipt_entry = asdict(row) + # JSON persists tuples as arrays. Normalize the expected binding + # to that durable representation so a migrated v2 receipt is + # accepted on the next restart instead of comparing list vs tuple. + receipt_entry["allowed_body_etags"] = list(row.allowed_body_etags) receipt_binding = { - "entry": asdict(row), + "entry": receipt_entry, "final_path": str(final), "expected_bytes": row.byte_length, "bytes": row.byte_length, "resolved_commit": row.revision, "expected_etag": row.accepted_etag, + "expected_git_blob_id": row.git_blob_id, + "expected_lfs_oid_sha256": row.lfs_oid_sha256, + "expected_xet_file_hash": row.xet_file_hash, + "allowed_body_etags": list(row.allowed_body_etags), "expected_lfs_oid": row.lfs_oid_sha256, "observed_lfs_oid": row.lfs_oid_sha256, "observed_header_length": row.accepted_header_length, @@ -593,25 +935,45 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: try: with receipt.open("r", encoding="utf-8") as handle: prior = json.load(handle) - if prior.get("completion") != "SOURCE_COMPLETE" or any(prior.get(key) != value for key, value in receipt_binding.items()): - raise ExecutorError(f"{row.filename}: source receipt binding mismatch") - accepted_observed = {str(value).strip('"').lower() for value in (row.accepted_etag, row.lfs_oid_sha256) if value} - if str(prior.get("observed_etag", "")).strip('"').lower() not in accepted_observed: - raise ExecutorError(f"{row.filename}: source receipt ETag binding mismatch") - if row.lfs_oid_sha256 and str(prior.get("observed_xet_file_hash", "")).strip('"').lower() != str(row.accepted_etag).strip('"').lower(): - raise ExecutorError(f"{row.filename}: source receipt Xet binding mismatch") - resumed = prior.get("resumed_from") - if resumed is not None and not 0 <= int(resumed) <= row.byte_length: - raise ExecutorError(f"{row.filename}: invalid source receipt resume offset") - if int(prior.get("body_bytes", -1)) < 0: - raise ExecutorError(f"{row.filename}: invalid source receipt body byte count") + if prior.get("receipt_version") == SOURCE_RECEIPT_VERSION: + prior_binding = dict(receipt_binding) + if prior.get("completion") != "SOURCE_COMPLETE" or any(prior.get(key) != value for key, value in prior_binding.items()): + raise ExecutorError(f"{row.filename}: source receipt binding mismatch") + accepted_observed = {_clean_identity(value) for value in row.allowed_body_etags} + if _clean_identity(prior.get("observed_body_etag") or prior.get("observed_etag")) not in accepted_observed: + raise ExecutorError(f"{row.filename}: source receipt body ETag binding mismatch") + if _clean_identity(prior.get("observed_metadata_etag") or prior.get("metadata_etag")) != _clean_identity(row.metadata_etag): + raise ExecutorError(f"{row.filename}: source receipt metadata identity mismatch") + if row.xet_file_hash and _clean_identity(prior.get("observed_xet_file_hash")) != _clean_identity(row.xet_file_hash): + raise ExecutorError(f"{row.filename}: source receipt Xet binding mismatch") + else: + # v1 receipt migration is metadata-only. Resolve the + # immutable source metadata before rewriting its receipt. + if self.execute: + if isinstance(self.transport, UrllibTransport): + remote = self.resolve_hf_metadata(row) + else: + head = self.transport.head(self._url(row), headers={}) + try: + remote = self.validate_metadata(row, head) + finally: + head.close() + else: + remote = {} + migrate_source_receipt_v1_to_v2( + receipt, + row=row, + source_inventory_fingerprint=self.manifest.source_inventory_fingerprint, + observed_metadata_etag=remote.get("metadata_etag") or remote.get("etag"), + observed_xet_file_hash=remote.get("xet_file_hash"), + observed_body_etag=prior.get("observed_body_etag") or prior.get("observed_etag"), + body_bytes=int(prior.get("body_bytes", 0) or 0), + ) except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: raise ExecutorError(f"{row.filename}: invalid source receipt") from exc else: - _atomic_json(receipt, {**receipt_binding, "observed_etag": row.lfs_oid_sha256 or row.accepted_etag, "observed_xet_file_hash": row.accepted_etag, "resumed_from": None, "body_bytes": 0, "completion": "SOURCE_COMPLETE", "recovered_after_promotion": True}) + _atomic_json(receipt, {**receipt_binding, "receipt_version": SOURCE_RECEIPT_VERSION, "identity_version": row.identity_version, "observed_metadata_etag": row.metadata_etag, "observed_etag": row.metadata_etag, "observed_xet_file_hash": row.xet_file_hash, "observed_body_etag": (row.allowed_body_etags[0] if row.allowed_body_etags else None), "resumed_from": None, "body_bytes": 0, "completion": "SOURCE_COMPLETE", "recovered_after_promotion": True}) return {"filename": row.filename, **result, "state": "SKIP_VALID_FINAL"} - if self.execute and not self.allow_network_body: - raise BodyTransferDisabled("execution mode requires explicit network-body authorization") if not self.execute: if partial.exists() or identity.exists(): raise BodyTransferDisabled("dry run cannot inspect/resume body partials") @@ -645,11 +1007,26 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: ) return {"filename": row.filename, **result} headers: dict[str, str] = {} + persisted_body_etag: str | None = None if current: - validated_etag = str(remote_meta.get("etag") or "") - if not validated_etag: - raise ResumeRejected("validated remote ETag unavailable for If-Range") - headers = {"Range": f"bytes={current}-", "If-Range": validated_etag} + try: + with identity.open("r", encoding="utf-8") as handle: + persisted = json.load(handle) + persisted_value = persisted.get("observed_body_etag") or (persisted.get("observed_etag") if row.identity_version < SOURCE_IDENTITY_VERSION else None) + persisted_body_etag = str(persisted_value) if persisted_value is not None else None + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ResumeRejected("invalid partial identity sidecar") from exc + if row.identity_version >= SOURCE_IDENTITY_VERSION and not persisted_body_etag: + raise ResumeRejected("validated observed_body_etag unavailable for If-Range") + if row.identity_version >= SOURCE_IDENTITY_VERSION: + headers = {"Range": f"bytes={current}-", "If-Range": persisted_body_etag} + else: + # v1 compatibility: its observed_etag represented metadata + # identity, so retain the historical If-Range value while v2 + # uses the exact transport-body ETag above. + headers = {"Range": f"bytes={current}-", "If-Range": str(remote_meta.get("etag") or persisted_body_etag or "")} + if not self.allow_network_body: + raise BodyTransferDisabled("execution mode requires explicit network-body authorization") self._semaphore.acquire() with self._active_lock: self.active += 1 @@ -687,15 +1064,18 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: elif int(rh.get("content-length", "-1")) != row.byte_length: raise ExecutorError("initial Content-Length mismatch") etag = rh.get("etag") - accepted = {str(value).strip('"').lower() for value in (row.accepted_etag, row.lfs_oid_sha256) if value} + accepted = {_clean_identity(value) for value in (row.allowed_body_etags if row.identity_version >= SOURCE_IDENTITY_VERSION else (row.accepted_etag, row.lfs_oid_sha256)) if value} if accepted and not etag: raise ExecutorError("body response omitted required ETag identity") if etag and accepted and etag.strip('"').lower() not in accepted: - raise ExecutorError("body ETag/Xet identity changed") + raise ExecutorError("body ETag identity changed") + if current and row.identity_version >= SOURCE_IDENTITY_VERSION and persisted_body_etag and _clean_identity(etag) != _clean_identity(persisted_body_etag): + raise ResumeRejected("body ETag changed across resume") + body_etag = str(etag) if etag is not None else None mode = "ab" if current else "wb" with partial.open(mode) as target: if not current: - _atomic_json(identity, self._identity(row, 0, remote_meta)) + _atomic_json(identity, self._identity(row, 0, remote_meta, observed_body_etag=body_etag)) for chunk in response.iter_bytes(): if not chunk: continue @@ -707,7 +1087,7 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: raise ExecutorError("response body exceeds expected length") target.write(chunk) received += len(chunk) - _atomic_json(identity, self._identity(row, current + received, remote_meta)) + _atomic_json(identity, self._identity(row, current + received, remote_meta, observed_body_etag=body_etag)) target.flush() os.fsync(target.fileno()) if current + received != row.byte_length: @@ -721,6 +1101,7 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: remote_meta, resumed_from=current, body_bytes_this_run=received, + observed_body_etag=body_etag, ) return {"filename": row.filename, **result} finally: @@ -738,7 +1119,10 @@ class Step9BExecutor: def __init__(self, manifest_path: str | os.PathLike[str], source_root: str | os.PathLike[str], target_root: str | os.PathLike[str], scratch_root: str | os.PathLike[str], logs_root: str | os.PathLike[str], *, builder_commit: str, runtime_worktree: str | os.PathLike[str], runtime_commit: str, source_inventory_fingerprint: str | None = None, source_revision: str = PINNED_REVISION, transfer_cap: int = MAX_TRANSFER_BYTES, toolchain_root: str | os.PathLike[str] | None = None, execute: bool = False, allow_network_body: bool = False, source_retirement_authorized: bool = False, min_disk_free: int = MIN_DISK_RESERVE_BYTES, min_host_free: int = MIN_HOST_FREE_BYTES, max_concurrent_downloads: int = MAX_DOWNLOADS, transport: Transport | None = None): self.manifest_path = Path(manifest_path) - self.manifest = AcquisitionManifest.load(manifest_path) + # Payload execution is canonical only with explicit v2 identities; + # dry-run remains able to inspect the frozen v1 plan for compatibility + # and migration planning. + self.manifest = AcquisitionManifest.load(manifest_path, require_v2=bool(execute)) if str(source_revision) != self.manifest.revision: raise ExecutorError("source revision does not match acquisition manifest") if int(transfer_cap) != self.manifest.transfer_cap: @@ -768,8 +1152,6 @@ def __init__(self, manifest_path: str | os.PathLike[str], source_root: str | os. self.state: dict[str, Any] = {"mode": "EXECUTE" if execute else "DRY_RUN", "source_retirement_authorized": self.source_retirement_authorized, "stages": {}} def preflight(self) -> dict[str, Any]: - if self.execute and not self.allow_network_body: - raise BodyTransferDisabled("--execute requires --allow-network-body") if self.source_retirement_authorized: # Real Step 9 handoff intentionally sets this false. A caller may # test true only with an explicit, separately reviewed controller. @@ -983,8 +1365,11 @@ def _source_bindings(self, rows: Iterable[SourceEntry]) -> list[dict[str, Any]]: value = json.load(handle) except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: raise ExecutorError(f"{row.filename}: source receipt unavailable for component binding") from exc + receipt_entry = dict(value.get("entry") or {}) + if isinstance(receipt_entry.get("allowed_body_etags"), list): + receipt_entry["allowed_body_etags"] = tuple(receipt_entry["allowed_body_etags"]) if (value.get("completion") != "SOURCE_COMPLETE" - or value.get("entry") != asdict(row) + or receipt_entry != asdict(row) or value.get("source_inventory_fingerprint") != self.source_inventory_fingerprint or value.get("resolved_commit") != self.manifest.revision or int(value.get("bytes", -1)) != row.byte_length @@ -1271,7 +1656,25 @@ def main(argv: list[str] | None = None) -> int: return 0 -__all__ = ["AcquisitionManifest", "BodyTransferDisabled", "Downloader", "ExecutorError", "JsonlLogger", "MAX_TRANSFER_BYTES", "SourceEntry", "Step9BExecutor", "TransferBudget", "UrllibTransport", "main"] +__all__ = [ + "ACQUISITION_MANIFEST_V1", + "ACQUISITION_MANIFEST_V2", + "AcquisitionManifest", + "BodyTransferDisabled", + "Downloader", + "ExecutorError", + "JsonlLogger", + "MAX_TRANSFER_BYTES", + "SourceEntry", + "Step9BExecutor", + "TransferBudget", + "UrllibTransport", + "generate_acquisition_manifest_v2", + "migrate_acquisition_manifest_v1_to_v2", + "migrate_source_receipt_v1_to_v2", + "normalize_source_entry_mapping", + "main", +] if __name__ == "__main__": diff --git a/tests/checkpoint/test_step9b_identity_v2.py b/tests/checkpoint/test_step9b_identity_v2.py new file mode 100644 index 000000000..257432e17 --- /dev/null +++ b/tests/checkpoint/test_step9b_identity_v2.py @@ -0,0 +1,310 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import shutil + +import pytest + +from freetoken.checkpoint.step9b_executor import ( + ACQUISITION_MANIFEST_V1, + ACQUISITION_MANIFEST_V2, + ACCEPTED_SOURCE_INVENTORY, + ExecutorError, + BodyTransferDisabled, + Downloader, + AcquisitionManifest, + ResumeRejected, + _atomic_json, + SourceEntry, + generate_acquisition_manifest_v2, + migrate_source_receipt_v1_to_v2, +) + + +PIN = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" +REPO = "RadixArk/Qwen3.8-Flash-Next-NVFP4" + + +def _hex(char: str, length: int) -> str: + return (char * length)[:length] + + +def _v2_row(*, name: str = "model.safetensors", git: str | None = None) -> SourceEntry: + lfs = _hex("a", 64) + xet = _hex("b", 64) + if git is None: + git = _hex("c", 40) + return SourceEntry( + filename=name, + byte_length=4, + source_class="BF16", + acquisition_order=1, + repository=REPO, + revision=PIN, + git_blob_id=git, + lfs_oid_sha256=lfs, + xet_file_hash=xet, + allowed_body_etags=(lfs, xet), + identity_version=2, + ) + + +def test_v2_rows_keep_git_and_lfs_xet_body_identities_distinct(): + row = _v2_row() + assert row.metadata_etag == row.lfs_oid_sha256 + assert set(row.allowed_body_etags) == {row.lfs_oid_sha256, row.xet_file_hash} + assert row.git_blob_id not in row.allowed_body_etags + with pytest.raises(ExecutorError, match="allowed_body_etags"): + SourceEntry(**{**row.__dict__, "allowed_body_etags": (row.git_blob_id, row.lfs_oid_sha256)}) + + +def test_v2_weight_requires_git_lfs_and_xet_and_v2_input_cannot_infer_xet(): + row = _v2_row() + with pytest.raises(ExecutorError, match="git_blob_id provenance"): + SourceEntry(**{**row.__dict__, "git_blob_id": None}) + with pytest.raises(ExecutorError, match="weight row requires"): + SourceEntry(**{**row.__dict__, "lfs_oid_sha256": None, "xet_file_hash": None, "allowed_body_etags": (row.git_blob_id,)}) + raw = {**row.__dict__, "xet_file_hash": None, "allowed_body_etags": None, "accepted_etag": row.xet_file_hash} + with pytest.raises(ExecutorError, match="Xet hash"): + SourceEntry.from_mapping(raw, schema_version=2) + + +def test_git_only_v2_row_allows_only_git_body_etag(): + row = SourceEntry( + filename="config.json", + byte_length=4, + source_class="METADATA", + acquisition_order=1, + repository=REPO, + revision=PIN, + git_blob_id=_hex("d", 40), + allowed_body_etags=(_hex("d", 40),), + identity_version=2, + ) + assert row.semantic_kind == "GIT" + with pytest.raises(ExecutorError, match="allowed_body_etags"): + SourceEntry(**{**row.__dict__, "allowed_body_etags": (_hex("e", 40),)}) + + +def test_manifest_generator_requires_frozen_xet_for_lfs_metadata(): + lfs = _hex("a", 64) + raw = { + "schema": ACQUISITION_MANIFEST_V1, + "repository": REPO, + "revision": PIN, + "required_small_metadata": [{ + "filename": "tokenizer.json", + "byte_length": 4, + "git_blob_id": _hex("d", 40), + "lfs_oid_sha256": lfs, + }], + "source_weight_shards": [], + } + with pytest.raises(ExecutorError, match="xet_file_hash"): + generate_acquisition_manifest_v2(raw) + generated = generate_acquisition_manifest_v2(raw, identity_overrides={"tokenizer.json": {"xet_file_hash": _hex("b", 64)}}) + assert generated["schema"] == ACQUISITION_MANIFEST_V2 + row = generated["required_small_metadata"][0] + assert row["allowed_body_etags"] == [lfs, _hex("b", 64)] + assert _hex("d", 40) not in row["allowed_body_etags"] + + +def test_receipt_migration_preserves_v1_bytes_and_records_predecessor_sha(tmp_path: Path): + body = {"completion": "SOURCE_COMPLETE", "body_bytes": 0, "observed_etag": _hex("b", 64)} + receipt = tmp_path / "001-model.receipt.json" + original = json.dumps(body, sort_keys=True).encode("utf-8") + receipt.write_bytes(original) + row = _v2_row() + migrated = migrate_source_receipt_v1_to_v2( + receipt, + row=row, + source_inventory_fingerprint=ACCEPTED_SOURCE_INVENTORY, + observed_metadata_etag=row.lfs_oid_sha256, + observed_xet_file_hash=row.xet_file_hash, + observed_body_etag=row.xet_file_hash, + ) + predecessor = Path(str(receipt) + ".v1") + assert predecessor.read_bytes() == original + assert migrated["predecessor_receipt_sha256"] == hashlib.sha256(original).hexdigest() + assert json.loads(receipt.read_text())["receipt_version"] == 2 + assert json.loads(receipt.read_text())["body_bytes"] == 0 + assert migrated["new_body_bytes"] == 0 + assert migrated["lifetime_transfer_accounting"] == "unchanged" + assert migrated["migration_reason"] == "storage_identity_v2" + assert migrated["original_body_bytes"] == 0 + assert migrated["original_completion"] == "SOURCE_COMPLETE" + + +class _Response: + def __init__(self, status: int, headers: dict[str, str], body: bytes = b""): + self.status = status + self.headers = headers + self._body = body + + def iter_bytes(self, chunk_bytes: int = 8 << 20): + if self._body: + yield self._body + + def close(self): + return None + + +class _SyntheticTransport: + def __init__(self, body: bytes, metadata_etag: str, xet: str | None, body_etag: str | None): + self.body = body + self.metadata_etag = metadata_etag + self.xet = xet + self.body_etag = body_etag + self.get_headers: list[dict[str, str]] = [] + self.head_calls = 0 + + def head(self, url: str, *, headers=None): + self.head_calls += 1 + values = {"Content-Length": str(len(self.body)), "ETag": self.metadata_etag, "X-Repo-Commit": PIN} + if self.xet: + values["X-Xet-Hash"] = self.xet + return _Response(200, values) + + def get(self, url: str, *, headers=None, allow_body=False): + self.get_headers.append(dict(headers or {})) + headers = dict(headers or {}) + start = int(headers.get("Range", "bytes=0-").split("=")[1].split("-")[0]) + payload = self.body[start:] + response_headers = {"Content-Length": str(len(payload)), "ETag": self.body_etag or ""} + if start: + response_headers["Content-Range"] = f"bytes {start}-{len(self.body)-1}/{len(self.body)}" + return _Response(206, response_headers, payload) + return _Response(200, response_headers, payload) + + +def _manifest_for_v2(row: SourceEntry) -> AcquisitionManifest: + return AcquisitionManifest(REPO, PIN, (row,), (), ACCEPTED_SOURCE_INVENTORY, row.byte_length, 1_000_000, schema_version=2) + + +def test_strict_metadata_checks_git_lfs_xet_mismatch_and_absence(): + body = b"metadata" + lfs = hashlib.sha256(body).hexdigest() + xet = _hex("b", 64) + git = _hex("c", 40) + row = SourceEntry("meta.bin", len(body), "BF16", 1, REPO, PIN, git_blob_id=git, lfs_oid_sha256=lfs, xet_file_hash=xet, allowed_body_etags=(lfs, xet), identity_version=2) + downloader = Downloader(Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7f-test-fixtures/meta-check"), _manifest_for_v2(row), execute=True, allow_network_body=True) + good = _Response(200, {"Content-Length": str(len(body)), "ETag": lfs, "X-Xet-Hash": xet, "X-Repo-Commit": PIN}) + assert downloader.validate_metadata(row, good)["xet_file_hash"] == xet + for headers, message in [ + ({"Content-Length": str(len(body)), "ETag": git, "X-Xet-Hash": xet, "X-Repo-Commit": PIN}, "metadata ETag"), + ({"Content-Length": str(len(body)), "ETag": lfs, "X-Repo-Commit": PIN}, "Xet"), + ({"Content-Length": str(len(body)), "ETag": lfs, "X-Xet-Hash": _hex("d", 64), "X-Repo-Commit": PIN}, "Xet"), + ]: + with pytest.raises(ExecutorError, match=message): + downloader.validate_metadata(row, _Response(200, headers)) + git_row = SourceEntry("config.json", len(body), "METADATA", 1, REPO, PIN, git_blob_id=git, allowed_body_etags=(git,), identity_version=2) + with pytest.raises(ExecutorError, match="unexpectedly exposed Xet"): + downloader.validate_metadata(git_row, _Response(200, {"Content-Length": str(len(body)), "ETag": git, "X-Xet-Hash": xet, "X-Repo-Commit": PIN})) + + +@pytest.mark.parametrize("body_etag", ["lfs", "xet", "git", "arbitrary", None]) +def test_body_etag_allowlist_accepts_only_lfs_or_xet(body_etag: str | None, tmp_path: Path): + body = b"body-etag" + lfs = hashlib.sha256(body).hexdigest() + xet = _hex("b", 64) + git = _hex("c", 40) + values = {"lfs": lfs, "xet": xet, "git": git, "arbitrary": "transport"} + etag = values.get(body_etag) if body_etag else None + row = SourceEntry("body.bin", len(body), "BF16", 1, REPO, PIN, git_blob_id=git, lfs_oid_sha256=lfs, xet_file_hash=xet, allowed_body_etags=(lfs, xet), identity_version=2) + transport = _SyntheticTransport(body, lfs, xet, etag) + root = Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7f-test-fixtures") / f"body-etag-{body_etag or 'missing'}" + shutil.rmtree(root, ignore_errors=True) + root.mkdir(parents=True) + downloader = Downloader(root, _manifest_for_v2(row), transport=transport, execute=True, allow_network_body=True) + if body_etag in {"lfs", "xet"}: + assert downloader.acquire(row)["state"] == "SOURCE_COMPLETE" + else: + with pytest.raises(ExecutorError, match="ETag"): + downloader.acquire(row) + assert not (root / row.filename).exists() + + +def test_v2_partial_identity_contains_distinct_metadata_and_body_fields(tmp_path: Path): + body = b"partial" + lfs = hashlib.sha256(body).hexdigest() + xet = _hex("b", 64) + row = SourceEntry("partial.bin", len(body), "BF16", 1, REPO, PIN, git_blob_id=_hex("c", 40), lfs_oid_sha256=lfs, xet_file_hash=xet, allowed_body_etags=(lfs, xet), identity_version=2) + manifest = _manifest_for_v2(row) + downloader = Downloader(Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7f-test-fixtures/partial-fields"), manifest, execute=True, allow_network_body=True) + identity = downloader._identity(row, 2, {"commit": PIN, "metadata_etag": lfs, "xet_file_hash": xet}, observed_body_etag='"' + xet + '"') + assert identity["observed_metadata_etag"] == lfs + assert identity["observed_body_etag"] == '"' + xet + '"' + assert identity["expected_xet_file_hash"] == xet + + +@pytest.mark.parametrize("body_etag", ['"' + "a" * 64 + '"', "b" * 64]) +def test_resume_if_range_uses_exact_persisted_body_etag(body_etag: str, tmp_path: Path): + body = b"resume-body" + lfs = hashlib.sha256(body).hexdigest() + xet = "a" * 64 if body_etag.startswith('"') else "b" * 64 + row = SourceEntry("resume-v2.bin", len(body), "BF16", 1, REPO, PIN, git_blob_id=_hex("c", 40), lfs_oid_sha256=lfs, xet_file_hash=xet, allowed_body_etags=(lfs, xet), identity_version=2) + root = Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7f-test-fixtures") / f"resume-v2-{xet[0]}" + shutil.rmtree(root, ignore_errors=True) + root.mkdir(parents=True) + transport = _SyntheticTransport(body, lfs, xet, body_etag) + downloader = Downloader(root, _manifest_for_v2(row), transport=transport, execute=True, allow_network_body=True) + partial = root / f"{row.filename}.partial" + partial.write_bytes(body[:3]) + _atomic_json(partial.with_name(partial.name + ".meta.json"), downloader._identity(row, 3, {"commit": PIN, "metadata_etag": lfs, "xet_file_hash": xet}, observed_body_etag=body_etag)) + result = downloader.acquire(row) + assert result["state"] == "SOURCE_COMPLETE" + assert transport.get_headers[-1]["If-Range"] == body_etag + + +def test_receipt_migration_rejects_invalid_or_ambiguous_predecessor(tmp_path: Path): + row = SourceEntry("receipt.bin", 1, "METADATA", 1, REPO, PIN, git_blob_id=_hex("c", 40), allowed_body_etags=(_hex("c", 40),), identity_version=2) + receipt = tmp_path / "bad.receipt.json" + receipt.write_text("not-json", encoding="utf-8") + with pytest.raises(ExecutorError, match="invalid v1"): + migrate_source_receipt_v1_to_v2(receipt, row=row, source_inventory_fingerprint=ACCEPTED_SOURCE_INVENTORY) + receipt.write_text(json.dumps({"completion": "SOURCE_COMPLETE"}), encoding="utf-8") + predecessor = Path(str(receipt) + ".v1") + predecessor.write_text("different", encoding="utf-8") + with pytest.raises(ExecutorError, match="predecessor"): + migrate_source_receipt_v1_to_v2(receipt, row=row, source_inventory_fingerprint=ACCEPTED_SOURCE_INVENTORY) + + +def test_body_disabled_restart_migrates_final_then_stops_before_next_get(tmp_path: Path): + body = b"final-body" + lfs = hashlib.sha256(body).hexdigest() + xet = _hex("b", 64) + row = SourceEntry("final.bin", len(body), "BF16", 1, REPO, PIN, git_blob_id=_hex("c", 40), lfs_oid_sha256=lfs, xet_file_hash=xet, allowed_body_etags=(lfs, xet), identity_version=2) + root = Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7f-test-fixtures/body-disabled-restart") + shutil.rmtree(root, ignore_errors=True) + root.mkdir(parents=True) + (root / row.filename).write_bytes(body) + receipt_root = root / "receipts" + receipt_root.mkdir() + receipt = receipt_root / "001-final.bin.receipt.json" + receipt.write_text(json.dumps({ + "completion": "SOURCE_COMPLETE", + "observed_etag": xet, + "body_bytes": 0, + "final_path": str(root / row.filename), + "expected_bytes": len(body), + "bytes": len(body), + "expected_lfs_oid": lfs, + "observed_lfs_oid": lfs, + "sha256": lfs, + }), encoding="utf-8") + transport = _SyntheticTransport(body, lfs, xet, '"' + xet + '"') + downloader = Downloader(root, _manifest_for_v2(row), receipt_root=receipt_root, transport=transport, execute=True, allow_network_body=False) + assert downloader.acquire(row)["state"] == "SKIP_VALID_FINAL" + assert (Path(str(receipt) + ".v1")).is_file() + # A second restart must accept the durable JSON list representation in the + # migrated v2 receipt without rewriting or redownloading anything. + migrated_bytes = receipt.read_bytes() + assert downloader.acquire(row)["state"] == "SKIP_VALID_FINAL" + assert receipt.read_bytes() == migrated_bytes + missing = SourceEntry("next.bin", len(body), "BF16", 2, REPO, PIN, git_blob_id=_hex("d", 40), lfs_oid_sha256=lfs, xet_file_hash=xet, allowed_body_etags=(lfs, xet), identity_version=2) + with pytest.raises(BodyTransferDisabled): + downloader.acquire(missing) + assert transport.head_calls >= 2 + assert len(transport.get_headers) == 0 From 396833e51f5c7ba0c5f56200d50406ec50b64e34 Mon Sep 17 00:00:00 2001 From: Your Name Date: Fri, 28 Aug 2026 23:07:38 -0400 Subject: [PATCH 14/17] fix(qwen4): recover Windows step9b atomic checkpoints --- docs/plans/FREETOKEN-QWEN4-001-STAGE7H.md | 31 ++ .../freetoken/checkpoint/step9b_executor.py | 235 ++++++++++- .../checkpoint/test_step9b_windows_atomic.py | 377 ++++++++++++++++++ 3 files changed, 640 insertions(+), 3 deletions(-) create mode 100644 docs/plans/FREETOKEN-QWEN4-001-STAGE7H.md create mode 100644 tests/checkpoint/test_step9b_windows_atomic.py diff --git a/docs/plans/FREETOKEN-QWEN4-001-STAGE7H.md b/docs/plans/FREETOKEN-QWEN4-001-STAGE7H.md new file mode 100644 index 000000000..5540b83aa --- /dev/null +++ b/docs/plans/FREETOKEN-QWEN4-001-STAGE7H.md @@ -0,0 +1,31 @@ +# FREETOKEN-QWEN4-001 Stage 7H execution plan + +Task: `STAGE7H-STEP9B-WINDOWS-ATOMIC-SIDECAR-RECOVERY` + +Definition of done: the accepted Stage 7G executor is extended only with bounded Windows atomic-JSON replacement retries and fail-closed orphan checkpoint recovery; the existing first-PLE partial is preserved byte-for-byte and its exact orphan identity is adopted without any network response-body bytes; the lifetime ledger, manifest, runtime, and B1 state remain unchanged; the body-disabled real restart derives the exact future range plan; all focused regressions pass; and a clean local Stage 7H commit plus pinned resume handoff and evidence are produced. + +Dependencies: + +- Historical Stage 7G executor at `313c043861df5c57dfd7c2f98ec168dc7a631d28` remains clean. +- Manifest v2 SHA-256 remains `8e4074cd1a8950bfb19ebdfdd4c5154b66db3ed538ba99fe41221d1be9361e74`. +- PR257 runtime remains `0307a6114c57b0efc61bc17688f3288fe0bf1dc7` and clean. +- Existing B1, partial body, sidecars, and transfer ledger remain available and unmodified until preservation evidence is captured. +- No body GET, Range request, Q3 conversion, or subsequent source acquisition is permitted. + +Validation commands: + +- Focused Stage 7F/7G/7H executor tests under `tests/checkpoint`. +- Z:-backed real Windows lock integration and repeated-publication stress tests. +- Body-disabled execution against the actual retained source state. +- `python -m compileall -q python/freetoken`. +- `git diff --check`. +- Post-recovery hashes/sizes for the partial, ledger, manifest, B1 receipt, runtime authority, and recovered sidecar. + +Steps: + +1. `DONE` — Freeze authorities and preserve pre-recovery evidence, including a full partial-body recovery fingerprint. +2. `DONE` — Audit executor-owned handles and map every orphan-adoption invariant. +3. `DONE` — Implement bounded Windows replace retry and fail-closed orphan recovery. +4. `DONE` — Run synthetic retry, Windows lock, stress, recovery, identity/resume, receipt, and transfer-budget regressions. +5. `IN PROGRESS` — Adopt the independently validated real orphan and run the body-disabled actual-state restart rehearsal. +6. `PENDING` — Produce Stage 7H evidence, commit locally, regenerate the Step 9B resume handoff, and verify clean closeout. diff --git a/python/freetoken/checkpoint/step9b_executor.py b/python/freetoken/checkpoint/step9b_executor.py index 36f71a8fe..f4e9e7a44 100644 --- a/python/freetoken/checkpoint/step9b_executor.py +++ b/python/freetoken/checkpoint/step9b_executor.py @@ -58,6 +58,10 @@ class ResumeRejected(ExecutorError): """Raised when the server cannot prove an identity-safe range response.""" +class AtomicJsonReplaceError(ExecutorError): + """A bounded Windows atomic-publication retry exhausted its deadline.""" + + def _sha256(path: Path, chunk: int = 8 << 20) -> str: digest = hashlib.sha256() with path.open("rb") as handle: @@ -91,7 +95,30 @@ def _z_path(path: str | os.PathLike[str], *, must_exist: bool = False) -> Path: return resolved -def _atomic_json(path: Path, value: Mapping[str, Any]) -> None: +_WINDOWS_REPLACE_TRANSIENT_ERRORS = frozenset({5, 32, 33}) + + +def _atomic_json( + path: Path, + value: Mapping[str, Any], + *, + replace_attempts: int = 8, + replace_deadline_seconds: float = 2.0, + replace_backoff_seconds: float = 0.025, +) -> None: + """Durably publish JSON with bounded Windows share-lock recovery. + + The temporary file is closed and fsynced before ``os.replace``. On + Windows, a transient share/access lock can make replace fail with WinError + 5, 32, or 33. Retry only those errors for a short bounded window; all + other errors fail immediately and the temporary file is intentionally + preserved as recovery evidence. The existing canonical destination is + never removed before a successful replace. + """ + if replace_attempts < 1: + raise ValueError("replace_attempts must be positive") + if replace_deadline_seconds < 0 or replace_backoff_seconds < 0: + raise ValueError("replace retry timing must be non-negative") path.parent.mkdir(parents=True, exist_ok=True) partial = path.with_name(f".{path.name}.partial-{os.getpid()}-{threading.get_ident()}") with partial.open("w", encoding="utf-8", newline="\n") as handle: @@ -102,7 +129,27 @@ def _atomic_json(path: Path, value: Mapping[str, Any]) -> None: os.fsync(handle.fileno()) except OSError: pass - os.replace(partial, path) + started = time.monotonic() + attempt = 0 + while True: + attempt += 1 + try: + os.replace(partial, path) + return + except OSError as exc: + winerror = getattr(exc, "winerror", None) + transient = os.name == "nt" and winerror in _WINDOWS_REPLACE_TRANSIENT_ERRORS + within_budget = attempt < replace_attempts and (time.monotonic() - started) < replace_deadline_seconds + if not transient or not within_budget: + if transient: + raise AtomicJsonReplaceError( + f"atomic JSON replace failed after {attempt} attempts; " + f"destination={path}; preserved_orphan={partial}; winerror={winerror}" + ) from exc + raise + delay = min(replace_backoff_seconds * (2 ** (attempt - 1)), max(0.0, replace_deadline_seconds - (time.monotonic() - started))) + if delay: + time.sleep(delay) def _publish_component_receipt(receipt_path: Path, value: Mapping[str, Any]) -> dict[str, Any]: @@ -655,6 +702,9 @@ def __init__(self, source_root: str | os.PathLike[str], manifest: AcquisitionMan self._cancel = threading.Event() self._file_locks: dict[str, threading.Lock] = {} self._file_locks_guard = threading.Lock() + # Last body-disabled/resume plans are retained for audit and tests; + # they contain no payload bytes. + self.resume_plans: dict[str, dict[str, Any]] = {} def _url(self, row: SourceEntry) -> str: return f"https://huggingface.co/{row.repository}/resolve/{row.revision}/{row.filename}" @@ -803,6 +853,158 @@ def _validate_partial_identity(self, row: SourceEntry, meta: Path, length: int, if length > 0 and observed_body not in allowed: raise ResumeRejected("partial identity missing allowed observed_body_etag") + @staticmethod + def _atomic_identity_candidates(identity: Path) -> tuple[Path, ...]: + """Return only this executor's atomic-temp siblings, deterministically.""" + prefix = f".{identity.name}.partial-" + return tuple(sorted((item for item in identity.parent.iterdir() if item.is_file() and item.name.startswith(prefix)), key=lambda item: item.name)) + + def _completed_source_bytes(self, *, excluding: str | None = None) -> int: + """Count exact-length final source files for ledger consistency checks.""" + total = 0 + for item in self.manifest.all_entries: + if item.filename == excluding: + continue + final = self.root / item.filename + if not final.is_file(): + continue + if final.stat().st_size != item.byte_length: + raise ResumeRejected(f"completed source has wrong length: {item.filename}") + try: + self.validate_existing(item, final) + except ExecutorError as exc: + raise ResumeRejected(f"completed source failed validation: {item.filename}") from exc + total += item.byte_length + return total + + def _validate_orphan_identity( + self, + row: SourceEntry, + candidate: Path, + partial: Path, + remote: Mapping[str, Any], + ) -> dict[str, Any]: + """Validate an atomic-temp checkpoint against the complete v2 contract.""" + try: + with candidate.open("r", encoding="utf-8") as handle: + value = json.load(handle) + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ResumeRejected(f"invalid orphan identity JSON: {candidate.name}") from exc + if not isinstance(value, Mapping): + raise ResumeRejected(f"orphan identity is not an object: {candidate.name}") + if int(value.get("identity_version", -1)) != SOURCE_IDENTITY_VERSION: + raise ResumeRejected("orphan identity schema is not v2") + actual_length = partial.stat().st_size + expected = self._identity(row, actual_length, remote) + required = ( + "repository", "revision", "resolved_commit", "filename", "expected_length", + "source_inventory_fingerprint", "acquisition_order", "identity_version", + "expected_git_blob_id", "expected_lfs_oid_sha256", "expected_xet_file_hash", + "observed_metadata_etag", "observed_xet_file_hash", + ) + for key in required: + if value.get(key) != expected.get(key): + raise ResumeRejected(f"orphan identity mismatch: {key}") + if int(value.get("partial_length", -1)) != actual_length: + raise ResumeRejected("orphan partial length does not match physical partial") + declared_allowed = {_clean_identity(item) for item in value.get("allowed_body_etags", ())} + expected_allowed = {_clean_identity(item) for item in row.allowed_body_etags} + if declared_allowed != expected_allowed: + raise ResumeRejected("orphan allowed body ETag set mismatch") + observed_body = _clean_identity(value.get("observed_body_etag")) + if actual_length and observed_body not in expected_allowed: + raise ResumeRejected("orphan observed body ETag is not allowed") + # The ledger must account for every validated final source plus this + # partial. It may include retransmitted bytes, hence the >= relation. + required_ledger = self._completed_source_bytes(excluding=row.filename) + actual_length + if self.budget.transferred < required_ledger: + raise ResumeRejected( + f"orphan transfer ledger is inconsistent: {self.budget.transferred} < {required_ledger}" + ) + return dict(value) + + def recover_partial_identity_checkpoint( + self, + row: SourceEntry, + *, + partial: Path | None = None, + identity: Path | None = None, + remote: Mapping[str, Any] | None = None, + ) -> dict[str, Any] | None: + """Fail-closed adoption of a valid executor atomic-temp sidecar. + + Only siblings matching ``..partial-*`` are inspected. A + candidate is adopted when it exactly describes the current physical + partial and the canonical sidecar is absent or strictly behind it. + Ambiguous or invalid candidates are rejected rather than guessed. + """ + partial = partial or (self.root / f"{row.filename}.partial") + identity = identity or partial.with_name(partial.name + ".meta.json") + if not partial.is_file(): + return None + candidates = self._atomic_identity_candidates(identity) + if not candidates: + return None + if remote is None: + if isinstance(self.transport, UrllibTransport): + remote = self.resolve_hf_metadata(row) + else: + head = self.transport.head(self._url(row), headers={}) + try: + remote = self.validate_metadata(row, head) + finally: + head.close() + validated: list[tuple[Path, dict[str, Any]]] = [] + for candidate in candidates: + validated.append((candidate, self._validate_orphan_identity(row, candidate, partial, remote))) + # All candidates must be semantically identical; otherwise fail closed. + baseline = validated[0][1] + identity_keys = ( + "identity_version", "repository", "revision", "resolved_commit", "filename", + "expected_length", "expected_git_blob_id", "expected_lfs_oid_sha256", + "expected_xet_file_hash", "allowed_body_etags", "observed_metadata_etag", + "observed_xet_file_hash", "observed_body_etag", "partial_length", + "source_inventory_fingerprint", "acquisition_order", + ) + for _, value in validated[1:]: + if any(value.get(key) != baseline.get(key) for key in identity_keys): + raise ResumeRejected("AMBIGUOUS ORPHAN CHECKPOINT") + canonical_value: dict[str, Any] | None = None + if identity.exists(): + try: + with identity.open("r", encoding="utf-8") as handle: + parsed = json.load(handle) + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ResumeRejected("canonical partial identity is malformed") from exc + if not isinstance(parsed, Mapping): + raise ResumeRejected("canonical partial identity is not an object") + canonical_value = dict(parsed) + canonical_length = int(canonical_value.get("partial_length", -1)) + if canonical_length > partial.stat().st_size: + raise ResumeRejected("canonical partial identity is ahead of physical partial") + for key in identity_keys: + if key == "partial_length": + continue + if canonical_value.get(key) != baseline.get(key): + raise ResumeRejected(f"canonical/orphan checkpoint conflict: {key}") + if canonical_length == partial.stat().st_size: + return {"state": "ALREADY_CURRENT", "candidate": str(validated[0][0]), "partial_length": partial.stat().st_size} + if canonical_length >= int(baseline.get("partial_length", -1)): + raise ResumeRejected("canonical partial identity is not strictly behind orphan") + # Publish the validated orphan contents through the hardened helper. + _atomic_json(identity, baseline) + with identity.open("r", encoding="utf-8") as handle: + adopted = json.load(handle) + if int(adopted.get("partial_length", -1)) != partial.stat().st_size: + raise ResumeRejected("adopted canonical partial identity changed unexpectedly") + return { + "state": "ADOPTED", + "candidate": str(validated[0][0]), + "partial_length": partial.stat().st_size, + "canonical_was_present": canonical_value is not None, + "body_bytes": 0, + } + def _validate_safetensors_header(self, row: SourceEntry, path: Path) -> None: if row.accepted_header_length is None or row.accepted_header_sha256 is None: return @@ -990,7 +1192,18 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: head.close() current = partial.stat().st_size if partial.exists() else 0 if partial.exists() != identity.exists(): - raise ResumeRejected(f"partial and identity sidecar must exist together for {row.filename}") + # A process crash can leave a valid atomic-temp identity sidecar + # beside the body partial while the canonical replace is pending. + # Recover it before enforcing the pair invariant. + if partial.exists() and not identity.exists(): + self.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote_meta) + if partial.exists() != identity.exists(): + raise ResumeRejected(f"partial and identity sidecar must exist together for {row.filename}") + elif partial.exists(): + # The canonical sidecar may lag the body after a sharing-lock + # failure. This call is a no-op when no orphan is present. + self.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote_meta) + current = partial.stat().st_size if current or partial.exists(): self._validate_partial_identity(row, identity, current, remote_meta) if current == row.byte_length: @@ -1025,6 +1238,16 @@ def _acquire_locked(self, row: SourceEntry) -> dict[str, Any]: # identity, so retain the historical If-Range value while v2 # uses the exact transport-body ETag above. headers = {"Range": f"bytes={current}-", "If-Range": str(remote_meta.get("etag") or persisted_body_etag or "")} + plan = { + "filename": row.filename, + "range_start": current, + "remaining_bytes": row.byte_length - current, + "if_range": persisted_body_etag, + "body_request_authorized": bool(self.allow_network_body), + } + self.resume_plans[row.filename] = plan + if self.logger: + self.logger.event("resume_plan", **plan) if not self.allow_network_body: raise BodyTransferDisabled("execution mode requires explicit network-body authorization") self._semaphore.acquire() @@ -1260,6 +1483,12 @@ def _validate_source_workspace(self) -> None: continue if name.endswith(".partial.meta.json") and name[:-18] in allowed: continue + # Atomic JSON publication uses an executor-owned hidden sibling + # (``..partial--``). Keep these + # restartable checkpoints in scope for orphan recovery; arbitrary + # hidden JSON remains rejected below. + if any(name.startswith(f".{item}.partial.meta.json.partial-") for item in allowed): + continue raise ExecutorError(f"unexpected source workspace file: {name}") def _remaining_target_bytes(self) -> int: diff --git a/tests/checkpoint/test_step9b_windows_atomic.py b/tests/checkpoint/test_step9b_windows_atomic.py new file mode 100644 index 000000000..43431297c --- /dev/null +++ b/tests/checkpoint/test_step9b_windows_atomic.py @@ -0,0 +1,377 @@ +from __future__ import annotations + +import hashlib +import json +import os +import threading +import time +import uuid +import ctypes +from ctypes import wintypes +from pathlib import Path + +import pytest + +import freetoken.checkpoint.step9b_executor as module +from freetoken.checkpoint.step9b_executor import ( + ACCEPTED_SOURCE_INVENTORY, + AcquisitionManifest, + AtomicJsonReplaceError, + Downloader, + ResumeRejected, + SourceEntry, + TransferBudget, + _atomic_json, +) + + +ROOT = Path("Z:/Qwen38-FlashNext-Cluster/artifacts/stage7h-test-fixtures") +PIN = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" +REPO = "RadixArk/Qwen3.8-Flash-Next-NVFP4" + + +def zroot(name: str) -> Path: + path = ROOT / f"{name}-{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" + path.mkdir(parents=True, exist_ok=True) + return path + + +def v2_row(name: str = "partial.safetensors", length: int = 32) -> SourceEntry: + lfs = hashlib.sha256(b"lfs-" + name.encode()).hexdigest() + xet = hashlib.sha256(b"xet-" + name.encode()).hexdigest() + return SourceEntry( + filename=name, + byte_length=length, + source_class="PLE", + acquisition_order=1, + repository=REPO, + revision=PIN, + git_blob_id="a" * 40, + lfs_oid_sha256=lfs, + xet_file_hash=xet, + allowed_body_etags=(lfs, xet), + identity_version=2, + ) + + +def downloader_for(root: Path, row: SourceEntry, transferred: int) -> Downloader: + manifest = AcquisitionManifest(REPO, PIN, (row,), (), ACCEPTED_SOURCE_INVENTORY, row.byte_length, 10_000, schema_version=2) + budget = TransferBudget(10_000, transferred=transferred, state_path=root / "budget.json") + return Downloader(root, manifest, budget=budget, execute=True, allow_network_body=False) + + +def _deny_delete_handle(path: Path): + """Open *path* without FILE_SHARE_DELETE so Windows replace is denied.""" + if os.name != "nt": + pytest.skip("Windows share-lock integration is Windows-only") + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + create = kernel32.CreateFileW + create.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, wintypes.LPVOID, wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE] + create.restype = wintypes.HANDLE + handle = create(str(path), 0x80000000, 0x00000001, None, 3, 0x80, None) # GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING + if handle == wintypes.HANDLE(-1).value: + raise ctypes.WinError(ctypes.get_last_error()) + return kernel32, handle + + +def _close_handle(kernel32, handle): + kernel32.CloseHandle(handle) + + +def test_windows_deny_delete_lock_releases_within_retry_window(): + if os.name != "nt": + pytest.skip("Windows share-lock integration is Windows-only") + root = zroot("atomic-lock-release") + destination = root / "state.json" + _atomic_json(destination, {"sequence": 1}) + kernel32, handle = _deny_delete_handle(destination) + + def release(): + time.sleep(0.20) + _close_handle(kernel32, handle) + + thread = threading.Thread(target=release) + thread.start() + try: + _atomic_json(destination, {"sequence": 2}, replace_deadline_seconds=1.5, replace_backoff_seconds=0.03) + finally: + thread.join(timeout=2) + assert json.loads(destination.read_text()) == {"sequence": 2} + assert not list(root.glob(".state.json.partial-*")) + + +def test_windows_deny_delete_lock_beyond_deadline_preserves_old_and_orphan(): + if os.name != "nt": + pytest.skip("Windows share-lock integration is Windows-only") + root = zroot("atomic-lock-timeout") + destination = root / "state.json" + _atomic_json(destination, {"sequence": 1}) + kernel32, handle = _deny_delete_handle(destination) + try: + with pytest.raises(AtomicJsonReplaceError, match="preserved_orphan=.*state.json"): + _atomic_json(destination, {"sequence": 2}, replace_attempts=4, replace_deadline_seconds=0.10, replace_backoff_seconds=0.02) + assert json.loads(destination.read_text()) == {"sequence": 1} + assert list(root.glob(".state.json.partial-*")) + finally: + _close_handle(kernel32, handle) + + +def test_repeated_atomic_publication_has_no_malformed_reads(): + root = zroot("atomic-stress") + destination = root / "state.json" + _atomic_json(destination, {"sequence": 0}) + stop = threading.Event() + malformed: list[str] = [] + observed: list[int] = [] + + def reader(): + while not stop.is_set(): + try: + value = json.loads(destination.read_text()) + sequence = int(value["sequence"]) + if observed and sequence < observed[-1]: + malformed.append("non-monotonic") + observed.append(sequence) + except OSError: + # Windows may briefly deny/open-race the pathname while the + # directory entry is replaced. That is not a malformed read. + pass + except (ValueError, TypeError, KeyError, json.JSONDecodeError): + malformed.append("malformed") + # Yield a bounded replacement window instead of continuously + # reacquiring a deny-delete CRT handle and starving the writer. + time.sleep(0.002) + + thread = threading.Thread(target=reader, daemon=True) + thread.start() + try: + for sequence in range(1, 101): + _atomic_json(destination, {"sequence": sequence}, replace_backoff_seconds=0.001) + finally: + stop.set() + thread.join(timeout=2) + assert not thread.is_alive() + assert malformed == [] + assert json.loads(destination.read_text())["sequence"] == 100 + + +def _write_orphan_case(root: Path, row: SourceEntry, *, mutate=None, canonical=None, extra=None): + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(row.byte_length))) + downloader = downloader_for(root, row, transferred=row.byte_length) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + identity = partial.with_name(partial.name + ".meta.json") + orphan = identity.with_name(f".{identity.name}.partial-invalid") + orphan_value = downloader._identity(row, row.byte_length, remote) + if mutate: + mutate(orphan_value) + orphan.write_text(json.dumps(orphan_value)) + if extra is not None: + extra_path = identity.with_name(f".{identity.name}.partial-second") + extra_path.write_text(json.dumps(extra)) + if canonical is not None: + identity.write_text(json.dumps(canonical)) + return downloader, row, partial, identity, remote + + +@pytest.mark.parametrize( + "field,value,match", + [ + ("revision", "deadbeef", "revision"), + ("expected_git_blob_id", "b" * 40, "expected_git_blob_id"), + ("expected_lfs_oid_sha256", "c" * 64, "expected_lfs_oid_sha256"), + ("expected_xet_file_hash", "d" * 64, "expected_xet_file_hash"), + ("observed_body_etag", "e" * 64, "observed body ETag"), + ("source_inventory_fingerprint", "f" * 64, "source_inventory_fingerprint"), + ("acquisition_order", 9, "acquisition_order"), + ], +) +def test_orphan_identity_negative_matrix(field, value, match): + root = zroot(f"orphan-negative-{field}") + row = v2_row(length=64) + mutate = lambda identity: identity.__setitem__(field, value) + downloader, row, partial, identity, remote = _write_orphan_case(root, row, mutate=mutate) + with pytest.raises(ResumeRejected, match=match): + downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote) + assert partial.stat().st_size == 64 + assert downloader.budget.transferred == 64 + + +def test_orphan_canonical_ahead_rejected(): + root = zroot("orphan-canonical-ahead") + row = v2_row(length=64) + downloader = downloader_for(root, row, transferred=64) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(64))) + identity = partial.with_name(partial.name + ".meta.json") + identity.write_text(json.dumps(downloader._identity(row, 65, remote))) + orphan = identity.with_name(f".{identity.name}.partial-ahead") + orphan.write_text(json.dumps(downloader._identity(row, 64, remote))) + with pytest.raises(ResumeRejected, match="ahead"): + downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote) + + +def test_conflicting_multiple_orphans_fail_closed(): + root = zroot("orphan-conflict") + row = v2_row(length=64) + downloader = downloader_for(root, row, transferred=64) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(64))) + identity = partial.with_name(partial.name + ".meta.json") + first = downloader._identity(row, 64, remote) + second = dict(first) + second["observed_body_etag"] = row.lfs_oid_sha256 + identity.with_name(f".{identity.name}.partial-a").write_text(json.dumps(first)) + identity.with_name(f".{identity.name}.partial-b").write_text(json.dumps(second)) + with pytest.raises(ResumeRejected, match="AMBIGUOUS"): + downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote) + + +def test_malformed_orphan_rejected(): + root = zroot("orphan-malformed") + row = v2_row(length=64) + downloader = downloader_for(root, row, transferred=64) + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(64))) + identity = partial.with_name(partial.name + ".meta.json") + identity.with_name(f".{identity.name}.partial-malformed").write_text("not-json") + with pytest.raises(ResumeRejected, match="invalid orphan"): + downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote={"commit": PIN}) + + +def test_canonical_missing_valid_orphan_is_adopted(): + root = zroot("orphan-canonical-missing") + row = v2_row(length=64) + downloader = downloader_for(root, row, transferred=64) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(64))) + identity = partial.with_name(partial.name + ".meta.json") + identity.with_name(f".{identity.name}.partial-orphan").write_text(json.dumps(downloader._identity(row, 64, remote))) + result = downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote) + assert result and result["state"] == "ADOPTED" + assert json.loads(identity.read_text())["partial_length"] == 64 + + +@pytest.mark.parametrize("winerror", [5, 32, 33]) +def test_atomic_json_retries_known_windows_transient(monkeypatch, tmp_path: Path, winerror: int): + path = tmp_path / "state.json" + _atomic_json(path, {"sequence": 1}) + original = os.replace + attempts = {"count": 0} + + def flaky(source, destination): + attempts["count"] += 1 + if attempts["count"] <= 2: + error = PermissionError(winerror, "transient lock") + error.winerror = winerror + raise error + return original(source, destination) + + monkeypatch.setattr(module.os, "replace", flaky) + _atomic_json(path, {"sequence": 2}, replace_backoff_seconds=0, replace_deadline_seconds=1) + assert json.loads(path.read_text()) == {"sequence": 2} + assert attempts["count"] == 3 + + +def test_atomic_json_nontransient_is_immediate_and_preserves_orphan(monkeypatch, tmp_path: Path): + path = tmp_path / "state.json" + _atomic_json(path, {"sequence": 1}) + original = os.replace + calls = {"count": 0} + + def fail(source, destination): + calls["count"] += 1 + error = OSError("invalid path") + error.winerror = 87 + raise error + + monkeypatch.setattr(module.os, "replace", fail) + with pytest.raises(OSError, match="invalid path"): + _atomic_json(path, {"sequence": 2}) + assert calls["count"] == 1 + assert json.loads(path.read_text()) == {"sequence": 1} + assert list(tmp_path.glob(".state.json.partial-*")) + monkeypatch.setattr(module.os, "replace", original) + + +def test_atomic_json_permanent_transient_preserves_canonical_and_orphan(monkeypatch, tmp_path: Path): + path = tmp_path / "state.json" + _atomic_json(path, {"sequence": 1}) + + def fail(source, destination): + error = PermissionError(5, "locked") + error.winerror = 5 + raise error + + monkeypatch.setattr(module.os, "replace", fail) + with pytest.raises(AtomicJsonReplaceError, match="preserved_orphan=.*state.json"): + _atomic_json(path, {"sequence": 2}, replace_attempts=2, replace_backoff_seconds=0) + assert json.loads(path.read_text()) == {"sequence": 1} + assert list(tmp_path.glob(".state.json.partial-*")) + + +def test_orphan_adoption_is_exact_and_body_ledger_unchanged(tmp_path: Path): + root = zroot("orphan-adopt") + row = v2_row(length=64) + partial = root / f"{row.filename}.partial" + body = bytes(range(64)) + partial.write_bytes(body) + downloader = downloader_for(root, row, transferred=len(body)) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + identity = partial.with_name(partial.name + ".meta.json") + _atomic_json(identity, downloader._identity(row, 16, remote)) + orphan = identity.with_name(f".{identity.name}.partial-test") + orphan.write_text(json.dumps(downloader._identity(row, len(body), remote))) + before = hashlib.sha256(body).hexdigest() + result = downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote) + assert result and result["state"] == "ADOPTED" + assert json.loads(identity.read_text())["partial_length"] == len(body) + assert partial.stat().st_size == len(body) + assert hashlib.sha256(partial.read_bytes()).hexdigest() == before + assert downloader.budget.transferred == len(body) + + +def test_orphan_wrong_length_rejected_without_body_or_ledger_mutation(tmp_path: Path): + root = zroot("orphan-invalid") + row = v2_row(length=64) + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(64))) + downloader = downloader_for(root, row, transferred=64) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + identity = partial.with_name(partial.name + ".meta.json") + orphan = identity.with_name(f".{identity.name}.partial-invalid") + orphan.write_text(json.dumps({**downloader._identity(row, 63, remote), "partial_length": 63})) + with pytest.raises(ResumeRejected, match="physical partial"): + downloader.recover_partial_identity_checkpoint(row, partial=partial, identity=identity, remote=remote) + assert downloader.budget.transferred == 64 + assert partial.stat().st_size == 64 + + +def test_resume_plan_is_recorded_before_body_disabled(tmp_path: Path): + root = zroot("resume-plan") + row = v2_row(length=64) + partial = root / f"{row.filename}.partial" + partial.write_bytes(bytes(range(64))) + downloader = downloader_for(root, row, transferred=64) + remote = {"commit": PIN, "metadata_etag": row.lfs_oid_sha256, "xet_file_hash": row.xet_file_hash, "observed_body_etag": row.xet_file_hash} + identity = partial.with_name(partial.name + ".meta.json") + _atomic_json(identity, downloader._identity(row, 16, remote, observed_body_etag=row.xet_file_hash)) + # Partial length is intentionally made consistent with the body so the + # body-disabled path reaches the exact resume-plan seam. + _atomic_json(identity, downloader._identity(row, 64, remote, observed_body_etag=row.xet_file_hash)) + # A complete partial would be promoted; use a shorter body instead. + partial.write_bytes(bytes(range(32))) + _atomic_json(identity, downloader._identity(row, 32, remote, observed_body_etag=row.xet_file_hash)) + downloader.resolve_hf_metadata = lambda _row: remote + with pytest.raises(module.BodyTransferDisabled): + downloader.acquire(row) + assert downloader.resume_plans[row.filename] == { + "filename": row.filename, + "range_start": 32, + "remaining_bytes": 32, + "if_range": row.xet_file_hash, + "body_request_authorized": False, + } From 106ce5ba5bc9892cc909849b741c17e7280d03ba Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 29 Aug 2026 01:33:33 -0400 Subject: [PATCH 15/17] fix(qwen4): bind step9b q3 to source ple layer --- docs/plans/FREETOKEN-QWEN4-001-STAGE7I.md | 37 ++++++ .../freetoken/checkpoint/step9b_executor.py | 109 +++++++++++++++- tests/checkpoint/test_step9b_ple_selector.py | 116 ++++++++++++++++++ 3 files changed, 260 insertions(+), 2 deletions(-) create mode 100644 docs/plans/FREETOKEN-QWEN4-001-STAGE7I.md create mode 100644 tests/checkpoint/test_step9b_ple_selector.py diff --git a/docs/plans/FREETOKEN-QWEN4-001-STAGE7I.md b/docs/plans/FREETOKEN-QWEN4-001-STAGE7I.md new file mode 100644 index 000000000..0abf1d7ab --- /dev/null +++ b/docs/plans/FREETOKEN-QWEN4-001-STAGE7I.md @@ -0,0 +1,37 @@ +# FREETOKEN-QWEN4-001 Stage 7I plan + +Task: `STAGE7I-STEP9B-PLE-LAYER-SELECTOR-REMEDIATION` + +## Definition of done + +- Preserve executor commit `d391b4a6e6a31ffa206f7c8339920e2130e174f3`, runtime commit `0307a6114c57b0efc61bc17688f3288fe0bf1dc7`, acquisition manifest v2, all 19 acquired source files, and lifetime transfer ledger `51,257,444,615 / 135,252,480,565`. +- Make the Step 9B executor validate and select the frozen production PLE source layer represented by the immutable local source index instead of requesting nonexistent layer 2. +- Fail closed unless the index contains exactly `shard_0..shard_127` plus the global `weight_scale` under the expected production layer. +- Prove the correction with synthetic/index-only tests and a body-disabled real-state rehearsal. Do not construct the real Q3 artifact. +- Create one local executor commit and regenerate the resume handoff. Do not resume network body acquisition or B3. + +## Dependencies + +- Parent executor commit: `d391b4a6e6a31ffa206f7c8339920e2130e174f3`. +- Immutable source revision: `7b719225242aacd3dbd3f9407468c2ee9a9d2594`. +- Acquisition manifest SHA-256: `8e4074cd1a8950bfb19ebdfdd4c5154b66db3ed538ba99fe41221d1be9361e74`. +- Existing local B1/B2 source and receipts; no new response-body transfer. + +## Steps + +1. `DONE` — Verify authorities, manifest, ledger, and acquired source state; create isolated Stage 7I worktree. +2. `DONE` — Implemented the smallest explicit PLE selector validation and focused regressions. +3. `DONE` — Ran focused executor/Q3 tests, relevant regression suite, compileall, and diff-check. +4. `DONE` — Completed independent read-only review and index/header-only real-state rehearsal. +5. `DONE` — Prepared evidence and the local commit for the regenerated resume handoff and clean closeout. + +## Validation commands + +```text +python -m pytest -q tests/checkpoint/test_step9b_ple_selector.py +python -m pytest -q tests/checkpoint/test_step9b_executor.py tests/checkpoint/test_step9b_executor_contract.py tests/checkpoint/test_step9b_identity_v2.py tests/checkpoint/test_step9b_windows_atomic.py tests/checkpoint/test_q3_ple_writer.py +python -m compileall -q python/freetoken +git diff --check +``` + +The real-state rehearsal must omit `--allow-network-body` and must stop before Q3 construction. diff --git a/python/freetoken/checkpoint/step9b_executor.py b/python/freetoken/checkpoint/step9b_executor.py index f4e9e7a44..9f5b4fac7 100644 --- a/python/freetoken/checkpoint/step9b_executor.py +++ b/python/freetoken/checkpoint/step9b_executor.py @@ -11,6 +11,7 @@ import hashlib import json import os +import re import shutil import struct import subprocess @@ -39,6 +40,8 @@ EXPERT_BYTES = 1_419_776_000 MAX_DOWNLOADS = 2 MAX_SAFETENSORS_HEADER_BYTES = 256 << 20 +PRODUCTION_PLE_SOURCE_LAYER_ID = 1 +PRODUCTION_PLE_SEGMENT_COUNT = 128 ACCEPTED_SOURCE_INVENTORY = "8572d200e31b344faff0fda f0dc72aa4726c1f062443d4109531b62ca63f66eb".replace(" ", "") ACQUISITION_MANIFEST_V1 = "freetoken-step9-acquisition-v1" ACQUISITION_MANIFEST_V2 = "freetoken-step9-acquisition-v2" @@ -62,6 +65,91 @@ class AtomicJsonReplaceError(ExecutorError): """A bounded Windows atomic-publication retry exhausted its deadline.""" +_PLE_SHARD_KEY = re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.ple\.ple_embedding\." + r"ngram_embedding\.shard_(?P\d+)\.weight$" +) +_PLE_SHARD_PREFIX = re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.ple\.ple_embedding\." + r"ngram_embedding\.shard_" +) +_PLE_SCALE_KEY = re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.ple\.ple_embedding\." + r"ngram_embedding\.weight_scale$" +) + + +def _resolve_production_ple_source_layer(index_path: Path) -> int: + """Validate and return the frozen production PLE source layer. + + The target model exposes one PLE table under source layer 1. Its 128 + logical tensors are distinct from the ten physical Safetensors files. Do + not infer a different layer from whichever matching key happens to appear: + require the exact frozen namespace and reject competing PLE shard sets. + """ + try: + with index_path.open("r", encoding="utf-8") as handle: + document = json.load(handle) + weight_map = document["weight_map"] + except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError) as exc: + raise ExecutorError(f"cannot read PLE source index: {index_path}") from exc + if not isinstance(weight_map, Mapping): + raise ExecutorError("source index weight_map must be an object") + + by_layer: dict[int, set[int]] = {} + scale_layers: set[int] = set() + for key, source_file in weight_map.items(): + if not isinstance(key, str): + continue + match = _PLE_SHARD_KEY.fullmatch(key) + if match is None: + malformed = _PLE_SHARD_PREFIX.match(key) + if malformed is not None: + raise ExecutorError(f"malformed PLE source tensor key: {key}") + scale = _PLE_SCALE_KEY.fullmatch(key) + if scale is not None: + scale_layers.add(int(scale.group("layer"))) + continue + if not isinstance(source_file, str) or not source_file: + raise ExecutorError(f"invalid PLE source file mapping for {key}") + layer = int(match.group("layer")) + index = int(match.group("index")) + if index in by_layer.setdefault(layer, set()): + raise ExecutorError(f"duplicate PLE source tensor index {index} for layer {layer}") + by_layer[layer].add(index) + + expected_indices = set(range(PRODUCTION_PLE_SEGMENT_COUNT)) + if set(by_layer) != {PRODUCTION_PLE_SOURCE_LAYER_ID}: + raise ExecutorError( + "production PLE source layer mismatch: " + f"expected only layer {PRODUCTION_PLE_SOURCE_LAYER_ID}, found {sorted(by_layer)}" + ) + observed = by_layer[PRODUCTION_PLE_SOURCE_LAYER_ID] + if observed != expected_indices: + missing = sorted(expected_indices - observed) + unexpected = sorted(observed - expected_indices) + raise ExecutorError( + "production PLE logical segment mismatch: " + f"missing={missing}; unexpected={unexpected}" + ) + + prefix = ( + f"model.language_model.layers.{PRODUCTION_PLE_SOURCE_LAYER_ID}." + "ple.ple_embedding.ngram_embedding" + ) + scale_key = prefix + ".weight_scale" + if not scale_layers: + raise ExecutorError(f"production PLE global scale is missing: {scale_key}") + if scale_layers != {PRODUCTION_PLE_SOURCE_LAYER_ID}: + raise ExecutorError( + "production PLE scale layer mismatch: " + f"expected only layer {PRODUCTION_PLE_SOURCE_LAYER_ID}, found {sorted(scale_layers)}" + ) + if not isinstance(weight_map[scale_key], str) or not weight_map[scale_key]: + raise ExecutorError(f"invalid PLE source file mapping for {scale_key}") + return PRODUCTION_PLE_SOURCE_LAYER_ID + + def _sha256(path: Path, chunk: int = 8 << 20) -> str: digest = hashlib.sha256() with path.open("rb") as handle: @@ -1648,7 +1736,12 @@ def convert_and_validate_q3(self) -> dict[str, Any]: if plan["segment_count"] != 128 or plan["total_bytes"] != Q3_BYTES: raise ExecutorError("Q3 production plan mismatch") if not self.execute: - return {"format": "q3_ple_32", **plan, "state": "PLANNED"} + return { + "format": "q3_ple_32", + **plan, + "source_layer_id": PRODUCTION_PLE_SOURCE_LAYER_ID, + "state": "PLANNED", + } self._disk_gate() self._host_gate() self.target_root.mkdir(parents=True, exist_ok=True) @@ -1660,7 +1753,19 @@ def convert_and_validate_q3(self) -> dict[str, Any]: if data_path.exists() != manifest_path.exists(): raise ExecutorError("incomplete Q3 final target pair") if not data_path.exists(): - result = write_q3_ple_from_safetensors(self.source_root, data_path, manifest_path, layer_id=2, split_parts=128, source_fingerprint=self.source_inventory_fingerprint, rows_per_segment=2_500_012, processing_chunk_rows=8192) + ple_source_layer_id = _resolve_production_ple_source_layer( + self.source_root / "model.safetensors.index.json" + ) + result = write_q3_ple_from_safetensors( + self.source_root, + data_path, + manifest_path, + layer_id=ple_source_layer_id, + split_parts=PRODUCTION_PLE_SEGMENT_COUNT, + source_fingerprint=self.source_inventory_fingerprint, + rows_per_segment=2_500_012, + processing_chunk_rows=8192, + ) if data_path.stat().st_size != Q3_BYTES: raise ExecutorError("Q3 target extent mismatch") with Q3PLEReader(manifest_path) as reader: diff --git a/tests/checkpoint/test_step9b_ple_selector.py b/tests/checkpoint/test_step9b_ple_selector.py new file mode 100644 index 000000000..55292cb6c --- /dev/null +++ b/tests/checkpoint/test_step9b_ple_selector.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from freetoken.checkpoint.step9b_executor import ( + ExecutorError, + PRODUCTION_PLE_SEGMENT_COUNT, + PRODUCTION_PLE_SOURCE_LAYER_ID, + _resolve_production_ple_source_layer, +) + + +def _key(layer: int, index: int) -> str: + return ( + f"model.language_model.layers.{layer}.ple.ple_embedding." + f"ngram_embedding.shard_{index}.weight" + ) + + +def _scale_key(layer: int) -> str: + return ( + f"model.language_model.layers.{layer}.ple.ple_embedding." + "ngram_embedding.weight_scale" + ) + + +def _write_index(path: Path, weight_map: dict[str, str]) -> Path: + path.write_text(json.dumps({"weight_map": weight_map}), encoding="utf-8") + return path + + +def _production_map() -> dict[str, str]: + result = { + _key(PRODUCTION_PLE_SOURCE_LAYER_ID, index): f"model-plefp8-{index // 13:05d}.safetensors" + for index in range(PRODUCTION_PLE_SEGMENT_COUNT) + } + result[_scale_key(PRODUCTION_PLE_SOURCE_LAYER_ID)] = "model-plefp8-00009.safetensors" + return result + + +def test_resolves_exact_layer_one_128_segment_contract(tmp_path: Path) -> None: + index = _write_index(tmp_path / "model.safetensors.index.json", _production_map()) + assert _resolve_production_ple_source_layer(index) == 1 + + +@pytest.mark.parametrize("missing", [0, 63, 127]) +def test_rejects_missing_logical_segment(tmp_path: Path, missing: int) -> None: + weight_map = _production_map() + del weight_map[_key(1, missing)] + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match="logical segment mismatch"): + _resolve_production_ple_source_layer(index) + + +def test_rejects_missing_global_scale(tmp_path: Path) -> None: + weight_map = _production_map() + del weight_map[_scale_key(1)] + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match="global scale is missing"): + _resolve_production_ple_source_layer(index) + + +def test_rejects_stale_layer_two_namespace(tmp_path: Path) -> None: + weight_map = { + _key(2, index): "ple.safetensors" + for index in range(PRODUCTION_PLE_SEGMENT_COUNT) + } + weight_map[_scale_key(2)] = "ple.safetensors" + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match=r"expected only layer 1, found \[2\]"): + _resolve_production_ple_source_layer(index) + + +def test_rejects_competing_layer_namespace(tmp_path: Path) -> None: + weight_map = _production_map() + weight_map[_key(2, 0)] = "other.safetensors" + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match=r"expected only layer 1, found \[1, 2\]"): + _resolve_production_ple_source_layer(index) + + +def test_rejects_malformed_ple_shard_suffix(tmp_path: Path) -> None: + weight_map = _production_map() + weight_map[ + "model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_seventeen.weight" + ] = "ple.safetensors" + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match="malformed PLE source tensor key"): + _resolve_production_ple_source_layer(index) + + +def test_rejects_competing_scale_namespace(tmp_path: Path) -> None: + weight_map = _production_map() + weight_map[_scale_key(2)] = "other.safetensors" + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match="scale layer mismatch"): + _resolve_production_ple_source_layer(index) + + +def test_rejects_invalid_source_file_mapping(tmp_path: Path) -> None: + weight_map = _production_map() + weight_map[_key(1, 17)] = "" + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + with pytest.raises(ExecutorError, match="invalid PLE source file mapping"): + _resolve_production_ple_source_layer(index) + + +def test_ignores_unrelated_model_weights(tmp_path: Path) -> None: + weight_map = _production_map() + weight_map["model.language_model.layers.1.self_attn.q_proj.weight"] = "active.safetensors" + index = _write_index(tmp_path / "model.safetensors.index.json", weight_map) + assert _resolve_production_ple_source_layer(index) == 1 From b7c1fa2b3ba25a8ce7c2008e645a4f77b0e192ef Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 29 Aug 2026 02:37:09 -0400 Subject: [PATCH 16/17] perf(qwen4): accelerate byte-exact q3 conversion --- docs/FREETOKEN_STAGE7J_Q3_ACCELERATION.md | 51 ++++++ ...TOKEN-QWEN4-001-STAGE7J-Q3-ACCELERATION.md | 35 ++++ python/freetoken/checkpoint/q3_ple.py | 164 ++++++++++++++++-- .../freetoken/checkpoint/step9b_executor.py | 8 +- tests/checkpoint/test_q3_ple_writer.py | 54 ++++++ 5 files changed, 292 insertions(+), 20 deletions(-) create mode 100644 docs/FREETOKEN_STAGE7J_Q3_ACCELERATION.md create mode 100644 docs/plans/FREETOKEN-QWEN4-001-STAGE7J-Q3-ACCELERATION.md diff --git a/docs/FREETOKEN_STAGE7J_Q3_ACCELERATION.md b/docs/FREETOKEN_STAGE7J_Q3_ACCELERATION.md new file mode 100644 index 000000000..85c8157a4 --- /dev/null +++ b/docs/FREETOKEN_STAGE7J_Q3_ACCELERATION.md @@ -0,0 +1,51 @@ +# Stage 7J — Byte-Exact Q3 Acceleration + +Stage 7J replaces only the production Q3 encoder's scalar row loop. The +Q3_PLE_32 format, 128 logical segment boundaries, per-segment hashes, global +weight scale, manifest schema, source inventory, and PR257 runtime remain +unchanged. + +## Reason for remediation + +The first real conversion measured approximately 198 kB/s and projected about +31 hours for the 22,400,107,520-byte Q3 extent while using roughly one CPU +core. The run was stopped before its first component receipt. Its incomplete +target partial is not a valid component; all ten verified PLE source files and +the lifetime transfer ledger remain unchanged. + +## Selected implementation + +The production Safetensors path now reads bounded 131,072-row chunks and +executes the existing two-pass codec as batched Torch operations. CUDA is +selected when available and the same batched arithmetic has a CPU fallback. +The scalar `quantize_block()` and `quantize_row()` functions remain the +reference authority. + +Exactness protections include: + +- the historical FP8-to-FP32 source conversion boundary; +- float64 scale and refinement arithmetic; +- sequential `cumsum` reductions matching the scalar left-fold order; +- per-block early-convergence state; +- the existing integer round-to-nearest-even BF16 scale conversion; +- final requantization against the stored BF16 scale; +- unchanged low-bit-first 3-bit packing; +- ordered writes and unchanged logical segment/hash construction. + +## Evidence + +- 50,000 deterministic FP8-origin rows: scalar and CUDA output byte-identical; +- SHA-256 of the 3,500,000-byte differential output: + `dfe54d8a9d122d356cb46b21a7fda6e2daea7a91b404bcd4a808a9ab9e012f39`; +- measured differential speedup: 125.9x; +- 1,048,576-row sustained synthetic stream: 425,479 rows/s including source + generation, host/device transfer, output materialization, and hashing; +- sustained projected encoding time: 12.53 minutes; +- 131,072-row warmed batch: 1,320,719 rows/s, projected 4.04 minutes for codec + execution alone; +- peak measured GPU allocation: 1,329,070,080 bytes; +- focused executor/Q3/resume regression matrix: 107 passed; +- Q3 writer suite: 20 passed. + +The real Q3 conversion is not restarted until this isolated branch is committed +and the resume handoff pins that exact commit. diff --git a/docs/plans/FREETOKEN-QWEN4-001-STAGE7J-Q3-ACCELERATION.md b/docs/plans/FREETOKEN-QWEN4-001-STAGE7J-Q3-ACCELERATION.md new file mode 100644 index 000000000..b34d49144 --- /dev/null +++ b/docs/plans/FREETOKEN-QWEN4-001-STAGE7J-Q3-ACCELERATION.md @@ -0,0 +1,35 @@ +# FREETOKEN-QWEN4-001 Stage 7J — Q3 Acceleration + +## Definition of done + +Replace the day-scale scalar Q3 conversion loop with the smallest bounded batched implementation that is byte-identical to the frozen Q3_PLE_32 codec, materially uses the available CPU or RTX 5070, and can safely resume Step 9B without changing the artifact format, source manifest, runtime, or lifetime transfer ledger. + +## Frozen dependencies + +- Parent executor commit: `32217a407cbf2c829ecc01f2e26ea391660da278` +- PR257 runtime: `0307a6114c57b0efc61bc17688f3288fe0bf1dc7` +- Source inventory: `8572d200e31b344faff0fdaf0dc72aa4726c1f062443d4109531b62ca63f66eb` +- Lifetime transfer ledger: `51,257,444,615 / 135,252,480,565` bytes +- Binary contract: 128 segments, 70 bytes/row, 22,400,107,520 total Q3 bytes + +## Plan + +1. **DONE** — Benchmark byte-identical batched CPU and CUDA candidates on deterministic synthetic rows. +2. **DONE** — Select the smallest implementation meeting exact byte equality and bounded-memory requirements. +3. **DONE** — Implement the selected batched codec without changing scalar reference behavior or binary layout. +4. **DONE** — Add differential tests covering random rows, edge values, BF16 rounding boundaries, refinement convergence, chunk invariance, and 128-segment writer identity. +5. **DONE** — Benchmark sustained throughput and verify projected production duration and CPU/GPU/RAM bounds. +6. **DONE** — Run executor, Q3, identity-v2, Windows recovery, and manifest regressions plus `compileall` and `git diff --check`. +7. **DONE** — Commit the isolated executor change and regenerate the resume handoff. Resume the real Q3 component only under the new exact commit. + +## Validation commands + +- `python -m pytest -q tests/checkpoint/test_q3_ple_writer.py tests/checkpoint/test_step9b_executor.py tests/checkpoint/test_step9b_executor_contract.py tests/checkpoint/test_step9b_identity_v2.py tests/checkpoint/test_step9b_windows_atomic.py tests/checkpoint/test_step9b_ple_selector.py` +- `python -m compileall -q python/freetoken` +- `git diff --check` +- Synthetic scalar-versus-batched byte differential with deterministic random and adversarial BF16-scale rows +- Bounded sustained throughput benchmark on the selected backend + +## Safety boundary + +No real source download, no source deletion, no real Q3 restart, no expert acquisition, no runtime modification, and no inference occur until the accelerated implementation passes byte-exact review and receives a pinned local commit. diff --git a/python/freetoken/checkpoint/q3_ple.py b/python/freetoken/checkpoint/q3_ple.py index 84eec13b5..d81767672 100644 --- a/python/freetoken/checkpoint/q3_ple.py +++ b/python/freetoken/checkpoint/q3_ple.py @@ -34,6 +34,7 @@ ALIGN = 4096 REFINEMENT_PASSES = 2 DEFAULT_SEGMENT_ROWS = 128 +DEFAULT_PROCESSING_CHUNK_ROWS = 131_072 # Production Qwen4 PLE geometry. ``segment_count`` is a logical source # tensor count; it is deliberately independent from the bounded row chunk used @@ -201,6 +202,105 @@ def quantize_row(values: Sequence[float], *, refinement_passes: int = REFINEMENT ) +def quantize_rows_batched( + rows: torch.Tensor, + *, + device: str | torch.device, + refinement_passes: int = REFINEMENT_PASSES, +) -> bytes: + """Encode a bounded ``[rows, 160]`` batch with scalar-codec-exact arithmetic. + + The production PLE source is FP8, so every source value and every product + used by the least-squares refinement is exactly representable in float64. + ``cumsum(...)[..., -1]`` deliberately preserves the reference encoder's + left-to-right Python ``sum`` order. The stored BF16 scale is rounded by the + same integer-bit recipe as :func:`_bf16_bits`, followed by the same final + requantization against that stored value. + + The returned bytes retain row-major, five-block-per-row order. Choosing a + CUDA device changes only where the bounded arithmetic executes; it does not + change the on-disk format or logical segment identity. + """ + + if refinement_passes < 0: + raise ValueError("refinement_passes must be non-negative") + if not isinstance(rows, torch.Tensor) or rows.ndim != 2 or rows.shape[1] != ROW_VALUES: + shape = tuple(rows.shape) if isinstance(rows, torch.Tensor) else None + raise ValueError(f"expected a [rows, {ROW_VALUES}] tensor, got {shape}") + if rows.shape[0] <= 0: + raise ValueError("Q3_PLE_32 batch must contain at least one row") + + target = torch.device(device) + # The historical source iterator converted FP8 to float32 before Python + # materialized each value. Keep that conversion boundary explicit before + # moving to float64 so the batched path has identical source semantics. + source = rows.to(dtype=torch.float32).to(device=target, dtype=torch.float64).reshape( + -1, BLOCKS_PER_ROW, BLOCK_VALUES + ) + if not bool(torch.isfinite(source).all().item()): + raise ValueError("Q3_PLE_32 cannot encode non-finite values") + + minimum = source.amin(dim=-1) + maximum = source.amax(dim=-1) + scale = torch.maximum(-minimum / 4.0, maximum / 3.0) + zero = scale == 0.0 + + def codes_for_scale(candidate: torch.Tensor) -> torch.Tensor: + safe = torch.where(candidate == 0.0, torch.ones_like(candidate), candidate) + codes = torch.round(source / safe.unsqueeze(-1)).clamp_(-4, 3).to(torch.int16) + 4 + return torch.where((candidate == 0.0).unsqueeze(-1), 4, codes) + + codes = codes_for_scale(scale) + done = zero.clone() + for _ in range(refinement_passes): + quants = codes.to(torch.int16) - 4 + denominator = ( + quants.to(torch.int64) * quants.to(torch.int64) + ).cumsum(dim=-1)[..., -1] + products = source * quants.to(torch.float64) + numerator = products.cumsum(dim=-1)[..., -1] + refined = torch.where( + denominator != 0, + numerator / denominator.to(torch.float64), + torch.zeros_like(numerator), + ) + valid = (~done) & (denominator != 0) & (refined > 0.0) & torch.isfinite(refined) + candidate = torch.where(valid, refined, torch.ones_like(refined)) + new_codes = codes_for_scale(candidate) + same = valid & (new_codes == codes).all(dim=-1) + codes = torch.where(valid.unsqueeze(-1), new_codes, codes) + scale = torch.where(valid, refined, scale) + done |= (~valid) | same + + # Match struct.pack('> 16) & 1)) >> 16 + ) & 0xFFFF + scale_bits = torch.where((scale > 0.0) & (scale_bits == 0), 1, scale_bits).to(torch.int32) + stored_scale = (scale_bits << 16).contiguous().view(torch.float32).to(torch.float64) + codes = codes_for_scale(stored_scale) + + output = torch.empty( + (source.shape[0], BLOCKS_PER_ROW, BLOCK_BYTES), + dtype=torch.uint8, + device=target, + ) + output[..., 0] = (scale_bits & 0xFF).to(torch.uint8) + output[..., 1] = ((scale_bits >> 8) & 0xFF).to(torch.uint8) + packed_codes = codes.to(torch.int64) + for group in range(4): + packed = torch.zeros_like(scale_bits, dtype=torch.int64) + for code_index in range(8): + packed |= packed_codes[..., group * 8 + code_index] << (3 * code_index) + byte_offset = 2 + group * 3 + output[..., byte_offset] = (packed & 0xFF).to(torch.uint8) + output[..., byte_offset + 1] = ((packed >> 8) & 0xFF).to(torch.uint8) + output[..., byte_offset + 2] = ((packed >> 16) & 0xFF).to(torch.uint8) + return output.contiguous().cpu().numpy().tobytes() + + def _unpack_codes(payload: bytes) -> list[int]: if len(payload) != 12: raise ValueError(f"Q3_PLE_32 code payload must be 12 bytes, got {len(payload)}") @@ -663,6 +763,8 @@ def write_q3_ple_segmented_sidecar( weight_scale: float, segment_count: int, rows_per_segment: int | None = None, + batched: bool = False, + quantization_device: str | torch.device | None = None, ) -> dict: """Write Q3 data with explicit logical source-segment boundaries. @@ -717,18 +819,33 @@ def write_q3_ple_segmented_sidecar( segment_offset = file_offset segment_digest = hashlib.sha256() segment_rows = 0 - for source_row in source_segment: - row_values = _materialize_row(source_row) - encoded_row = quantize_row(row_values, refinement_passes=REFINEMENT_PASSES) - if len(encoded_row) != ROW_BYTES: - raise AssertionError(f"Q3_PLE_32 row has wrong size: {len(encoded_row)}") - output.write(encoded_row) - whole_digest.update(encoded_row) - payload_digest.update(encoded_row) - segment_digest.update(encoded_row) - file_offset += len(encoded_row) - rows_written += 1 - segment_rows += 1 + for source_item in source_segment: + if batched: + if not isinstance(source_item, torch.Tensor): + raise ValueError("batched Q3_PLE_32 input items must be tensors") + batch_rows = int(source_item.shape[0]) if source_item.ndim == 2 else 0 + encoded = quantize_rows_batched( + source_item, + device=quantization_device or "cpu", + refinement_passes=REFINEMENT_PASSES, + ) + expected_bytes = batch_rows * ROW_BYTES + else: + row_values = _materialize_row(source_item) + encoded = quantize_row(row_values, refinement_passes=REFINEMENT_PASSES) + batch_rows = 1 + expected_bytes = ROW_BYTES + if len(encoded) != expected_bytes: + raise AssertionError( + f"Q3_PLE_32 batch has wrong size: {len(encoded)} != {expected_bytes}" + ) + output.write(encoded) + whole_digest.update(encoded) + payload_digest.update(encoded) + segment_digest.update(encoded) + file_offset += len(encoded) + rows_written += batch_rows + segment_rows += batch_rows if segment_rows <= 0: raise ValueError(f"Q3_PLE_32 logical segment {segment_index} is empty") if expected_rows is not None and segment_rows != expected_rows: @@ -804,11 +921,12 @@ def write_q3_ple_from_safetensors( layer_id: int, split_parts: int, source_fingerprint: str, - rows_per_chunk: int = 8192, + rows_per_chunk: int = DEFAULT_PROCESSING_CHUNK_ROWS, segment_rows: int | None = None, processing_chunk_rows: int | None = None, segment_count: int | None = None, rows_per_segment: int | None = None, + quantization_device: str | torch.device | None = None, ) -> dict: """Stream the official FP8 PLE shards into the native Q3 sidecar. @@ -822,7 +940,7 @@ def write_q3_ple_from_safetensors( if not folder.is_dir(): raise ValueError(f"Q3 PLE source must be a local checkpoint directory: {folder}") if processing_chunk_rows is not None: - if rows_per_chunk != 8192: + if rows_per_chunk != DEFAULT_PROCESSING_CHUNK_ROWS: raise ValueError("specify only one of rows_per_chunk or processing_chunk_rows") rows_per_chunk = processing_chunk_rows if rows_per_chunk <= 0 or split_parts <= 0: @@ -879,7 +997,7 @@ def write_q3_ple_from_safetensors( scale = handle.get_tensor(scale_key).reshape(()) weight_scale = float(scale.float().item()) - def iter_segment_rows(key: str): + def iter_segment_batches(key: str): source_file = folder / weight_map[key] with safetensors.safe_open(source_file, framework="pt", device="cpu") as handle: sliced = handle.get_slice(key) @@ -894,8 +1012,11 @@ def iter_segment_rows(key: str): chunk = sliced[start : min(start + int(rows_per_chunk), shape[0])] if chunk.dtype != torch.float8_e4m3fn: raise ValueError(f"unexpected PLE source dtype for {key}: {chunk.dtype}") - for row in chunk.float(): - yield row + yield chunk + + def iter_segment_rows(key: str): + for chunk in iter_segment_batches(key): + yield from chunk.float() # Explicit segmented mode is the production contract. Legacy callers can # request flat row segmentation by passing ``segment_rows`` explicitly. @@ -914,7 +1035,10 @@ def iter_rows(): logical_count = int(segment_count if segment_count is not None else split_parts) if logical_count != int(split_parts): raise ValueError("segment_count must equal split_parts for PLE Safetensors conversion") - ordered_segments = (iter_segment_rows(indexed_keys[index]) for index in range(logical_count)) + selected_device = quantization_device + if selected_device is None: + selected_device = "cuda" if torch.cuda.is_available() else "cpu" + ordered_segments = (iter_segment_batches(indexed_keys[index]) for index in range(logical_count)) return write_q3_ple_segmented_sidecar( ordered_segments, data_path, @@ -923,6 +1047,8 @@ def iter_rows(): weight_scale=weight_scale, segment_count=logical_count, rows_per_segment=rows_per_segment, + batched=True, + quantization_device=selected_device, ) @@ -931,6 +1057,7 @@ def iter_rows(): "BLOCK_BYTES", "BLOCK_VALUES", "DEFAULT_SEGMENT_ROWS", + "DEFAULT_PROCESSING_CHUNK_ROWS", "FORMAT", "REFINEMENT_PASSES", "Q3PLEReader", @@ -946,6 +1073,7 @@ def iter_rows(): "plan_q3_ple_production", "quantize_block", "quantize_row", + "quantize_rows_batched", "write_q3_ple_sidecar", "write_q3_ple_segmented_sidecar", "write_q3_ple_from_safetensors", diff --git a/python/freetoken/checkpoint/step9b_executor.py b/python/freetoken/checkpoint/step9b_executor.py index 9f5b4fac7..444c8dffb 100644 --- a/python/freetoken/checkpoint/step9b_executor.py +++ b/python/freetoken/checkpoint/step9b_executor.py @@ -1745,7 +1745,11 @@ def convert_and_validate_q3(self) -> dict[str, Any]: self._disk_gate() self._host_gate() self.target_root.mkdir(parents=True, exist_ok=True) - from freetoken.checkpoint.q3_ple import Q3PLEReader, write_q3_ple_from_safetensors + from freetoken.checkpoint.q3_ple import ( + DEFAULT_PROCESSING_CHUNK_ROWS, + Q3PLEReader, + write_q3_ple_from_safetensors, + ) data_path = self.target_root / "ple-q3-000.bin" manifest_path = self.target_root / "ple-q3.json" receipt_path = self.scratch_root / "receipts" / "B2-q3.json" @@ -1764,7 +1768,7 @@ def convert_and_validate_q3(self) -> dict[str, Any]: split_parts=PRODUCTION_PLE_SEGMENT_COUNT, source_fingerprint=self.source_inventory_fingerprint, rows_per_segment=2_500_012, - processing_chunk_rows=8192, + processing_chunk_rows=DEFAULT_PROCESSING_CHUNK_ROWS, ) if data_path.stat().st_size != Q3_BYTES: raise ExecutorError("Q3 target extent mismatch") diff --git a/tests/checkpoint/test_q3_ple_writer.py b/tests/checkpoint/test_q3_ple_writer.py index 2b9835913..0d90f0719 100644 --- a/tests/checkpoint/test_q3_ple_writer.py +++ b/tests/checkpoint/test_q3_ple_writer.py @@ -23,6 +23,8 @@ PRODUCTION_TOTAL_BYTES, PRODUCTION_TOTAL_ROWS, quantize_block, + quantize_row, + quantize_rows_batched, plan_q3_ple_production, write_q3_ple_from_safetensors, write_q3_ple_sidecar, @@ -117,6 +119,58 @@ def test_writer_payload_is_byte_identical_to_authoritative_reference(z_fixture_d assert manifest["file_bytes"] == len(expected) +def test_batched_quantizer_is_byte_identical_to_scalar_reference() -> None: + generator = torch.Generator().manual_seed(20260829) + random_rows = (torch.randn((8192, ROW_VALUES), generator=generator) * 12.0).to( + torch.float8_e4m3fn + ) + adversarial = torch.stack( + [ + torch.zeros(ROW_VALUES), + torch.arange(-80, 80, dtype=torch.float32) / 8.0, + torch.tensor(([0.5, -0.5, 1.5, -1.5] * 40), dtype=torch.float32), + torch.full((ROW_VALUES,), 448.0), + torch.full((ROW_VALUES,), -448.0), + ] + ).to(torch.float8_e4m3fn) + rows = torch.cat((random_rows, adversarial), dim=0) + expected = b"".join(quantize_row(row.float()) for row in rows) + + assert quantize_rows_batched(rows, device="cpu") == expected + if torch.cuda.is_available(): + assert quantize_rows_batched(rows, device="cuda") == expected + + +def test_batched_segmented_writer_matches_scalar_bytes(z_fixture_dir: Path) -> None: + rows = (torch.arange(7 * ROW_VALUES, dtype=torch.float32).reshape(7, ROW_VALUES) / 32.0).to( + torch.float8_e4m3fn + ) + scalar_manifest = write_q3_ple_segmented_sidecar( + ([row.float() for row in rows[:3]], [row.float() for row in rows[3:]]), + z_fixture_dir / "scalar.bin", + z_fixture_dir / "scalar.json", + source_fingerprint="9" * 64, + weight_scale=1.0, + segment_count=2, + ) + batched_manifest = write_q3_ple_segmented_sidecar( + ((rows[:2], rows[2:3]), (rows[3:6], rows[6:])), + z_fixture_dir / "batched.bin", + z_fixture_dir / "batched.json", + source_fingerprint="9" * 64, + weight_scale=1.0, + segment_count=2, + batched=True, + quantization_device="cuda" if torch.cuda.is_available() else "cpu", + ) + + assert (z_fixture_dir / "batched.bin").read_bytes() == ( + z_fixture_dir / "scalar.bin" + ).read_bytes() + for key in ("rows", "payload_bytes", "file_bytes", "sha256", "payload_sha256", "segments"): + assert batched_manifest[key] == scalar_manifest[key] + + def test_writer_consumes_rows_once_and_rejects_nonfinite_without_finalizing(z_fixture_dir: Path) -> None: data_path = z_fixture_dir / "ple-q3.bin" manifest_path = z_fixture_dir / "ple-q3.json" From 893f63df9e45e0ed4a67811d8788bfbeb210ee58 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 29 Aug 2026 05:25:00 -0400 Subject: [PATCH 17/17] fix(qwen4): preserve frozen active ftw extent --- ...KEN_STAGE7K_ACTIVE_FTW_TERMINAL_PADDING.md | 46 +++ docs/plans/FREETOKEN-QWEN4-001-STAGE7K.md | 22 ++ python/freetoken/checkpoint/__init__.py | 3 +- python/freetoken/checkpoint/ftw.py | 297 ++++++++++++++++++ .../freetoken/checkpoint/step9b_executor.py | 20 +- tests/checkpoint/test_ftw_terminal_padding.py | 251 +++++++++++++++ 6 files changed, 635 insertions(+), 4 deletions(-) create mode 100644 docs/FREETOKEN_STAGE7K_ACTIVE_FTW_TERMINAL_PADDING.md create mode 100644 docs/plans/FREETOKEN-QWEN4-001-STAGE7K.md create mode 100644 tests/checkpoint/test_ftw_terminal_padding.py diff --git a/docs/FREETOKEN_STAGE7K_ACTIVE_FTW_TERMINAL_PADDING.md b/docs/FREETOKEN_STAGE7K_ACTIVE_FTW_TERMINAL_PADDING.md new file mode 100644 index 000000000..d6a53a8a6 --- /dev/null +++ b/docs/FREETOKEN_STAGE7K_ACTIVE_FTW_TERMINAL_PADDING.md @@ -0,0 +1,46 @@ +# Stage 7K — active FTW terminal-padding remediation + +## Outcome + +Stage 7K preserves the frozen active FTW v1 extent of `4,804,403,200` bytes without changing a tensor or introducing artifact v2. The production tensor stream naturally ends at `4,804,399,104` bytes. The exact `4,096`-byte difference is represented as one zero-filled terminal compatibility page. + +## Root cause + +The accepted analytical layout included the two-byte global PLE scale in the active resident-byte calculation. FTW alignment turns that scalar into one 4-KiB page. Production correctly binds the global PLE scale once in `ple-q3.json`, so the active tensor stream contains no duplicate PLE-scale tensor and ends one aligned page earlier. + +The real candidate contains 1,698 tensor entries. Its final tensor ends exactly at byte `4,804,399,104`. Existing tensor payloads reconcile as: + +```text +packed NVFP4 weights 1,772,748,800 +FP8 block scales 221,593,600 +FP16 row globals 4,162,176 +protected tensors 2,804,402,200 +raw tensor bytes 4,802,906,776 +inter-entry alignment 1,492,328 +tensor-stream extent 4,804,399,104 +terminal compatibility page 4,096 +frozen active extent 4,804,403,200 +``` + +The page is not a tensor and does not represent a second PLE scale. FTW readers traverse indexed tensor entries and ignore the reserved tail. + +## Recovery contract + +`ensure_ftw_terminal_padding()` accepts only FTW v1 with 4-KiB alignment and only these two states: + +- the exact unpadded production extent, followed by appending and fsyncing one zero page; +- the exact frozen extent with an already published or recoverable zero page. + +It rejects nonterminal shard disagreement, nonzero tail bytes, partial pages, oversize tails, unknown extents, malformed geometry, and incompatible format/alignment. Publication order is shard append and fsync, then bounded Windows-safe atomic index replacement. A crash after append but before index replacement is recovered by validating and adopting the exact zero tail. Repeated recovery is idempotent. + +Only the final shard `nbytes` and index `total_bytes` change. Tensor keys, order, offsets, lengths, dtypes, shapes, kinds, and all pre-tail bytes remain unchanged. + +## Gate boundaries + +- Source bodies remain retained and are not downloaded again. +- The lifetime transfer ledger remains `135,252,480,565 / 135,252,480,565`. +- Q3 and all 48 expert sidecars remain unchanged. +- PR257 runtime commit `0307a6114c57b0efc61bc17688f3288fe0bf1dc7` remains unchanged. +- B5 and isolated C6 may proceed only after B4 publishes a valid receipt at the frozen extent. +- No full model, inference, generation, serving, or Step 10 work is authorized. + diff --git a/docs/plans/FREETOKEN-QWEN4-001-STAGE7K.md b/docs/plans/FREETOKEN-QWEN4-001-STAGE7K.md new file mode 100644 index 000000000..ad595671d --- /dev/null +++ b/docs/plans/FREETOKEN-QWEN4-001-STAGE7K.md @@ -0,0 +1,22 @@ +# FREETOKEN-QWEN4-001 Stage 7K plan + +Task: `STAGE7K-STEP9B-ACTIVE-FTW-TERMINAL-PADDING-REMEDIATION` + +Definition of done: preserve all acquired source files, Q3, and 48 expert sidecars; prove the active FTW discrepancy is exactly one terminal aligned page; make FTW v1 publication deterministically reach the frozen 4,804,403,200-byte extent without changing any tensor key, offset, length, dtype, or payload byte; recover B4; finalize B5; and pass isolated C6 with no network body transfer, full-model construction, inference, source deletion, runtime modification, or artifact v2. + +Dependencies: + +- Executor parent `3153f5e8f39a22aeb3d4283dba75336e988b1ba7`. +- Runtime `0307a6114c57b0efc61bc17688f3288fe0bf1dc7` remains clean and byte-identical. +- Acquisition manifest SHA-256 `8e4074cd1a8950bfb19ebdfdd4c5154b66db3ed538ba99fe41221d1be9361e74` remains unchanged. +- Lifetime transfer ledger remains `135252480565 / 135252480565`; all continuation work is body-disabled. + +| Step | Status | Validation | +|---|---|---| +| Freeze and audit the current 4,804,399,104-byte FTW candidate | DONE | Index, shard, tensor inventory, hashes, and 4-KiB reconciliation | +| Implement the smallest FTW-v1 terminal-padding contract | DONE | Focused unit tests; no tensor metadata or payload changes | +| Prove recovery/adoption of the existing active candidate | PENDING | Before/after tensor inventory and prefix hashes; exact 4,804,403,200 bytes | +| Run relevant executor, FTW, artifact, active, and regression tests | DONE | `148 passed`; `compileall`; `git diff --check` | +| Commit the Stage 7K executor and regenerate the resume handoff | PENDING | Exact parent/commit, clean worktree, pinned zero-body command | +| Resume body-disabled B4 recovery, B5 finalization, and isolated C6 | PENDING | Receipts, known-total reconciliation, C6 static-only PASS | +| Close out machine, source-retention, ledger, cache, and evidence audits | PENDING | 215 retained sources, unchanged ledger/runtime/manifest, reserve checks | diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index 60dec40be..46f86b481 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -7,6 +7,7 @@ from .ftw import ( FTWReader, FTWWriter, + ensure_ftw_terminal_padding, is_ftw_checkpoint, iter_ftw_weights, load_ftw_banks, @@ -22,7 +23,7 @@ ) __all__ = [ - "FTWReader", "FTWWriter", "is_ftw_checkpoint", + "FTWReader", "FTWWriter", "ensure_ftw_terminal_padding", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", "Q3PLEReader", "Q3PLESegment", "write_q3_ple_sidecar", "write_q3_ple_segmented_sidecar", "write_q3_ple_from_safetensors", diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index dcd6fe5be..d907e2dc3 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -39,6 +39,8 @@ import re import sys import threading +import time +import uuid from concurrent.futures import ThreadPoolExecutor import torch @@ -58,6 +60,7 @@ # Per-layer expert-bank entry name (converter streaming path, see checkpoint/convert.py): # each layer of a bank is its own FTW tensor instead of one flat [num_layers*E, ...] region. _LAYER_ENTRY_RE = re.compile(r"^(?P.+)#L(?P\d{5})$") +_WINDOWS_REPLACE_TRANSIENT_ERRORS = frozenset({5, 32, 33}) def layer_bank_entry_name(bank_name: str, layer_id: int) -> str: @@ -106,6 +109,300 @@ def is_ftw_checkpoint(path: str) -> bool: return os.path.isfile(os.path.join(path, INDEX_NAME)) +def _atomic_json_document(path: str, document: dict, *, attempts: int = 8, + deadline_seconds: float = 2.0, + backoff_seconds: float = 0.025) -> None: + """Durably publish one JSON document beside its destination. + + FTW terminal-padding recovery updates the index only after the shard tail is + durable. The replacement is therefore the final publication step. Keep the + previous index intact if replacement fails, and retain the temporary document so + a later recovery pass can inspect it. The bounded Windows retry mirrors the + executor's atomic receipt path without coupling this low-level format module to + the orchestration module. + """ + if attempts < 1: + raise ValueError("attempts must be positive") + if deadline_seconds < 0 or backoff_seconds < 0: + raise ValueError("retry timing must be non-negative") + parent = os.path.dirname(path) or "." + os.makedirs(parent, exist_ok=True) + temporary = os.path.join( + parent, + f".{os.path.basename(path)}.padding-{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}", + ) + with open(temporary, "w", encoding="utf-8", newline="\n") as handle: + json.dump(document, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + try: + os.fsync(handle.fileno()) + except OSError: + pass + + started = time.monotonic() + attempt = 0 + while True: + attempt += 1 + try: + os.replace(temporary, path) + return + except OSError as exc: + winerror = getattr(exc, "winerror", None) + if winerror is None: + winerror = getattr(exc, "errno", None) + transient = os.name == "nt" and winerror in _WINDOWS_REPLACE_TRANSIENT_ERRORS + if not transient: + raise + if attempt >= attempts or time.monotonic() - started >= deadline_seconds: + raise OSError( + f"FTW index atomic replacement failed after {attempt} attempts; " + f"destination={path}; preserved_temp={temporary}; winerror={winerror}" + ) from exc + delay = min( + backoff_seconds * (2 ** (attempt - 1)), + max(0.0, deadline_seconds - (time.monotonic() - started)), + ) + if delay: + time.sleep(delay) + + +def _read_ftw_index(path: str) -> tuple[str, dict]: + root = os.fspath(path) + if not os.path.isdir(root): + raise ValueError(f"FTW checkpoint directory is missing: {root}") + index_path = os.path.join(root, INDEX_NAME) + try: + with open(index_path, "r", encoding="utf-8") as handle: + index = json.load(handle) + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read FTW index: {index_path}") from exc + if not isinstance(index, dict): + raise ValueError("FTW index must be a JSON object") + try: + version = int(index.get("version", -1)) + except (TypeError, ValueError) as exc: + raise ValueError("FTW index version is invalid") from exc + if index.get("format") != FORMAT_TAG or version != FORMAT_VERSION: + raise ValueError("terminal padding requires FTW format version 1") + return index_path, index + + +def _remove_redundant_padding_indexes(index_path: str, published: dict) -> None: + """Remove only exact stale copies produced by terminal-padding publication. + + A failed Windows replace intentionally preserves its complete temporary JSON. Once + a later recovery publishes the same document, leaving that temporary file inside + the FTW directory would make it an unintended artifact component. Conflicting or + malformed candidates are never guessed away; they fail closed for operator review. + """ + parent = os.path.dirname(index_path) or "." + prefix = f".{os.path.basename(index_path)}.padding-" + redundant: list[str] = [] + for name in sorted(os.listdir(parent)): + if not name.startswith(prefix): + continue + candidate = os.path.join(parent, name) + try: + with open(candidate, "r", encoding="utf-8") as handle: + document = json.load(handle) + except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: + raise ValueError(f"conflicting FTW terminal-padding temp: {candidate}") from exc + if document != published: + raise ValueError(f"conflicting FTW terminal-padding temp: {candidate}") + redundant.append(candidate) + # Validate the complete candidate set before mutating it. A conflict therefore + # preserves every sibling document for deterministic operator inspection. + for candidate in redundant: + os.unlink(candidate) + + +def _validate_ftw_geometry(root: str, index: dict, *, align: int) -> tuple[list[dict], list[int], int]: + """Validate shard/tensor geometry without reading tensor payloads.""" + try: + indexed_align = int(index.get("align", -1)) + total = int(index.get("total_bytes", -1)) + except (TypeError, ValueError) as exc: + raise ValueError("FTW alignment/total geometry is invalid") from exc + if indexed_align != align: + raise ValueError("terminal padding requires the FTW alignment") + if total < 0 or total % align: + raise ValueError("FTW total_bytes must be a non-negative aligned integer") + shards = index.get("shards") + if not isinstance(shards, list) or not shards: + raise ValueError("FTW index has no shards") + if any(not isinstance(item, dict) for item in shards): + raise ValueError("FTW shard entry must be an object") + try: + ordered = sorted(shards, key=lambda item: int(item.get("global_off", -1))) + except (TypeError, ValueError) as exc: + raise ValueError("FTW shard global offset is invalid") from exc + if ordered != shards: + raise ValueError("FTW shards are not in global offset order") + expected_offset = 0 + physical_sizes: list[int] = [] + for shard in ordered: + if not isinstance(shard, dict): + raise ValueError("FTW shard entry must be an object") + name = shard.get("file") + try: + global_off = int(shard.get("global_off", -1)) + nbytes = int(shard.get("nbytes", -1)) + except (TypeError, ValueError) as exc: + raise ValueError("FTW shard geometry is invalid") from exc + if not isinstance(name, str) or not name or global_off != expected_offset: + raise ValueError("FTW shards are not contiguous") + if nbytes <= 0 or nbytes % align: + raise ValueError("FTW shard length is not aligned") + shard_path = os.path.join(root, name) + try: + physical = os.path.getsize(shard_path) + except OSError as exc: + raise ValueError(f"FTW shard is missing: {shard_path}") from exc + physical_sizes.append(physical) + expected_offset += nbytes + if expected_offset != total: + raise ValueError("FTW shard geometry does not reconcile with total_bytes") + + tensors = index.get("tensors") + if not isinstance(tensors, list): + raise ValueError("FTW index tensors must be a list") + names: set[str] = set() + for tensor in tensors: + if not isinstance(tensor, dict): + raise ValueError("FTW tensor entry must be an object") + name = tensor.get("name") + try: + global_off = int(tensor.get("global_off", -1)) + nbytes = int(tensor.get("nbytes", -1)) + except (TypeError, ValueError) as exc: + raise ValueError("FTW tensor geometry is invalid") from exc + if not isinstance(name, str) or not name or name in names: + raise ValueError("FTW tensor names must be unique") + names.add(name) + if global_off < 0 or global_off % align or nbytes < 0 or global_off + nbytes > total: + raise ValueError("FTW tensor geometry is invalid") + return ordered, physical_sizes, total + + +def _assert_zero_tail(path: str, pad_bytes: int) -> None: + with open(path, "rb") as handle: + handle.seek(-pad_bytes, os.SEEK_END) + tail = handle.read(pad_bytes) + if len(tail) != pad_bytes or any(tail): + raise ValueError("FTW terminal padding is not an all-zero page") + + +def ensure_ftw_terminal_padding( + path: str, + *, + expected_unpadded_bytes: int, + target_bytes: int, + pad_bytes: int = ALIGN, +) -> dict: + """Adopt or append one deterministic terminal FTW-v1 padding page. + + ``expected_unpadded_bytes`` and ``target_bytes`` are explicit so callers cannot + silently pad an arbitrary extent. The only accepted discrepancy is exactly one + ``pad_bytes`` page. Existing tensor entries and all bytes before the terminal + page are left untouched. If a process crashed after appending the page but before + publishing the index, the physical all-zero tail is adopted idempotently. + """ + root = os.fspath(path) + if pad_bytes <= 0 or pad_bytes % ALIGN: + raise ValueError("pad_bytes must be a positive FTW alignment multiple") + if expected_unpadded_bytes < 0 or expected_unpadded_bytes % pad_bytes: + raise ValueError("expected_unpadded_bytes must be non-negative and aligned") + if target_bytes < 0 or target_bytes % pad_bytes: + raise ValueError("target_bytes must be non-negative and aligned") + if target_bytes != expected_unpadded_bytes + pad_bytes: + raise ValueError("terminal padding target must be exactly one page above source") + index_path, index = _read_ftw_index(root) + shards, physical_sizes, indexed_total = _validate_ftw_geometry(root, index, align=pad_bytes) + if indexed_total not in (expected_unpadded_bytes, target_bytes): + raise ValueError( + f"unexpected FTW extent {indexed_total}; expected {expected_unpadded_bytes} " + f"or {target_bytes}" + ) + last = shards[-1] + last_name = last["file"] + last_path = os.path.join(root, last_name) + indexed_last_bytes = int(last["nbytes"]) + physical_last_bytes = physical_sizes[-1] + + # Every non-terminal shard must agree exactly with its index entry. Only the + # final shard may carry the one-page orphan tail while the old index is still + # published. + for shard, physical in zip(shards[:-1], physical_sizes[:-1]): + if physical != int(shard["nbytes"]): + raise ValueError("non-terminal FTW shard length does not match its index") + + if indexed_total == target_bytes: + if physical_last_bytes != indexed_last_bytes: + raise ValueError("padded FTW index does not match physical shard length") + _assert_zero_tail(last_path, pad_bytes) + _remove_redundant_padding_indexes(index_path, index) + return { + "state": "ALREADY_PADDED", + "index": index_path, + "target_bytes": target_bytes, + "pad_bytes": pad_bytes, + } + + # The old index may describe the original shard or a shard whose terminal page + # was durably appended immediately before a crash. Any other size is ambiguous + # and must fail closed instead of truncating or retransferring data. + adopted_orphan = physical_last_bytes == indexed_last_bytes + pad_bytes + if physical_last_bytes == indexed_last_bytes: + with open(last_path, "ab") as handle: + remaining = pad_bytes + zero_page = bytes(min(1 << 20, pad_bytes)) + while remaining: + chunk = min(remaining, len(zero_page)) + handle.write(zero_page[:chunk]) + remaining -= chunk + handle.flush() + try: + os.fsync(handle.fileno()) + except OSError: + pass + physical_last_bytes = os.path.getsize(last_path) + elif adopted_orphan: + _assert_zero_tail(last_path, pad_bytes) + else: + raise ValueError("FTW shard has a non-terminal or non-aligned extent discrepancy") + + if physical_last_bytes != indexed_last_bytes + pad_bytes: + raise ValueError("FTW terminal padding append did not reach the expected extent") + _assert_zero_tail(last_path, pad_bytes) + updated = dict(index) + updated_shards = [dict(shard) for shard in shards] + updated_last = dict(updated_shards[-1]) + updated_last["nbytes"] = indexed_last_bytes + pad_bytes + updated_shards[-1] = updated_last + updated["shards"] = updated_shards + updated["total_bytes"] = target_bytes + _atomic_json_document(index_path, updated) + + # Reopen the published index and verify that only the terminal metadata changed. + _published_path, published = _read_ftw_index(root) + _validate_ftw_geometry(root, published, align=pad_bytes) + if int(published.get("total_bytes", -1)) != target_bytes: + raise ValueError("FTW terminal padding index publication did not reach target") + published_shards = published.get("shards") + if published_shards[:-1] != shards[:-1] or published_shards[-1].get("global_off") != last.get("global_off"): + raise ValueError("FTW terminal padding changed non-terminal shard metadata") + if int(published_shards[-1].get("nbytes", -1)) != indexed_last_bytes + pad_bytes: + raise ValueError("FTW terminal padding shard length mismatch") + _remove_redundant_padding_indexes(index_path, published) + return { + "state": "RECOVERED" if adopted_orphan else "PADDED", + "index": index_path, + "target_bytes": target_bytes, + "pad_bytes": pad_bytes, + } + + # ============================== writer ============================== class FTWWriter: """Stream tensors into the FTW, rolling shard files at ``shard_limit``. diff --git a/python/freetoken/checkpoint/step9b_executor.py b/python/freetoken/checkpoint/step9b_executor.py index 444c8dffb..b82a6613b 100644 --- a/python/freetoken/checkpoint/step9b_executor.py +++ b/python/freetoken/checkpoint/step9b_executor.py @@ -1820,18 +1820,32 @@ def convert_and_validate_active(self) -> dict[str, Any]: result: dict[str, Any] = {} if not active.exists(): result = convert_checkpoint(str(self.source_root), str(self.target_root), artifact_format="qwen4_modular_v1", source_inventory_sha256=self.source_inventory_fingerprint) - from freetoken.checkpoint.ftw import INDEX_NAME + from freetoken.checkpoint.ftw import INDEX_NAME, ensure_ftw_terminal_padding index = active / INDEX_NAME if not active.is_dir() or not index.is_file(): raise ExecutorError("active FTW was not created") with index.open("r", encoding="utf-8") as handle: index_data = json.load(handle) - if int(index_data.get("total_bytes", -1)) != ACTIVE_BYTES: + active_bytes = int(index_data.get("total_bytes", -1)) + padding_result = None + if active_bytes in (ACTIVE_BYTES - 4096, ACTIVE_BYTES): + try: + padding_result = ensure_ftw_terminal_padding( + str(active), + expected_unpadded_bytes=ACTIVE_BYTES - 4096, + target_bytes=ACTIVE_BYTES, + ) + except (OSError, ValueError) as exc: + raise ExecutorError(f"active FTW terminal padding failed: {exc}") from exc + with index.open("r", encoding="utf-8") as handle: + index_data = json.load(handle) + active_bytes = int(index_data.get("total_bytes", -1)) + if active_bytes != ACTIVE_BYTES: raise ExecutorError("active FTW extent mismatch") tree_bytes, tree_sha = _tree_sha256(active) expected = {"stage": "B4", "format": "nvfp4_w4a16_v1", "target": str(active), "target_bytes": ACTIVE_BYTES, "target_tree_bytes": tree_bytes, "target_sha256": tree_sha, "source_inventory_fingerprint": self.source_inventory_fingerprint, "source_revision": self.manifest.revision, "builder_commit": self.builder_commit, "runtime_commit": self.runtime_commit, "source_inputs": self._source_bindings(self.manifest.rows_for_stage("B4"))} if not _receipt_matches(receipt_path, expected): - _publish_component_receipt(receipt_path, {**expected, "copied_metadata": list(result.get("copied_metadata", ())), "validation": {"ftw_index": True}, "recovered_after_promotion": not bool(result)}) + _publish_component_receipt(receipt_path, {**expected, "copied_metadata": list(result.get("copied_metadata", ())), "validation": {"ftw_index": True, "terminal_padding": padding_result}, "recovered_after_promotion": not bool(result)}) return {**expected, "state": "COMPLETE", "recovered_after_promotion": not bool(result)} def finalize_artifact(self) -> dict[str, Any]: diff --git a/tests/checkpoint/test_ftw_terminal_padding.py b/tests/checkpoint/test_ftw_terminal_padding.py new file mode 100644 index 000000000..612b5bfb7 --- /dev/null +++ b/tests/checkpoint/test_ftw_terminal_padding.py @@ -0,0 +1,251 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.checkpoint.ftw as ftw_module +import freetoken.checkpoint.step9b_executor as executor_module +from freetoken.checkpoint.ftw import ( + FTWReader, + FTWWriter, + ensure_ftw_terminal_padding, + iter_ftw_weights, +) +from freetoken.checkpoint.step9b_executor import Step9BExecutor + + +ALIGN = 4096 + + +def _fixture(root: Path, *, shard_limit: int = 1 << 20) -> tuple[Path, dict, bytes]: + root.mkdir() + writer = FTWWriter(str(root), shard_limit=shard_limit) + tensors = { + "first": torch.arange(32, dtype=torch.int32), + "second": torch.arange(15, dtype=torch.float32), + "third": torch.tensor([3, 1, 4, 1, 5], dtype=torch.bfloat16), + } + for name, tensor in tensors.items(): + writer.add_tensor(name, tensor) + index = writer.finalize({"source_inventory_sha256": "a" * 64}) + shard = root / index["shards"][-1]["file"] + return root, index, shard.read_bytes() + + +def _index(root: Path) -> dict: + return json.loads((root / "freetoken_weight.json").read_text(encoding="utf-8")) + + +def test_terminal_padding_preserves_tensor_inventory_prefix_and_reader(tmp_path: Path) -> None: + root, before, before_bytes = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + before_tensors = before["tensors"] + before_prefix = before_bytes + + result = ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + + after = _index(root) + shard = root / after["shards"][-1]["file"] + after_bytes = shard.read_bytes() + assert result["state"] == "PADDED" + assert after["total_bytes"] == target + assert after["tensors"] == before_tensors + assert after["shards"][:-1] == before["shards"][:-1] + assert after["shards"][-1]["global_off"] == before["shards"][-1]["global_off"] + assert after["shards"][-1]["nbytes"] == before["shards"][-1]["nbytes"] + ALIGN + assert after_bytes[: len(before_prefix)] == before_prefix + assert after_bytes[len(before_prefix) :] == b"\0" * ALIGN + + loaded = dict(iter_ftw_weights(str(root), workers=1)) + assert torch.equal(loaded["first"], torch.arange(32, dtype=torch.int32)) + assert torch.equal(loaded["second"], torch.arange(15, dtype=torch.float32)) + assert torch.equal(loaded["third"], torch.tensor([3, 1, 4, 1, 5], dtype=torch.bfloat16)) + reader = FTWReader(str(root)) + assert [entry["name"] for entry in reader.entries()] == ["first", "second", "third"] + reader.close() + + +def test_terminal_padding_repeated_recovery_is_idempotent(tmp_path: Path) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + first = ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + first_digest = hashlib.sha256((root / "freetoken-00000.ftw").read_bytes()).hexdigest() + second = ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + second_digest = hashlib.sha256((root / "freetoken-00000.ftw").read_bytes()).hexdigest() + assert first["state"] == "PADDED" + assert second["state"] == "ALREADY_PADDED" + assert first_digest == second_digest + assert _index(root)["total_bytes"] == target + + +def test_terminal_padding_adopts_durable_orphan_tail(tmp_path: Path) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + shard = root / before["shards"][-1]["file"] + with shard.open("ab") as handle: + handle.write(b"\0" * ALIGN) + handle.flush() + result = ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + assert result["state"] == "RECOVERED" + assert _index(root)["total_bytes"] == target + assert shard.stat().st_size == target + + +@pytest.mark.parametrize("extra", [1, ALIGN - 1, ALIGN + 1, ALIGN * 2]) +def test_terminal_padding_rejects_non_exact_or_oversize_tail(tmp_path: Path, extra: int) -> None: + root, before, _ = _fixture(tmp_path / f"active-{extra}") + expected = int(before["total_bytes"]) + shard = root / before["shards"][-1]["file"] + with shard.open("ab") as handle: + handle.write(b"\0" * extra) + with pytest.raises(ValueError, match="extent discrepancy"): + ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=expected + ALIGN + ) + assert _index(root)["total_bytes"] == expected + + +def test_terminal_padding_rejects_nonzero_orphan_tail(tmp_path: Path) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + shard = root / before["shards"][-1]["file"] + with shard.open("ab") as handle: + handle.write(b"\0" * (ALIGN - 1) + b"x") + with pytest.raises(ValueError, match="all-zero"): + ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=expected + ALIGN + ) + assert _index(root)["total_bytes"] == expected + + +def test_terminal_padding_rejects_corrupt_target_tail(tmp_path: Path) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + ensure_ftw_terminal_padding(str(root), expected_unpadded_bytes=expected, target_bytes=target) + shard = root / _index(root)["shards"][-1]["file"] + with shard.open("r+b") as handle: + handle.seek(-1, 2) + handle.write(b"x") + with pytest.raises(ValueError, match="all-zero"): + ensure_ftw_terminal_padding(str(root), expected_unpadded_bytes=expected, target_bytes=target) + + +def test_terminal_padding_index_publication_failure_preserves_old_index_and_recovers( + tmp_path: Path, monkeypatch +) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + original_replace = ftw_module.os.replace + + def denied_replace(_source, _destination): + error = OSError(5, "sharing/access lock") + error.winerror = 5 + raise error + + monkeypatch.setattr(ftw_module.os, "replace", denied_replace) + with pytest.raises(OSError, match="atomic replacement failed"): + ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + assert _index(root)["total_bytes"] == expected + assert (root / "freetoken-00000.ftw").stat().st_size == target + assert list(root.glob(".freetoken_weight.json.padding-*")) + + monkeypatch.setattr(ftw_module.os, "replace", original_replace) + recovered = ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + assert recovered["state"] == "RECOVERED" + assert _index(root)["total_bytes"] == target + assert not list(root.glob(".freetoken_weight.json.padding-*")) + + +def test_terminal_padding_rejects_conflicting_stale_index_temp(tmp_path: Path) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + ensure_ftw_terminal_padding(str(root), expected_unpadded_bytes=expected, target_bytes=target) + stale = root / ".freetoken_weight.json.padding-conflict" + stale.write_text('{"format":"not-the-published-index"}\n', encoding="utf-8") + with pytest.raises(ValueError, match="conflicting FTW terminal-padding temp"): + ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + assert stale.is_file() + + +def test_terminal_padding_conflict_preserves_all_stale_candidates(tmp_path: Path) -> None: + root, before, _ = _fixture(tmp_path / "active") + expected = int(before["total_bytes"]) + target = expected + ALIGN + ensure_ftw_terminal_padding(str(root), expected_unpadded_bytes=expected, target_bytes=target) + published = _index(root) + valid = root / ".freetoken_weight.json.padding-aaa-valid" + valid.write_text(json.dumps(published), encoding="utf-8") + conflict = root / ".freetoken_weight.json.padding-zzz-conflict" + conflict.write_text('{"format":"conflict"}\n', encoding="utf-8") + with pytest.raises(ValueError, match="conflicting FTW terminal-padding temp"): + ensure_ftw_terminal_padding( + str(root), expected_unpadded_bytes=expected, target_bytes=target + ) + assert valid.is_file() + assert conflict.is_file() + + +def test_executor_b4_receipt_recovery_seam_pads_candidate(tmp_path: Path, monkeypatch) -> None: + source = tmp_path / "source" + source.mkdir() + target = tmp_path / "target" + scratch = tmp_path / "scratch" + base = ALIGN + frozen = base + ALIGN + + def fake_convert(_source: str, out_dir: str, **_kwargs): + active = Path(out_dir) / "qwen4-active-v1.ftw" + writer = FTWWriter(str(active), shard_limit=1 << 20) + writer.add_tensor("active.weight", torch.arange(8, dtype=torch.float32)) + index = writer.finalize({"source_inventory_sha256": "b" * 64}) + assert index["total_bytes"] == base + return {"copied_metadata": []} + + monkeypatch.setattr(executor_module, "ACTIVE_BYTES", frozen) + monkeypatch.setattr("freetoken.checkpoint.convert.convert_checkpoint", fake_convert) + executor = object.__new__(Step9BExecutor) + executor.execute = True + executor.source_root = source + executor.target_root = target + executor.scratch_root = scratch + executor.source_inventory_fingerprint = "b" * 64 + executor.builder_commit = "c" * 40 + executor.runtime_commit = "d" * 40 + executor.manifest = SimpleNamespace(revision="e" * 40, rows_for_stage=lambda _stage: ()) + executor._disk_gate = lambda: None + executor._host_gate = lambda: None + executor._source_bindings = lambda _rows: [] + + result = Step9BExecutor.convert_and_validate_active(executor) + assert result["state"] == "COMPLETE" + assert result["target_bytes"] == frozen + index = _index(target / "qwen4-active-v1.ftw") + assert index["total_bytes"] == frozen + receipt = json.loads((scratch / "receipts" / "B4-active.json").read_text()) + assert receipt["completion"] == "COMPONENT_COMPLETE" + assert receipt["validation"]["terminal_padding"]["state"] == "PADDED"