Running AITextGen on Apple Silicon Mac

Question

I'm using a MacBook Pro with an M1 chip, and I am having trouble getting aitextgen to recognise its GPU -- it looks for a CUDA GPU, which Apple Silicon Macs don't use. It works without using the GPU but is of course slower.

Below is the code I ran:

# Info on GPT Neo Models: https://huggingface.co/models?other=gpt_neo

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


# This code is from Apple and it invokes MPS : https://developer.apple.com/metal/pytorch/

import torch
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.")    

ai = aitextgen(model="EleutherAI/gpt-neo-1.3b", to_gpu=True)

Below is the output from Python:

tensor([1.], device='mps:0')Generate config GenerationConfig {  "_from_model_config": true,  "bos_token_id": 50256,  "eos_token_id": 50256,  "transformers_version": "4.28.1"} loading file vocab.json from cache at aitextgen/models--EleutherAI--gpt-neo-1.3b/snapshots/8282180b53cba30a1575e49de1530019e5931739/vocab.jsonloading file merges.txt from cache at aitextgen/models--EleutherAI--gpt-neo-1.3b/snapshots/8282180b53cba30a1575e49de1530019e5931739/merges.txtloading file tokenizer.json from cache at Noneloading file added_tokens.json from cache at Noneloading file special_tokens_map.json from cache at aitextgen/models--EleutherAI--gpt-neo-1.3b/snapshots/8282180b53cba30a1575e49de1530019e5931739/special_tokens_map.json loading file tokenizer_config.json from cache at aitextgen/models--EleutherAI--gpt-neo-1.3b/snapshots/8282180b53cba30a1575e49de1530019e5931739/tokenizer_config.json AssertionError: CUDA is not installed.

It looks like aitextgen doesn't support the MPS GPU framework from Apple. Has anyone had any luck using aitextgen and Huggingface models on an Apple Silicon Mac?

Answer 1

You can see in the code that there's no way to explicitly set an arbitrary device:

https://github.com/minimaxir/aitextgen/blob/master/aitextgen/aitextgen.py

If you want it to work, you could modify the code yourself (and maybe submit a PR to them).