loading file tokenizer.json from cache at /kaggle/input/llama/other/default/1/models/models--meta-llama--Llama-3.2-1B-Instruct/snapshots/9213176726f574b556790deb65791e0c5aa438b6/tokenizer.json
loading file tokenizer.model from cache at None
loading file added_tokens.json from cache at None
loading file special_tokens_map.json from cache at /kaggle/input/llama/other/default/1/models/models--meta-llama--Llama-3.2-1B-Instruct/snapshots/9213176726f574b556790deb65791e0c5aa438b6/special_tokens_map.json
loading file tokenizer_config.json from cache at /kaggle/input/llama/other/default/1/models/models--meta-llama--Llama-3.2-1B-Instruct/snapshots/9213176726f574b556790deb65791e0c5aa438b6/tokenizer_config.json
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
loading configuration file config.json from cache at /kaggle/input/llama/other/default/1/models/models--meta-llama--Llama-3.2-1B-Instruct/snapshots/9213176726f574b556790deb65791e0c5aa438b6/config.json
Model config LlamaConfig {
  "_name_or_path": "meta-llama/Llama-3.2-1B-Instruct",
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": [
    128001,
    128008,
    128009
  ],
  "head_dim": 64,
  "hidden_act": "silu",
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "intermediate_size": 8192,
  "max_position_embeddings": 131072,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 16,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": {
    "factor": 32.0,
    "high_freq_factor": 4.0,
    "low_freq_factor": 1.0,
    "original_max_position_embeddings": 8192,
    "rope_type": "llama3"
  },
  "rope_theta": 500000.0,
  "tie_word_embeddings": true,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.46.2",
  "use_cache": true,
  "vocab_size": 128256
}

loading weights file model.safetensors from cache at /kaggle/input/llama/other/default/1/models/models--meta-llama--Llama-3.2-1B-Instruct/snapshots/9213176726f574b556790deb65791e0c5aa438b6/model.safetensors
Number of GPUs available: 2
All model checkpoint weights were used when initializing LlamaForSequenceClassification.

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: 858,112 || all params: 1,236,678,656 || trainable%: 0.0694
/opt/conda/lib/python3.10/site-packages/transformers/training_args.py:1568: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead
  warnings.warn(
PyTorch: setting up devices
/tmp/ipykernel_30/3441435683.py:182: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `CustomTrainer.__init__`. Use `processing_class` instead.
  trainer = CustomTrainer(
You have loaded a model on multiple GPUs. `is_model_parallel` attribute will be force-set to `True` to avoid any unexpected behavior such as device placement mismatching.
Using auto half precision backend
***** Running training *****
  Num examples = 24,000
  Num Epochs = 3
  Instantaneous batch size per device = 2
  Total train batch size (w. parallel, distributed & accumulation) = 32
  Gradient Accumulation steps = 16
  Total optimization steps = 2,250
  Number of trainable parameters = 858,112
/opt/conda/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  self.pid = os.fork()
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
 [ 8/2250 00:34 < 3:32:29, 0.18 it/s, Epoch 0.01/3]

```
import os
# Configure environment variables BEFORE importing torch
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
os.environ["PYTORCH_NO_CUDA_MEMORY_CACHING"] = "1"

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoConfig,
    Trainer,
    TrainingArguments,
    DataCollatorWithPadding,
)
from torch.utils.data import Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from peft import LoraConfig, get_peft_model
import numpy as np
from transformers.utils import logging

# Set up logging
logging.set_verbosity_info()

# Model and Tokenizer loading
model_name = "meta-llama/Llama-3.2-1B-Instruct"
cache_dir = "/kaggle/input/llama/other/default/1/models/"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    cache_dir=cache_dir,
    local_files_only=True,
    padding_side="right",
)
tokenizer.pad_token = tokenizer.eos_token

# Load the model configuration
config = AutoConfig.from_pretrained(
    model_name,
    cache_dir=cache_dir,
    local_files_only=True,
)
config.pad_token_id = tokenizer.pad_token_id
config.num_labels = 3

# Check if GPUs are available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    n_gpus = torch.cuda.device_count()
    print(f"Number of GPUs available: {n_gpus}")
else:
    n_gpus = 0
    print("No GPUs available.")

# Load the model
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    config=config,
    cache_dir=cache_dir,
    local_files_only=True,
    device_map={"": device} if device.type == "cuda" else None,  # Load on the specified 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="SEQ_CLS",
)

# Apply LoRA to the model
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Ensure LoRA modules are on the correct devices
def move_lora_to_device(model):
    for name, module in model.named_modules():
        if hasattr(module, 'lora_A') or hasattr(module, 'lora_B'):
            device = next(module.parameters()).device
            if hasattr(module, 'lora_A'):
                module.lora_A.to(device)
            if hasattr(module, 'lora_B'):
                module.lora_B.to(device)

move_lora_to_device(model)

# 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.encodings = tokenizer(
            texts,
            truncation=True,
            padding=True,
            max_length=512,
            return_tensors="pt",
        )
        self.labels = torch.tensor(labels)

    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item['labels'] = self.labels[idx]
        return item

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

train_dataset = SentimentDataset(train_texts, train_labels)
val_dataset = SentimentDataset(val_texts, val_labels)

# Data collator
data_collator = DataCollatorWithPadding(tokenizer, padding='longest', return_tensors='pt')

# Define the metrics function
def compute_metrics(pred):
    labels = pred.label_ids
    preds = np.argmax(pred.predictions, axis=-1)
    accuracy = accuracy_score(labels, preds)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted', zero_division=0)
    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1
    }

# Training arguments
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=16,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    fp16=True,
    fp16_full_eval=True,
    save_total_limit=2,
    load_best_model_at_end=True,
    logging_steps=50,
    dataloader_num_workers=2,
    report_to="none",
)

# Custom Trainer to handle device placement of labels
from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        # Move inputs to the correct device
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        # Forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # Move labels to the same device as logits
        labels = labels.to(logits.device)
        # Compute loss
        loss_fct = torch.nn.CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

# Initialize the Custom Trainer
trainer = CustomTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# Train the model
trainer.train()

# Evaluate the model
eval_results = trainer.evaluate()
print(f"Evaluation results: {eval_results}")

# Save the model and tokenizer
trainer.save_model("fine_tuned_llama_sentiment_model")
tokenizer.save_pretrained("fine_tuned_llama_sentiment_model")
```

i have two gpus out of which this is using only 1 (Vram and Cores). Please fix this (kaggle notebook). And return full code. Thanks!