vllm-omni/probe/issue-5/marlin_isolate.py

#!/usr/bin/env python
"""Issue 5 diagnostic: isolate the compressed-tensors W8A16 serve path per layer.

Builds vLLM parallel linear layers exactly as the qwen_image_21 transformer
does (fused QKV, row-parallel to_out, img_mlp proj/out), loads the int8
artifact's packed tensors through the production weight-loader machinery,
runs the Marlin repack (process_weights_after_loading) on the GPU, and
compares each layer's forward output against a float32 reference computed
from the artifact's own unpacked+dequantized weights.

Expected rel-L2 when the load->repack->apply chain is correct: ~1-2%
(int8 quantization error plus bf16 activation rounding). A large error
pins corruption inside the serve path with the checkpoint already
verified correct (verify_artifact.py).

Usage: marlin_isolate.py --artifact ARTIFACT_DIR [--block N]
"""

from __future__ import annotations

import argparse
import json
import pathlib

import torch
from compressed_tensors.compressors.pack_quantized.helpers import unpack_from_int32
from safetensors.torch import load_file


def dequantized_weight(art: dict, base: str) -> torch.Tensor:
    shape = tuple(art[f"{base}_shape"].tolist())
    q = unpack_from_int32(art[f"{base}_packed"].cpu(), 8, shape)
    return q.to(torch.float32) * art[f"{base}_scale"].cpu().to(torch.float32)


def load_tensor(layer: torch.nn.Module, param_name: str, tensor: torch.Tensor, shard: str | None) -> None:
    param = getattr(layer, param_name)
    if shard is None:
        param.weight_loader(param, tensor)
    else:
        param.weight_loader(param, tensor, shard)


def check(name: str, layer: torch.nn.Module, art: dict, base: str, shards: list[tuple[str, str]], in_features: int, dtype: torch.dtype) -> dict:
    device = torch.device("cuda")
    for suffix in ("weight_packed", "weight_scale", "weight_shape"):
        for ckpt_base, shard in shards:
            load_tensor(layer, suffix, art[f"{ckpt_base}.{suffix}"].to(device), shard)

    # Bitwise check of the loaded packed weights against the artifact.
    loaded = layer.weight_packed.detach().cpu()
    if len(shards) > 1:
        ref_packed = torch.cat([art[f"{b}.weight_packed"] for b, _ in shards], dim=0)
    else:
        ref_packed = art[f"{base}.weight_packed"]
    packed_ok = bool(torch.equal(loaded, ref_packed))
    scale_loaded = layer.weight_scale.detach().float()
    scale_ok = bool(torch.isfinite(scale_loaded).all())

    layer.quant_method.process_weights_after_loading(layer)
    repacked_scale = getattr(layer, "weight_scale", None)
    scale_post = (
        {
            "finite": bool(torch.isfinite(repacked_scale.detach().float()).all()),
            "min": float(repacked_scale.detach().float().min()),
            "max": float(repacked_scale.detach().float().max()),
        }
        if repacked_scale is not None
        else {}
    )

    torch.manual_seed(2301)
    ref_w = torch.cat([dequantized_weight(art, f"{b}.weight") for b, _ in shards], dim=0) if len(shards) > 1 else dequantized_weight(art, f"{base}.weight")
    outs = {}
    for act_dtype in (torch.bfloat16, torch.float16):
        x = torch.randn(512, in_features, dtype=act_dtype, device=device)
        result = layer(x)
        out = (result[0] if isinstance(result, tuple) else result).float()
        ref = x.float() @ ref_w.T.to(device)
        finite = bool(torch.isfinite(out).all())
        rel = float((out - ref).norm() / ref.norm())
        outs[str(act_dtype).replace("torch.", "")] = rel if finite else "NaN"

    return {
        "layer": name,
        "rel_l2_by_act_dtype": outs,
        "packed_bitwise_ok": packed_ok,
        "scale_finite_pre_repack": scale_ok,
        "scale_post_repack": scale_post,
    }


