from transformers import AutoConfig

# 替换为你的模型名称或本地路径
config = AutoConfig.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")

# 获取总 Attention 头数 (Query Heads)
num_heads = getattr(config, "num_attention_heads", None)

# 获取 KV 头数 (Key/Value Heads, GQA 模型关键参数)
num_kv_heads = getattr(config, "num_key_value_heads", num_heads)

print(f"总注意力头数 (num_attention_heads): {num_heads}")
print(f"KV 注意力头数 (num_key_value_heads): {num_kv_heads}")