190 lines
6.5 KiB
Python
190 lines
6.5 KiB
Python
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import re
|
||
|
|
|
||
|
|
from collections import defaultdict
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
LAYER_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
|
||
|
|
CANDIDATE_INDEX_FILES = [
|
||
|
|
"model.safetensors.index.json",
|
||
|
|
"consolidated.safetensors.index.json",
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_numeric_segments(name: str) -> str:
|
||
|
|
"""
|
||
|
|
Replace pure numeric path segments with '*'.
|
||
|
|
Example: model.layers.12.mlp.experts.3.up_proj.weight
|
||
|
|
-> model.layers.*.mlp.experts.*.up_proj.weight
|
||
|
|
"""
|
||
|
|
return ".".join("*" if part.isdigit() else part for part in name.split("."))
|
||
|
|
|
||
|
|
|
||
|
|
def load_weight_map(input_json_path: str) -> dict[str, str]:
|
||
|
|
with open(input_json_path, "r") as f:
|
||
|
|
data = json.load(f)
|
||
|
|
|
||
|
|
if isinstance(data, dict) and isinstance(data.get("weight_map"), dict):
|
||
|
|
return data["weight_map"]
|
||
|
|
|
||
|
|
if isinstance(data, dict) and all(
|
||
|
|
isinstance(k, str) and isinstance(v, str) for k, v in data.items()
|
||
|
|
):
|
||
|
|
return data
|
||
|
|
|
||
|
|
raise ValueError(
|
||
|
|
"Input json must be either {'weight_map': {...}} or a plain key->file mapping."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def find_index_json(input_dir: str) -> str:
|
||
|
|
for name in CANDIDATE_INDEX_FILES:
|
||
|
|
path = os.path.join(input_dir, name)
|
||
|
|
if os.path.exists(path):
|
||
|
|
return path
|
||
|
|
raise ValueError(f"Can not find index file in {input_dir}")
|
||
|
|
|
||
|
|
|
||
|
|
def load_weight_map_from_input_dir(input_dir: str) -> tuple[dict[str, str], str]:
|
||
|
|
index_json_path = find_index_json(input_dir)
|
||
|
|
return load_weight_map(index_json_path), index_json_path
|
||
|
|
|
||
|
|
|
||
|
|
def build_structure_report(
|
||
|
|
weight_map: dict[str, str],
|
||
|
|
max_examples: int,
|
||
|
|
max_pattern_preview: int,
|
||
|
|
include_all_layer_patterns: bool,
|
||
|
|
) -> dict:
|
||
|
|
pattern_to_keys: dict[str, list[str]] = defaultdict(list)
|
||
|
|
layer_to_suffix_patterns: dict[int, set[str]] = defaultdict(set)
|
||
|
|
non_layer_patterns: set[str] = set()
|
||
|
|
|
||
|
|
for key in weight_map.keys():
|
||
|
|
normalized_full = normalize_numeric_segments(key)
|
||
|
|
pattern_to_keys[normalized_full].append(key)
|
||
|
|
|
||
|
|
layer_match = LAYER_KEY_PATTERN.match(key)
|
||
|
|
if layer_match is None:
|
||
|
|
non_layer_patterns.add(normalized_full)
|
||
|
|
continue
|
||
|
|
|
||
|
|
layer_id = int(layer_match.group(1))
|
||
|
|
layer_suffix = layer_match.group(2)
|
||
|
|
normalized_suffix = normalize_numeric_segments(layer_suffix)
|
||
|
|
layer_to_suffix_patterns[layer_id].add(normalized_suffix)
|
||
|
|
|
||
|
|
normalized_key_patterns = []
|
||
|
|
for pattern, keys in sorted(
|
||
|
|
pattern_to_keys.items(),
|
||
|
|
key=lambda item: (-len(item[1]), item[0]),
|
||
|
|
):
|
||
|
|
normalized_key_patterns.append(
|
||
|
|
{
|
||
|
|
"pattern": pattern,
|
||
|
|
"count": len(keys),
|
||
|
|
"example_keys": sorted(keys)[:max_examples],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
all_patterns = [item["pattern"] for item in normalized_key_patterns]
|
||
|
|
|
||
|
|
signature_to_layers: dict[tuple[str, ...], list[int]] = defaultdict(list)
|
||
|
|
for layer_id, suffix_patterns in layer_to_suffix_patterns.items():
|
||
|
|
signature = tuple(sorted(suffix_patterns))
|
||
|
|
signature_to_layers[signature].append(layer_id)
|
||
|
|
|
||
|
|
layer_structure_groups = []
|
||
|
|
sorted_signature_items = sorted(
|
||
|
|
signature_to_layers.items(),
|
||
|
|
key=lambda item: (len(item[1]), item[1][0]),
|
||
|
|
reverse=True,
|
||
|
|
)
|
||
|
|
for idx, (signature, layer_ids) in enumerate(sorted_signature_items, start=1):
|
||
|
|
layer_ids = sorted(layer_ids)
|
||
|
|
entry = {
|
||
|
|
"group_id": f"layer_struct_{idx:03d}",
|
||
|
|
"layer_ids": layer_ids,
|
||
|
|
"layer_count": len(layer_ids),
|
||
|
|
"pattern_count": len(signature),
|
||
|
|
"pattern_preview": list(signature)[:max_pattern_preview],
|
||
|
|
}
|
||
|
|
if include_all_layer_patterns:
|
||
|
|
entry["patterns"] = list(signature)
|
||
|
|
layer_structure_groups.append(entry)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"all_patterns": all_patterns,
|
||
|
|
"total_weight_keys": len(weight_map),
|
||
|
|
"normalized_key_pattern_count": len(normalized_key_patterns),
|
||
|
|
"normalized_key_patterns": normalized_key_patterns,
|
||
|
|
"layer_ids_seen": sorted(layer_to_suffix_patterns.keys()),
|
||
|
|
"layer_structure_group_count": len(layer_structure_groups),
|
||
|
|
"layer_structure_groups": layer_structure_groups,
|
||
|
|
"non_layer_pattern_count": len(non_layer_patterns),
|
||
|
|
"non_layer_patterns": sorted(non_layer_patterns),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def get_default_output_path(input_json_path: str) -> str:
|
||
|
|
input_path = Path(input_json_path)
|
||
|
|
return str(input_path.with_name(f"{input_path.stem}.layer_structures.json"))
|
||
|
|
|
||
|
|
|
||
|
|
def resolve_io_paths(
|
||
|
|
input_dir: str | None, input_json: str | None, output_json: str | None
|
||
|
|
) -> tuple[dict[str, str], str, str]:
|
||
|
|
if not input_dir and not input_json:
|
||
|
|
raise ValueError("Either --input-dir or --input-json is required.")
|
||
|
|
if input_dir and input_json:
|
||
|
|
raise ValueError("Please provide only one of --input-dir or --input-json.")
|
||
|
|
|
||
|
|
if input_dir:
|
||
|
|
weight_map, input_json_path = load_weight_map_from_input_dir(input_dir)
|
||
|
|
else:
|
||
|
|
input_json_path = input_json
|
||
|
|
weight_map = load_weight_map(input_json_path)
|
||
|
|
|
||
|
|
output_json_path = output_json or get_default_output_path(input_json_path)
|
||
|
|
return weight_map, input_json_path, output_json_path
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Group weight_map keys by normalized pattern and layer structure."
|
||
|
|
)
|
||
|
|
parser.add_argument("--input-dir", default=None, type=str)
|
||
|
|
parser.add_argument("--input-json", default=None, type=str)
|
||
|
|
parser.add_argument("--output-json", default=None, type=str)
|
||
|
|
parser.add_argument("--max-examples", default=3, type=int)
|
||
|
|
parser.add_argument("--max-pattern-preview", default=20, type=int)
|
||
|
|
parser.add_argument("--include-all-layer-patterns", action="store_true")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
weight_map, input_json_path, output_json = resolve_io_paths(
|
||
|
|
input_dir=args.input_dir,
|
||
|
|
input_json=args.input_json,
|
||
|
|
output_json=args.output_json,
|
||
|
|
)
|
||
|
|
report = build_structure_report(
|
||
|
|
weight_map=weight_map,
|
||
|
|
max_examples=args.max_examples,
|
||
|
|
max_pattern_preview=args.max_pattern_preview,
|
||
|
|
include_all_layer_patterns=args.include_all_layer_patterns,
|
||
|
|
)
|
||
|
|
|
||
|
|
with open(output_json, "w") as f:
|
||
|
|
json.dump(report, f, indent=2)
|
||
|
|
|
||
|
|
print(f"Saved grouped structure report to: {output_json}")
|
||
|
|
print(f"Loaded weight_map from: {input_json_path}")
|
||
|
|
print(f"Total keys: {report['total_weight_keys']}")
|
||
|
|
print(f"Layer structure groups: {report['layer_structure_group_count']}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|