[Python] How to run inference large size model in multi-GPU effeciently?

I'm trying to run only inference with large 70B sized model with multi-GPU env, but facing some issues.

1. The loading time takes so long, about 15mins.
2. I'm not sure this works properly to shard model over multi-GPU automatically and do their jobs in batch. Currently, it keeps error with such messages. I guess if it worked well, using only 4 GPU with this sized model should properly work.

[rank4]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 448.00 MiB. GPU 0 has a total capacity of 47.51 GiB of which 319.56 MiB is free. Including non-PyTorch memory, this process has 24.97 GiB memory in use. Process 1068220 has 22.22 GiB memory in use. Of the allocated memory 24.55 GiB is allocated by PyTorch, and 1.56 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)


Is there any suggestions to fix my code?

p.s. Also wonder what would be the good prompts for the Instruction models? I know there's tons of prompt engineering methods, but just curious how to control the output quality well for the Llama-Instruction model.

import os
import torch
import pandas as pd
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM


os.environ["HF_HOME"] = "/mystorage"
os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,5,6,7"


MODEL_ID = "meta-llama/Meta-Llama-3-70B-Instruct"
CSV_FILE = "my_csv.csv"
COLUMN_NAME = "column1"
ADDITIONAL_COLUMNS = ["column2"]
OUTPUT_FILE = "output.csv"
BATCH_SIZE = 32


print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
tokenizer.pad_token_id = tokenizer.eos_token_id

model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map="auto",
torch_dtype=torch.float16,
)


print("Loading data...")
data = pd.read_csv(CSV_FILE)
cells = data[COLUMN_NAME].dropna().tolist()
additional_data = data[ADDITIONAL_COLUMNS].loc[data[COLUMN_NAME].notna()]

PROMPTS = [
f"myprompt {cell}\n Given this cell, summarize in one sentence concisely."
for cell in cells
]


results = []


print("Starting inference...")
progress_bar = tqdm(total=len(PROMPTS))

for i in range(0, len(PROMPTS), BATCH_SIZE):
batch_prompts = PROMPTS[i:i + BATCH_SIZE]


inputs = tokenizer(batch_prompts, return_tensors="pt", padding=True, truncation=True).to("cuda")

with torch.no_grad():
output_ids = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=40,
pad_token_id=tokenizer.eos_token_id,
)


decoded_outputs = [
tokenizer.decode(output, skip_special_tokens=True).split(prompt)[-1].strip()
for output, prompt in zip(output_ids, batch_prompts)
]
results.extend(zip(batch_prompts, decoded_outputs))
progress_bar.update(len(batch_prompts))


print("Saving results...")
result_df = pd.DataFrame(results, columns=["Prompt", "Result"])
result_df = pd.concat([result_df, additional_data.reset_index(drop=True)], axis=1)
result_df.to_csv(OUTPUT_FILE, index=False)

print(f"Results saved to '{OUTPUT_FILE}'")

Continue reading...