qwen3.5-122B-A10B on DGX Spark: vLLM + DFlash + dense-bandwidth stack, one-shot installer
This commit is contained in:
ent
@@ -0,0 +1,463 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Build a hybrid GPTQ-INT4 + FP8 checkpoint for Qwen3.5-122B-A10B.
|
||||
|
||||
Takes MoE expert weights from the GPTQ-INT4 checkpoint (0.5 bytes/param),
|
||||
and dense layers (attention, shared experts, embeddings) from the official
|
||||
FP8 checkpoint (1 byte/param + calibrated block scales).
|
||||
|
||||
Result: a checkpoint that is ~9 GB smaller than the GPTQ-INT4 original
|
||||
while using properly calibrated FP8 scales (not naive cast), yielding
|
||||
better decode throughput on bandwidth-limited hardware.
|
||||
|
||||
NOTE: Requires the hybrid FP8 dispatch patch from https://github.com/rmstxrx/vllm/tree/v0.17.1-hybrid-fp8
|
||||
|
||||
Usage:
|
||||
python build-hybrid-checkpoint.py \
|
||||
--gptq-dir ~/inference/models/hf/qwen3.5-122b-a10b-gptq-int4 \
|
||||
--fp8-repo Qwen/Qwen3.5-122B-A10B-FP8 \
|
||||
--output ~/inference/models/hf/qwen3.5-122b-a10b-fp8hybrid
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_fp8_non_expert_manifest(fp8_repo: str) -> dict[str, str]:
|
||||
"""Get the weight map for non-expert tensors from the FP8 checkpoint.
|
||||
|
||||
Args:
|
||||
fp8_repo: Hugging Face repo ID for the FP8 checkpoint.
|
||||
|
||||
Returns:
|
||||
Mapping of tensor names to shard filenames for non-expert tensors.
|
||||
"""
|
||||
idx_path = hf_hub_download(fp8_repo, "model.safetensors.index.json")
|
||||
with open(idx_path, encoding="utf-8") as f:
|
||||
idx = json.load(f)
|
||||
|
||||
wm = idx["weight_map"]
|
||||
return {k: v for k, v in wm.items() if ".experts." not in k}
|
||||
|
||||
|
||||
def download_fp8_shards(fp8_repo: str, shards: set[str], cache_dir: Path) -> dict[str, Path]:
|
||||
"""Download only the needed FP8 shards.
|
||||
|
||||
Args:
|
||||
fp8_repo: Hugging Face repo ID for the FP8 checkpoint.
|
||||
shards: Shard filenames that contain non-expert tensors.
|
||||
cache_dir: Local cache directory for downloaded shards.
|
||||
|
||||
Returns:
|
||||
Mapping of shard filename to downloaded local path.
|
||||
"""
|
||||
shard_paths: dict[str, Path] = {}
|
||||
for shard in sorted(shards):
|
||||
logger.info(" Downloading %s...", shard)
|
||||
path = hf_hub_download(fp8_repo, shard, local_dir=str(cache_dir))
|
||||
shard_paths[shard] = Path(path)
|
||||
logger.info(" -> %s", path)
|
||||
return shard_paths
|
||||
|
||||
|
||||
def extract_fp8_tensors(shard_paths: dict[str, Path], wanted: dict[str, str]) -> dict[str, torch.Tensor]:
|
||||
"""Extract the requested FP8 tensors from downloaded shards.
|
||||
|
||||
Args:
|
||||
shard_paths: Mapping of shard filename to local path.
|
||||
wanted: Mapping of tensor name to shard filename.
|
||||
|
||||
Returns:
|
||||
Mapping of tensor name to loaded FP8 tensor.
|
||||
"""
|
||||
tensors: dict[str, torch.Tensor] = {}
|
||||
for shard_name, shard_path in shard_paths.items():
|
||||
keys_in_shard = [k for k, v in wanted.items() if v == shard_name]
|
||||
if not keys_in_shard:
|
||||
continue
|
||||
|
||||
logger.info(" Extracting %d tensors from %s...", len(keys_in_shard), shard_name)
|
||||
with safe_open(str(shard_path), framework="pt") as f:
|
||||
for k in keys_in_shard:
|
||||
tensors[k] = f.get_tensor(k)
|
||||
|
||||
return tensors
|
||||
|
||||
|
||||
def find_model_safetensors_files(directory: Path) -> list[Path]:
|
||||
"""Find model safetensors files in a GPTQ checkpoint directory.
|
||||
|
||||
Args:
|
||||
directory: Directory containing GPTQ checkpoint files.
|
||||
|
||||
Returns:
|
||||
Sorted list of shard paths, or a single-element list for single-file checkpoints.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If no supported model safetensors files are found.
|
||||
"""
|
||||
gptq_shards = sorted(directory.glob("model.safetensors-*"))
|
||||
if gptq_shards:
|
||||
return gptq_shards
|
||||
|
||||
# AutoRound / HF naming: model-NNNNN-of-NNNNN.safetensors
|
||||
gptq_shards = sorted(directory.glob("model-*-of-*.safetensors"))
|
||||
if gptq_shards:
|
||||
return gptq_shards
|
||||
|
||||
single_file = directory / "model.safetensors"
|
||||
if single_file.is_file():
|
||||
return [single_file]
|
||||
|
||||
raise FileNotFoundError(
|
||||
f"No model.safetensors files found in {directory}. Expected a sharded "
|
||||
"checkpoint (model.safetensors-NNNNN-of-NNNNN)."
|
||||
)
|
||||
|
||||
|
||||
def validate_gptq_input(gptq_dir: Path) -> None:
|
||||
"""Validate the GPTQ checkpoint directory before any downloads.
|
||||
|
||||
Args:
|
||||
gptq_dir: Path to the local GPTQ checkpoint directory.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the directory or required files do not exist.
|
||||
NotADirectoryError: If the path exists but is not a directory.
|
||||
"""
|
||||
if not gptq_dir.exists():
|
||||
raise FileNotFoundError(f"GPTQ directory does not exist: {gptq_dir}")
|
||||
if not gptq_dir.is_dir():
|
||||
raise NotADirectoryError(f"GPTQ path is not a directory: {gptq_dir}")
|
||||
if not any(path.is_file() for path in gptq_dir.glob("*.safetensors*")):
|
||||
raise FileNotFoundError(f"No .safetensors files found in {gptq_dir}")
|
||||
if not (gptq_dir / "config.json").is_file():
|
||||
raise FileNotFoundError(f"Missing config.json in {gptq_dir}")
|
||||
|
||||
|
||||
def validate_output_dir(output_dir: Path, force: bool) -> None:
|
||||
"""Validate the output directory before building.
|
||||
|
||||
Args:
|
||||
output_dir: Path where the hybrid checkpoint will be written.
|
||||
force: Whether destructive cleanup is allowed.
|
||||
|
||||
Raises:
|
||||
FileExistsError: If the output directory is non-empty without `force`.
|
||||
NotADirectoryError: If the output path exists but is not a directory.
|
||||
"""
|
||||
if output_dir.exists() and not output_dir.is_dir():
|
||||
raise NotADirectoryError(f"Output path is not a directory: {output_dir}")
|
||||
|
||||
if output_dir.exists() and any(output_dir.iterdir()) and not force:
|
||||
raise FileExistsError(
|
||||
f"Output directory {output_dir} exists and is not empty. Use --force "
|
||||
"to remove existing model.safetensors* and config.json files before building."
|
||||
)
|
||||
|
||||
|
||||
def prepare_output_dir(output_dir: Path, force: bool) -> None:
|
||||
"""Create or clean the output directory before writing files.
|
||||
|
||||
Args:
|
||||
output_dir: Path where the hybrid checkpoint will be written.
|
||||
force: Whether destructive cleanup is allowed.
|
||||
"""
|
||||
if not output_dir.exists():
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
return
|
||||
|
||||
if not force:
|
||||
return
|
||||
|
||||
delete_targets = sorted(
|
||||
{
|
||||
path
|
||||
for pattern in ("model.safetensors*", "config.json")
|
||||
for path in output_dir.glob(pattern)
|
||||
if path.is_file()
|
||||
}
|
||||
)
|
||||
for path in delete_targets:
|
||||
path.unlink()
|
||||
logger.info(" Deleted %s", path)
|
||||
|
||||
|
||||
def build_hybrid_checkpoint(
|
||||
gptq_dir: Path,
|
||||
fp8_tensors: dict[str, torch.Tensor],
|
||||
output_dir: Path,
|
||||
force: bool,
|
||||
) -> tuple[int, int, int]:
|
||||
"""Build the hybrid checkpoint from GPTQ and FP8 tensors.
|
||||
|
||||
Args:
|
||||
gptq_dir: Path to the local GPTQ checkpoint directory.
|
||||
fp8_tensors: FP8 tensors keyed by tensor name.
|
||||
output_dir: Output directory for rewritten checkpoint shards.
|
||||
force: Whether to continue if many unexpected FP8 tensors are unmatched.
|
||||
|
||||
Returns:
|
||||
Tuple of replaced tensor count, added scale tensor count, and bytes saved.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If no supported GPTQ model files are found.
|
||||
RuntimeError: If too many FP8 tensors cannot be matched and `force` is not set.
|
||||
ValueError: If a matched FP8 tensor has an incompatible shape.
|
||||
"""
|
||||
# 1. Copy non-safetensors files
|
||||
for f in gptq_dir.iterdir():
|
||||
if f.suffix != ".safetensors" and not f.name.startswith("."):
|
||||
if f.is_file():
|
||||
shutil.copy2(f, output_dir / f.name)
|
||||
|
||||
# 2. Process each GPTQ shard
|
||||
gptq_shards = find_model_safetensors_files(gptq_dir)
|
||||
total_replaced = 0
|
||||
total_added = 0
|
||||
total_saved_bytes = 0
|
||||
|
||||
# Track which FP8 tensors have been placed
|
||||
placed_fp8: set[str] = set()
|
||||
|
||||
for i, shard_path in enumerate(gptq_shards):
|
||||
shard_name = shard_path.name
|
||||
logger.info(" [%d/%d] %s...", i + 1, len(gptq_shards), shard_name)
|
||||
|
||||
gptq_tensors = load_file(str(shard_path))
|
||||
output_tensors: dict[str, torch.Tensor] = {}
|
||||
replaced = 0
|
||||
|
||||
for name, tensor in gptq_tensors.items():
|
||||
if name in fp8_tensors:
|
||||
# Replace BF16 tensor with FP8 version
|
||||
fp8_tensor = fp8_tensors[name]
|
||||
if tensor.shape != fp8_tensor.shape:
|
||||
raise ValueError(
|
||||
f"Shape mismatch for {name}: GPTQ={tensor.shape}, FP8={fp8_tensor.shape}"
|
||||
)
|
||||
old_bytes = tensor.numel() * tensor.element_size()
|
||||
new_bytes = fp8_tensor.numel() * fp8_tensor.element_size()
|
||||
total_saved_bytes += old_bytes - new_bytes
|
||||
output_tensors[name] = fp8_tensor
|
||||
placed_fp8.add(name)
|
||||
replaced += 1
|
||||
|
||||
# Also add the scale tensor if it exists
|
||||
scale_name = name.replace(".weight", ".weight_scale_inv")
|
||||
if scale_name in fp8_tensors and scale_name != name:
|
||||
output_tensors[scale_name] = fp8_tensors[scale_name]
|
||||
placed_fp8.add(scale_name)
|
||||
total_added += 1
|
||||
else:
|
||||
output_tensors[name] = tensor
|
||||
|
||||
total_replaced += replaced
|
||||
|
||||
save_file(output_tensors, str(output_dir / shard_name))
|
||||
logger.info(" replaced=%d, tensors=%d", replaced, len(output_tensors))
|
||||
|
||||
# 3. Check for any FP8 tensors not yet placed (e.g. scale tensors
|
||||
# for weights that exist in shards we already processed)
|
||||
unplaced = set(fp8_tensors.keys()) - placed_fp8
|
||||
if unplaced:
|
||||
expected_unplaced = {
|
||||
name
|
||||
for name in unplaced
|
||||
if name.endswith(".weight_scale_inv")
|
||||
and f"{name.removesuffix('.weight_scale_inv')}.weight" in placed_fp8
|
||||
}
|
||||
unexpected_unplaced = sorted(unplaced - expected_unplaced)
|
||||
|
||||
for name in unexpected_unplaced:
|
||||
logger.warning("WARNING: unexpected unmatched FP8 tensor %s", name)
|
||||
|
||||
if len(unexpected_unplaced) > 5 and not force:
|
||||
message = (
|
||||
f"{len(unexpected_unplaced)} FP8 tensors could not be matched to GPTQ tensors. "
|
||||
"This likely indicates a naming mismatch between the GPTQ and FP8 checkpoints. "
|
||||
"Use --force to proceed anyway."
|
||||
)
|
||||
logger.error("ERROR: %s", message)
|
||||
raise RuntimeError(message)
|
||||
if len(unexpected_unplaced) > 5 and force:
|
||||
logger.warning(
|
||||
"WARNING: proceeding despite %d unexpected unmatched FP8 tensors because --force was provided",
|
||||
len(unexpected_unplaced),
|
||||
)
|
||||
|
||||
if expected_unplaced:
|
||||
logger.info(" Adding %d expected unplaced FP8 scale tensors to final shard...", len(expected_unplaced))
|
||||
|
||||
if expected_unplaced:
|
||||
# Load last shard, add expected scale tensors, re-save
|
||||
last_shard = output_dir / gptq_shards[-1].name
|
||||
existing = load_file(str(last_shard))
|
||||
for name in sorted(expected_unplaced):
|
||||
existing[name] = fp8_tensors[name]
|
||||
total_added += 1
|
||||
save_file(existing, str(last_shard))
|
||||
|
||||
return total_replaced, total_added, total_saved_bytes
|
||||
|
||||
|
||||
def update_safetensors_index(output_dir: Path) -> None:
|
||||
"""Rebuild `model.safetensors.index.json` from actual shard contents.
|
||||
|
||||
Args:
|
||||
output_dir: Directory containing rewritten model shards.
|
||||
"""
|
||||
weight_map: dict[str, str] = {}
|
||||
total_size = 0
|
||||
|
||||
for shard_path in find_model_safetensors_files(output_dir):
|
||||
with safe_open(str(shard_path), framework="pt") as f:
|
||||
for key in f.keys():
|
||||
weight_map[key] = shard_path.name
|
||||
tensor = f.get_tensor(key)
|
||||
total_size += tensor.numel() * tensor.element_size()
|
||||
|
||||
index = {
|
||||
"metadata": {"total_size": total_size},
|
||||
"weight_map": weight_map
|
||||
}
|
||||
|
||||
with open(output_dir / "model.safetensors.index.json", "w", encoding="utf-8") as f:
|
||||
json.dump(index, f, indent=2, sort_keys=True)
|
||||
|
||||
logger.info(" Index rebuilt: %d tensors, %.2f GB", len(weight_map), total_size / 1e9)
|
||||
|
||||
|
||||
def update_config(output_dir: Path) -> None:
|
||||
"""Update `config.json` with hybrid quantization metadata.
|
||||
|
||||
Args:
|
||||
output_dir: Directory containing the hybrid checkpoint.
|
||||
"""
|
||||
config_path = output_dir / "config.json"
|
||||
with open(config_path, encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
config["_hybrid_quant_info"] = {
|
||||
"description": "Hybrid GPTQ-INT4 + FP8 checkpoint for single-GPU deployment",
|
||||
"moe_experts": "GPTQ INT4 (group_size=128, sym=True, Marlin kernels)",
|
||||
"dense_layers": "FP8 E4M3 block-128 (from official Qwen/Qwen3.5-122B-A10B-FP8, calibrated scales)",
|
||||
"norms_gates_embeddings": "Preserved at source dtype (BF16 for norms/gates, FP8 for others)",
|
||||
"source_gptq": "Qwen/Qwen3.5-122B-A10B-GPTQ-Int4",
|
||||
"source_fp8": "Qwen/Qwen3.5-122B-A10B-FP8",
|
||||
"vllm_patch": "https://github.com/rmstxrx/vllm/tree/v0.17.1-hybrid-fp8",
|
||||
"target_hardware": "NVIDIA DGX Spark (GB10, 128GB unified, 273 GB/s)",
|
||||
"converter": "build-hybrid-checkpoint.py"
|
||||
}
|
||||
|
||||
with open(config_path, "w", encoding="utf-8") as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Build a hybrid GPTQ-INT4 + FP8 checkpoint.
|
||||
|
||||
Returns:
|
||||
Process exit code.
|
||||
"""
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
parser = argparse.ArgumentParser(description="Build hybrid GPTQ-INT4 + FP8 checkpoint")
|
||||
parser.add_argument("--gptq-dir", required=True, help="Path to GPTQ-INT4 model")
|
||||
parser.add_argument("--fp8-repo", default="Qwen/Qwen3.5-122B-A10B-FP8", help="HF repo for FP8 model")
|
||||
parser.add_argument("--output", required=True, help="Output directory")
|
||||
parser.add_argument("--dry-run", action="store_true")
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Allow output cleanup and continue despite many unmatched FP8 tensors",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
gptq_dir = Path(args.gptq_dir)
|
||||
output_dir = Path(args.output)
|
||||
|
||||
validate_gptq_input(gptq_dir)
|
||||
find_model_safetensors_files(gptq_dir)
|
||||
validate_output_dir(output_dir, args.force)
|
||||
|
||||
logger.info("=== Hybrid GPTQ-INT4 + FP8 Checkpoint Builder ===")
|
||||
logger.info(" GPTQ source: %s", gptq_dir)
|
||||
logger.info(" FP8 source: %s", args.fp8_repo)
|
||||
logger.info(" Output: %s", output_dir)
|
||||
logger.info("")
|
||||
|
||||
# Step 1: Get manifest of non-expert tensors from FP8 checkpoint
|
||||
logger.info("[1/5] Fetching FP8 tensor manifest...")
|
||||
fp8_manifest = get_fp8_non_expert_manifest(args.fp8_repo)
|
||||
shards_needed = set(fp8_manifest.values())
|
||||
logger.info(" Non-expert tensors: %d", len(fp8_manifest))
|
||||
logger.info(" Shards to download: %s", sorted(shards_needed))
|
||||
|
||||
if args.dry_run:
|
||||
logger.info("")
|
||||
logger.info("[DRY RUN] Would download shards and build hybrid. Exiting.")
|
||||
return 0
|
||||
|
||||
prepare_output_dir(output_dir, args.force)
|
||||
cache_dir = output_dir / ".fp8_cache"
|
||||
cache_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Step 2: Download needed FP8 shards
|
||||
logger.info("")
|
||||
logger.info("[2/5] Downloading %d FP8 shards...", len(shards_needed))
|
||||
shard_paths = download_fp8_shards(args.fp8_repo, shards_needed, cache_dir)
|
||||
|
||||
# Step 3: Extract non-expert FP8 tensors
|
||||
logger.info("")
|
||||
logger.info("[3/5] Extracting FP8 tensors...")
|
||||
fp8_tensors = extract_fp8_tensors(shard_paths, fp8_manifest)
|
||||
logger.info(" Extracted: %d tensors", len(fp8_tensors))
|
||||
|
||||
# Show dtype breakdown
|
||||
dtypes: dict[str, int] = {}
|
||||
for name, t in fp8_tensors.items():
|
||||
d = str(t.dtype)
|
||||
dtypes[d] = dtypes.get(d, 0) + 1
|
||||
logger.info(" Dtypes: %s", dtypes)
|
||||
|
||||
# Step 4: Build hybrid checkpoint
|
||||
logger.info("")
|
||||
logger.info("[4/5] Building hybrid checkpoint...")
|
||||
replaced, added, saved = build_hybrid_checkpoint(
|
||||
gptq_dir,
|
||||
fp8_tensors,
|
||||
output_dir,
|
||||
args.force,
|
||||
)
|
||||
|
||||
# Step 5: Update index and config
|
||||
logger.info("")
|
||||
logger.info("[5/5] Updating index and config...")
|
||||
update_safetensors_index(output_dir)
|
||||
update_config(output_dir)
|
||||
|
||||
# Cleanup downloaded FP8 shards
|
||||
shutil.rmtree(cache_dir, ignore_errors=True)
|
||||
|
||||
logger.info("")
|
||||
logger.info("=== Complete ===")
|
||||
logger.info(" Tensors replaced (BF16→FP8): %d", replaced)
|
||||
logger.info(" Scale tensors added: %d", added)
|
||||
logger.info(" Bytes saved: %.2f GB", saved / 1e9)
|
||||
logger.info(" Output: %s", output_dir)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,115 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Inspect the Intel INT4 AutoRound checkpoint to decide if albond's hybrid
|
||||
INT4+FP8 build will work: it swaps BF16 *dense* (non-expert) tensors for FP8 by
|
||||
name. If AutoRound quantized the dense linears to INT4 (.qweight), the swap is a
|
||||
near no-op. We need attention/shared_expert dense weights stored as BF16 (.weight).
|
||||
|
||||
Also pulls the FP8 repo's index to confirm it exists and that names line up.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
INT4_DIR = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(
|
||||
"/home/ent/.cache/huggingface/hub/models--Intel--Qwen3.5-122B-A10B-int4-AutoRound/"
|
||||
"snapshots/3045d02bb737effc4581da91bddbad3be02934e4")
|
||||
FP8_REPO = sys.argv[2] if len(sys.argv) > 2 else "Qwen/Qwen3.5-122B-A10B-FP8"
|
||||
|
||||
idx = json.loads((INT4_DIR / "model.safetensors.index.json").read_text())
|
||||
wm = idx["weight_map"]
|
||||
names = list(wm.keys())
|
||||
|
||||
# Non-expert = no ".experts." (routed experts stay INT4 in the hybrid)
|
||||
nonexp = [n for n in names if ".experts." not in n]
|
||||
|
||||
def bucket(n):
|
||||
if ".experts." in n:
|
||||
return "ROUTED-EXPERT"
|
||||
if "shared_expert" in n:
|
||||
return "shared_expert"
|
||||
if "self_attn" in n or "linear_attn" in n or ".attn" in n:
|
||||
return "attention"
|
||||
if "embed_tokens" in n or "lm_head" in n:
|
||||
return "embed/head"
|
||||
if "mtp" in n:
|
||||
return "mtp"
|
||||
return "other"
|
||||
|
||||
print(f"=== Intel INT4 checkpoint: {INT4_DIR.name} ===")
|
||||
print(f"total tensors: {len(names)} non-expert: {len(nonexp)}")
|
||||
|
||||
# Suffix histogram tells us quant scheme: .qweight/.scales/.qzeros => INT4; .weight => dense
|
||||
suffix = Counter(n.rsplit(".", 1)[-1] for n in nonexp)
|
||||
print("\nnon-expert tensor SUFFIX histogram (qweight/scales/qzeros = INT4-packed; weight = dense):")
|
||||
for s, c in suffix.most_common():
|
||||
print(f" .{s:20s} {c}")
|
||||
|
||||
# For each functional group, does it have .weight (BF16 dense) or .qweight (INT4)?
|
||||
print("\nper-group quant scheme (sample names):")
|
||||
groups = {}
|
||||
for n in nonexp:
|
||||
g = bucket(n)
|
||||
groups.setdefault(g, {"weight": 0, "qweight": 0, "scale": 0, "other": 0, "ex": None})
|
||||
suf = n.rsplit(".", 1)[-1]
|
||||
if suf == "weight":
|
||||
groups[g]["weight"] += 1
|
||||
elif suf == "qweight":
|
||||
groups[g]["qweight"] += 1
|
||||
elif "scale" in suf or suf in ("qzeros",):
|
||||
groups[g]["scale"] += 1
|
||||
else:
|
||||
groups[g]["other"] += 1
|
||||
if groups[g]["ex"] is None and suf in ("weight", "qweight"):
|
||||
groups[g]["ex"] = n
|
||||
for g, d in sorted(groups.items()):
|
||||
scheme = "INT4(.qweight)" if d["qweight"] else ("DENSE(.weight)" if d["weight"] else "?")
|
||||
print(f" {g:16s} weight={d['weight']:4d} qweight={d['qweight']:4d} scale={d['scale']:4d} -> {scheme}")
|
||||
print(f" e.g. {d['ex']}")
|
||||
|
||||
# Dtypes of a few non-expert .weight tensors (open the shard header only).
|
||||
print("\ndtypes of sample non-expert '.weight' tensors (FP8 swap needs BF16 here):")
|
||||
from safetensors import safe_open # noqa: E402
|
||||
sample = [n for n in nonexp if n.endswith(".weight")
|
||||
and ("self_attn" in n or "shared_expert" in n or "embed" in n or "lm_head" in n)]
|
||||
seen_shards = {}
|
||||
shown = 0
|
||||
for n in sample:
|
||||
shard = wm[n]
|
||||
f = seen_shards.get(shard)
|
||||
if f is None:
|
||||
f = safe_open(str(INT4_DIR / shard), framework="pt")
|
||||
seen_shards[shard] = f
|
||||
try:
|
||||
t = f.get_slice(n)
|
||||
print(f" {n:60s} {t.get_dtype()} {tuple(t.get_shape())}")
|
||||
except Exception as e:
|
||||
print(f" {n:60s} <err {e}>")
|
||||
shown += 1
|
||||
if shown >= 12:
|
||||
break
|
||||
|
||||
# FP8 repo: confirm exists + list its non-expert names/dtypes for name-match sanity.
|
||||
print(f"\n=== FP8 repo manifest: {FP8_REPO} ===")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
p = hf_hub_download(FP8_REPO, "model.safetensors.index.json")
|
||||
fidx = json.loads(Path(p).read_text())
|
||||
fwm = fidx["weight_map"]
|
||||
fnon = [n for n in fwm if ".experts." not in n]
|
||||
print(f"FP8 total tensors: {len(fwm)} non-expert: {len(fnon)}")
|
||||
fsuf = Counter(n.rsplit('.', 1)[-1] for n in fnon)
|
||||
print("FP8 non-expert suffix histogram:")
|
||||
for s, c in fsuf.most_common(10):
|
||||
print(f" .{s:20s} {c}")
|
||||
# how many FP8 non-expert .weight names also exist in INT4 as .weight?
|
||||
int4_weight = {n for n in nonexp if n.endswith('.weight')}
|
||||
fp8_weight = {n for n in fnon if n.endswith('.weight')}
|
||||
match = int4_weight & fp8_weight
|
||||
print(f"\nname overlap (FP8 '.weight' that also exist as '.weight' in INT4): {len(match)} / {len(fp8_weight)} FP8 weights")
|
||||
only_fp8 = sorted(fp8_weight - int4_weight)[:8]
|
||||
print(f"FP8 '.weight' NOT present as '.weight' in INT4 (would not swap): {len(fp8_weight - int4_weight)}")
|
||||
for n in only_fp8:
|
||||
print(f" {n} (INT4 has: {'qweight' if n[:-7]+'.qweight' in set(names) else 'MISSING'})")
|
||||
except Exception as e:
|
||||
print(f"<FP8 repo fetch failed: {e}>")
|
||||
@@ -0,0 +1,96 @@
|
||||
#!/usr/bin/env python3
|
||||
"""v2 harness: proper bf16 GEMV baseline + a single BATCHED int8 w8a16 GEMM kernel
|
||||
(one launch for any B, dot-based, pads B to >=16) vs albond's per-row loop, at the
|
||||
real lm-head shape and B in {1,5,13} (base + DFlash verify-batch sizes).
|
||||
"""
|
||||
import time
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
V, H = 248320, 3072
|
||||
DEV = "cuda"
|
||||
torch.manual_seed(0)
|
||||
|
||||
W_bf16 = (torch.randn(V, H, device=DEV, dtype=torch.float32) * 0.02).to(torch.bfloat16)
|
||||
scales = (W_bf16.float().abs().amax(dim=1) / 127.0).clamp(min=1e-12)
|
||||
W_int8 = (W_bf16.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
|
||||
scales_f16 = scales.to(torch.float16)
|
||||
INT8_BYTES = V * H
|
||||
BF16_BYTES = V * H * 2
|
||||
|
||||
|
||||
def bench(fn, iters=30, warmup=5):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(iters):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
return (time.perf_counter() - t0) / iters
|
||||
|
||||
|
||||
# ---- batched int8 w8a16 GEMM: out[B,N] = (x[B,K] @ (W_int8[N,K]*s[N]).T) ----
|
||||
@triton.jit
|
||||
def _k_batched(x_ptr, w_ptr, s_ptr, o_ptr, B, N, K,
|
||||
sxb, sxk, swn, swk, sob, son,
|
||||
BLOCK_B: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
|
||||
pid_n = tl.program_id(0)
|
||||
offs_b = tl.arange(0, BLOCK_B)
|
||||
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
offs_k = tl.arange(0, BLOCK_K)
|
||||
x_ptrs = x_ptr + offs_b[:, None] * sxb + offs_k[None, :] * sxk
|
||||
w_ptrs = w_ptr + offs_n[:, None] * swn + offs_k[None, :] * swk
|
||||
acc = tl.zeros((BLOCK_B, BLOCK_N), dtype=tl.float32)
|
||||
for k in range(0, K, BLOCK_K):
|
||||
km = (offs_k[None, :] + k) < K
|
||||
x = tl.load(x_ptrs, mask=(offs_b[:, None] < B) & km, other=0.0).to(tl.float16)
|
||||
w = tl.load(w_ptrs, mask=(offs_n[:, None] < N) & km, other=0).to(tl.float16)
|
||||
acc += tl.dot(x, w.T) # [BB,BK] @ [BK,BN] -> [BB,BN]
|
||||
x_ptrs += BLOCK_K * sxk
|
||||
w_ptrs += BLOCK_K * swk
|
||||
s = tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32)
|
||||
acc = acc * s[None, :]
|
||||
o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son
|
||||
tl.store(o_ptrs, acc.to(tl.float16), mask=(offs_b[:, None] < B) & (offs_n[None, :] < N))
|
||||
|
||||
|
||||
def run_batched(x, BLOCK_N=128, BLOCK_K=64, num_warps=4, num_stages=3):
|
||||
B = x.shape[0]
|
||||
BLOCK_B = max(16, triton.next_power_of_2(B))
|
||||
out = torch.empty(B, V, dtype=torch.float16, device=DEV)
|
||||
xf = x.to(torch.float16)
|
||||
grid = ((V + BLOCK_N - 1) // BLOCK_N,)
|
||||
_k_batched[grid](xf, W_int8, scales_f16, out, B, V, H,
|
||||
xf.stride(0), xf.stride(1), W_int8.stride(0), W_int8.stride(1),
|
||||
out.stride(0), out.stride(1),
|
||||
BLOCK_B=BLOCK_B, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
|
||||
num_warps=num_warps, num_stages=num_stages)
|
||||
return out
|
||||
|
||||
|
||||
for B in (1, 5, 13):
|
||||
print(f"\n=== B={B} ===")
|
||||
x = torch.randn(B, H, device=DEV, dtype=torch.bfloat16) * 0.1
|
||||
ref_bf16 = (x.float() @ W_bf16.float().T)
|
||||
ref_deq = (x.float() @ (W_int8.float() * scales.unsqueeze(1)).T)
|
||||
am_floor = (ref_bf16.argmax(-1) == ref_deq.argmax(-1)).float().mean().item()
|
||||
print(f" quant floor: argmax_match={am_floor*100:.1f}% maxerr={(ref_deq-ref_bf16).abs().max():.4f}")
|
||||
|
||||
# real bf16 GEMV baseline (stays bf16, reads 1.5GB)
|
||||
dt = bench(lambda: torch.matmul(x, W_bf16.t()))
|
||||
print(f" {'bf16 GEMV (real baseline)':30s} {'':33s}{dt*1e3:7.3f} ms {BF16_BYTES/dt/1e9:6.1f} GB/s")
|
||||
|
||||
# batched int8 kernel, a few configs
|
||||
for (bn, bk, nw, ns) in [(128, 64, 4, 3), (256, 64, 8, 3), (128, 128, 4, 3), (64, 128, 4, 3)]:
|
||||
try:
|
||||
out = run_batched(x, bn, bk, nw, ns)
|
||||
err = (out.float() - ref_deq.float()).abs().max().item()
|
||||
am = (out.float().argmax(-1) == ref_bf16.argmax(-1)).float().mean().item()
|
||||
dt = bench(lambda: run_batched(x, bn, bk, nw, ns))
|
||||
print(f" batched N{bn}/K{bk}/w{nw}/s{ns:<2d} maxerr={err:8.4f} argmax={am*100:5.1f}% "
|
||||
f"{dt*1e3:7.3f} ms {INT8_BYTES/dt/1e9:6.1f} GB/s")
|
||||
except Exception as e:
|
||||
print(f" batched N{bn}/K{bk}/w{nw}/s{ns}: ERR {str(e)[:70]}")
|
||||
@@ -0,0 +1,27 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Validate patch_inc_hybrid.py landed on AEON 0.23's inc.py. Run AFTER the patch
|
||||
in the same container: python3 /host/patch_inc_hybrid.py && python3 /host/validate_inc_patch.py
|
||||
"""
|
||||
import inspect
|
||||
|
||||
import vllm.model_executor.layers.quantization.inc as m
|
||||
|
||||
C = m.INCConfig
|
||||
src = inspect.getsource(m)
|
||||
print("import OK")
|
||||
print("sentinel in source :", "spark-dflash-hybrid-fp8" in src)
|
||||
print("maybe_update_config OWN :", "maybe_update_config" in C.__dict__)
|
||||
print("_is_layer_fp8 OWN :", "_is_layer_fp8" in C.__dict__)
|
||||
print("maybe_update_config hasattr:", hasattr(C, "maybe_update_config"))
|
||||
print("_is_layer_fp8 hasattr :", hasattr(C, "_is_layer_fp8"))
|
||||
# signature must accept hf_config kw (config/vllm.py calls it that way)
|
||||
try:
|
||||
sig = inspect.signature(C.maybe_update_config)
|
||||
print("maybe_update_config sig :", str(sig))
|
||||
print("accepts hf_config kw :", "hf_config" in sig.parameters)
|
||||
except Exception as e:
|
||||
print("sig err:", e)
|
||||
# count FP8 dispatch sites
|
||||
print("Fp8LinearMethod dispatch ct:", src.count("return Fp8LinearMethod(self.fp8_config)"))
|
||||
# byte-compile sanity already implied by import; show line count
|
||||
print("inc.py lines :", len(src.splitlines()))
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Validate patch_int8_lmhead_v3 landed + the helper runs on GPU. Run AFTER the
|
||||
patch in the same container."""
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.logits_processor as lp
|
||||
|
||||
print("import OK | helpers:",
|
||||
hasattr(lp, "_spark_int8_gemm"),
|
||||
hasattr(lp, "_spark_int8_lmhead_apply"),
|
||||
hasattr(lp, "_spark_k_int8"))
|
||||
src = open(lp.__file__).read()
|
||||
print("sentinel in _get_logits:", "DGX_SPARK_INT8_LMHEAD_V3: int8 w8a16" in src)
|
||||
|
||||
V, H = 4096, 512
|
||||
torch.manual_seed(0)
|
||||
W = torch.randn(V, H, device="cuda") * 0.02
|
||||
s = (W.abs().amax(1) / 127).clamp(min=1e-12)
|
||||
wi = (W / s.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8).contiguous()
|
||||
sf = s.to(torch.float16)
|
||||
for B in (1, 5, 13):
|
||||
x = torch.randn(B, H, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
out = lp._spark_int8_gemm(x, wi, sf)
|
||||
ref = x.float() @ (wi.float() * s.unsqueeze(1)).T
|
||||
am = (out.argmax(-1) == ref.argmax(-1)).float().mean().item()
|
||||
print(f" B={B}: out{tuple(out.shape)} {out.dtype} argmax={am*100:.0f}% "
|
||||
f"maxerr={(out - ref).abs().max().item():.4f}")
|
||||
|
||||
# Exercise the FULL apply path (print + quantize-once + gemm) with a mock lm_head
|
||||
# (vocab > 100k to trigger the int8 path). self is unused -> None.
|
||||
print("--- full _spark_int8_lmhead_apply path (mock lm_head, V=131072) ---")
|
||||
Vbig = 131072
|
||||
|
||||
class _MockLMHead:
|
||||
pass
|
||||
|
||||
mh = _MockLMHead()
|
||||
mh.weight = (torch.randn(Vbig, H, device="cuda", dtype=torch.float32) * 0.02).to(torch.bfloat16)
|
||||
hs = torch.randn(3, H, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
o1 = lp._spark_int8_lmhead_apply(None, mh, hs, None) # first call: quantizes + prints
|
||||
o2 = lp._spark_int8_lmhead_apply(None, mh, hs, None) # second call: reuses int8
|
||||
refb = hs.float() @ mh.weight.float().T
|
||||
am = (o2.argmax(-1) == refb.argmax(-1)).float().mean().item()
|
||||
print(f" apply: out{tuple(o2.shape)} {o2.dtype} argmax_vs_bf16={am*100:.0f}% "
|
||||
f"int8_ready={getattr(mh, '_spark_int8_ready', None)} weight_kept={mh.weight.numel() > 0}")
|
||||
Reference in New Issue
Block a user