# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

import torch
from peft import get_peft_model, LoraConfig, TaskType, PeftModel
from accelerate import Accelerator, FullyShardedDataParallelPlugin

torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2
)
accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B")
max_seq_length = 2048

dtype = torch.bfloat16
device_index = torch.cuda.current_device()
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B", # Note, warning about sliding window attn is not real: https://github.com/huggingface/transformers/pull/36316
    torch_dtype=dtype,
)

lora_config = LoraConfig(
    r = 32,
    lora_alpha = 128,
    target_modules = "all-linear",
    lora_dropout = 0,
    bias = "none",
    task_type = TaskType.CAUSAL_LM,
)

print(model.dtype)
# Get LoRA and setup model
model = get_peft_model(model, lora_config)

# Make sure we don't require gradients on non-lora params
with torch.no_grad():
    for name, param in model.named_parameters():
        if ".lora_A." in name or ".lora_B." in name: param.requires_grad_(True)
        else: param.requires_grad_(False)

model.gradient_checkpointing_enable()
model.enable_input_require_grads()

optimizer = torch.optim.Adam(model.parameters())

model.train()

print(f"[Device {device_index}] Memory summary before accelerator prepare")
print(torch.cuda.memory_summary())

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

torch.cuda.empty_cache()
print(f"[Device {device_index}] Memory summary after accelerator prepare")
print(torch.cuda.memory_summary())

print(f"[Device {device_index}] Is accelerator FSDP2? {accelerator.is_fsdp2}")

sample = "Lorem ipsum dolor sit amet consectetur adipiscing elit. Quisque faucibus ex sapien vitae pellentesque sem placerat."


enc = tokenizer(sample)
tokens = torch.tensor(enc.input_ids).to(model.device)
attn_mask = torch.tensor(enc.attention_mask).to(model.device)
input_ids = tokens[None, :-1]
labels = tokens[None, 1:]

out = model(input_ids, labels=labels)
loss = out['loss']

accelerator.backward(loss)
optimizer.step()
optimizer.zero_grad()

torch.cuda.empty_cache()
print(f"[Device {device_index}] Memory summary after one backward pass")
print(torch.cuda.memory_summary())