transformer_parameters.py

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b',
                                              trust_remote_code=True,
                                              torch_dtype=torch.bfloat16,
                                            )
model.eval()
model.cuda()

for name, param in model.named_parameters():
    print(f"{name}   Modelsize: {param.numel()/1000**2:.1f}M parameters")
    if "31" not in name:
        param.requires_grad = False
    print(name, param.requires_grad)