From bc13dd709b4b0fe20aa13211ba5630014ac24a2e Mon Sep 17 00:00:00 2001 From: "chucai.dzq" Date: Fri, 18 Sep 2026 13:44:15 +0800 Subject: [PATCH 1/5] fix(qwen): preserve PLE export precision and bound QSA memory --- src/mcore_bridge/model/gpts/qwen4_exp.py | 12 ++++++- .../model/modules/kernels/ple_kernels.py | 6 ++++ src/mcore_bridge/model/modules/ple.py | 6 ++-- src/mcore_bridge/model/modules/qsa_indexer.py | 27 ++++++++------ tests/test_ple_checkpoint.py | 35 +++++++++++++------ tests/test_qsa_indexer.py | 18 ++++++++++ 6 files changed, 79 insertions(+), 25 deletions(-) diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index aa0c4dc..fbd5ab7 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -176,7 +176,7 @@ def _qsa_select(self, hidden_states, attn_kwargs, position_ids=None): 'Use --padding_free false with context_parallel_size 1 to take the bool-mask path, ' f'or set {QSA_SPARSE_KERNEL_ENV}=0 to fall back to full attention.') if cp_size > 1 and getattr(self.config, 'cp_comm_type', None) != 'all_gather': - raise RuntimeError(f"QSA sparse selection with context_parallel_size={cp_size} requires " + raise RuntimeError(f'QSA sparse selection with context_parallel_size={cp_size} requires ' f"cp_comm_type='all_gather' (got {getattr(self.config, 'cp_comm_type', None)!r}): the " 'selection has to see every key before attention runs, which ring/p2p cannot provide.') rotary_pos_emb = attn_kwargs.get('rotary_pos_emb') @@ -287,6 +287,16 @@ def __init__(self, *args, **kwargs): class Qwen4ExpBridge(Qwen3NextBridge): hf_mixer_prefix = 'model.' + def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> None: + # PLE export emits every shard. If it omits the scale, these are already + # dequantized parameters; copying the source FP8 scale would corrupt them. + accounted_keys = set(saved_keys) + suffix = 'ngram_embedding.shard_0.weight' + for key in saved_keys: + if key.endswith(f'ple.ple_embedding.{suffix}'): + accounted_keys.add(key[:-len(suffix)] + 'ngram_embedding.weight_scale') + super()._save_missing_weights(saver, accounted_keys, source_model_dir) + def _get_hf_experts_attr(self, is_mtp: bool = False): # The checkpoint stores experts as packed per-layer tensors # (`mlp.experts.gate_up_proj` / `mlp.experts.down_proj`). diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py index ef1f0c6..c66c75a 100644 --- a/src/mcore_bridge/model/modules/kernels/ple_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py @@ -477,13 +477,19 @@ def backward(ctx, dout): DIL=dilation, BLOCK_W=BW) + # The convolution backward has consumed the recomputed norm output. + del normed + # norm_conv backward: dwc on host, dx via kernel (fp32). x_hat = (gated.view(T, n, C) * rstdc.unsqueeze(-1)).view(T, W) dwc = (dnormed * x_hat).sum(dim=0).to(wc.dtype) + del x_hat dgated_norm = torch.empty(T, W, dtype=torch.float32, device=dev) if T > 0: _ple_norm_bwd_kernel[(T * n, )](gated, wc, rstdc, dnormed, dgated_norm, T, N=n, C=C, BLOCK_C=block_c) dgated += dgated_norm + # Release token-sized FP32 temporaries before gate gradient buffers. + del gated, dnormed, dgated_norm dkey = torch.empty_like(key) dquery = torch.empty_like(hc_state) diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index b1f67db..5cb75a6 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -265,8 +265,10 @@ def iter_export_table_to_hf(self, prefix=''): # scalar `weight_scale`, so divide by the scale stashed during loading # and cast back to fp8. Without a known scale the values cannot be # represented as fp8 + scale; keep the current dtype and warn. - scale = getattr(self, '_ngram_weight_scale', None) - if scale is None: + # Updated BF16 parameters cannot be losslessly requantized with the + # original checkpoint scale. Only immutable host tables reuse it. + scale = getattr(self, '_ngram_weight_scale', None) if self.cpu_offload else None + if scale is None and self.cpu_offload: get_logger().warning(f'`{self._NGRAM_SCALE_KEY}` was not seen during loading; exporting the PLE ngram ' 'embedding without re-quantizing to fp8.') # Reduce on GPU: NCCL has no CPU backend, and the host table is pinned diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py index 85595b2..d8db862 100644 --- a/src/mcore_bridge/model/modules/qsa_indexer.py +++ b/src/mcore_bridge/model/modules/qsa_indexer.py @@ -4,7 +4,7 @@ from megatron.core.extensions.transformer_engine import TELinear from torch import nn -# Byte budget for the transient score tile in select_token_indices_thd: the (token, block) +# Byte budget for the transient score tile in packed and unpacked selection: the (token, block) # scoring is chunked over queries to avoid OOM from the full [T, n_heads, NB] fp32 tensor. # Tests monkeypatch this to compare one chunk (un-chunked reference) vs many chunks. _QSA_INDEX_SCORE_CHUNK_BYTES = 1024 * 1024 * 1024 @@ -180,18 +180,23 @@ def apply_rope(t, cos_, sin_): starts = torch.arange(max_blocks, device=device) * R block_keys = apply_rope(pooled, cos[:, starts], sin[:, starts]) # [b, nb, d] - # ---- score all (query, block) pairs ---- - scores = torch.einsum('bqhd,bkd->bqhk', q.float(), block_keys.float()) - scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim) # [b, s, nb] - - # ---- restrict to blocks fully inside the causal prefix ---- - n_blocks = (torch.arange(s, device=device) + 1) // R # [s] + # Bound score workspace by queries, preserving all candidate blocks. + chunk_size = max(1, min(s, _QSA_INDEX_SCORE_CHUNK_BYTES // max(1, b * self.index_n_heads * max_blocks * 4))) + n_blocks = (torch.arange(s, device=device) + 1) // R block_ids = torch.arange(max_blocks, device=device) - scores = scores.masked_fill((block_ids[None, :] >= n_blocks[:, None])[None], float('-inf')) - k = min(self.block_topk, max_blocks) - top_blocks = scores.topk(k, dim=-1).indices # [b, s, k] - keep = top_blocks < n_blocks[None, :, None] # drop the -inf padding slots + top_blocks = torch.empty((b, s, k), dtype=torch.long, device=device) + keep = torch.empty((b, s, k), dtype=torch.bool, device=device) + keys_float = block_keys.float() + for start in range(0, s, chunk_size): + end = min(start + chunk_size, s) + scores = torch.einsum('bqhd,bkd->bqhk', q[:, start:end].float(), keys_float) + scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim) + scores = scores.masked_fill((block_ids[None, :] >= n_blocks[start:end, None])[None], float('-inf')) + selected = scores.topk(k, dim=-1).indices + top_blocks[:, start:end] = selected + keep[:, start:end] = selected < n_blocks[None, start:end, None] + del scores, selected return top_blocks, keep, n_blocks @torch.no_grad() diff --git a/tests/test_ple_checkpoint.py b/tests/test_ple_checkpoint.py index f131f64..47fcba5 100644 --- a/tests/test_ple_checkpoint.py +++ b/tests/test_ple_checkpoint.py @@ -64,11 +64,11 @@ def __init__(self, table, dim): @pytest.mark.parametrize('tp_size', [1, 2, 8]) @pytest.mark.parametrize('with_scale', [False, True]) -def test_export_updated_table(monkeypatch, tp_group, tp_size, with_scale): +def test_export_updated_table(monkeypatch, tmp_path, tp_group, tp_size, with_scale): """Combine actual contributions from every simulated rank into HF shards.""" total, dim, parts = 32, 4, 3 full_weight = (torch.arange(total * dim).reshape(total, dim) % 13 - 6).to(torch.bfloat16) - expected = full_weight + 1 + expected = full_weight + 0.0078125 tp_group.size.return_value = tp_size contributions = [[] for _ in range(parts)] shard_index = 0 @@ -90,9 +90,9 @@ def all_reduce(tensor, group): table = _Table(full_weight, tp_group, parts) if with_scale: table._ngram_weight_scale = torch.tensor(0.5) - # Stand in for a training update; values remain exactly representable in FP8. + # Preserve a small BF16 update that the original FP8 scale would round away. with torch.no_grad(): - table.ngram_embedding.weight.add_(1) + table.ngram_embedding.weight.add_(0.0078125) shard_index = 0 exported = {} table.export_table_to_hf(exported, prefix=prefix) @@ -101,17 +101,30 @@ def all_reduce(tensor, group): assert exported == {} weights = [exported[f'{prefix}ple.ple_embedding.ngram_embedding.shard_{i}.weight'] for i in range(parts)] - assert len(exported) == parts + int(with_scale) + assert len(exported) == parts assert [w.shape[0] for w in weights] == [11, 11, 10] - assert all(w.dtype == (torch.float8_e4m3fn if with_scale else torch.bfloat16) for w in weights) + assert all(w.dtype == torch.bfloat16 for w in weights) restored = torch.cat([w.float() for w in weights]) - if with_scale: - scale = exported[f'{prefix}{table._NGRAM_SCALE_KEY}'] - assert scale.shape == () - assert scale.item() == 0.5 - restored *= scale torch.testing.assert_close(restored, expected.float(), rtol=0, atol=0) + # Exercise the real missing-weight copy path against an FP8 source scale. + from safetensors.torch import save_file + + import mcore_bridge.bridge.gpt_bridge as gb + from mcore_bridge.model.gpts.qwen4_exp import Qwen4ExpBridge + + source = {'visual.weight': torch.ones(2)} + if with_scale: + source[f'{prefix}{table._NGRAM_SCALE_KEY}'] = torch.tensor(0.5) + save_file(source, str(tmp_path / 'model.safetensors')) + monkeypatch.setattr(gb, 'is_master', lambda: True) + bridge = object.__new__(Qwen4ExpBridge) + saver = Mock() + saver.add_tensor.side_effect = exported.__setitem__ + bridge._save_missing_weights(saver, set(exported), str(tmp_path)) + assert f'{prefix}{table._NGRAM_SCALE_KEY}' not in exported + torch.testing.assert_close(exported['visual.weight'], source['visual.weight']) + @pytest.mark.parametrize('tp_size', [8, 2]) def test_ple_sharded_state_dict_tp_layouts(tp_group, tp_size): diff --git a/tests/test_qsa_indexer.py b/tests/test_qsa_indexer.py index 419aa50..886495b 100644 --- a/tests/test_qsa_indexer.py +++ b/tests/test_qsa_indexer.py @@ -148,3 +148,21 @@ def test_rotate_half_matches_reference(): x1, x2 = x[..., :4], x[..., 4:] torch.testing.assert_close(got, torch.cat((-x2, x1), dim=-1)) + + +@pytest.mark.parametrize('ties', [False, True]) +def test_qsa_unpacked_chunked_matches_single_chunk_bitwise(monkeypatch, ties): + import mcore_bridge.model.modules.qsa_indexer as qi + device = 'cuda' if torch.cuda.is_available() else 'cpu' + torch.manual_seed(37) + idx, cfg = _make_idx(compress_ratio=4, budget=32, device=device) + hidden = torch.randn(73, 2, cfg.hidden_size, device=device) + if ties: + hidden.zero_() + freqs = torch.randn(73, 1, 1, cfg.indexer_head_dim, device=device) + monkeypatch.setattr(qi, '_QSA_INDEX_SCORE_CHUNK_BYTES', 1 << 62) + expected = idx._score_and_topk_blocks(hidden, freqs) + monkeypatch.setattr(qi, '_QSA_INDEX_SCORE_CHUNK_BYTES', 1024) + actual = idx._score_and_topk_blocks(hidden, freqs) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, rtol=0, atol=0) From 4374a56cc38dff9fb64b22c267831ef641be1ba2 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 18 Sep 2026 17:52:10 +0800 Subject: [PATCH 2/5] refactor(qwen): own checkpoint validation and export completion --- .../utils/qwen4_exp_checkpoint.py | 392 +++++++++++ tests/test_qwen4_exp_checkpoint.py | 631 ++++++++++++++++++ 2 files changed, 1023 insertions(+) create mode 100644 src/mcore_bridge/utils/qwen4_exp_checkpoint.py create mode 100644 tests/test_qwen4_exp_checkpoint.py diff --git a/src/mcore_bridge/utils/qwen4_exp_checkpoint.py b/src/mcore_bridge/utils/qwen4_exp_checkpoint.py new file mode 100644 index 0000000..04b73a8 --- /dev/null +++ b/src/mcore_bridge/utils/qwen4_exp_checkpoint.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate Qwen4-Exp PLE assets and complete HF text-training checkpoints.""" +from __future__ import annotations + +import hashlib +import json +import os +import re +import struct +import tempfile +import torch +from collections.abc import Collection +from contextlib import ExitStack +from copy import deepcopy +from pathlib import Path +from safetensors import safe_open +from safetensors.torch import save_file +from typing import Any + +_VISION_PREFIX = 'model.visual.' +_MTP_PREFIXES = ('mtp.', ) + + +def validate_ple_checkpoint( + path: str, + config: Any, + layers_prefix: str, + models: list[torch.nn.Module] | None = None, +) -> dict[str, dict[str, Any]]: + """Check PLE assets without materializing the large embedding tables. + + Hashes cover the small lookup buffers, not the table contents. Full checkpoint + file hashes belong to the experiment's source manifest. + """ + model_dir = Path(path) + index_path = model_dir / 'model.safetensors.index.json' + if index_path.is_file(): + with index_path.open() as stream: + weight_map = json.load(stream)['weight_map'] + else: + with safe_open(model_dir / 'model.safetensors', framework='pt', device='cpu') as handle: + weight_map = dict.fromkeys(handle.keys(), 'model.safetensors') + + heads = (config.ngram_size - 1) * config.heads_per_ngram + parts = config.split_ngram_parts + divisor = config.make_ngram_vocab_size_divisible_by + if heads <= 0 or parts <= 0 or divisor <= 0 or config.ple_embed_dim % heads: + raise ValueError('Invalid Qwen4-Exp PLE head, shard, or embedding dimensions.') + head_dim = config.ple_embed_dim // heads + manifest = {} + with ExitStack() as stack: + files = {} + + def tensor_handle(key: str): + if key not in weight_map: + raise ValueError(f"Missing required PLE checkpoint tensor: {key}") + filename = weight_map[key] + if filename not in files: + files[filename] = stack.enter_context(safe_open(model_dir / filename, framework='pt', device='cpu')) + handle = files[filename] + if key not in handle.keys(): + raise ValueError(f"PLE checkpoint index points to an absent tensor: {key}") + return handle + + for layer_id in config.ple_layer_ids: + prefix = f"{layers_prefix}.{layer_id - 1}.ple.ple_embedding." + buffers = {} + for name, length in ( + ('layer_multipliers', config.ngram_size), + ('ngram_heads_offsets', heads), + ('ngram_heads_vocab_sizes', heads), + ): + key = prefix + name + handle = tensor_handle(key) + metadata = handle.get_slice(key) + if metadata.get_shape() != [length] or metadata.get_dtype() != 'I64': + raise ValueError(f"Invalid PLE hash-buffer shape or dtype: {key}") + buffer = handle.get_tensor(key) + buffers[name] = buffer + manifest[key] = { + 'shape': [length], + 'dtype': 'I64', + 'sha256': hashlib.sha256(buffer.numpy().tobytes()).hexdigest(), + } + sizes = buffers['ngram_heads_vocab_sizes'] + offsets = buffers['ngram_heads_offsets'] + expected_offsets = torch.cat((sizes.new_zeros(1), sizes.cumsum(0)[:-1])) + if not bool(torch.all(sizes > 0)) or not torch.equal(offsets, expected_offsets): + raise ValueError(f"Invalid PLE hash-table sizes or offsets: {prefix}") + total = ((int(sizes.sum()) + divisor - 1) // divisor) * divisor + shard_size = (total + parts - 1) // parts + shard_dtypes = set() + for part in range(parts): + key = f"{prefix}ngram_embedding.shard_{part}.weight" + metadata = tensor_handle(key).get_slice(key) + expected_shape = [ + max(0, min(shard_size, total - part * shard_size)), + head_dim, + ] + shape = metadata.get_shape() + dtype = metadata.get_dtype() + if shape != expected_shape: + raise ValueError(f"Invalid PLE shard shape for {key}: expected {expected_shape}, got {shape}") + if dtype not in ('BF16', 'F16', 'F32', 'F8_E4M3'): + raise ValueError(f"Unsupported PLE shard dtype for {key}: {dtype}") + shard_dtypes.add(dtype) + manifest[key] = {'shape': shape, 'dtype': dtype} + if len(shard_dtypes) != 1: + raise ValueError(f"Mixed PLE shard dtypes: {prefix}") + scale_key = f"{prefix}ngram_embedding.weight_scale" + if 'F8_E4M3' in shard_dtypes or scale_key in weight_map: + handle = tensor_handle(scale_key) + metadata = handle.get_slice(scale_key) + if metadata.get_shape() not in ([], [1]): + raise ValueError(f"PLE weight scale must be scalar: {scale_key}") + scale = handle.get_tensor(scale_key).float() + if not bool(torch.all(torch.isfinite(scale) & (scale > 0))): + raise ValueError(f"PLE weight scale must be finite and positive: {scale_key}") + manifest[scale_key] = { + 'shape': metadata.get_shape(), + 'dtype': metadata.get_dtype(), + } + # These buffers are deterministic functions of config/seed. Check the local + # model before loading can overwrite a mismatched hash function from disk. + for model in models or []: + for layer in model.modules(): + ple = getattr(layer, 'ple', None) + if ple is None or not hasattr(layer, 'layer_number'): + continue + prefix = f"{layers_prefix}.{layer.layer_number - 1}.ple.ple_embedding." + for name in ( + 'layer_multipliers', + 'ngram_heads_offsets', + 'ngram_heads_vocab_sizes', + ): + buffer = getattr(ple.ple_embedding, name).detach().cpu().contiguous() + fingerprint = hashlib.sha256(buffer.numpy().tobytes()).hexdigest() + if fingerprint != manifest[prefix + name]['sha256']: + raise ValueError(f"PLE hash buffer disagrees with the model configuration: {prefix}{name}") + return manifest + + +def qwen4_exp_export_config(hf_config: Any, *, mtp_enabled: bool) -> Any: + """Return an export-only config; keep the running model's config untouched.""" + exported = deepcopy(hf_config) + text_config = getattr(exported, 'text_config', exported) + layer_types = getattr(text_config, 'layer_types', None) + if layer_types is not None: + # Newer Transformers normalizes this checkpoint spelling to QSA at + # load time. Preserve the original portable spelling when exporting: + # the pinned SGLang Qwen4-Exp loader selects its QSA implementation + # through "full_attention", not "qwen_sparse_attention". + text_config.layer_types = [ + 'full_attention' if kind == 'qwen_sparse_attention' else kind for kind in layer_types + ] + if not mtp_enabled: + text_config.mtp = None + text_config.mtp_num_hidden_layers = 0 + for config in (exported, text_config): + if hasattr(config, 'mtp'): + config.mtp = None + for field in ( + 'mtp_num_hidden_layers', + 'mtp_num_layers', + 'num_nextn_predict_layers', + ): + if hasattr(config, field): + setattr(config, field, 0) + return exported + + +def _checkpoint_inventory(directory: Path, ) -> tuple[dict[str, str], dict[str, dict[str, Any]], dict[str, Any]]: + """Read tensor headers and verify the index against files without loading weights.""" + index_path = directory / 'model.safetensors.index.json' + index = None + if index_path.is_file(): + if not index_path.resolve().is_relative_to(directory.resolve()): + raise ValueError(f"Checkpoint index escapes its directory: {index_path}") + with index_path.open() as stream: + index = json.load(stream) + if not isinstance(index.get('weight_map'), dict): + raise ValueError(f"Invalid safetensors weight_map: {index_path}") + for filename in index['weight_map'].values(): + if (not isinstance(filename, str) or Path(filename).is_absolute() or '..' in Path(filename).parts + or Path(filename).suffix != '.safetensors'): + raise ValueError(f"Invalid checkpoint shard path: {filename!r}") + filenames = {path.name for path in directory.glob('*.safetensors')} + if index is not None: + filenames.update(index['weight_map'].values()) + if not filenames: + raise ValueError(f"No safetensors weight files in {directory}") + weight_map = {} + tensors = {} + for filename in sorted(filenames): + path = directory / filename + if not path.resolve().is_relative_to(directory.resolve()): + raise ValueError(f"Checkpoint shard escapes its directory: {path}") + # safe_open validates the file format before offsets are used for the + # same byte-accounting convention as MegatronEngine's index rebuild. + with safe_open(path, framework='pt', device='cpu') as handle: + with path.open('rb') as stream: + header_size = struct.unpack(' dict[str, Any]: + """Validate an HF export and restore only missing ``model.visual.*`` tensors. + + Call on the saving rank after the bridge finished writing all tensor shards. + Callers own distributed error propagation and config/tokenizer preservation. + Each bucket is bounded by ``max_shard_size_bytes``; an indivisible larger + tensor gets its own shard. Existing tensors are never copied over. With MTP + disabled, only top-level ``mtp.*`` keys are allowed to be omitted; callers + may narrow that whitelist by providing exact ``omitted_mtp_keys``. + """ + source = Path(source_path).resolve() + output = Path(output_path).resolve() + if source == output: + raise ValueError('Source checkpoint and output checkpoint must differ.') + if max_shard_size_bytes <= 0: + raise ValueError('max_shard_size_bytes must be positive.') + with (source / 'config.json').open() as stream: + config = json.load(stream) + if config.get('model_type') != 'qwen4_exp': + raise ValueError('Fixed-asset restoration requires a qwen4_exp source checkpoint.') + source_map, source_tensors, _ = _checkpoint_inventory(source) + output_map, output_tensors, output_metadata = _checkpoint_inventory(output) + source_keys = set(source_map) + output_keys = set(output_map) + if omitted_mtp_keys is None: + omitted_mtp_keys = (set() if mtp_enabled else {key for key in source_keys if key.startswith(_MTP_PREFIXES)}) + else: + omitted_mtp_keys = set(omitted_mtp_keys) + if mtp_enabled and omitted_mtp_keys: + raise ValueError('Enabled MTP cannot have omitted checkpoint keys.') + if any(not key.startswith(_MTP_PREFIXES) for key in omitted_mtp_keys): + raise ValueError('The omission whitelist accepts exact MTP tensor keys only.') + if omitted_mtp_keys - source_keys: + raise ValueError(f"Omitted MTP keys are absent from the source: {sorted(omitted_mtp_keys - source_keys)}") + if omitted_mtp_keys & output_keys: + raise ValueError('The output contains MTP tensors explicitly declared omitted.') + unknown = output_keys - source_keys + if unknown: + raise ValueError(f"Export contains unknown checkpoint tensors: {sorted(unknown)}") + # The bridge exports dequantized PLE tables as current BF16 parameters. + # Their original FP8 scale must disappear, but only after every source + # shard has been exported with the same shape in the new representation. + ple_groups: dict[str, dict[int, str]] = {} + for key in source_keys: + match = re.fullmatch( + r'(model\.language_model\.layers\.\d+\.ple\.ple_embedding' + r'\.ngram_embedding)\.shard_(\d+)\.weight', + key, + ) + if match: + ple_groups.setdefault(match[1], {})[int(match[2])] = key + converted_ple_keys: set[str] = set() + omitted_ple_scales: set[str] = set() + for prefix, shards in ple_groups.items(): + if not any(source_tensors[key]['dtype'] == 'F8_E4M3' and output_tensors.get(key, {}).get('dtype') == 'BF16' + for key in shards.values()): + continue + if set(shards) != set(range(len(shards))) or any( + source_tensors[key]['dtype'] != 'F8_E4M3' or output_tensors.get(key, {}).get('dtype') != 'BF16' + or output_tensors[key]['shape'] != source_tensors[key]['shape'] for key in shards.values()): + raise ValueError(f"Incomplete or invalid BF16 PLE conversion: {prefix}") + scale_key = f"{prefix}.weight_scale" + if scale_key not in source_keys or scale_key in output_keys: + raise ValueError(f"BF16 PLE conversion requires removing the source scale: {scale_key}") + converted_ple_keys.update(shards.values()) + omitted_ple_scales.add(scale_key) + vision_keys = {key for key in source_keys if key.startswith(_VISION_PREFIX)} + missing = source_keys - output_keys - omitted_mtp_keys - omitted_ple_scales + copy_keys = missing & vision_keys if language_model_only else set() + missing_required = missing - copy_keys + if missing_required: + raise ValueError(f"Export is missing required non-restorable tensors: {sorted(missing_required)}") + for key in output_keys: + if output_tensors[key]['shape'] != source_tensors[key]['shape']: + raise ValueError(f"Export tensor shape differs from source: {key}") + if (output_tensors[key]['dtype'] != source_tensors[key]['dtype'] and key not in converted_ple_keys): + raise ValueError(f"Export tensor dtype differs from source: {key}: " + f"{source_tensors[key]['dtype']} -> {output_tensors[key]['dtype']}") + if language_model_only and not vision_keys: + raise ValueError('The source contains no recognized model.visual.* fixed assets.') + + bucket: dict[str, torch.Tensor] = {} + bucket_size = 0 + next_shard = 1 + new_shards = [] + pending_shards = [] + + def flush_bucket(cleanup: ExitStack) -> None: + nonlocal bucket_size, next_shard + if not bucket: + return + filename = f"model-fixed-visual-{next_shard:05d}.safetensors" + while (output / filename).exists(): + next_shard += 1 + filename = f"model-fixed-visual-{next_shard:05d}.safetensors" + with tempfile.NamedTemporaryFile( + dir=output, prefix='.fixed-visual-', suffix='.pending', delete=False) as stream: + temporary_path = Path(stream.name) + cleanup.callback(temporary_path.unlink, missing_ok=True) + save_file(bucket, temporary_path, metadata={'format': 'pt'}) + pending_shards.append((temporary_path, output / filename)) + for key in bucket: + output_map[key] = filename + output_tensors[key] = source_tensors[key] + new_shards.append(filename) + bucket.clear() + bucket_size = 0 + next_shard += 1 + + with ExitStack() as cleanup: + # Group reads by source shard. get_tensor touches only selected vision + # tensors, even when that file also stores trainable text weights. + for filename in sorted({source_map[key] for key in copy_keys}): + with safe_open(source / filename, framework='pt', device='cpu') as handle: + for key in sorted(key for key in copy_keys if source_map[key] == filename): + nbytes = source_tensors[key]['nbytes'] + if bucket and bucket_size + nbytes > max_shard_size_bytes: + flush_bucket(cleanup) + bucket[key] = handle.get_tensor(key) + bucket_size += nbytes + if bucket_size >= max_shard_size_bytes: + flush_bucket(cleanup) + flush_bucket(cleanup) + for temporary_path, shard_path in pending_shards: + os.replace(temporary_path, shard_path) + cleanup.callback(shard_path.unlink, missing_ok=True) + + # HF prefers model.safetensors over an index, so it must become a shard + # when extra files are added. The trained tensor bytes remain unchanged. + if ('model.safetensors' in output_map.values() and len(set(output_map.values())) > 1): + shard_number = 1 + filename = f"model-exported-{shard_number:05d}.safetensors" + while (output / filename).exists(): + shard_number += 1 + filename = f"model-exported-{shard_number:05d}.safetensors" + os.replace(output / 'model.safetensors', output / filename) + cleanup.callback(os.replace, output / filename, output / 'model.safetensors') + output_map = {key: filename if value == 'model.safetensors' else value for key, value in output_map.items()} + + total_size = sum(tensor['nbytes'] for tensor in output_tensors.values()) + output_metadata = dict(output_metadata, total_size=total_size) + index = { + 'metadata': output_metadata, + 'weight_map': dict(sorted(output_map.items())), + } + with tempfile.NamedTemporaryFile( + mode='w', + dir=output, + prefix='.model-index-', + suffix='.pending', + delete=False, + ) as stream: + temporary_path = Path(stream.name) + cleanup.callback(temporary_path.unlink, missing_ok=True) + json.dump(index, stream, indent=2) + stream.write('\n') + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_path, output / 'model.safetensors.index.json') + cleanup.pop_all() + return { + 'restored_keys': sorted(copy_keys), + 'omitted_mtp_keys': sorted(omitted_mtp_keys), + 'new_shards': new_shards, + 'total_size': total_size, + } diff --git a/tests/test_qwen4_exp_checkpoint.py b/tests/test_qwen4_exp_checkpoint.py new file mode 100644 index 0000000..89f285b --- /dev/null +++ b/tests/test_qwen4_exp_checkpoint.py @@ -0,0 +1,631 @@ +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import importlib.util +import json +import pytest +import torch +from pathlib import Path +from safetensors import safe_open +from safetensors.torch import save_file +from transformers import PretrainedConfig +from types import SimpleNamespace + +from mcore_bridge.utils import qwen4_exp_checkpoint as mcore_bridge_checkpoint +from mcore_bridge.utils.qwen4_exp_checkpoint import (qwen4_exp_export_config, restore_qwen4_exp_fixed_assets, + validate_ple_checkpoint) + +LAYERS_PREFIX = 'model.language_model.layers' +PLE_PREFIX = f"{LAYERS_PREFIX}.1.ple.ple_embedding." + + +def _write_weights(directory: Path, tensors: dict, *, indexed: bool = False) -> None: + directory.mkdir(exist_ok=True) + save_file(tensors, directory / 'model.safetensors', metadata={'format': 'pt'}) + if indexed: + (directory / 'model.safetensors.index.json').write_text( + json.dumps({ + 'metadata': { + 'total_size': -1, + 'test_metadata': 'preserved' + }, + 'weight_map': dict.fromkeys(tensors, 'model.safetensors'), + })) + + +@pytest.fixture +def checkpoints(tmp_path): + source = tmp_path / 'source' + output = tmp_path / 'output' + source_tensors = { + 'model.language_model.layers.0.mlp.experts.gate_up_proj': torch.ones(2, 2, dtype=torch.bfloat16), + 'lm_head.weight': torch.full((2, 2), 2.0, dtype=torch.bfloat16), + 'model.visual.blocks.0.attn.qkv.weight': torch.arange(4, dtype=torch.bfloat16).reshape(2, 2), + 'model.visual.patch_embed.proj.weight': torch.arange(12, dtype=torch.bfloat16).reshape(3, 4), + 'model.visual.pos_embed.weight': torch.arange(6, dtype=torch.float32), + 'model.visual.merger.bias': torch.tensor(2.0), + 'mtp.fc_hidden.weight': torch.ones(2, 2), + } + exported_tensors = { + 'model.language_model.layers.0.mlp.experts.gate_up_proj': torch.full((2, 2), 9.0, dtype=torch.bfloat16), + 'lm_head.weight': torch.full((2, 2), 7.0, dtype=torch.bfloat16), + } + _write_weights(source, source_tensors) + (source / 'config.json').write_text(json.dumps({'model_type': 'qwen4_exp'})) + _write_weights(output, exported_tensors) + return source, output, source_tensors, exported_tensors + + +def _read_indexed_weights(directory: Path) -> tuple[dict, dict]: + index = json.loads((directory / 'model.safetensors.index.json').read_text()) + tensors = {} + for filename in set(index['weight_map'].values()): + with safe_open(directory / filename, framework='pt', device='cpu') as handle: + for key in handle.keys(): + tensors[key] = handle.get_tensor(key) + return index, tensors + + +@pytest.mark.parametrize('indexed', [False, True]) +def test_restore_only_visual_assets_preserves_trained_text_and_index(checkpoints, indexed): + source, output, source_tensors, exported_tensors = checkpoints + _write_weights(source, source_tensors, indexed=indexed) + _write_weights(output, exported_tensors, indexed=indexed) + text_shard_hash = hashlib.sha256((output / 'model.safetensors').read_bytes()).hexdigest() + + report = restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True, max_shard_size_bytes=32) + + index, tensors = _read_indexed_weights(output) + assert set(tensors) == set(source_tensors) - {'mtp.fc_hidden.weight'} + assert report['omitted_mtp_keys'] == ['mtp.fc_hidden.weight'] + assert report['restored_keys'] == sorted(key for key in source_tensors if key.startswith('model.visual.')) + assert len(report['new_shards']) >= 2 + for filename in report['new_shards']: + shard_bytes = sum(tensors[key].numel() * tensors[key].element_size() + for key, shard_name in index['weight_map'].items() if shard_name == filename) + assert shard_bytes <= 32 + for key, value in tensors.items(): + expected = (exported_tensors[key] if key in exported_tensors else source_tensors[key]) + torch.testing.assert_close(value, expected, rtol=0, atol=0) + assert report['total_size'] == sum(tensor.numel() * tensor.element_size() for tensor in tensors.values()) + assert index['metadata']['total_size'] == report['total_size'] + if indexed: + assert index['metadata']['test_metadata'] == 'preserved' + # Transformers prioritizes the monolithic filename over an index. + assert not (output / 'model.safetensors').exists() + exported_filename = index['weight_map']['lm_head.weight'] + assert (hashlib.sha256((output / exported_filename).read_bytes()).hexdigest() == text_shard_hash) + assert not list(output.glob('*.pending')) + + +def test_restore_repeated_call_is_idempotent(checkpoints): + source, output, _, _ = checkpoints + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + first_files = {path.name: path.read_bytes() for path in output.iterdir()} + + report = restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + assert report['restored_keys'] == [] + assert report['new_shards'] == [] + assert {path.name: path.read_bytes() for path in output.iterdir()} == first_files + + +def test_restore_hf_sharded_loader_reads_text_and_visual_assets(checkpoints): + from transformers.trainer_utils import load_sharded_checkpoint + + source, output, source_tensors, exported_tensors = checkpoints + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + model = torch.nn.Module() + for key, tensor in source_tensors.items(): + if key.startswith('mtp.'): + continue + module = model + components = key.split('.') + for component in components[:-1]: + if not hasattr(module, component): + module.add_module(component, torch.nn.Module()) + module = getattr(module, component) + module.register_parameter(components[-1], torch.nn.Parameter(torch.zeros_like(tensor))) + + result = load_sharded_checkpoint(model, str(output), strict=True, prefer_safe=True) + + assert result.missing_keys == [] + assert result.unexpected_keys == [] + for key, tensor in model.state_dict().items(): + expected = (exported_tensors[key] if key in exported_tensors else source_tensors[key]) + torch.testing.assert_close(tensor, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + 'missing_key', + ['lm_head.weight', 'model.language_model.layers.0.mlp.experts.gate_up_proj'], +) +def test_restore_rejects_missing_text_before_writing(checkpoints, missing_key): + source, output, _, exported_tensors = checkpoints + exported_tensors.pop(missing_key) + _write_weights(output, exported_tensors) + before = {path.name: path.read_bytes() for path in output.iterdir()} + + with pytest.raises(ValueError, match='missing required non-restorable tensors'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + assert {path.name: path.read_bytes() for path in output.iterdir()} == before + + +def test_restore_does_not_copy_missing_ple_or_nested_mtp(checkpoints): + source, output, source_tensors, _ = checkpoints + source_tensors['model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight'] = torch.ones(2, 2) + source_tensors['model.language_model.mtp.fc_hidden.weight'] = torch.ones(2, 2) + _write_weights(source, source_tensors) + + with pytest.raises(ValueError, match='missing required non-restorable tensors'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + assert not list(output.glob('model-fixed-visual-*.safetensors')) + + +def test_restore_rejects_unknown_output_keys(checkpoints): + source, output, _, exported_tensors = checkpoints + exported_tensors['unexpected.weight'] = torch.ones(2) + _write_weights(output, exported_tensors) + + with pytest.raises(ValueError, match='unknown checkpoint tensors'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + +def test_restore_rejects_text_shape_change(checkpoints): + source, output, _, exported_tensors = checkpoints + exported_tensors['lm_head.weight'] = torch.ones(1, 2) + _write_weights(output, exported_tensors) + + with pytest.raises(ValueError, match='shape differs from source'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + +@pytest.mark.parametrize('dtype', [torch.int64, torch.bool, torch.float16]) +def test_restore_rejects_changed_tensor_dtype(checkpoints, dtype): + source, output, _, exported_tensors = checkpoints + exported_tensors['lm_head.weight'] = exported_tensors['lm_head.weight'].to(dtype) + _write_weights(output, exported_tensors) + + with pytest.raises(ValueError, match='dtype differs from source'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + assert not list(output.glob('model-fixed-visual-*.safetensors')) + + +def test_restore_preserves_source_fp32_buffer_dtype(checkpoints): + source, output, source_tensors, exported_tensors = checkpoints + key = 'model.language_model.layers.0.linear_attn.A_log' + source_tensors[key] = torch.ones(2, dtype=torch.float32) + exported_tensors[key] = torch.full((2, ), 2.0, dtype=torch.float32) + _write_weights(source, source_tensors) + _write_weights(output, exported_tensors) + + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + _, tensors = _read_indexed_weights(output) + assert tensors[key].dtype == torch.float32 + torch.testing.assert_close(tensors[key], exported_tensors[key], rtol=0, atol=0) + + +@pytest.mark.parametrize('bad_path', ['../source/model.safetensors', '/outside/model.safetensors']) +@pytest.mark.parametrize('side', ['source', 'output']) +def test_restore_rejects_nonlocal_index_shard_paths(checkpoints, bad_path, side): + source, output, source_tensors, exported_tensors = checkpoints + directory = source if side == 'source' else output + tensors = source_tensors if side == 'source' else exported_tensors + _write_weights(directory, tensors, indexed=True) + index_path = directory / 'model.safetensors.index.json' + index = json.loads(index_path.read_text()) + index['weight_map']['lm_head.weight'] = bad_path + index_path.write_text(json.dumps(index)) + + with pytest.raises(ValueError, match='Invalid checkpoint shard path'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + +@pytest.mark.parametrize('side', ['source', 'output']) +def test_restore_rejects_external_symlink_shard(checkpoints, side): + source, output, _, _ = checkpoints + directory, external = (source, output) if side == 'source' else (output, source) + (directory / 'external.safetensors').symlink_to(external / 'model.safetensors') + + with pytest.raises(ValueError, match='escapes its directory'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + +def test_restore_rejects_external_index_symlink(checkpoints): + source, output, source_tensors, _ = checkpoints + _write_weights(source, source_tensors, indexed=True) + (output / 'model.safetensors.index.json').symlink_to(source / 'model.safetensors.index.json') + + with pytest.raises(ValueError, match='Checkpoint index escapes'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + +def test_restore_visual_mode_requires_exporter_to_include_vision(checkpoints): + source, output, _, _ = checkpoints + + with pytest.raises(ValueError, match='missing required non-restorable tensors'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=False) + + +@pytest.mark.parametrize('mtp_enabled', [False, True]) +def test_restore_mtp_contract_is_explicit(checkpoints, mtp_enabled): + source, output, source_tensors, exported_tensors = checkpoints + if not mtp_enabled: + exported_tensors['mtp.fc_hidden.weight'] = source_tensors['mtp.fc_hidden.weight'] + _write_weights(output, exported_tensors) + message = ('declared omitted' if not mtp_enabled else 'missing required non-restorable') + + with pytest.raises(ValueError, match=message): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True, mtp_enabled=mtp_enabled) + + +def test_restore_mtp_whitelist_cannot_hide_missing_text(checkpoints): + source, output, _, _ = checkpoints + + with pytest.raises(ValueError, match='exact MTP tensor keys only'): + restore_qwen4_exp_fixed_assets( + str(source), + str(output), + language_model_only=True, + omitted_mtp_keys=['lm_head.weight'], + ) + + +def test_restore_rejects_index_that_disagrees_with_tensor_files(checkpoints): + source, output, _, exported_tensors = checkpoints + _write_weights(output, exported_tensors, indexed=True) + index_path = output / 'model.safetensors.index.json' + index = json.loads(index_path.read_text()) + index['weight_map']['ghost.weight'] = 'model.safetensors' + index_path.write_text(json.dumps(index)) + + with pytest.raises(ValueError, match='index disagrees'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + +def test_restore_oversized_tensor_gets_its_own_shard(checkpoints): + source, output, source_tensors, _ = checkpoints + source_tensors['model.visual.patch_embed.proj.weight'] = torch.ones(20) + _write_weights(source, source_tensors) + + report = restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True, max_shard_size_bytes=32) + + for filename in report['new_shards']: + with safe_open(output / filename, framework='pt', device='cpu') as handle: + sizes = [handle.get_tensor(key).numel() * handle.get_tensor(key).element_size() for key in handle.keys()] + assert sum(sizes) <= 32 or len(sizes) == 1 + + +def test_restore_staging_failure_preserves_existing_export(checkpoints, monkeypatch): + source, output, _, _ = checkpoints + original_save = mcore_bridge_checkpoint.save_file + before = {path.name: path.read_bytes() for path in output.iterdir()} + calls = 0 + + def fail_second_shard(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise OSError('disk full') + return original_save(*args, **kwargs) + + monkeypatch.setattr(mcore_bridge_checkpoint, 'save_file', fail_second_shard) + + with pytest.raises(OSError, match='disk full'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True, max_shard_size_bytes=32) + + assert {path.name: path.read_bytes() for path in output.iterdir()} == before + + +def test_restore_index_publish_failure_rolls_back_new_shards(checkpoints, monkeypatch): + source, output, _, _ = checkpoints + original_replace = mcore_bridge_checkpoint.os.replace + before = {path.name: path.read_bytes() for path in output.iterdir()} + + def fail_index_publish(src, dst): + if Path(dst).name == 'model.safetensors.index.json': + raise OSError('cannot publish index') + return original_replace(src, dst) + + monkeypatch.setattr(mcore_bridge_checkpoint.os, 'replace', fail_index_publish) + + with pytest.raises(OSError, match='cannot publish index'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + assert {path.name: path.read_bytes() for path in output.iterdir()} == before + + +@pytest.mark.parametrize('invalid', [None, 'missing', 'mixed', 'shape', 'stale_scale', 'non_ple']) +def test_restore_fp8_ple_conversion_requires_complete_bf16_export(checkpoints, invalid): + source, output, source_tensors, exported_tensors = checkpoints + prefix = 'model.language_model.layers.0.ple.ple_embedding.ngram_embedding' + shards = [f"{prefix}.shard_{index}.weight" for index in range(2)] + scale = f"{prefix}.weight_scale" + for key in shards: + source_tensors[key] = torch.ones(2, 2).to(torch.float8_e4m3fn) + exported_tensors[key] = torch.full((2, 2), 7.0, dtype=torch.bfloat16) + source_tensors[scale] = torch.tensor(0.5) + if invalid == 'missing': + del exported_tensors[shards[1]] + elif invalid == 'mixed': + exported_tensors[shards[1]] = source_tensors[shards[1]] + elif invalid == 'shape': + exported_tensors[shards[1]] = torch.ones(1, 2, dtype=torch.bfloat16) + elif invalid == 'stale_scale': + exported_tensors[scale] = source_tensors[scale] + elif invalid == 'non_ple': + source_tensors['lm_head.weight'] = source_tensors['lm_head.weight'].to(torch.float8_e4m3fn) + _write_weights(source, source_tensors) + _write_weights(output, exported_tensors) + + if invalid is not None: + with pytest.raises(ValueError, match='PLE conversion|dtype differs'): + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + assert not list(output.glob('model-fixed-visual-*.safetensors')) + return + + restore_qwen4_exp_fixed_assets(str(source), str(output), language_model_only=True) + + index, tensors = _read_indexed_weights(output) + assert scale not in index['weight_map'] + assert set(tensors) == set(source_tensors) - {scale, 'mtp.fc_hidden.weight'} + for key in shards: + assert tensors[key].dtype == torch.bfloat16 + torch.testing.assert_close(tensors[key], exported_tensors[key], rtol=0, atol=0) + + +@pytest.fixture +def ple_checkpoint(): + config = SimpleNamespace( + hf_model_type='qwen4_exp', + ngram_size=3, + heads_per_ngram=1, + ple_embed_dim=4, + split_ngram_parts=3, + make_ngram_vocab_size_divisible_by=4, + ple_layer_ids=[2], + ) + tensors = { + PLE_PREFIX + 'layer_multipliers': torch.tensor([11, 13, 17]), + PLE_PREFIX + 'ngram_heads_offsets': torch.tensor([0, 5]), + PLE_PREFIX + 'ngram_heads_vocab_sizes': torch.tensor([5, 7]), + PLE_PREFIX + 'ngram_embedding.weight_scale': torch.tensor(0.25), + } + for part in range(config.split_ngram_parts): + tensors[f"{PLE_PREFIX}ngram_embedding.shard_{part}.weight"] = ( + torch.arange(8, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn)) + return config, tensors + + +@pytest.mark.parametrize('indexed', [False, True]) +def test_ple_checkpoint_complete_assets_validate_without_loading_tables(tmp_path, ple_checkpoint, indexed): + config, tensors = ple_checkpoint + save_file(tensors, tmp_path / 'model.safetensors') + if indexed: + (tmp_path / 'model.safetensors.index.json').write_text( + json.dumps({'weight_map': dict.fromkeys(tensors, 'model.safetensors')})) + + manifest = validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX) + + assert set(manifest) == set(tensors) + buffer_key = PLE_PREFIX + 'layer_multipliers' + assert (manifest[buffer_key]['sha256'] == hashlib.sha256(tensors[buffer_key].numpy().tobytes()).hexdigest()) + shard_key = PLE_PREFIX + 'ngram_embedding.shard_0.weight' + assert manifest[shard_key] == {'shape': [4, 2], 'dtype': 'F8_E4M3'} + + +@pytest.mark.parametrize( + 'missing_key', + [ + 'ngram_embedding.shard_1.weight', + 'ngram_embedding.weight_scale', + 'layer_multipliers', + ], +) +def test_ple_checkpoint_missing_required_asset_raises(tmp_path, ple_checkpoint, missing_key): + config, tensors = ple_checkpoint + del tensors[PLE_PREFIX + missing_key] + save_file(tensors, tmp_path / 'model.safetensors') + + with pytest.raises(ValueError, match='Missing required PLE checkpoint tensor'): + validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX) + + +def test_ple_checkpoint_index_ghost_tensor_raises(tmp_path, ple_checkpoint): + config, tensors = ple_checkpoint + weight_map = dict.fromkeys(tensors, 'model.safetensors') + del tensors[PLE_PREFIX + 'ngram_embedding.shard_1.weight'] + save_file(tensors, tmp_path / 'model.safetensors') + (tmp_path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': weight_map})) + + with pytest.raises(ValueError, match='index points to an absent tensor'): + validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX) + + +@pytest.mark.parametrize( + ('key', 'replacement', 'message'), + [ + ('ngram_embedding.shard_1.weight', torch.ones(3, 2), 'Invalid PLE shard shape'), + ('ngram_heads_offsets', torch.tensor([0, 4]), 'Invalid PLE hash-table'), + ( + 'layer_multipliers', + torch.tensor([11, 13, 17], dtype=torch.int32), + 'hash-buffer', + ), + ( + 'ngram_embedding.weight_scale', + torch.tensor(float('nan')), + 'finite and positive', + ), + ('ngram_embedding.weight_scale', torch.tensor(0.0), 'finite and positive'), + ('ngram_embedding.weight_scale', torch.ones(2), 'must be scalar'), + ], +) +def test_ple_checkpoint_corrupt_metadata_raises(tmp_path, ple_checkpoint, key, replacement, message): + config, tensors = ple_checkpoint + tensors[PLE_PREFIX + key] = replacement + save_file(tensors, tmp_path / 'model.safetensors') + + with pytest.raises(ValueError, match=message): + validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX) + + +def test_ple_checkpoint_dequantized_table_without_scale_is_valid(tmp_path, ple_checkpoint): + config, tensors = ple_checkpoint + del tensors[PLE_PREFIX + 'ngram_embedding.weight_scale'] + for key in tensors: + if '.shard_' in key: + tensors[key] = tensors[key].to(torch.bfloat16) + save_file(tensors, tmp_path / 'model.safetensors') + + manifest = validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX) + + assert manifest[PLE_PREFIX + 'ngram_embedding.shard_0.weight']['dtype'] == 'BF16' + + +def test_ple_checkpoint_hash_must_match_model_configuration(tmp_path, ple_checkpoint): + config, tensors = ple_checkpoint + save_file(tensors, tmp_path / 'model.safetensors') + layer = torch.nn.Module() + layer.layer_number = 2 + layer.ple = torch.nn.Module() + layer.ple.ple_embedding = torch.nn.Module() + for key, value in tensors.items(): + name = key.removeprefix(PLE_PREFIX) + if '.' not in name: + layer.ple.ple_embedding.register_buffer(name, value.clone()) + validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX, [layer]) + layer.ple.ple_embedding.layer_multipliers[0] += 1 + + with pytest.raises(ValueError, match='disagrees with the model configuration'): + validate_ple_checkpoint(str(tmp_path), config, LAYERS_PREFIX, [layer]) + + +class _SourceConfig(PretrainedConfig): + model_type = 'qwen4_exp' + + +def _source_config(): + return _SourceConfig( + text_config=PretrainedConfig( + mtp={ + 'hybrid': True, + 'num_hidden_layers': 1, + 'layer_types': ['full_attention'], + }, + mtp_num_hidden_layers=1, + mtp_use_dedicated_embeddings=False, + )) + + +def test_export_config_removes_mtp_on_a_copy_only(tmp_path): + config = _source_config() + before = config.to_dict() + + exported = qwen4_exp_export_config(config, mtp_enabled=False) + exported.save_pretrained(tmp_path) + + saved = json.loads((tmp_path / 'config.json').read_text()) + assert saved['text_config']['mtp'] is None + assert saved['text_config']['mtp_num_hidden_layers'] == 0 + assert config.to_dict() == before + assert exported.text_config is not config.text_config + + +def test_enabled_mtp_export_preserves_configuration_without_aliasing(): + config = _source_config() + + exported = qwen4_exp_export_config(config, mtp_enabled=True) + + assert exported.to_dict() == config.to_dict() + exported.text_config.mtp['num_hidden_layers'] = 9 + assert config.text_config.mtp['num_hidden_layers'] == 1 + + +@pytest.mark.parametrize('mtp_enabled', [False, True]) +def test_export_qsa_uses_portable_layer_names_without_mutating_runtime(tmp_path, mtp_enabled): + config = _source_config() + config.text_config.layer_types = ['linear_attention', 'qwen_sparse_attention'] + before = config.to_dict() + + exported = qwen4_exp_export_config(config, mtp_enabled=mtp_enabled) + exported.save_pretrained(tmp_path) + + saved = json.loads((tmp_path / 'config.json').read_text()) + assert saved['text_config']['layer_types'] == [ + 'linear_attention', + 'full_attention', + ] + assert config.to_dict() == before + + +def test_exact_qwen4_exp_hf_tiny_export_config_has_no_mtp(tmp_path): + if importlib.util.find_spec('transformers.models.qwen4_exp') is None: + pytest.skip('Qwen4-Exp requires the pinned Transformers runtime') + from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpConfig + from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpForConditionalGeneration + + config = Qwen4ExpConfig( + text_config={ + 'vocab_size': 32, + 'hidden_size': 16, + 'num_hidden_layers': 4, + 'num_attention_heads': 2, + 'num_key_value_heads': 1, + 'head_dim': 8, + 'linear_key_head_dim': 8, + 'linear_value_head_dim': 8, + 'linear_num_key_heads': 1, + 'linear_num_value_heads': 1, + 'num_experts': 2, + 'num_experts_per_tok': 1, + 'moe_intermediate_size': 8, + 'shared_expert_intermediate_size': 8, + 'hc_count': 2, + 'hc_lowrank': 4, + 'ple_layer_ids': [2], + 'ple_embed_dim': 4, + 'heads_per_ngram': 1, + 'ngram_vocab_size_base': 11, + 'make_ngram_vocab_size_divisible_by': 4, + 'split_ngram_parts': 2, + 'eos_token_id': 2, + 'indexer_n_heads': 1, + 'indexer_kv_heads': 1, + 'indexer_head_dim': 8, + 'indexer_budget': 4, + 'indexer_compress_ratio': 2, + 'mtp': { + 'num_hidden_layers': 1 + }, + 'mtp_num_hidden_layers': 1, + 'rope_parameters': { + 'rope_type': 'default', + 'rope_theta': 10000.0 + }, + }, + vision_config={ + 'depth': 1, + 'hidden_size': 16, + 'intermediate_size': 16, + 'num_heads': 4, + 'patch_size': 2, + 'temporal_patch_size': 1, + 'out_hidden_size': 16, + 'num_position_embeddings': 16, + }, + ) + before = config.to_dict() + exported = qwen4_exp_export_config(config, mtp_enabled=False) + exported.save_pretrained(tmp_path) + reloaded = Qwen4ExpConfig.from_pretrained(tmp_path) + + model = Qwen4ExpForConditionalGeneration._from_config(reloaded, attn_implementation='eager') + + assert not any('mtp' in name.split('.') for name, _ in model.named_modules()) + assert reloaded.text_config.mtp is None + assert reloaded.text_config.mtp_num_hidden_layers == 0 + assert config.to_dict() == before From 4964a6fdcb3db9bb6bfa35c577e9ca6f917e6775 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 18 Sep 2026 22:40:04 +0800 Subject: [PATCH 3/5] fix: validate dequantized PLE exports across floating dtypes --- .../utils/qwen4_exp_checkpoint.py | 19 +++++++++++-------- tests/test_qwen4_exp_checkpoint.py | 14 +++++++++----- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/src/mcore_bridge/utils/qwen4_exp_checkpoint.py b/src/mcore_bridge/utils/qwen4_exp_checkpoint.py index 04b73a8..76a46d1 100644 --- a/src/mcore_bridge/utils/qwen4_exp_checkpoint.py +++ b/src/mcore_bridge/utils/qwen4_exp_checkpoint.py @@ -263,7 +263,7 @@ def restore_qwen4_exp_fixed_assets( unknown = output_keys - source_keys if unknown: raise ValueError(f"Export contains unknown checkpoint tensors: {sorted(unknown)}") - # The bridge exports dequantized PLE tables as current BF16 parameters. + # The bridge exports dequantized PLE tables in the parameter dtype. # Their original FP8 scale must disappear, but only after every source # shard has been exported with the same shape in the new representation. ple_groups: dict[str, dict[int, str]] = {} @@ -277,17 +277,20 @@ def restore_qwen4_exp_fixed_assets( ple_groups.setdefault(match[1], {})[int(match[2])] = key converted_ple_keys: set[str] = set() omitted_ple_scales: set[str] = set() + dequantized_dtypes = {'BF16', 'F16', 'F32'} for prefix, shards in ple_groups.items(): - if not any(source_tensors[key]['dtype'] == 'F8_E4M3' and output_tensors.get(key, {}).get('dtype') == 'BF16' - for key in shards.values()): + if not any(source_tensors[key]['dtype'] == 'F8_E4M3' + and output_tensors.get(key, {}).get('dtype') in dequantized_dtypes for key in shards.values()): continue - if set(shards) != set(range(len(shards))) or any( - source_tensors[key]['dtype'] != 'F8_E4M3' or output_tensors.get(key, {}).get('dtype') != 'BF16' - or output_tensors[key]['shape'] != source_tensors[key]['shape'] for key in shards.values()): - raise ValueError(f"Incomplete or invalid BF16 PLE conversion: {prefix}") + output_dtypes = {output_tensors.get(key, {}).get('dtype') for key in shards.values()} + if (len(output_dtypes) != 1 or not output_dtypes <= dequantized_dtypes + or set(shards) != set(range(len(shards))) + or any(source_tensors[key]['dtype'] != 'F8_E4M3' + or output_tensors[key]['shape'] != source_tensors[key]['shape'] for key in shards.values())): + raise ValueError(f"Incomplete or invalid floating-point PLE conversion: {prefix}") scale_key = f"{prefix}.weight_scale" if scale_key not in source_keys or scale_key in output_keys: - raise ValueError(f"BF16 PLE conversion requires removing the source scale: {scale_key}") + raise ValueError(f"Floating-point PLE conversion requires removing the source scale: {scale_key}") converted_ple_keys.update(shards.values()) omitted_ple_scales.add(scale_key) vision_keys = {key for key in source_keys if key.startswith(_VISION_PREFIX)} diff --git a/tests/test_qwen4_exp_checkpoint.py b/tests/test_qwen4_exp_checkpoint.py index 89f285b..6f5d2ba 100644 --- a/tests/test_qwen4_exp_checkpoint.py +++ b/tests/test_qwen4_exp_checkpoint.py @@ -338,22 +338,26 @@ def fail_index_publish(src, dst): assert {path.name: path.read_bytes() for path in output.iterdir()} == before -@pytest.mark.parametrize('invalid', [None, 'missing', 'mixed', 'shape', 'stale_scale', 'non_ple']) -def test_restore_fp8_ple_conversion_requires_complete_bf16_export(checkpoints, invalid): +@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize('invalid', [None, 'missing', 'mixed', 'mixed_float', 'shape', 'stale_scale', 'non_ple']) +def test_restore_fp8_ple_conversion_requires_complete_float_export(checkpoints, invalid, dtype): source, output, source_tensors, exported_tensors = checkpoints prefix = 'model.language_model.layers.0.ple.ple_embedding.ngram_embedding' shards = [f"{prefix}.shard_{index}.weight" for index in range(2)] scale = f"{prefix}.weight_scale" for key in shards: source_tensors[key] = torch.ones(2, 2).to(torch.float8_e4m3fn) - exported_tensors[key] = torch.full((2, 2), 7.0, dtype=torch.bfloat16) + exported_tensors[key] = torch.full((2, 2), 7.0, dtype=dtype) source_tensors[scale] = torch.tensor(0.5) if invalid == 'missing': del exported_tensors[shards[1]] elif invalid == 'mixed': exported_tensors[shards[1]] = source_tensors[shards[1]] + elif invalid == 'mixed_float': + other_dtype = torch.float32 if dtype != torch.float32 else torch.float16 + exported_tensors[shards[1]] = exported_tensors[shards[1]].to(other_dtype) elif invalid == 'shape': - exported_tensors[shards[1]] = torch.ones(1, 2, dtype=torch.bfloat16) + exported_tensors[shards[1]] = torch.ones(1, 2, dtype=dtype) elif invalid == 'stale_scale': exported_tensors[scale] = source_tensors[scale] elif invalid == 'non_ple': @@ -373,7 +377,7 @@ def test_restore_fp8_ple_conversion_requires_complete_bf16_export(checkpoints, i assert scale not in index['weight_map'] assert set(tensors) == set(source_tensors) - {scale, 'mtp.fc_hidden.weight'} for key in shards: - assert tensors[key].dtype == torch.bfloat16 + assert tensors[key].dtype == dtype torch.testing.assert_close(tensors[key], exported_tensors[key], rtol=0, atol=0) From 61517221aad0f3c3ca2f666f69132a781c616c02 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 20 Sep 2026 15:26:56 +0800 Subject: [PATCH 4/5] fix(qwen): use 64-bit PLE offsets for long sequences --- .../model/modules/kernels/ple_kernels.py | 15 +- tests/test_ple_kernels.py | 131 ++++++++++++++++++ 2 files changed, 139 insertions(+), 7 deletions(-) create mode 100644 tests/test_ple_kernels.py diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py index c66c75a..28fb9ac 100644 --- a/src/mcore_bridge/model/modules/kernels/ple_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py @@ -33,7 +33,7 @@ def _gather_ple_rows_from_pinned( # One program per flattened (token, hash-head) id. Rows outside this TP # rank's [row_start, row_end) are written as zero; the caller sums the # per-rank results across TP to reassemble the full embedding. - row_id = tl.program_id(0) + row_id = tl.program_id(0).to(tl.int64) global_idx = tl.load(ids_ptr + row_id) in_range = (global_idx >= row_start) & (global_idx < row_end) local_idx = tl.where(in_range, global_idx - row_start, 0) @@ -110,7 +110,8 @@ def _ple_gate_fwd_kernel( ): # Fused: grouped RMSNorm(key) * grouped RMSNorm(query) -> per-group score, # gate = sigmoid(sign(s)*sqrt(max(|s|,1e-6))), out = gate * value. - pid = tl.program_id(0) + # Promote before multiplying: 256K * (4 * 2560) exceeds int32. + pid = tl.program_id(0).to(tl.int64) t = pid // N c = pid % N if t >= T: @@ -162,7 +163,7 @@ def _ple_gate_bwd_kernel( SQRTC: tl.constexpr, BLOCK_C: tl.constexpr, ): - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) t = pid // N c = pid % N if t >= T: @@ -222,7 +223,7 @@ def _ple_norm_fwd_kernel( BLOCK_C: tl.constexpr, ): # Grouped zero-centered RMSNorm: out = x * rstd * (1 + w), fp32 out. - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) t = pid // N c = pid % N if t >= T: @@ -248,7 +249,7 @@ def _ple_norm_bwd_kernel( C: tl.constexpr, BLOCK_C: tl.constexpr, ): - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) t = pid // N c = pid % N if t >= T: @@ -281,7 +282,7 @@ def _ple_conv_fwd_kernel( ): # Causal dilated depthwise conv; rows never read across their segment # start. out = gated + silu(conv(normed)). - t = tl.program_id(0) + t = tl.program_id(0).to(tl.int64) wb = tl.program_id(1) if t >= T: return @@ -319,7 +320,7 @@ def _ple_conv_bwd_kernel( DIL: tl.constexpr, BLOCK_W: tl.constexpr, ): - t = tl.program_id(0) + t = tl.program_id(0).to(tl.int64) wb = tl.program_id(1) if t >= T: return diff --git a/tests/test_ple_kernels.py b/tests/test_ple_kernels.py new file mode 100644 index 0000000..1c0a1f0 --- /dev/null +++ b/tests/test_ple_kernels.py @@ -0,0 +1,131 @@ +"""PLE fused numerics and opt-in GPU regressions beyond the int32 boundary.""" +import importlib.util +import math +import os +import pytest +import torch +import torch.nn.functional as F +from pathlib import Path + +_KERNEL_PATH = Path(__file__).parents[1] / 'src/mcore_bridge/model/modules/kernels/ple_kernels.py' +_spec = importlib.util.spec_from_file_location('ple_kernels_under_test', _KERNEL_PATH) +kernels = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(kernels) +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not kernels.HAVE_TRITON, reason='PLE kernels require CUDA and Triton') + + +@pytest.mark.parametrize('rows,seq_len', [(1, 19), (2, 13)]) +def test_ple_forward_backward_match_torch(rows, seq_len): + torch.manual_seed(42) + n, c, k, dilation, eps = 2, 16, 4, 3, 1e-6 + t, w = rows * seq_len, n * c + shapes = [(t, w), (t, w), (t, c), (w, ), (w, ), (w, ), (w, 1, k)] + inputs = [(torch.randn(shape, device='cuda') * 0.2).requires_grad_() for shape in shapes] + query, key, value, wk, wq, wc, conv = inputs + + def norm(x, weight): + groups = x.reshape(t, n, c) + normalized = groups * torch.rsqrt(groups.square().mean(-1, keepdim=True) + eps) + return normalized * (1 + weight.reshape(n, c)) + + score = (norm(key, wk) * norm(query, wq)).sum(-1) / math.sqrt(c) + gate = torch.sigmoid(score.sign() * score.abs().clamp_min(1e-6).sqrt()) + gated = (gate[..., None] * value[:, None, :]).reshape(t, w) + normalized = norm(gated, wc).reshape(rows, seq_len, w).transpose(1, 2) + convolved = F.conv1d( + F.pad(normalized, (dilation * (k - 1), 0)), conv, dilation=dilation, groups=w).transpose(1, 2).reshape(t, w) + expected = gated + F.silu(convolved) + actual = kernels.ple_gate_conv_triton(*inputs, n, eps, dilation, seq_len) + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-4) + grad = torch.randn_like(actual) + actual_grads = torch.autograd.grad(actual, inputs, grad) + expected_grads = torch.autograd.grad(expected, inputs, grad) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close(actual_grad, expected_grad, atol=2e-4, rtol=2e-3) + + +@pytest.mark.skipif( + os.environ.get('MCORE_BRIDGE_TEST_LONG_PLE') != '1', + reason='Set MCORE_BRIDGE_TEST_LONG_PLE=1 on a GPU with at least 80 GiB free') +@pytest.mark.parametrize('stage', ['gate_fwd', 'gate_bwd', 'norm_fwd', 'norm_bwd', 'conv_fwd', 'conv_bwd']) +def test_ple_256k_offsets_cross_int32_boundary(stage): + # Actual model width: the first overflow is token 209715, channel 2048. + t, n, c, w = 262144, 4, 2560, 10240 + torch.cuda.empty_cache() + if torch.cuda.mem_get_info()[0] < 80 * 1024**3: + pytest.skip('Long PLE regression requires at least 80 GiB free GPU memory') + + def zeros(shape, dtype=torch.float32): + return torch.zeros(shape, dtype=dtype, device='cuda') + + def empty(shape, dtype=torch.float32): + return torch.empty(shape, dtype=dtype, device='cuda') + + def check(tensor, value=0): + torch.cuda.synchronize() + for index in [0, 209714, 209715, t - 1]: + row = tensor[index] + torch.testing.assert_close(row, torch.full_like(row, value), atol=1e-6, rtol=1e-6) + + weights = zeros(w, torch.bfloat16) + x = zeros((t, w), torch.bfloat16) + block = kernels.triton.next_power_of_2(c) + if stage == 'gate_fwd': + value = torch.ones((t, c), device='cuda', dtype=torch.bfloat16) + out, gate, rk, rq = empty((t, w)), empty((t, n)), empty((t, n)), empty((t, n)) + kernels._ple_gate_fwd_kernel[(t * n, )]( + x, x, value, weights, weights, out, gate, rk, rq, t, N=n, C=c, EPS=1e-6, SQRTC=math.sqrt(c), BLOCK_C=block) + check(out, 1 / (1 + math.exp(-0.001))) + check(gate, 1 / (1 + math.exp(-0.001))) + elif stage == 'gate_bwd': + dg, value, stats = zeros((t, w)), zeros((t, c), torch.bfloat16), torch.ones((t, n), device='cuda') + dk, dq = empty((t, w), torch.bfloat16), empty((t, w), torch.bfloat16) + dv, dwk, dwq = empty((t, w)), empty((t, w)), empty((t, w)) + kernels._ple_gate_bwd_kernel[(t * n, )]( + dg, + x, + x, + value, + weights, + weights, + stats, + stats, + stats, + dk, + dq, + dv, + dwk, + dwq, + t, + N=n, + C=c, + SQRTC=math.sqrt(c), + BLOCK_C=block) + for out in [dk, dq, dv, dwk, dwq]: + check(out) + elif stage == 'norm_fwd': + out, stats = empty((t, w)), empty((t, n)) + kernels._ple_norm_fwd_kernel[(t * n, )](x, weights, out, stats, t, N=n, C=c, EPS=1e-6, BLOCK_C=block) + check(out) + check(stats, 1000) + elif stage == 'norm_bwd': + stats, out = torch.ones((t, n), device='cuda'), empty((t, w)) + kernels._ple_norm_bwd_kernel[(t * n, )](x, weights, stats, x, out, t, N=n, C=c, BLOCK_C=block) + check(out) + else: + lo, hi = kernels._uniform_seg_bounds(t, t, 'cuda') + conv = zeros((w, 4), torch.bfloat16) + out, pre = empty((t, w)), zeros((t, w)) + grid = (t, kernels.triton.cdiv(w, 256)) + if stage == 'conv_fwd': + kernels._ple_conv_fwd_kernel[grid](x, x, conv, lo, out, pre, t, w, K=4, DIL=3, BLOCK_W=256) + check(out) + check(pre) + else: + dw, residual = zeros((w, 4)), empty((t, w)) + kernels._ple_conv_bwd_kernel[grid]( + x, pre, x, conv, lo, hi, out, dw, residual, t, w, K=4, DIL=3, BLOCK_W=256) + check(out) + check(residual) + torch.testing.assert_close(dw, torch.zeros_like(dw), atol=0, rtol=0) From 38c675ed356c3ab316da7ba9a88857952591bbc6 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 20 Sep 2026 15:46:20 +0800 Subject: [PATCH 5/5] fix(qwen): use 64-bit offsets for sparse attention --- .../modules/kernels/qsa_block_sparse_attn.py | 13 ++-- .../model/modules/kernels/qsa_kernels.py | 2 +- tests/test_qsa_sparse_kernels.py | 70 +++++++++++++++++++ 3 files changed, 78 insertions(+), 7 deletions(-) create mode 100644 tests/test_qsa_sparse_kernels.py diff --git a/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py index f0b8586..6231c0b 100644 --- a/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py +++ b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py @@ -78,7 +78,8 @@ def _qsa_bs_fwd_kernel( BK: tl.constexpr, BLK: tl.constexpr, ): - pid_t = tl.program_id(0) + # The selection bitmap exceeds 2 GiB at long context lengths. + pid_t = tl.program_id(0).to(tl.int64) pid_h = tl.program_id(1) kv_head = pid_h // GROUP @@ -105,7 +106,7 @@ def _qsa_bs_fwd_kernel( # over the tile's queries; per-query exactness comes from the mask below. n_tiles = tl.load(KCNT + pid_t) for i in range(0, n_tiles): - kt = tl.load(KLIST + pid_t * stride_kl + i) + kt = tl.load(KLIST + pid_t * stride_kl + i).to(tl.int64) offs_k = kt * BK + tl.arange(0, BK) k_in = offs_k < T @@ -193,7 +194,7 @@ def _qsa_bs_dq_kernel( BK: tl.constexpr, BLK: tl.constexpr, ): - pid_t = tl.program_id(0) + pid_t = tl.program_id(0).to(tl.int64) pid_h = tl.program_id(1) kv_head = pid_h // GROUP @@ -217,7 +218,7 @@ def _qsa_bs_dq_kernel( dq = tl.zeros((BQ, D), tl.float32) n_tiles = tl.load(KCNT + pid_t) for i in range(0, n_tiles): - kt = tl.load(KLIST + pid_t * stride_kl + i) + kt = tl.load(KLIST + pid_t * stride_kl + i).to(tl.int64) offs_k = kt * BK + tl.arange(0, BK) k_in = offs_k < T @@ -298,7 +299,7 @@ def _qsa_bs_dkdv_kernel( from separate programs would have them overwrite each other (the gather kernel got away with it only because it used atomic_add). """ - pid_k = tl.program_id(0) + pid_k = tl.program_id(0).to(tl.int64) kv_head = tl.program_id(1) offs_k = pid_k * BK + tl.arange(0, BK) @@ -314,7 +315,7 @@ def _qsa_bs_dkdv_kernel( n_q = tl.load(QCNT + pid_k) for i in range(0, n_q): - qt = tl.load(QLIST + pid_k * stride_ql + i) + qt = tl.load(QLIST + pid_k * stride_ql + i).to(tl.int64) offs_q = qt * BQ + tl.arange(0, BQ) q_mask = offs_q < T lo = tl.load(LO + offs_q, mask=q_mask, other=0) diff --git a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py index 485137c..fa57d75 100644 --- a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """QSA sparse attention wrappers around the vendored tensor-core triton kernel. -The kernel itself lives in ``qsa_block_sparse_attn.py``, vendored verbatim from +The kernel itself lives in ``qsa_block_sparse_attn.py``, adapted from miles PR #2777 (commit 0f5dff4). This file owns only the glue mcore needs: sbhd<->thd flattening, context parallelism, and the ``core_attention`` shim. diff --git a/tests/test_qsa_sparse_kernels.py b/tests/test_qsa_sparse_kernels.py new file mode 100644 index 0000000..f051301 --- /dev/null +++ b/tests/test_qsa_sparse_kernels.py @@ -0,0 +1,70 @@ +"""QSA sparse attention numerical and long-context address regressions.""" +import importlib.util +import os +import pytest +import torch +from pathlib import Path + +pytest.importorskip('triton') +_spec = importlib.util.spec_from_file_location( + 'qsa_sparse_under_test', + Path(__file__).parents[1] / 'src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py') +kernels = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(kernels) +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason='QSA kernels require CUDA') + + +@pytest.mark.parametrize('packed', [False, True]) +def test_qsa_forward_backward_match_dense_attention(packed): + torch.manual_seed(123) + t, h, d = 39, 2, 32 + q = (torch.randn(t, h, d, device='cuda', dtype=torch.bfloat16) * 0.2).requires_grad_() + k = (torch.randn(t, 1, d, device='cuda', dtype=torch.bfloat16) * 0.2).requires_grad_() + v = torch.randn_like(k, requires_grad=True) + pos = torch.arange(t, device='cuda') + sel = torch.zeros(t, (t + 3) // 4, device='cuda', dtype=torch.uint8) + sel[:, 0] = 1 + sel[pos, pos // 4] = 1 + lo = torch.where(pos >= 19, 19, 0).int() if packed else torch.zeros_like(pos, dtype=torch.int32) + hi = pos.int() + zero = torch.zeros_like(lo) + actual = kernels.qsa_block_sparse_attention_triton(q, k, v, sel, lo, hi, zero, zero, d**-0.5, 4) + allowed = sel[:, pos // 4].bool() & (pos[None, :] >= lo[:, None]) & (pos[None, :] <= hi[:, None]) + scores = torch.einsum('thd,shd->hts', q.float(), k.expand(-1, h, -1).float()) * d**-0.5 + probs = scores.masked_fill(~allowed[None], float('-inf')).softmax(-1) + expected = torch.einsum('hts,shd->thd', probs, v.expand(-1, h, -1).float()).to(q.dtype) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.008) + grad = torch.randn_like(actual) + actual_grads = torch.autograd.grad(actual, (q, k, v), grad) + expected_grads = torch.autograd.grad(expected, (q, k, v), grad) + for a, e in zip(actual_grads, expected_grads): + torch.testing.assert_close(a, e, rtol=0.04, atol=0.015) + + +@pytest.mark.skipif( + os.environ.get('MCORE_BRIDGE_TEST_LONG_QSA') != '1', + reason='Set MCORE_BRIDGE_TEST_LONG_QSA=1 for the 16 GiB bitmap regression') +def test_qsa_256k_bitmap_forward_backward_cross_int32_boundary(): + torch.cuda.empty_cache() + if torch.cuda.mem_get_info()[0] < 32 * 1024**3: + pytest.skip('Long QSA regression requires at least 32 GiB free GPU memory') + t, d = 262144, 16 + q = torch.zeros(t, 1, d, device='cuda', dtype=torch.bfloat16, requires_grad=True) + k = torch.zeros_like(q, requires_grad=True) + v = torch.ones_like(q, requires_grad=True) + # Every query selects the first four keys. The bitmap row stride remains + # the real 256K model stride, so rows >=32768 require >int32 addressing. + sel = torch.zeros(t, t // 4, device='cuda', dtype=torch.uint8) + sel[:, 0] = 1 + lo = torch.zeros(t, device='cuda', dtype=torch.int32) + hi = torch.full_like(lo, 3) + out = kernels.qsa_block_sparse_attention_triton(q, k, v, sel, lo, hi, lo, lo, d**-0.5, 4) + torch.cuda.synchronize() + torch.testing.assert_close(out, torch.ones_like(out), rtol=0, atol=0) + out.sum().backward() + torch.cuda.synchronize() + torch.testing.assert_close(q.grad, torch.zeros_like(q), rtol=0, atol=0) + torch.testing.assert_close(k.grad, torch.zeros_like(k), rtol=0, atol=0) + expected = torch.zeros_like(v) + expected[:4] = t / 4 + torch.testing.assert_close(v.grad, expected, rtol=0, atol=0)