LLM ingores max_memory in inference
It seems, like my model ingores max_memory parametre. It crashes withtorch.cuda.OutOfMemoryError: CUDA out of memory.. And the problem is that first 3 generations goes well. Just model allocates 25GB, GPU Memory is 80GB, but than it’s just grow. I’ve tried to delete all objects, but it doesn’t work for me.UPD: Forgot to mention. In error said, that 79 GB allocated and used, by I’ve specified only 70 in the model.Model init code:base_model = 'google/flan-t5-xxl'
ckpt = './results/checkpoints_t5_1/checkpoint-4600'
device = 'cuda:3'

tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)

model = AutoModelForSeq2SeqLM.from_pretrained(
    base_model,
    device_map=device,
    max_memory={3:'70GB'},
    trust_remote_code=True,
    torch_dtype=torch.float16,
)

model = PeftModel.from_pretrained(
    model,
    ckpt
)

model.config.use_cache = False
model.eval()
model = torch.compile(model)Generation config and code:generation_config = GenerationConfig(
    do_sample=True,
    temperature=0.8,
    top_p=0.75,
    top_k=40,
    num_beams=4,
    max_new_tokens=256,
    stream_output=False,
    model=model,
)
inputs = create_prompt(example)
inputs = inputs.to(device)
top1 = example['input'][0]
gt = example['output']

with torch.inference_mode():
    generation_output = model.generate(
    input_ids=inputs["input_ids"],
    generation_config=generation_config,
    return_dict_in_generate=False,
    output_scores=False,
)