def synthetic_artifact(block: int) -> dict[str, torch.Tensor]:
    """Random int8-quantized tensors in the artifact's layout, for stack-only isolation."""
    from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32

    def quant(name: str, out: int, inp: int) -> None:
        w = torch.randn(out, inp, dtype=torch.float32)
        scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
        q = torch.round(w / scale).clamp(-127, 127).to(torch.int8)
        art[f"{name}.weight_packed"] = pack_to_int32(q.contiguous(), 8).contiguous()
        art[f"{name}.weight_scale"] = scale.to(torch.float32).contiguous()
        art[f"{name}.weight_shape"] = torch.tensor([out, inp], dtype=torch.int64)

    art: dict[str, torch.Tensor] = {}
    p = f"transformer_blocks.{block}"
    for part in ("to_q", "to_k", "to_v"):
        quant(f"{p}.attn.{part}", 4096, 4096)
    quant(f"{p}.attn.to_out.0", 4096, 4096)
    quant(f"{p}.img_mlp.proj", 12288, 4096)
    quant(f"{p}.img_mlp.out", 4096, 12288)
    return art


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--artifact", type=pathlib.Path, required=True)
    parser.add_argument("--block", type=int, default=0)
    parser.add_argument(
        "--synthetic",
        action="store_true",
        help="ignore --artifact; use random int8 weights (isolates the stack from the checkpoint)",
    )
    args = parser.parse_args()

    import vllm.distributed
    from vllm.config import VllmConfig, set_current_vllm_config
    from vllm.distributed import (
        destroy_distributed_environment,
        destroy_model_parallel,
        initialize_model_parallel,
    )
    from vllm.model_executor.layers.linear import (
        ColumnParallelLinear,
        QKVParallelLinear,
        RowParallelLinear,
    )

    vllm.distributed.init_distributed_environment(
        backend="gloo",
        world_size=1,
        rank=0,
        local_rank=0,
        distributed_init_method="tcp://127.0.0.1:29601",
    )
    with set_current_vllm_config(VllmConfig()):
        initialize_model_parallel(tensor_model_parallel_size=1)

    config = json.loads((args.artifact / "transformer" / "config.json").read_text())
    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config

    quant_config = resolve_component_quant_config(
        build_quant_config(config["quantization_config"]), "transformer"
    )
    if args.synthetic:
        art = synthetic_artifact(args.block)
    else:
        art = load_file(str(args.artifact / "transformer" / "diffusion_pytorch_model.safetensors"))
    b = args.block
    p = f"transformer.transformer_blocks.{b}"

    results = []
    with torch.device("cuda"):
        to_qkv = QKVParallelLinear(
            hidden_size=4096, head_size=128, total_num_heads=32, bias=False,
            quant_config=quant_config, prefix=f"{p}.attn.to_qkv",
        )
        results.append(
            check(
                "attn.to_qkv (fused q/k/v)",
                to_qkv,
                art,
                f"transformer_blocks.{b}.attn.to_qkv",
                [
                    (f"transformer_blocks.{b}.attn.to_q", "q"),
                    (f"transformer_blocks.{b}.attn.to_k", "k"),
                    (f"transformer_blocks.{b}.attn.to_v", "v"),
                ],
                4096,
                torch.bfloat16,
            )
        )

        to_out = RowParallelLinear(
            4096, 4096, bias=False, input_is_parallel=True, return_bias=False,
            quant_config=quant_config, prefix=f"{p}.attn.to_out.0",
        )
        results.append(
            check(
                "attn.to_out.0 (row)",
                to_out,
                art,
                f"transformer_blocks.{b}.attn.to_out.0",
                [(f"transformer_blocks.{b}.attn.to_out.0", None)],
                4096,
                torch.bfloat16,
            )
        )

        proj = ColumnParallelLinear(
            4096, 12288, bias=False, gather_output=False, return_bias=False,
            quant_config=quant_config, prefix=f"{p}.img_mlp.proj",
        )
        results.append(
            check(
                "img_mlp.proj (column)",
                proj,
                art,
                f"transformer_blocks.{b}.img_mlp.proj",
                [(f"transformer_blocks.{b}.img_mlp.proj", None)],
                4096,
                torch.bfloat16,
            )
        )

        mlp_out = RowParallelLinear(
            12288, 4096, bias=False, input_is_parallel=True, return_bias=False,
            quant_config=quant_config, prefix=f"{p}.img_mlp.out",
        )
        results.append(
            check(
                "img_mlp.out (row, wide-in)",
                mlp_out,
                art,
                f"transformer_blocks.{b}.img_mlp.out",
                [(f"transformer_blocks.{b}.img_mlp.out", None)],
                12288,
                torch.bfloat16,
            )
        )

    destroy_model_parallel()
    destroy_distributed_environment()
    print(json.dumps(results, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())