first commit

dev/tools/group_weight_map_structure.py

@@ -0,0 +1,189 @@
+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()