import torch, os, gc
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig, AdamW, get_scheduler, BitsAndBytesConfig
from torch.utils.data import DataLoader, Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from tqdm import tqdm
from torch.cuda.amp import GradScaler, autocast
from peft import LoraConfig, get_peft_model
from accelerate import Accelerator

# Crucial: Configure environment variables BEFORE importing torch
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"  # Reduce allocation block size
os.environ["PYTORCH_NO_CUDA_MEMORY_CACHING"] = "1"  # Disable CUDA memory caching

# Initialize Accelerator
accelerator = Accelerator()
device = accelerator.device

# Quantization Configuration
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Model and Tokenizer loading
model_name = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/kaggle/input/llama/other/default/1/models/")
tokenizer.pad_token = tokenizer.eos_token
config = AutoConfig.from_pretrained(model_name, cache_dir="/kaggle/input/llama/other/default/1/models/")
config.pad_token_id = tokenizer.pad_token_id
config.num_labels = 3

model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    config=config,
    cache_dir="/kaggle/input/llama/other/default/1/models/",
    quantization_config=bnb_config,
    device_map={"": device} if device.type == "cuda" else None  # Load on the accelerator's device
)

# LoRA configuration
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Data Loading and Preprocessing
df = pd.read_csv("/kaggle/input/amazon-reviews/limited_reviews.csv")
train_texts, val_texts, train_labels, val_labels = train_test_split(
    df['text'].astype(str).tolist(), df['sentiment'].tolist(), test_size=0.2, random_state=42
)

label_map = {'positive': 2, 'negative': 0, 'neutral': 1}
train_labels = [label_map.get(label, -1) for label in train_labels]
val_labels = [label_map.get(label, -1) for label in val_labels]

class SentimentDataset(Dataset):
    def __init__(self, texts, labels):
        self.texts = texts
        self.labels = labels

    def __getitem__(self, idx):
        return {'text': self.texts[idx], 'labels': self.labels[idx]}

    def __len__(self):
        return len(self.labels)

def custom_collate_fn(batch):
    texts = [item['text'] for item in batch]
    labels = [item['labels'] for item in batch]
    encodings = tokenizer(texts, truncation=True, padding=True, return_tensors='pt')
    encodings['labels'] = torch.tensor(labels)
    return encodings

train_dataset = SentimentDataset(train_texts, train_labels)
val_dataset = SentimentDataset(val_texts, val_labels)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=custom_collate_fn)
val_loader = DataLoader(val_dataset, batch_size=4, collate_fn=custom_collate_fn)

