RuntimeError: CUDA error: device-side assert triggered

[Mahhos]

Hello! I have created a pipeline using the torchfastText library which works fine with 1M data. However, with a larger dataset ~5M, I keep getting RuntimeError: CUDA error: device-side assert triggered error, which I cannot understand why it is happening. The labelEncoder part is fine as below. Do you have any idea what might cause this issue?
I am running this on a Databricks notebook with a cluster using g5.16xlarge with 256GB, 64 core.

Number of unique classes: 10285
Min label index: 0
Max label index: 10284


X_train shape: (4000000, 1)
X_train dtype: object

Here is how I am reading and training the model:

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from torchFastText import torchFastText
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
import torch
import os


# Load data
train_df = pd.read_parquet("/Volumes/train_5M.parquet")
test_df = pd.read_parquet('/Volumes/test.parquet')

# Fit LabelEncoder on training labels
encoder = LabelEncoder()
encoder.fit(train_df["label"])

# Memory-efficient split using pandas sample
val_df = train_df.sample(frac=0.2, random_state=42)
train_df_small = train_df.drop(val_df.index)

# Prepare training data
X_train = train_df_small["text"].astype(str).values.reshape(-1, 1)
y_train = encoder.transform(train_df_small["label"])

# Prepare validation data
X_val = val_df["text"].astype(str).values.reshape(-1, 1)
y_val = encoder.transform(val_df["label"])

# Filter test set to only include known labels
test_df_filtered = test_df[test_df["label"].isin(encoder.classes_)]

# Prepare test data
X_test = test_df_filtered["text"].astype(str).values.reshape(-1, 1)
y_test = encoder.transform(test_df_filtered["label"])



# Initialize model
model = torchFastText(
    num_tokens=100000,
    embedding_dim=100,
    min_count=5,
    min_n=3,
    max_n=6,
    len_word_ngrams=2,
    sparse=False
)



# Train model

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
gpu_count = torch.cuda.device_count()

# print(gpu_count)
# num_workers=min(4 * gpu_count if gpu_count > 0 else 2, os.cpu_count() - 1)
# oss= os.cpu_count() - 1
# print("num_workers", num_workers, "oss", oss)

model.train(
X_train=X_train,
y_train=y_train,
X_val=X_val,
y_val=y_val,
num_epochs=14,
batch_size=64,
lr=4e-3,
verbose=True,
num_workers=min(4 * gpu_count if gpu_count > 0 else 2, os.cpu_count() - 1)
)

[micedre]

However, with a larger dataset ~5M, I keep getting RuntimeError: CUDA error: device-side assert triggered error, 

This error can happen when the GPU can't allocate enough memory. You should have another error in all your stack trace besides the assert error.

[Mahhos]

../aten/src/ATen/native/cuda/Loss.cu:250: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [1,0,0] Assertion t >= 0 && t < n_classesfailed.

[micedre]

Do you have at least one observation in X_train for each label in y?

Sorry, I'm sure you already made sure of it for it, but we want to make sure we did not miss something obvious before troubleshooting this in depth.

[Mahhos]

thanks for your response.

I checked my data:

num_unique_labels = len(np.unique(y_train))
Number of unique labels in y_train: 10281 


unique_labels = np.unique(y_train)
print(f"Unique labels in y_train: {unique_labels}")
Unique labels in y_train: [    0     1     2 ... 10282 10283 10284]

Label range: 0 to 10284
Expected count in range: 10285
Actual unique labels: 10281
Missing label numbers: [680, 2262, 4805, 7038]

I explicitly passed the num_classes to the model:

model = torchFastText(
    num_tokens=100000,
    embedding_dim=100,
    min_count=5,
    min_n=3,
    max_n=6,
    len_word_ngrams=2,
    sparse=False,
    num_classes=len(encoder.classes_) 
)

Model change it to 10281:
WARNING:torchFastText.torchFastText:Old num_classes value is 10285. New num_classes value is 10281.

Even after removing missing labels from val and test, with or without passing num_classes to the model, the CUDA error is there:

Validation samples to remove: 9 out of 1000000
Test samples to remove: 1 out of 428973
Some examples of unseen labels in validation set: [7038 2262 7038 2262 7038]
Some examples of unseen labels in test set: [7038]
Number of unique classes in training set: 10281
New label range: 0 to 10280

./aten/src/ATen/native/cuda/Loss.cu:250: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [10,0,0] Assertion `t >= 0 && t < n_classes` failed.
RuntimeError: CUDA error: device-side assert triggered

[meilame-tayebjee]

Temporarily getting back to cpu would give you better insights of the error.

I suspect that at some point in your training data you have a label higher than 10280 (10281 or 10282 for instance), but you can not access it because the last layer has only a size of 10281 (accepting labels in 0,...,10280).

The model automatically changes this size as you saw, taking len(np.unique(y_train)) : see here. We could update this in a next version to set it to np.max(y_train) +1: that would be more robust to errors such as yours, but less frugal, creating a bigger layer than necessary (to be discussed - @micedre).

For now, you need to make sure that np.max(y_train) == len(np.unique(y_train))-1.

A next version could also simply raise an error if the above condition is verified, for better clarity on what is expected in the preprocessing step.

[Mahhos]

Thank you, np.max(y_train) == len(np.unique(y_train))-1 resolved the issue.
Does the model have a default early stopping? With more data (5M), it stops at epoch 3, while I set num_epochs=14, and surprisingly the performance drops drastically when training with more data.

[meilame-tayebjee]

Yes, the patience for early stopping can be controlled by the patience_train parameter in model.train(...) function. Default value is 3 (see here), you can increase it.

However, the fact that it stops exactly after 3 epochs means that your loss is not decreasing at all over the 3 first epochs, which is weird. You should consider changing the learning rate and see if this behavior changes.