1-torch模型的相关操作

import json
from safetensors import safe_open

root = "/home/ctbots/llm/Shanghai_AI_Laboratory/internlm2-chat-7b"  # 放 shard 的目录
index_path = f"{root}/model.safetensors.index.json"

with open(index_path, "r") as f:
    index = json.load(f)

weight_map = index["weight_map"]   # key -> shard file

# 按 shard 分组 key（减少反复 open）
shard2keys = {}
for k, shard in weight_map.items():
    shard2keys.setdefault(shard, []).append(k)

total = 0

for shard, keys in shard2keys.items():
    shard_path = f"{root}/{shard}"
    print(f"\n=== Open shard: {shard} ===")

    with safe_open(shard_path, framework="pt", device="cpu") as f:
        for k in keys:
            t = f.get_tensor(k)
            numel = t.numel()
            total += numel
            print(
                f"{k:60s} | "
                f"dtype={t.dtype} | "
                f"shape={list(t.shape)} | "
                f"numel={numel}"
            )

print(f"\nTotal numel: {total}")