# Training with Gradient Accumulation, Mixed Precision, and LoRA
optimizer = AdamW(model.parameters(), lr=5e-5)
num_epochs = 3
accumulation_steps = 32
num_training_steps = (len(train_loader) // accumulation_steps) * num_epochs
lr_scheduler = get_scheduler("linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps)
scaler = torch.amp.GradScaler()  # Updated Scaler instantiation

# Prepare with accelerator (This will ensure data loaders are on the right device)
train_loader, lr_scheduler = accelerator.prepare(train_loader, lr_scheduler)

model.train()
for epoch in range(num_epochs):
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs}", disable=not accelerator.is_local_main_process)
    optimizer.zero_grad()
    for i, batch in enumerate(progress_bar):
        batch = {k: v.to(device) for k, v in batch.items()}  # Move batch to the device explicitly
        with autocast():
            outputs = model(**batch)
            loss = outputs.loss.mean() / accumulation_steps
        accelerator.backward(scaler.scale(loss))

        if (i + 1) % accumulation_steps == 0:
            scaler.step(optimizer)
            scaler.update()
            lr_scheduler.step()
            optimizer.zero_grad()

            gc.collect()
            torch.cuda.empty_cache()

        total_loss += loss.item() * accumulation_steps
        progress_bar.set_postfix({"loss": total_loss / (i + 1)})

    avg_train_loss = total_loss / len(train_loader)
    accelerator.print(f"Epoch {epoch + 1}/{num_epochs}, Average Training Loss: {avg_train_loss}")

# Evaluation
model.eval()
predictions = []
true_labels = []
with torch.no_grad():
    for batch in tqdm(val_loader, desc="Evaluating", disable=not accelerator.is_local_main_process):
        batch = {k: v.to(device) for k, v in batch.items()}  # Move batch to the device explicitly
        outputs = model(**batch)
        logits = outputs.logits
        preds = torch.argmax(logits, dim=1)
        predictions.extend(preds.cpu().tolist())
        true_labels.extend(batch['labels'].cpu().tolist())

accuracy = accuracy_score(true_labels, predictions)
precision, recall, f1, _ = precision_recall_fscore_support(true_labels, predictions, average='weighted')

accelerator.print(f"Accuracy: {accuracy}")
accelerator.print(f"Precision: {precision}")
accelerator.print(f"Recall: {recall}")
accelerator.print(f"F1 Score: {f1}")

# Saving the model
accelerator.wait_for_everyone()
if accelerator.is_main_process:
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained("fine_tuned_llama_sentiment_model")
    tokenizer.save_pretrained("fine_tuned_llama_sentiment_model")

Some weights of LlamaForSequenceClassification were not initialized from the model checkpoint at meta-llama/Llama-3.2-1B-Instruct and are newly initialized: ['score.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
trainable params: 851,968 || all params: 1,236,672,512 || trainable%: 0.0689
/opt/conda/lib/python3.10/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning
  warnings.warn(
Epoch 1/3:   0%|          | 0/6000 [00:00<?, ?it/s]/tmp/ipykernel_30/131544246.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with autocast():
Epoch 1/3:  61%|██████▏   | 3679/6000 [54:34<34:25,  1.12it/s, loss=1.03]  
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[2], line 109
    107 batch = {k: v.to(device) for k, v in batch.items()}  # Move batch to the device explicitly
    108 with autocast():
--> 109     outputs = model(**batch)
    110     loss = outputs.loss.mean() / accumulation_steps
    111 accelerator.backward(scaler.scale(loss))

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/peft/peft_model.py:1644, in PeftModelForCausalLM.forward(self, input_ids, attention_mask, inputs_embeds, labels, output_attentions, output_hidden_states, return_dict, task_ids, **kwargs)
   1642     with self._enable_peft_forward_hooks(**kwargs):
   1643         kwargs = {k: v for k, v in kwargs.items() if k not in self.special_peft_forward_args}
-> 1644         return self.base_model(
   1645             input_ids=input_ids,
   1646             attention_mask=attention_mask,
   1647             inputs_embeds=inputs_embeds,
   1648             labels=labels,
   1649             output_attentions=output_attentions,
   1650             output_hidden_states=output_hidden_states,
   1651             return_dict=return_dict,
   1652             **kwargs,
   1653         )
   1655 batch_size = _get_batch_size(input_ids, inputs_embeds)
   1656 if attention_mask is not None:
   1657     # concat prompt attention mask

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/peft/tuners/tuners_utils.py:197, in BaseTuner.forward(self, *args, **kwargs)
    196 def forward(self, *args: Any, **kwargs: Any):
--> 197     return self.model.forward(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/accelerate/hooks.py:170, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    168         output = module._old_forward(*args, **kwargs)
    169 else:
--> 170     output = module._old_forward(*args, **kwargs)
    171 return module._hf_hook.post_forward(module, output)

File /opt/conda/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:1282, in LlamaForSequenceClassification.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1274 r"""
   1275 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
   1276     Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
   1277     config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
   1278     `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
   1279 """
   1280 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-> 1282 transformer_outputs = self.model(
   1283     input_ids,
   1284     attention_mask=attention_mask,
   1285     position_ids=position_ids,
   1286     past_key_values=past_key_values,
   1287     inputs_embeds=inputs_embeds,
   1288     use_cache=use_cache,
   1289     output_attentions=output_attentions,
   1290     output_hidden_states=output_hidden_states,
   1291     return_dict=return_dict,
   1292 )
   1293 hidden_states = transformer_outputs[0]
   1294 logits = self.score(hidden_states)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/accelerate/hooks.py:170, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    168         output = module._old_forward(*args, **kwargs)
    169 else:
--> 170     output = module._old_forward(*args, **kwargs)
    171 return module._hf_hook.post_forward(module, output)

File /opt/conda/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:945, in LlamaModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
    933     layer_outputs = self._gradient_checkpointing_func(
    934         decoder_layer.__call__,
    935         hidden_states,
   (...)
    942         position_embeddings,
    943     )
    944 else:
--> 945     layer_outputs = decoder_layer(
    946         hidden_states,
    947         attention_mask=causal_mask,
    948         position_ids=position_ids,
    949         past_key_value=past_key_values,
    950         output_attentions=output_attentions,
    951         use_cache=use_cache,
    952         cache_position=cache_position,
    953         position_embeddings=position_embeddings,
    954     )
    956 hidden_states = layer_outputs[0]
    958 if use_cache:

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/accelerate/hooks.py:170, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    168         output = module._old_forward(*args, **kwargs)
    169 else:
--> 170     output = module._old_forward(*args, **kwargs)
    171 return module._hf_hook.post_forward(module, output)

File /opt/conda/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:692, in LlamaDecoderLayer.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position, position_embeddings, **kwargs)
    690 residual = hidden_states
    691 hidden_states = self.post_attention_layernorm(hidden_states)
--> 692 hidden_states = self.mlp(hidden_states)
    693 hidden_states = residual + hidden_states
    695 outputs = (hidden_states,)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/accelerate/hooks.py:170, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    168         output = module._old_forward(*args, **kwargs)
    169 else:
--> 170     output = module._old_forward(*args, **kwargs)
    171 return module._hf_hook.post_forward(module, output)

File /opt/conda/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py:258, in LlamaMLP.forward(self, x)
    256     down_proj = sum(down_proj)
    257 else:
--> 258     down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
    260 return down_proj

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File /opt/conda/lib/python3.10/site-packages/accelerate/hooks.py:170, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    168         output = module._old_forward(*args, **kwargs)
    169 else:
--> 170     output = module._old_forward(*args, **kwargs)
    171 return module._hf_hook.post_forward(module, output)

File /opt/conda/lib/python3.10/site-packages/bitsandbytes/nn/modules.py:484, in Linear4bit.forward(self, x)
    481     x = x.to(self.compute_dtype)
    483 bias = None if self.bias is None else self.bias.to(self.compute_dtype)
--> 484 out = bnb.matmul_4bit(x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state)
    486 out = out.to(inp_dtype)
    488 return out

File /opt/conda/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py:579, in matmul_4bit(A, B, quant_state, out, bias)
    577         return out
    578 else:
--> 579     return MatMul4Bit.apply(A, B, out, bias, quant_state)

File /opt/conda/lib/python3.10/site-packages/torch/autograd/function.py:574, in Function.apply(cls, *args, **kwargs)
    571 if not torch._C._are_functorch_transforms_active():
    572     # See NOTE: [functorch vjp and autograd interaction]
    573     args = _functorch.utils.unwrap_dead_wrappers(args)
--> 574     return super().apply(*args, **kwargs)  # type: ignore[misc]
    576 if not is_setup_ctx_defined:
    577     raise RuntimeError(
    578         "In order to use an autograd.Function with functorch transforms "
    579         "(vmap, grad, jvp, jacrev, ...), it must override the setup_context "
    580         "staticmethod. For more details, please see "
    581         "https://pytorch.org/docs/main/notes/extending.func.html"
    582     )

File /opt/conda/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py:509, in MatMul4Bit.forward(ctx, A, B, out, bias, quant_state)
    505         return torch.empty(A.shape[:-1] + B_shape[:1], dtype=A.dtype, device=A.device)
    507 # 1. Dequantize
    508 # 2. MatmulnN
--> 509 output = torch.nn.functional.linear(A, F.dequantize_4bit(B, quant_state).to(A.dtype).t(), bias)
    511 # 3. Save state
    512 ctx.state = quant_state

RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

fix and return full code nothing else.