Mistral-7B-v0.1模型参数计数异常及显存占用合理性咨询

问题背景 #

加载mistralai/Mistral-7B-v0.1模型时，通过遍历model.parameters()计数得到约3.7B参数，但model.num_parameters()返回预期的约7B参数。同时对4bit量化后的显存占用有疑问，具体问题如下：

1. 遍历参数计数的操作存在哪些错误？模型分片是否会影响计算结果？
2. 7B参数在4bit量化下，model.get_memory_footprint()返回4.55GB是否合理？

代码示例 #

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

base_model_id = "mistralai/Mistral-7B-v0.1"

# 创建量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

# 加载量化后的模型
model = AutoModelForCausalLM.from_pretrained(
        base_model_id, quantization_config=bnb_config
)

# 自定义参数计数函数
def print_parameters(model):
    all_param = 0
    for param in model.parameters():
        all_param += param.numel()
    print(
        f"all params: {all_param} ||"
    )
    
print_parameters(model)
>>> 3752071168
print(model.num_parameters())
>>> 7241732096

依赖库版本 #

transformers==4.36.1
torch==2.0.1
bitsandbytes==0.41.3.post2

问题解答 #

1. 参数计数差异的原因 #

• 遍历model.parameters()的局限性：使用BitsAndBytes做4bit量化时，模型核心的线性层权重会被转为4bit专用张量格式，这类量化后的参数不会出现在model.parameters()的迭代列表中。该方法仅返回未被量化的参数（比如偏置、层归一化参数）和量化辅助参数，所以计数结果远低于实际总参数。
• model.num_parameters()的准确性：这个方法会统计模型的全部参数，包括被量化的权重，因此返回的是正确的7B左右总参数数。
• 模型分片的影响：模型分片只是为了拆分权重到不同设备以实现大模型加载，不会影响参数计数结果。无论是否分片，model.num_parameters()都会统计全量参数，model.parameters()的计数逻辑也依然只返回未量化部分。

2. 4bit量化后显存占用的合理性 #

理论上7B参数4bit量化的显存占用为：7e9 * 4bit / 8 = 3.5GB，但实际返回4.55GB是合理的，原因包括：

• 量化辅助参数开销：BitsAndBytes的4bit量化会额外存储缩放因子、偏移量等辅助参数，这些参数通常以bfloat16精度存储，会占用额外显存。
• 计算 dtype 相关缓存：配置中bnb_4bit_compute_dtype=torch.bfloat16意味着计算时会把4bit参数反量化为bfloat16，model.get_memory_footprint()会统计这类反量化相关的缓存和中间结构。
• 未量化组件的显存：模型的嵌入层、层归一化、偏置等参数并未被量化，仍以原精度（bfloat16或float32）存储，这部分也会占用额外显存。

内容的提问来源于stack exchange，提问作者Rgkpdx