repro.py

from transformers import AutoModelForCausalLM
from accelerate import Accelerator
import torch

torch.cuda.memory._record_memory_history()

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
accelerator = Accelerator()

model = AutoModelForCausalLM.from_pretrained(model_id)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)

model, optimizer = accelerator.prepare(model, optimizer)

print(model)

torch.cuda.memory._dump_snapshot("pre-patch.pkl")