debug_w4a8.py

# debug_layer17_load_fix.py
import json
from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
from vllm.config import LoadConfig, ModelConfig

# Load the index
with open("./Llama-3.1-8B-Instruct-W4A8H100/model.safetensors.index.json") as f:
    index = json.load(f)

layer17_keys = [k for k in index["weight_map"].keys() if "model.layers.17" in k]
print("Layer 17 keys in index:")
for k in layer17_keys:
    print(" ", k)

# Setup loader and config
load_config = LoadConfig(
    load_format="fastsafetensors",  # or "safetensors"
    download_dir="./Llama-3.1-8B-Instruct-W4A8H100",
    use_tqdm_on_load=False
)
model_config = ModelConfig(
    model="./Llama-3.1-8B-Instruct-W4A8H100",
    quantization="w4a8"
)

loader = DefaultModelLoader(load_config)

# Iterate through all weights and print those in layer 17
print("\nTrying to load Layer 17 weights:")
for name, tensor in loader.get_all_weights(model_config, None):
    if "layers.17" in name:
        print("Loaded:", name, "->", tensor.shape)