40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Patch _allocate_kv_cache_tensors to print the layer name mismatch."""
|
|
import sys
|
|
|
|
def patch(path):
|
|
with open(path, 'r') as f:
|
|
content = f.read()
|
|
|
|
if "CLAWMINE_DEBUG_LAYERS" in content:
|
|
print("Already patched, skipping")
|
|
return
|
|
|
|
old = """ assert layer_names == set(kv_cache_raw_tensors.keys()), (
|
|
"Some layers are not correctly initialized"
|
|
)"""
|
|
|
|
new = """ # CLAWMINE_DEBUG_LAYERS: print mismatch instead of asserting
|
|
missing = layer_names - set(kv_cache_raw_tensors.keys())
|
|
extra = set(kv_cache_raw_tensors.keys()) - layer_names
|
|
if missing or extra:
|
|
print(f"CLAWMINE DEBUG: missing layers ({len(missing)}): {sorted(missing)[:20]}")
|
|
print(f"CLAWMINE DEBUG: extra layers ({len(extra)}): {sorted(extra)[:20]}")
|
|
print(f"CLAWMINE DEBUG: expected ({len(layer_names)}), got ({len(kv_cache_raw_tensors.keys())})")
|
|
assert layer_names == set(kv_cache_raw_tensors.keys()), (
|
|
"Some layers are not correctly initialized"
|
|
)"""
|
|
|
|
if old not in content:
|
|
print("ERROR: Could not find the code to patch")
|
|
sys.exit(1)
|
|
|
|
content = content.replace(old, new)
|
|
|
|
with open(path, 'w') as f:
|
|
f.write(content)
|
|
print("Patched gpu_model_runner.py for debug layer names")
|
|
|
|
if __name__ == "__main__":
|
|
patch(sys.argv[1])
|