Cuda OOM
Author: apoalquaary
Pull request: no
Labels: name: bug

### System Info

I'm trying to pre-train a small LLM with Llama architecture with 1.4B parameters. I'm using a node which is 8 * H100. Every thing is working fine. I noticed a small problem. when I load a big dataset in the code it gives an OOM error regardless the batch size after some steps. At the beginning I thought it would be some samples are too long so I did some experiments when all the samples have the max length (2048) and the code worked fine. so the problem probably not from the size of the  batch and the inputs.

when I'm adding select(range(10000)) to the end of wikitext dataset object to get sub dataset it works fine, but when I did add select(range(500000)) it gives OOM error after some steps.

Where exactly the OOM comes from?

```
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import LlamaForCausalLM, LlamaConfig, Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoTokenizer
from datasets import config
from datasets import load_dataset, load_from_disk, IterableDataset
from utils import count_parameters
from tokenizers import processors
from sim_data import create_test_dataset, filter_long_sequences

# Environment settings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = False


# Initialize distributed training
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# Tokenizer setup
tokenizer = AutoTokenizer.from_pretrained('/mnt/fs/huggingface/hub/tokenizers/Aranizer-PBE-86k')

tokenizer._tokenizer.post_processor = processors.TemplateProcessing(
    single=["
", "$A", "
"],
    special_tokens=[("
", tokenizer.bos_token_id), ("
", tokenizer.eos_token_id)]
)

# Special tokens setup
tokenizer.bos_token = '
'
tokenizer.bos_token_id = 0
tokenizer.eos_token = '
'
tokenizer.eos_token_id = 2
tokenizer.pad_token = '<pad>'
tokenizer.pad_token_id = 1



# Model configuration
model_config = LlamaConfig(
    vocab_size=len(tokenizer),
    hidden_size=2048,
    intermediate_size=8192,
    num_hidden_layers=16,
    num_attention_heads=16,
    max_position_embeddings=2048,
    dropout_rate=0.1,
    layer_norm_eps=1e-6,
    pad_token_id=tokenizer.pad_token_id,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    use_cache=False,
)

# Initialize model and move to appropriate GPU
model = LlamaForCausalLM(model_config)

# Enable gradient checkpointing BEFORE wrapping with DDP

print("great")
# Move model to GPU
model = model.to(local_rank)
# 2. Enable gradient checkpointing before DDP wrapping
model.gradient_checkpointing_enable()

# Wrap model with DDP
model = DDP(
    model, 
    device_ids=[local_rank],
    output_device=local_rank,
    find_unused_parameters=False,
)


dataset = load_dataset("wikitext", "wikitext-103-v1", split='train')


eval_dataset = load_dataset("wikitext", "wikitext-103-v1", split='validation')



def tokenize_function(examples):

    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=2046,  # 2046 - 2 to leave room for BOS and EOS
        padding=True,
        return_tensors=None,
        add_special_tokens=True,
    )


# Tokenize datasets
tokenized_datasets = dataset.map(
    tokenize_function, 
    batched=True, 
    remove_columns=["text"],
    # num_proc=32,
    # load_from_cache_file=False
)

tokenized_eval_dataset = eval_dataset.map(
    tokenize_function, 
    batched=True, 
    remove_columns=["text"],
    # num_proc=32,
    # load_from_cache_file=False
)

steps_var = 100

#torchrun --nproc_per_node=8 train_distributed.py

# interleave_datasets

# Training arguments
training_args = TrainingArguments(

    accelerator_config={"dispatch_batches": False},

    # deepspeed=deepspeed_config,
    output_dir="./Mulhem-1.4B",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=14,
    gradient_accumulation_steps=2,
    save_steps=steps_var,
    save_total_limit=2,
    
    # Resume training settings
    resume_from_checkpoint=True, # Enable resuming from checkpoint
    # dataloader_pin_memory=True,

    # Distributed training settings
    local_rank=local_rank,

    # Optimization settings
    learning_rate=5e-5,
    warmup_steps=500,
    weight_decay=0.01,
    adam_beta1=0.9,
    adam_beta2=0.999,
    max_grad_norm=1.0,
    
    # Evaluation and logging
    eval_strategy="steps",
    eval_steps=steps_var,
    eval_delay=0.1, # 
    logging_dir="./runs",
    logging_strategy="steps",
    logging_steps=steps_var,
    report_to=["tensorboard"],
    
    # should be checked:
    remove_unused_columns=False,

    # Distributed specific
    ddp_find_unused_parameters=False,
    ddp_bucket_cap_mb=25,
    
    # Precision settings
    fp16=False,
    bf16=False,
    
)

# Initialize trainer
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
    train_dataset=tokenized_datasets,
    eval_dataset=tokenized_eval_dataset,
)

if __name__ == "__main__":
    # Start training
    trainer.train()
    
    # Cleanup
    dist.destroy_process_group()

    
```



### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import LlamaForCausalLM, LlamaConfig, Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoTokenizer
from datasets import config
from datasets import load_dataset, load_from_disk, IterableDataset
from utils import count_parameters
from tokenizers import processors
from sim_data import create_test_dataset, filter_long_sequences

# Environment settings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = False


