#!/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())