CPU环境下torch.float16调用memory_efficient_attention_forward报错求助

问题 #

在搭载Intel Core i7处理器的MacBook Pro上，使用transformers库的AutoModelForCausalLM加载PyTorch模型THUDM/cogvlm-chat-hf时，出现memory_efficient_attention_forward算子不支持CPU与torch.float16组合的报错。

运行代码 #

import torch
import requests
from PIL import Image
from IPython.display import display
from transformers import AutoModelForCausalLM, LlamaTokenizer

# Load tokenizer and model
tokenizer = LlamaTokenizer.from_pretrained('lmsys/vicuna-7b-v1.5')
model = AutoModelForCausalLM.from_pretrained(
    'THUDM/cogvlm-chat-hf',
    torch_dtype=torch.float16,  # Using torch.float16
    low_cpu_mem_usage=True,
    trust_remote_code=True
).eval()

def generate(query: str, img_url: str, max_length: int = 2048) -> str:
    image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
    display(image)

    # Generate token inputs
    inputs = model.build_conversation_input_ids(tokenizer, query=query, history=[], images=[image], template_version='vqa')

    # Convert tensors to appropriate types
    input_ids = inputs['input_ids'].unsqueeze(0).to(torch.long)
    token_type_ids = inputs['token_type_ids'].unsqueeze(0).to(torch.long)
    attention_mask = inputs['attention_mask'].unsqueeze(0).to(torch.float16)
    images = [[inputs['images'][0].to(torch.float16)]]

    inputs = {
        'input_ids': input_ids,
        'token_type_ids': token_type_ids,
        'attention_mask': attention_mask,
        'images': images,
    }
    
    gen_kwargs = {"max_length": max_length, "do_sample": False}
    
    with torch.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, input_ids.shape[1]:]
        return tokenizer.decode(outputs[0])

query = 'Describe this image in detail'
img_url = 'https://i.ibb.co/x1nH9vr/Slide1.jpg'
generate(query, img_url)

报错信息 #

NotImplementedError: No operator found for `memory_efficient_attention_forward` with inputs:
     query       : shape=(1, 1226, 16, 112) (torch.float16)
     key         : shape=(1, 1226, 16, 112) (torch.float16)
     value       : shape=(1, 1226, 16, 112) (torch.float16)
     attn_bias   : <class 'NoneType'>
     p           : 0.0
`ck_decoderF` is not supported because:
    device=cpu (supported: {'cuda'})
    operator wasn't built - see `python -m xformers.info` for more info
`ckF` is not supported because:
    device=cpu (supported: {'cuda'})
    operator wasn't built - see `python -m xformers.info` for more info

解决方案 #

方案1：改用float32数据类型 #

xformers的内存高效注意力算子不支持CPU上的float16计算，切换为torch.float32即可避开该问题，修改点如下：

• 模型加载时指定torch_dtype=torch.float32
• 调整输入张量的类型为float32

修改后的代码片段：

# 加载模型时修改dtype
model = AutoModelForCausalLM.from_pretrained(
    'THUDM/cogvlm-chat-hf',
    torch_dtype=torch.float32,  # 改为float32
    low_cpu_mem_usage=True,
    trust_remote_code=True
).eval()

# generate函数中修改张量类型
attention_mask = inputs['attention_mask'].unsqueeze(0).to(torch.float32)
images = [[inputs['images'][0].to(torch.float32)]]

方案2：禁用xformers，使用PyTorch原生注意力实现 #

通过指定attn_implementation="eager"，强制模型使用PyTorch原生的注意力计算逻辑，无需修改数据类型：

model = AutoModelForCausalLM.from_pretrained(
    'THUDM/cogvlm-chat-hf',
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    attn_implementation="eager"  # 添加该参数禁用xformers
).eval()

方案3：内存优化补充（可选） #

如果使用float32时内存不足，可尝试调整内存参数：

model = AutoModelForCausalLM.from_pretrained(
    'THUDM/cogvlm-chat-hf',
    torch_dtype=torch.float32,
    low_cpu_mem_usage=False,
    trust_remote_code=True,
    gradient_checkpointing=True  # 开启梯度检查点节省内存
).eval()

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