# Initialize distributed training
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# Tokenizer setup
tokenizer = AutoTokenizer.from_pretrained('/mnt/fs/huggingface/hub/tokenizers/Aranizer-PBE-86k')

tokenizer._tokenizer.post_processor = processors.TemplateProcessing(
    single=["
", "$A", "
"],
    special_tokens=[("
", tokenizer.bos_token_id), ("
", tokenizer.eos_token_id)]
)

# Special tokens setup
tokenizer.bos_token = '
'
tokenizer.bos_token_id = 0
tokenizer.eos_token = '
'
tokenizer.eos_token_id = 2
tokenizer.pad_token = '<pad>'
tokenizer.pad_token_id = 1



# Model configuration
model_config = LlamaConfig(
    vocab_size=len(tokenizer),
    hidden_size=2048,
    intermediate_size=8192,
    num_hidden_layers=16,
    num_attention_heads=16,
    max_position_embeddings=2048,
    dropout_rate=0.1,
    layer_norm_eps=1e-6,
    pad_token_id=tokenizer.pad_token_id,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    use_cache=False,
)

# Initialize model and move to appropriate GPU
model = LlamaForCausalLM(model_config)

# Enable gradient checkpointing BEFORE wrapping with DDP

print("great")
# Move model to GPU
model = model.to(local_rank)
# 2. Enable gradient checkpointing before DDP wrapping
model.gradient_checkpointing_enable()

# Wrap model with DDP
model = DDP(
    model, 
    device_ids=[local_rank],
    output_device=local_rank,
    find_unused_parameters=False,
)


dataset = load_dataset("wikitext", "wikitext-103-v1", split='train')


eval_dataset = load_dataset("wikitext", "wikitext-103-v1", split='validation')



def tokenize_function(examples):

    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=2046,  # 2046 - 2 to leave room for BOS and EOS
        padding=True,
        return_tensors=None,
        add_special_tokens=True,
    )


# Tokenize datasets
tokenized_datasets = dataset.map(
    tokenize_function, 
    batched=True, 
    remove_columns=["text"],
    # num_proc=32,
    # load_from_cache_file=False
)

tokenized_eval_dataset = eval_dataset.map(
    tokenize_function, 
    batched=True, 
    remove_columns=["text"],
    # num_proc=32,
    # load_from_cache_file=False
)

steps_var = 100

#torchrun --nproc_per_node=8 train_distributed.py

# interleave_datasets

# Training arguments
training_args = TrainingArguments(

    accelerator_config={"dispatch_batches": False},

    # deepspeed=deepspeed_config,
    output_dir="./Mulhem-1.4B",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=14,
    gradient_accumulation_steps=2,
    save_steps=steps_var,
    save_total_limit=2,
    
    # Resume training settings
    resume_from_checkpoint=True, # Enable resuming from checkpoint
    # dataloader_pin_memory=True,

    # Distributed training settings
    local_rank=local_rank,

    # Optimization settings
    learning_rate=5e-5,
    warmup_steps=500,
    weight_decay=0.01,
    adam_beta1=0.9,
    adam_beta2=0.999,
    max_grad_norm=1.0,
    
    # Evaluation and logging
    eval_strategy="steps",
    eval_steps=steps_var,
    eval_delay=0.1, # 
    logging_dir="./runs",
    logging_strategy="steps",
    logging_steps=steps_var,
    report_to=["tensorboard"],
    
    # should be checked:
    remove_unused_columns=False,

    # Distributed specific
    ddp_find_unused_parameters=False,
    ddp_bucket_cap_mb=25,
    
    # Precision settings
    fp16=False,
    bf16=False,
    
)

# Initialize trainer
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
    train_dataset=tokenized_datasets,
    eval_dataset=tokenized_eval_dataset,
)

if __name__ == "__main__":
    # Start training
    trainer.train()
    
    # Cleanup
    dist.destroy_process_group()

    

### Expected behavior

to answer my question, where the error comes from?

Comments:
@apoalquaary 
I think I can help you with this issue. 
There are two ways to fix it.
**1** - try to reduce the batch size in training arguments, 
``` python
per_device_train_batch_size=4,
```

**2** - Train model in bf16 or in fp16, it can help to reduce model size.

@yash3056, Thank you for your response. my main question is why that happens? I want to keep the training with large batch.

I want it fp32 and I tried the max input with this batch size and it worked.

@apoalquaary 
that's a great question.
fp32 and bf16 are formats, which uses 32 bits and 16bits respectively. It is better for you to study in detail about this if you wish.

you know to do this it is better to used [FSDP](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html) or deepspeed.

@yash3056 
Thank you again for your response. I do know the difference, I just want to pre train it with 32 bits. The question If you can see the issue again, why when I load bigger datasets I gets an OOM error. It should be irrelevant. 

I'll give [FSDP] a look, 

@apoalquaary you should know that data is converted into tensors and put on the same device as the model.
And these tensors also take more space then regular text data which you usually see in datasets 

@yash3056 
Thank you for your insights again. but the data tensors should stay in the cpu and just the batch is sent to the GPU,


@apoalquaary yes, you are right. but according to my understanding and as I mentioned only batch is sent to gpu, while cpu converts it into tensor

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.