import os                       
import torch                    
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig
from trl import SFTTrainer


dataset = load_dataset("csv", data_files="dataset/data.csv")

base_model = "meta-llama/Llama-3.2-1B"
compute_dtype = getattr(torch, "float16")

# Configure memory-efficient quantization
compute_dtype = getattr(torch, "float16")
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=True,  # Enable double quantization
)

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=quant_config,
    device_map="auto",  # Let transformers handle device mapping
    torch_dtype=torch.float16,  # Use fp16 for model weights
    low_cpu_mem_usage=True,    # Enable memory optimization
)

torch.cuda.empty_cache()
model.config.use_cache = False
model.config.pretraining_tp = 1


# Configure PEFT using LoRA for efficient fine-tuning of the model.
peft_params = LoraConfig(
    lora_alpha=16,    
    lora_dropout=0.1, 
    r=8,              
    bias="none",       
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)

training_params = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=2,  
    gradient_accumulation_steps=2,   
    optim="paged_adamw_8bit",    
    save_steps=50,
    logging_steps=50,
    learning_rate=2e-4,
    weight_decay=0.001,
    fp16=True,                   
    bf16=False,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
    report_to="tensorboard",
    gradient_checkpointing=True,   
)


tokenizer = AutoTokenizer.from_pretrained(
    base_model,
    padding_side="right",
    truncation_side="right",
)
tokenizer.pad_token = tokenizer.eos_token

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset['train'],
    peft_config=peft_params,
    dataset_text_field="input_text",
    max_seq_length=512,
    tokenizer=tokenizer,
    args=training_params,
    packing=False,
)

trainer.train()

I am using the above code block to fine tune Llama 1-B parameters using my PC which has 128 GB RAM and a 4090 GPU. Accordingly the PC has all the requirements met for the model but while executing the trainer = SFTTrainer(....) line the RAM gets filled up and the kernel dies which is unexpected. The dataset size is only 10 GB with 7400 rows of data. Would be glad if anyone could help me solve this issue. The error traceback is like this (it uses up all the 128GB memory and the terminal is shut down)

Deprecated positional argument(s) used in SFTTrainer, please use the SFTConfig to set these arguments instead.
  warnings.warn(message, FutureWarning)
/home/.../python3.10/site-packages/trl/trainer/sft_trainer.py:300: UserWarning: You passed a `max_seq_length` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.
  warnings.warn(
/home/.../python3.10/site-packages/trl/trainer/sft_trainer.py:328: UserWarning: You passed a `dataset_text_field` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.
  warnings.warn(
Map:  14%|██████████████████▍                                                                                                                    | 1000/7346 [02:59<18:56,  5.58 examples/s]
Killed

I tried all the changes and configuration I could find on Google and previous relevant Stack Overflow blogs but none of those solved the issue.