from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from torch.profiler import ProfilerActivity
import os

from torchao.quantization import quantize_
from torchao.quantization import float8_dynamic_activation_float8_weight
from torchao.quantization.observer import PerRow

os.environ["TOKENIZERS_PARALLELISM"] = "false"
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="cuda:4",
)

model.generation_config.cache_implementation = "static"

quantize_(model, float8_dynamic_activation_float8_weight(granularity=PerRow()))
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

input_text = "Hello my name is "
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda:4")

torch.profiler._utils._init_for_cuda_graphs()
prof = torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    # with_stack=True,
    # record_shapes=True,
)

with prof:
    outputs = model.generate(**input_ids, max_new_tokens=2)

print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
prof.export_chrome_trace("llama_fp8_dynamic_trace.json.gz")