нехватка памяти при запуске модели llama2-13b-chat на машине с несколькими графическими процессорами

Question

Я пытаюсь запустить модель llama2 13b с веревочным масштабированием на машине AWS g4dn.12xlarge с 4 графическими процессорами по 16 ГБ видеопамяти каждый, но возникает ошибка cuda из памяти.

Код:

      from transformers import AutoModelForCausalLM, AutoTokenizer,pipeline
import transformers

tokenizer = AutoTokenizer.from_pretrained("TheBloke/Llama-2-13B-Chat-fp16", use_fast=False)
model = AutoModelForCausalLM.from_pretrained("TheBloke/Llama-2-13B-Chat-fp16", device_map = 'auto',
                                             **{"rope_scaling":{"factor": 2.0,"type": "linear"}}
    )

user_prompt = "..."

pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto",
)


sequences = pipeline(
   user_prompt,
    max_length=8000,
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
)

print(sequences)

Это ошибка, которую я получаю, когда размер приглашения превышает 4 КБ.

torch.cuda.OutOfMemoryError: CUDA не хватает памяти. Попыталось выделить 4,51 ГиБ (GPU 0; общая емкость 14,61 ГиБ; 11,92 ГиБ уже выделено; 1,76 ГиБ свободно; всего 12,14 ГиБ зарезервировано PyTorch). Если зарезервированная память >> выделена, попробуйте установить max_split_size_mb, чтобы избежать фрагментации. См. документацию по управлению памятью и PYTORCH_CUDA_ALLOC_CONF.

64 ГБ недостаточно для запуска модели с контекстом 8k или в моем коде ошибка?