如何在Colab的GPU显存中以8位精度加载google/flan-t5-xl模型

先补全必要依赖 #

Colab环境下先确保安装支持8位量化的核心库：

!pip install -q bitsandbytes transformers accelerate sentencepiece

正确的8位量化加载代码 #

不要和torch.set_default_tensor_type混用，直接用以下代码加载——device_map='auto'本身就会自动把模型分配到CUDA（只要GPU可用），之前加载失败大概率是依赖缺失或版本不兼容导致的：

from transformers import T5ForConditionalGeneration, AutoTokenizer
import torch

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-xl')
# 加载8位量化模型
model = T5ForConditionalGeneration.from_pretrained(
    'google/flan-t5-xl',
    load_in_8bit=True,
    device_map='auto',
    torch_dtype=torch.float16  # 配合8位量化进一步压缩显存占用
)

关键注意点 #

• 确认Colab分配的GPU支持8位量化：T4、A100等型号都支持，若分到K80则不兼容，可重启会话更换GPU。
• 别再用torch.set_default_tensor_type(torch.cuda.FloatTensor)，这会干扰bitsandbytes的8位张量处理逻辑。
• 要是还碰显存瓶颈，可加offload_folder='./offload'参数，把部分模型组件暂存到磁盘（会小幅降低速度）。

验证模型是否在CUDA上 #

加载完成后运行以下代码确认：

print(next(model.parameters()).device)
# 正常输出应为 cuda:0

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