CUDA Out of memory with large dataset, but OK with same batch_size and context_length for small test set

rachad_lakis: <p>I am getting an OOM with cuda when using a large dataset  (I am trying to fine tune Llama3.2 1B) but  when i use around 20 elements dataset for testing, with the same context length and same batch size I have no issues.</p>
<p>(max_context_length = 8192 and batch_size = 1)</p>
<p>This is the simple code I am using:</p>
<pre><code class="lang-auto">Llama32_1B_MAX_LENGTH = 8192
NUM_SAMPLES = 100 #or 10 when testing

thepile_10k = load_dataset("NeelNanda/pile-10k")
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, texts, tokenizer, max_length=None):
    self.texts = texts
    self.tokenizer = tokenizer
    self.max_length = max_length

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

    def __getitem__(self, idx):
        text = self.texts[idx]
        tokens = self.tokenizer(text, max_length=self.max_length, truncation=True, padding=False, return_tensors="pt")

        return {"input_ids": tokens['input_ids'].squeeze(0), "attention_mask": tokens['attention_mask'].squeeze(0)}

def collate_fn(batch):
    input_ids = [item['input_ids'] for item in batch]
    attention_masks = [item['attention_mask'] for item in batch]
    # Pad sequences to the same length
    input_ids_padded = pad_sequence(input_ids, batch_first=True, 
    padding_value=llama32_1B_tokenizer.pad_token_id)
    attention_masks_padded = pad_sequence(attention_masks, batch_first=True, 
    padding_value=0)
    return {
    'input_ids': input_ids_padded,
    'attention_mask': attention_masks_padded
     }



thepile_10k_train_dataset = CustomDataset(
thepile_10k_train_subset['text'],
tokenizer=llama32_1B_tokenizer,
max_length=Llama32_1B_MAX_LENGTH
)

train_dataloader = DataLoader(thepile_10k_train_dataset, batch_size=1, collate_fn=collate_fn ,shuffle=True, pin_memory=True)

val_dataloader = DataLoader(thepile_10k_train_dataset, batch_size=1, collate_fn=collate_fn, shuffle=True, pin_memory=True)

epochs = 6
lr_train = 5e-5

llama32_1B_losses = train_model(
llama32_1B_model,
optimizer=AdamW(llama32_1B_model.parameters(), lr=lr_train),
epochs=epochs,
train_loader=train_dataloader,
validation_loader=val_dataloader,
patience=2,
train_verbose=True,
eval_verbose=True,
with_autocast=False,
)


</code></pre>

ptrblck: <p>If I understand your code correctly, the <code>max_length</code> does not necessarily mean this length is used, so check the shapes of the samples in the small subset vs. the full dataset.</p>

rachad_lakis: <p>Tried to test max_length=32, and I got this (it seems it is truncating to this max_length):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/e/fe60e92ea28ca606ec76136cf83a7c1b4e6f7104.png" data-download-href="https://discuss.pytorch.org/uploads/default/fe60e92ea28ca606ec76136cf83a7c1b4e6f7104" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe60e92ea28ca606ec76136cf83a7c1b4e6f7104_2_690x338.png" alt="image" data-base62-sha1="AikUiCHI9Csw643OlAaJL42B1Ck" width="690" height="338" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe60e92ea28ca606ec76136cf83a7c1b4e6f7104_2_690x338.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe60e92ea28ca606ec76136cf83a7c1b4e6f7104_2_1035x507.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe60e92ea28ca606ec76136cf83a7c1b4e6f7104_2_1380x676.png 2x" data-dominant-color="232930"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1708×839 23.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

rachad_lakis: <p>and sorry, but I knew that You are the best one to help, so I allowed myself to send you a LinkedIn message about this question <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>Was the code working by reducing the <code>max_context_length</code> to <code>32</code>? If so, my previous guess might still be valid and the subset of your data might not have used any long sequences (i.e. it contained sequences smaller than the allowed <code>8192</code> length).</p>

rachad_lakis: <p>That is absolutely correct. Thank you for your help.<br>
What i was trying to do is to fine tune Llama3.2 1B without quantization nor PEFT on a 40 GB GPU on AWS, but I am getting an Out Of Memory when having an 8192 context length input.<br>
It seems that I should go for FSDP on AWS (I was trying to escape that, too complex) or use a shorter context length (I was also trying to escape that and stay with the 8K inintial context length of Llama3.2 1B).</p>