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..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..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..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..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..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.