meta-llama

/

Meta-Llama-3.1-8B-Instruct

like 2.51k

I am running the model on RTXA5000 with 24g memory, which should satisfy the need, yet when I run the code, it still output CUDA issue. How should I fix this?

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, pipeline
import torch

# Function to clean CUDA memory
def clean_cuda_memory():
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

# Clean CUDA memory before starting
clean_cuda_memory()

model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Load the model configuration and modify the rope_scaling parameter
model_config = AutoConfig.from_pretrained(model_id)
model_config.rope_scaling = {"type": "linear", "factor": 8.0}  # Adjust to the required format

# Load the model with the modified configuration
model = AutoModelForCausalLM.from_pretrained(model_id, config=model_config, torch_dtype=torch.float32)

# Initialize the text generation pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=1
)

system_role = '''
MY DESIGN FOR SYSTEM ROLE
'''

user_input = '''
MY INPUT
'''

# Define the messages
messages = [
    {"role": "system", "content": system_role},
    {"role": "user", "content": user_input},
]

# Concatenate messages into a single prompt
prompt = ""
for message in messages:
    if message["role"] == "system":
        prompt += f"System: {message['content']}\n"
    elif message["role"] == "user":
        prompt += f"User: {message['content']}\n"

# Generate text based on the prompt
output = pipe(
    prompt,
    max_new_tokens=512,
    eos_token_id=tokenizer.eos_token_id,
    do_sample=True,
    temperature=0.6,
    top_p=0.9,
)

# Extract the generated text
generated_text = output[0]["generated_text"]

# Print the response
print(generated_text)

The error message:
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB (GPU 1; 23.69 GiB total capacity; 22.86 GiB already allocated; 128.06 MiB free; 22.87 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF