[P1] Weird Memory issues during finetuning

[LuisMalhadas]

I have been trying to use this to finetune a llama3.2 3B instruct model.

After some adaptation from previous iterations, I have the following script:

import torch
torch.cuda.empty_cache()
import json
import transformers
import pyreft
from transformers import BitsAndBytesConfig, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model

device = "cuda" if torch.cuda.is_available() else "cpu"
model_max_length = 8192 
dtype = None 
load_in_4bit = True 
access_token = "....."

model_name_or_path = "meta-llama/Llama-3.2-3B-Instruct"
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_name_or_path, 
    torch_dtype=torch.bfloat16, 
    device_map=device
)

# get tokenizer
tokenizer = transformers.AutoTokenizer.from_pretrained(
    model_name_or_path, 
    model_max_length=model_max_length, 
    padding_side="right", 
    use_fast=False, 
    token=access_token,
)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.eos_token='<|eot_id|>'

terminators = [
    tokenizer.eos_token_id,
]

prompt_no_input_template = """<|start_header_id|>system<|end_header_id|>
You are a helpful AI assistant for summarizing text in easy to read format.
<|eot_id|>
<|start_header_id|>user<|end_header_id|>
%s
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
"""

peft_config = LoraConfig(
    r=4, 
    lora_alpha=32, 
    target_modules=["o_proj"], 
    layers_to_transform=[15],
    use_rslora=True, 
    lora_dropout=0.05, 
    bias="none", 
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, peft_config)

reft_config = pyreft.ReftConfig(representations=[{
    "layer": l, 
    "component": f"base_model.model.model.layers[{l}].output",
    "low_rank_dimension": 4,
    "intervention": pyreft.LoreftIntervention(embed_dim=model.config.hidden_size,
    low_rank_dimension=4)} for l in [15]])

reft_model = pyreft.get_reft_model(model, reft_config)

# you need to call this to re-enable lora grads!
reft_model.model.enable_adapter_layers()
reft_model.print_trainable_parameters()

# position info about the interventions
share_weights = False # whether the prefix and suffix interventions sharing weights.
positions="f3+l3"    # the intervening positions of prefix tokens (f[irst]1) and suffix tokens (l[ast]1).
first_n, last_n = pyreft.parse_positions(positions)

from datasets import load_dataset, Dataset
# read the txt file line by line
training_examples = []
with open("datasets/validated_easyreading.jsonl", "r") as f:
    for line in f:
        data = json.loads(line.strip())
        input_text = data.get("hard")
        output_text = data.get("easy")
        training_examples.append([input_text, output_text])

print("Trainning Examples: \n", training_examples[0][0], "\n", training_examples[0][1])

data_module = pyreft.make_multiple_position_supervised_data_module(
    tokenizer, model, 
    [prompt_no_input_template % e[0] for e in training_examples],
    [e[1] for e in training_examples], 
    positions=positions, 
    num_interventions=len(reft_config.representations), 
    share_weights=share_weights, 
    nonstop=False
)

# train
training_args = transformers.TrainingArguments(
    num_train_epochs=1, 
    output_dir="./tmp", 
    per_device_train_batch_size=1, 
    learning_rate=4e-3, 
    report_to=[], 
    logging_steps=10
)
trainer = pyreft.ReftTrainerForCausalLM(
    model=reft_model, 
    tokenizer=tokenizer,
    args=training_args, 
    **data_module
)
_ = trainer.train()

...

I'm having this extremely weird issue, where the train fails 'randomly' alternating between the following errors:

CUDA_VISIBLE_DEVICES=0 python train_ReFT.py 
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00,  1.06it/s]
trainable intervention params: 24,580 || trainable model params: 24,576
model params: 3,212,774,400 || trainable%: 0.0015300171714515655
Trainning Examples: 
 - * 86 profesionales en 50 centros* orientan con metodologías y recursos a profesionales de centros ordinarios para que aumente la inclusión en las aulas.  Una llave para avanzar hacia la educación inclusiva podría estar en Cataluña y tiene una sigla larga. CEEPSIR significa Centro de Educación Especial Proveedor de Servicios y Recursos nacidos para apoyar la inclusión. Se trata de profesionales de Centros de Educación Especial que están orientando a profesoras y profesores de centros de educación ordinaria para aumentar la inclusión en las aulas. De esta manera, se logra que niñas, niños y jóvenes con discapacidad puedan tener la opción a aprender en espacios ordinarios, con niños sin discapacidad y los apoyos necesarios. Los CEEPSIR pertenecen al Departamento de Educación de la Generalitat de Catalunya y se crearon tras el Decreto 150/2017 de Educación Inclusiva . “Con el desarrollo de este proyecto se ofrece una alternativa viable a la necesidad de extender la educación inclusiva sin perder el valor y la experiencia de los centros de educación especial”, señala Carmen Laucirica, presidenta de Plena inclusión. Y añade: “Son centros pioneros con los que estamos colaborando desde nuestro movimiento asociativo”. Precisamente, con la intención de conocer su trabajo y transmitir aprendizajes y retos, durante mayo tuvo lugar una visita a varios de estos centros en Barcelona en la que participó Francisco Javier Amaya Flores, subdirector de la Subdirección General de Cooperación Territorial e Innovación Educativa del Ministerio de Educación. También participó en la visita Beatriz Vega, responsable del área de Educación de Plena inclusión España, que destaca que “no se enfocan tanto en atender casos individuales de niñas y niños sino en ayudar al centro a contar con herramientas para solucionar no casos puntuales, sino situaciones diversas”. La labor de los CEEPSIR consiste en aportar propuestas para exponer las asignaturas de forma sencilla y adaptada a alumnado con y sin discapacidad, materiales en lectura fácil, cómo señalizar el centro para facilitar la orientación o cómo agrupar al alumnado. Dincat Plena inclusión Cataluña cuenta con 28 de estos centros y, desde el inicio de su creación, consideran que los CEEPSIR son una herramienta imprescindible en el camino hacia la inclusión. “Aquí en Cataluña se ha hecho una gran apuesta en este sentido los últimos años y el resultado es más que positivo. Aun sí, todavía falta mejorar muchos aspectos como la coordinación entre las y los diferentes agentes implicados, la planificación del recurso, la definición de los indicadores o la evaluación del impacto que tiene el CEEPSIR tanto para el alumnado como para el entorno ordinario. Por ello, esperamos que se siga apostando por este modelo, invirtiendo los recursos necesarios para su buen desarrollo”, explica Alba Cortina, representante de Educación en Dincat y directora de la Escola Jeroni de Moragas de AMPANS , uno de los centros CEEPSIR de Cataluña. A su vez, Plena inclusión España ha elaborado dos estudios con la Universidad Autónoma de Madrid -que se presentarán en octubre- para conocer cómo transformar centros de educación especial en centros de recursos para la educación inclusiva. Y también, con el objetivo de profundizar en la visión de los propios centros de educación especial sobre esta transformación. Plena inclusión España es la confederación agrupa a 950 asociaciones a través de federaciones en todas las comunidades y ciudades autónomas. Estas entidades apoyan a más de 150.0000 personas con discapacidad intelectual y del desarrollo y a sus familias, desde 4.000 centros y servicios, y gracias al compromiso de 46.000 profesionales y más de 8.000 personas voluntarias, promoviendo la igualdad en el acceso a derechos fundamentales como el empleo, la educación, el ocio, la cultura o la salud, entre otros. 
 Una llave para avanzar
 hacia la educación inclusiva
 podrían ser los centros CEEPSIR.
  ¿Qué es la educación inclusiva? 
 Son 50 centros con 86 profesionales.
 Su trabajo es ayudar
 a otras y otros profesores
 para que la educación sea más inclusiva.
 Por ejemplo, dan ideas como:

 
 -  Cómo explicar los temas
 que las niñas, niños y jóvenes
 deben aprender
 de forma fácil de entender. 
 
 -  Hacer materiales en lectura fácil. 
 
 -  Poner señales en los centros
 para que nadie se pierda. 
 
 -  Cómo hacer grupos
 con las alumnas y los alumnos
 para que aprendan mejor.  Estos centros están en Cataluña.
 Es la comunidad autónoma que antes
 ha puesto en marcha esta idea.
 Dincat Plena inclusión Cataluña
 cuenta con 28 de estos centros.
 Plena inclusión España ha hecho 2 estudios
 con la Universidad Autónoma de Madrid
 sobre estos centros.
 En octubre presentará los resultados.
{'loss': 2.0008, 'grad_norm': 11.575387001037598, 'learning_rate': 0.003963702359346642, 'epoch': 0.01}                                                                   
{'loss': 1.8515, 'grad_norm': 2.7920827865600586, 'learning_rate': 0.003927404718693285, 'epoch': 0.02}                                                                   
{'loss': 1.7237, 'grad_norm': 3.1544675827026367, 'learning_rate': 0.0038911070780399274, 'epoch': 0.03}                                                                  
{'loss': 1.6755, 'grad_norm': 6.611184597015381, 'learning_rate': 0.00385480943738657, 'epoch': 0.04}                                                                     
{'loss': 1.911, 'grad_norm': 20.23436164855957, 'learning_rate': 0.0038185117967332125, 'epoch': 0.05}                                                                    
{'loss': 1.8584, 'grad_norm': 1.7269971370697021, 'learning_rate': 0.003782214156079855, 'epoch': 0.05}                                                                   
{'loss': 1.9298, 'grad_norm': 2.129941940307617, 'learning_rate': 0.0037459165154264972, 'epoch': 0.06}                                                                   
{'loss': 1.9631, 'grad_norm': 4.218372821807861, 'learning_rate': 0.00370961887477314, 'epoch': 0.07}                                                                     
{'loss': 1.8263, 'grad_norm': 9.363015174865723, 'learning_rate': 0.0036733212341197824, 'epoch': 0.08}                                                                   
{'loss': 1.701, 'grad_norm': 2.8993072509765625, 'learning_rate': 0.0036370235934664245, 'epoch': 0.09}                                                                   
{'loss': 1.7242, 'grad_norm': 4.215240001678467, 'learning_rate': 0.0036007259528130675, 'epoch': 0.1}                                                                    
{'loss': 2.1187, 'grad_norm': 2.250020980834961, 'learning_rate': 0.0035644283121597097, 'epoch': 0.11}                                                                   
{'loss': 1.7096, 'grad_norm': 2.840819835662842, 'learning_rate': 0.0035281306715063522, 'epoch': 0.12}                                                                   
{'loss': 1.4442, 'grad_norm': 4.925025939941406, 'learning_rate': 0.003491833030852995, 'epoch': 0.13}                                                                    
{'loss': 1.8288, 'grad_norm': 4.087110996246338, 'learning_rate': 0.003455535390199637, 'epoch': 0.14}                                                                    
{'loss': 1.7688, 'grad_norm': 2.643681049346924, 'learning_rate': 0.0034192377495462795, 'epoch': 0.15}                                                                   
{'loss': 1.7245, 'grad_norm': 3.4367027282714844, 'learning_rate': 0.003382940108892922, 'epoch': 0.15}                                                                   
{'loss': 1.8289, 'grad_norm': 3.4109046459198, 'learning_rate': 0.0033466424682395647, 'epoch': 0.16}                                                                     
{'loss': 1.5225, 'grad_norm': 4.121783256530762, 'learning_rate': 0.003310344827586207, 'epoch': 0.17}                                                                    
{'loss': 1.7151, 'grad_norm': 3.042985200881958, 'learning_rate': 0.00327404718693285, 'epoch': 0.18}                                                                     
{'loss': 1.9969, 'grad_norm': 10.616415977478027, 'learning_rate': 0.003237749546279492, 'epoch': 0.19}                                                                   
{'loss': 1.8255, 'grad_norm': 3.870309352874756, 'learning_rate': 0.003201451905626134, 'epoch': 0.2}                                                                     
{'loss': 1.8267, 'grad_norm': 5.23873233795166, 'learning_rate': 0.003165154264972777, 'epoch': 0.21}                                                                     
{'loss': 1.689, 'grad_norm': 2.7276580333709717, 'learning_rate': 0.0031288566243194193, 'epoch': 0.22}                                                                   
{'loss': 1.7351, 'grad_norm': 2.4301626682281494, 'learning_rate': 0.003092558983666062, 'epoch': 0.23}                                                                   
{'loss': 1.7584, 'grad_norm': 3.251612424850464, 'learning_rate': 0.0030562613430127044, 'epoch': 0.24}                                                                   
{'loss': 1.6712, 'grad_norm': 2.4075534343719482, 'learning_rate': 0.003019963702359347, 'epoch': 0.25}                                                                   
{'loss': 1.6811, 'grad_norm': 1.869033932685852, 'learning_rate': 0.002983666061705989, 'epoch': 0.25}                                                                    
{'loss': 1.6329, 'grad_norm': 8.135303497314453, 'learning_rate': 0.0029473684210526317, 'epoch': 0.26}                                                                   
{'loss': 1.5052, 'grad_norm': 6.392016410827637, 'learning_rate': 0.0029110707803992743, 'epoch': 0.27}                                                                   
{'loss': 1.5604, 'grad_norm': 3.028606653213501, 'learning_rate': 0.0028747731397459164, 'epoch': 0.28}                                                                   
{'loss': 1.8425, 'grad_norm': 5.725487232208252, 'learning_rate': 0.0028384754990925594, 'epoch': 0.29}                                                                   
 29%|██████████████████████████████████████▎                                                                                           | 325/1102 [01:06<02:41,  4.82it/s]Traceback (most recent call last):
  File "/home/djfil/support-platform/experiments/llm-finetune-qlora/train_ReFT.py", line 115, in <module>
    _ = trainer.train()
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2052, in train
    return inner_training_loop(
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2388, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3518, in training_step
    self.accelerator.backward(loss, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/accelerate/accelerator.py", line 2196, in backward
    loss.backward(**kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/_tensor.py", line 581, in backward
    torch.autograd.backward(
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.17 GiB. GPU 0 has a total capacity of 23.58 GiB of which 1.33 GiB is free. Including non-PyTorch memory, this process has 22.24 GiB memory in use. Of the allocated memory 20.23 GiB is allocated by PyTorch, and 1.66 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
 29%|██▉       | 325/1102 [01:08<02:43,  4.77it/s]  

Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00,  1.05it/s]
trainable intervention params: 24,580 || trainable model params: 24,576
model params: 3,212,774,400 || trainable%: 0.0015300171714515655
Trainning Examples: 
 - * 86 profesionales en 50 centros* orientan con metodologías y recursos a profesionales de centros ordinarios para que aumente la inclusión en las aulas.  Una llave para avanzar hacia la educación inclusiva podría estar en Cataluña y tiene una sigla larga. CEEPSIR significa Centro de Educación Especial Proveedor de Servicios y Recursos nacidos para apoyar la inclusión. Se trata de profesionales de Centros de Educación Especial que están orientando a profesoras y profesores de centros de educación ordinaria para aumentar la inclusión en las aulas. De esta manera, se logra que niñas, niños y jóvenes con discapacidad puedan tener la opción a aprender en espacios ordinarios, con niños sin discapacidad y los apoyos necesarios. Los CEEPSIR pertenecen al Departamento de Educación de la Generalitat de Catalunya y se crearon tras el Decreto 150/2017 de Educación Inclusiva . “Con el desarrollo de este proyecto se ofrece una alternativa viable a la necesidad de extender la educación inclusiva sin perder el valor y la experiencia de los centros de educación especial”, señala Carmen Laucirica, presidenta de Plena inclusión. Y añade: “Son centros pioneros con los que estamos colaborando desde nuestro movimiento asociativo”. Precisamente, con la intención de conocer su trabajo y transmitir aprendizajes y retos, durante mayo tuvo lugar una visita a varios de estos centros en Barcelona en la que participó Francisco Javier Amaya Flores, subdirector de la Subdirección General de Cooperación Territorial e Innovación Educativa del Ministerio de Educación. También participó en la visita Beatriz Vega, responsable del área de Educación de Plena inclusión España, que destaca que “no se enfocan tanto en atender casos individuales de niñas y niños sino en ayudar al centro a contar con herramientas para solucionar no casos puntuales, sino situaciones diversas”. La labor de los CEEPSIR consiste en aportar propuestas para exponer las asignaturas de forma sencilla y adaptada a alumnado con y sin discapacidad, materiales en lectura fácil, cómo señalizar el centro para facilitar la orientación o cómo agrupar al alumnado. Dincat Plena inclusión Cataluña cuenta con 28 de estos centros y, desde el inicio de su creación, consideran que los CEEPSIR son una herramienta imprescindible en el camino hacia la inclusión. “Aquí en Cataluña se ha hecho una gran apuesta en este sentido los últimos años y el resultado es más que positivo. Aun sí, todavía falta mejorar muchos aspectos como la coordinación entre las y los diferentes agentes implicados, la planificación del recurso, la definición de los indicadores o la evaluación del impacto que tiene el CEEPSIR tanto para el alumnado como para el entorno ordinario. Por ello, esperamos que se siga apostando por este modelo, invirtiendo los recursos necesarios para su buen desarrollo”, explica Alba Cortina, representante de Educación en Dincat y directora de la Escola Jeroni de Moragas de AMPANS , uno de los centros CEEPSIR de Cataluña. A su vez, Plena inclusión España ha elaborado dos estudios con la Universidad Autónoma de Madrid -que se presentarán en octubre- para conocer cómo transformar centros de educación especial en centros de recursos para la educación inclusiva. Y también, con el objetivo de profundizar en la visión de los propios centros de educación especial sobre esta transformación. Plena inclusión España es la confederación agrupa a 950 asociaciones a través de federaciones en todas las comunidades y ciudades autónomas. Estas entidades apoyan a más de 150.0000 personas con discapacidad intelectual y del desarrollo y a sus familias, desde 4.000 centros y servicios, y gracias al compromiso de 46.000 profesionales y más de 8.000 personas voluntarias, promoviendo la igualdad en el acceso a derechos fundamentales como el empleo, la educación, el ocio, la cultura o la salud, entre otros. 
 Una llave para avanzar
 hacia la educación inclusiva
 podrían ser los centros CEEPSIR.
  ¿Qué es la educación inclusiva? 
 Son 50 centros con 86 profesionales.
 Su trabajo es ayudar
 a otras y otros profesores
 para que la educación sea más inclusiva.
 Por ejemplo, dan ideas como:

 
 -  Cómo explicar los temas
 que las niñas, niños y jóvenes
 deben aprender
 de forma fácil de entender. 
 
 -  Hacer materiales en lectura fácil. 
 
 -  Poner señales en los centros
 para que nadie se pierda. 
 
 -  Cómo hacer grupos
 con las alumnas y los alumnos
 para que aprendan mejor.  Estos centros están en Cataluña.
 Es la comunidad autónoma que antes
 ha puesto en marcha esta idea.
 Dincat Plena inclusión Cataluña
 cuenta con 28 de estos centros.
 Plena inclusión España ha hecho 2 estudios
 con la Universidad Autónoma de Madrid
 sobre estos centros.
 En octubre presentará los resultados.
{'loss': 1.7763, 'grad_norm': 3.616866111755371, 'learning_rate': 0.003963800904977376, 'epoch': 0.01}                                                                    
{'loss': 1.9981, 'grad_norm': 6.343026638031006, 'learning_rate': 0.003927601809954751, 'epoch': 0.02}                                                                    
{'loss': 2.0399, 'grad_norm': 3.420192241668701, 'learning_rate': 0.003891402714932127, 'epoch': 0.03}                                                                    
{'loss': 2.025, 'grad_norm': 22.433563232421875, 'learning_rate': 0.0038552036199095024, 'epoch': 0.04}                                                                   
{'loss': 2.0788, 'grad_norm': 5.083585262298584, 'learning_rate': 0.0038190045248868783, 'epoch': 0.05}                                                                   
{'loss': 1.7099, 'grad_norm': 7.139271259307861, 'learning_rate': 0.0037828054298642533, 'epoch': 0.05}                                                                   
{'loss': 2.0154, 'grad_norm': 5.275798320770264, 'learning_rate': 0.0037466063348416292, 'epoch': 0.06}                                                                   
{'loss': 1.8586, 'grad_norm': 3.426173686981201, 'learning_rate': 0.0037104072398190047, 'epoch': 0.07}                                                                   
{'loss': 1.8094, 'grad_norm': 5.704839706420898, 'learning_rate': 0.00367420814479638, 'epoch': 0.08}                                                                     
{'loss': 1.6718, 'grad_norm': 6.754937648773193, 'learning_rate': 0.0036380090497737556, 'epoch': 0.09}                                                                   
{'loss': 1.8681, 'grad_norm': 3.4839773178100586, 'learning_rate': 0.0036018099547511315, 'epoch': 0.1}                                                                   
{'loss': 1.7791, 'grad_norm': 2.5712482929229736, 'learning_rate': 0.003565610859728507, 'epoch': 0.11}                                                                   
{'loss': 1.874, 'grad_norm': 4.03897762298584, 'learning_rate': 0.0035294117647058825, 'epoch': 0.12}                                                                     
{'loss': 1.8406, 'grad_norm': 2.6527585983276367, 'learning_rate': 0.003493212669683258, 'epoch': 0.13}                                                                   
{'loss': 1.6617, 'grad_norm': 3.3806071281433105, 'learning_rate': 0.003457013574660634, 'epoch': 0.14}                                                                   
{'loss': 1.7789, 'grad_norm': 3.1879842281341553, 'learning_rate': 0.003420814479638009, 'epoch': 0.14}                                                                   
{'loss': 1.8814, 'grad_norm': 3.91975736618042, 'learning_rate': 0.003384615384615385, 'epoch': 0.15}                                                                     
{'loss': 1.9034, 'grad_norm': 4.3993000984191895, 'learning_rate': 0.0033484162895927603, 'epoch': 0.16}                                                                  
{'loss': 1.6992, 'grad_norm': 2.986009359359741, 'learning_rate': 0.0033122171945701357, 'epoch': 0.17}                                                                   
{'loss': 1.8968, 'grad_norm': 6.679447174072266, 'learning_rate': 0.003276018099547511, 'epoch': 0.18}                                                                    
{'loss': 1.6596, 'grad_norm': 5.53939151763916, 'learning_rate': 0.003239819004524887, 'epoch': 0.19}                                                                     
{'loss': 1.9562, 'grad_norm': 12.79240894317627, 'learning_rate': 0.0032036199095022626, 'epoch': 0.2}                                                                    
{'loss': 1.7928, 'grad_norm': 4.095736980438232, 'learning_rate': 0.003167420814479638, 'epoch': 0.21}                                                                    
{'loss': 1.6319, 'grad_norm': 3.0026490688323975, 'learning_rate': 0.0031312217194570135, 'epoch': 0.22}                                                                  
{'loss': 1.7756, 'grad_norm': 3.098665475845337, 'learning_rate': 0.0030950226244343894, 'epoch': 0.23}                                                                   
{'loss': 1.7177, 'grad_norm': 6.132549285888672, 'learning_rate': 0.0030588235294117644, 'epoch': 0.24}                                                                   
{'loss': 1.7903, 'grad_norm': 5.206614017486572, 'learning_rate': 0.0030226244343891403, 'epoch': 0.24}                                                                   
{'loss': 1.9207, 'grad_norm': 8.382951736450195, 'learning_rate': 0.002986425339366516, 'epoch': 0.25}                                                                    
{'loss': 1.7168, 'grad_norm': 3.0222861766815186, 'learning_rate': 0.0029502262443438917, 'epoch': 0.26}                                                                  
{'loss': 1.5926, 'grad_norm': 4.645003795623779, 'learning_rate': 0.0029140271493212667, 'epoch': 0.27}                                                                   
{'loss': 1.9577, 'grad_norm': 2.6779160499572754, 'learning_rate': 0.0028778280542986426, 'epoch': 0.28}                                                                  
{'loss': 1.7194, 'grad_norm': 4.0169596672058105, 'learning_rate': 0.002841628959276018, 'epoch': 0.29}                                                                   
{'loss': 1.5649, 'grad_norm': 2.4001996517181396, 'learning_rate': 0.0028054298642533936, 'epoch': 0.3}                                                                   
{'loss': 1.7419, 'grad_norm': 3.1197869777679443, 'learning_rate': 0.002769230769230769, 'epoch': 0.31}                                                                   
{'loss': 1.8444, 'grad_norm': 3.07463002204895, 'learning_rate': 0.002733031674208145, 'epoch': 0.32}                                                                     
{'loss': 1.5558, 'grad_norm': 2.684776544570923, 'learning_rate': 0.0026968325791855204, 'epoch': 0.33}                                                                   
{'loss': 1.7142, 'grad_norm': 2.771249771118164, 'learning_rate': 0.002660633484162896, 'epoch': 0.33}                                                                    
{'loss': 1.6536, 'grad_norm': 2.758201837539673, 'learning_rate': 0.0026244343891402714, 'epoch': 0.34}                                                                   
{'loss': 1.6682, 'grad_norm': 2.7577850818634033, 'learning_rate': 0.0025882352941176473, 'epoch': 0.35}                                                                  
{'loss': 1.6112, 'grad_norm': 1.2884979248046875, 'learning_rate': 0.0025520361990950223, 'epoch': 0.36}                                                                  
{'loss': 1.805, 'grad_norm': 4.771214962005615, 'learning_rate': 0.002515837104072398, 'epoch': 0.37}                                                                     
{'loss': 1.5651, 'grad_norm': 5.079401969909668, 'learning_rate': 0.0024796380090497737, 'epoch': 0.38}                                                                   
{'loss': 1.7925, 'grad_norm': 5.37261962890625, 'learning_rate': 0.0024434389140271496, 'epoch': 0.39}                                                                    
{'loss': 1.5472, 'grad_norm': 2.8810372352600098, 'learning_rate': 0.0024072398190045246, 'epoch': 0.4}                                                                   
{'loss': 1.63, 'grad_norm': 4.205113410949707, 'learning_rate': 0.0023710407239819005, 'epoch': 0.41}                                                                     
{'loss': 1.6719, 'grad_norm': 4.840310096740723, 'learning_rate': 0.0023348416289592764, 'epoch': 0.42}                                                                   
{'loss': 1.6965, 'grad_norm': 3.3432343006134033, 'learning_rate': 0.0022986425339366514, 'epoch': 0.43}                                                                  
{'loss': 1.6191, 'grad_norm': 4.2115960121154785, 'learning_rate': 0.0022624434389140274, 'epoch': 0.43}                                                                  
{'loss': 1.8101, 'grad_norm': 2.018566131591797, 'learning_rate': 0.002226244343891403, 'epoch': 0.44}                                                                    
{'loss': 1.6261, 'grad_norm': 2.128119468688965, 'learning_rate': 0.0021900452488687787, 'epoch': 0.45}                                                                   
 45%|██████████████████████████████████████████████████████████▊                                                                       | 500/1105 [01:43<01:47,  5.62it/s]Directory './tmp/checkpoint-500/intervenable_model' already exists.
{'loss': 1.4937, 'grad_norm': 3.4188449382781982, 'learning_rate': 0.0021538461538461538, 'epoch': 0.46}                                                                  
{'loss': 1.5994, 'grad_norm': 3.9260153770446777, 'learning_rate': 0.0021176470588235297, 'epoch': 0.47}                                                                  
{'loss': 1.5838, 'grad_norm': 4.29110050201416, 'learning_rate': 0.002081447963800905, 'epoch': 0.48}                                                                     
{'loss': 1.7872, 'grad_norm': 2.3661904335021973, 'learning_rate': 0.0020452488687782806, 'epoch': 0.49}                                                                  
{'loss': 1.5149, 'grad_norm': 1.819394588470459, 'learning_rate': 0.002009049773755656, 'epoch': 0.5}                                                                     
{'loss': 1.8831, 'grad_norm': 2.6721067428588867, 'learning_rate': 0.001972850678733032, 'epoch': 0.51}                                                                   
{'loss': 1.5401, 'grad_norm': 2.8439788818359375, 'learning_rate': 0.0019366515837104072, 'epoch': 0.52}                                                                  
{'loss': 1.6857, 'grad_norm': 1.7938305139541626, 'learning_rate': 0.0019004524886877827, 'epoch': 0.52}                                                                  
{'loss': 1.7157, 'grad_norm': 3.063899517059326, 'learning_rate': 0.0018642533936651584, 'epoch': 0.53}                                                                   
{'loss': 1.6262, 'grad_norm': 3.4250340461730957, 'learning_rate': 0.0018280542986425338, 'epoch': 0.54}                                                                  
{'loss': 2.0905, 'grad_norm': 2.200345993041992, 'learning_rate': 0.0017918552036199095, 'epoch': 0.55}                                                                   
{'loss': 1.7155, 'grad_norm': 3.3154056072235107, 'learning_rate': 0.001755656108597285, 'epoch': 0.56}                                                                   
{'loss': 1.6819, 'grad_norm': 2.451686143875122, 'learning_rate': 0.0017194570135746607, 'epoch': 0.57}                                                                   
{'loss': 1.5713, 'grad_norm': 2.4018261432647705, 'learning_rate': 0.0016832579185520362, 'epoch': 0.58}                                                                  
{'loss': 1.644, 'grad_norm': 3.3602707386016846, 'learning_rate': 0.0016470588235294116, 'epoch': 0.59}                                                                   
 59%|█████████████████████████████████████████████████████████████████████████████                                                     | 655/1105 [02:14<01:31,  4.93it/s]Traceback (most recent call last):
  File "/home/djfil/support-platform/experiments/llm-finetune-qlora/train_ReFT.py", line 115, in <module>
    _ = trainer.train()
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2052, in train
    return inner_training_loop(
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2388, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3485, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/djfil/.local/lib/python3.10/site-packages/pyreft/reft_trainer.py", line 92, in compute_loss
    base_outputs, cf_outputs = intervenable(
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/pyvene/models/intervenable_base.py", line 1947, in forward
    raise e
  File "/home/djfil/.local/lib/python3.10/site-packages/pyvene/models/intervenable_base.py", line 1931, in forward
    counterfactual_outputs = self.model(**base, **model_kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/peft/peft_model.py", line 1644, in forward
    return self.base_model(
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/peft/tuners/tuners_utils.py", line 197, in forward
    return self.model.forward(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1229, in forward
    loss = loss_fct(shift_logits, shift_labels)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 1293, in forward
    return F.cross_entropy(
  File "/home/djfil/.local/lib/python3.10/site-packages/torch/nn/functional.py", line 3479, in cross_entropy
    return torch._C._nn.cross_entropy_loss(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.17 GiB. GPU 0 has a total capacity of 23.58 GiB of which 2.76 GiB is free. Including non-PyTorch memory, this process has 20.81 GiB memory in use. Of the allocated memory 17.81 GiB is allocated by PyTorch, and 2.64 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
 59%|█████▉    | 655/1105 [02:16<01:33,  4.80it/s]

I started removing entries (up to 2/3 of all data) from the dataset but it fails again.
Here's an example: considering 1000 entries, if it fails at 500, and then i remove it, it may fail now at entry 100, which was ok before.
I have been like this, for a week now... :(

Is anyone able to give me some idea?