I am training an Unsloth model in a Google Colab notebook. When I reach the Trainer.train() step. And I run the cell, it throws this error:

PicklingError: Can't pickle : it's not the same object as trl.trainer.sft_config.SFTConfig

I have the Google Colab Pro Plus plan, I have tried it on all the heavy-duty GPUs (H100, A100, L4, T4, and High-RAM), none worked if you look at the code, I am even using Google's sample json data. I have even used data from Hugging Faceyahma/alpaca-cleaned.

This is the error

PicklingError                             Traceback (most recent call last)
/tmp/ipykernel_22154/2279315892.py in ()
----> 1 trainer_stats = trainer.train()

10 frames
/usr/local/lib/python3.12/dist-packages/torch/serialization.py in _save(obj, zip_file, pickle_module, pickle_protocol, _disable_byteorder_record)
   1225 
   1226     pickler = PyTorchPickler(data_buf, protocol=pickle_protocol)
-> 1227     pickler.dump(obj)
   1228 
   1229     # The class def keeps the persistent_id closure alive, leaking memory.

PicklingError: Can't pickle : it's not the same object as trl.trainer.sft_config.SFTConfig

This is my training configuration cell code before I run the trainer.train() in the next cell then I get the piclke error.

    #Train the model using HuggingFace TRLs wait for the trainer variable to be created
    import sys
    import importlib
    import torch
    from datasets import load_dataset
    # Force reload TRL components to sync memory references
    if "trl" in sys.modules:
        importlib.reload(sys.modules["trl"])
    
    from transformers import TrainingArguments
    from unsloth import is_bfloat16_supported
    from trl import SFTConfig, SFTTrainer
    
    trainer = SFTTrainer(
        output_dir = "/content/drive/MyDrive/outputDir",
        model = model,
        tokenizer = tokenizer,
        train_dataset = dataset,
        dataset_text_field = "text",
        max_seq_length = max_seq_length,
        dataset_num_proc = 2,
        packing = False, # Can make training 5x faster for short sequences.
        args = TrainingArguments(
            per_device_train_batch_size = 1,#it makes no difference when it is 2
            gradient_accumulation_steps = 1,#when i set the gradient_accumulation_steps to 1or23o4 the loss decreasa up to steps 7 and8, then it starts to increse again
            warmup_steps = 1,
             num_train_epochs = 1, # Set this for 1 full training run.
            gradient_checkpointing = True,
            max_steps = 60,
            learning_rate = 2e-4,
            fp16 = not is_bfloat16_supported(),
            bf16 = is_bfloat16_supported(),
            logging_steps = 1,
            optim = "adamw_8bit",
            weight_decay = 0.01,
            lr_scheduler_type = "linear",
            seed = 3407,
            
            report_to = "none", # Use this for WandB etc
        ),
    )

that is the training cell code below the one above

trainer_stats = trainer.train()