在Apple Silicon Mac上运行aitextgen的GPU识别问题求助

问题核心在于aitextgen的to_gpu=True参数仅针对CUDA设备做了适配，默认不识别Apple Silicon的MPS框架，因此会抛出CUDA未安装的错误。要利用MPS加速，需手动指定设备，而非依赖内置的to_gpu参数。

修改后的代码示例 #

# Info on GPT Neo Models
import numpy as np
import pandas as pd 
from aitextgen.TokenDataset import TokenDataset 
from aitextgen.tokenizers import train_tokenizer
from aitextgen.utils import GPT2ConfigCPU
from aitextgen import aitextgen

import torch
# 确认MPS设备可用
if torch.backends.mps.is_available():    
    mps_device = torch.device("mps")    
    x = torch.ones(1, device=mps_device)    
    print(x)
else:    
    print("MPS device not found.")    

# 初始化模型时不使用to_gpu参数
ai = aitextgen(model="EleutherAI/gpt-neo-1.3b")
# 手动将模型移到MPS设备
ai.model = ai.model.to(mps_device)
# 适配tokenizer上下文长度
ai.tokenizer.model_max_length = ai.config.n_ctx

生成文本时的配置 #

调用generate方法时，需明确指定使用MPS设备：

ai.generate(device=mps_device, max_length=200)

额外注意事项 #

• PyTorch版本要求：必须使用支持MPS的PyTorch版本（建议1.12及以上），确保环境已适配Apple Silicon。
• 内存限制：1.3B参数的GPT-Neo模型在M1/M2芯片上运行，建议至少16GB内存；内存不足时可尝试更小的模型（如125M版本）。
• 源码定制（可选）：若频繁使用，可修改aitextgen源码中的GPU检测逻辑，将torch.backends.mps.is_available()加入判断条件，替换原有的CUDA单一检测逻辑，实现to_gpu参数自动适配MPS。

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