`text-generation` `Pipeline` prohibitively slow to load, even with cached model
I used thetimeitmodule to test the difference between including and excluding thedevice=0argument when instantiating a pipeline forgpt2, and found an enormous performance benefit of addingdevice=0; over 50 repetitions, the best time for usingdevice=0was 184 seconds, while the development node I was working on killed my process after 3 repetitions.However, it looks like this performance benefit does not uniformly apply across different models. I am using the following code to load an LLM as atext-generationPipelineinstance:from transformers import pipeline
pipe = pipeline('text-generation', model='MetaIX/GPT4-X-Alpasta-30b', device=0)Even withdevice=0, my process gets killed before completion of a single load (processes on this node are killed when they reach 2 hours of CPU time).What else can I do to speed up the loading of LLMs into pipeline objects?EDIT: TheMetaIX/GPT4-X-Alpasta-30bmodel is already downloaded locally, and this still happens
2023-05-23T18:42:45Z
Here’s what I figured out in case it’s helpful to anyone else.It seems like the model I chose was just substantially larger & required a lot more memory than thegpt2model. In order to get both model loading and inference to work without OOM Errors, I used the following code to generate text for a givenprompt:# Load model and tokenizer
checkpoint = 'MetaIX/GPT4-X-Alpasta-30b'
model = AutoModelForCausalLM.from_pretrained(checkpoint,
        torch_dtype=torch.float16, device_map='auto', load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# Generate predictions
inputs = tokenizer(prompt, return_tensors='pt')
inputs = inputs.to(0)
output = model.generate(inputs['input_ids'], max_new_tokens=500)
response = tokenizer.decode(output[0].tolist())The key for me was theload_in_8bitparameter. While I could load the model just fine without it, I would get CUDA OOM errors when it came time to usemodel.generate()for inference. Inference is still pretty slow, but at least it doesn’t get killed!