Relational Training: CUDA OOM Out of Memory Error

[solitaryangler]

Hi,

Thanks for developing and releasing this codebase. I'm using it to train on a tabular data. I tried both in Tabular format and Relational format. But in the Relational format I'm getting CUDA OOM (Out of Memory Error).

Original Table (Raw data): The original table has only 8 cols x 10,000 rows (which I have subsampled for testing). The model in Tabular mode trains perfectly fine and I am able to generate synthetic samples.

Relational Table Format (Parent / Child): In the relational format the tables have the following statistics:

• Parent: 4 cols x 5359 rows
• Child: 6 cols x 10,000 rows, where one parent row has ~150 corresponding child rows (at most).

However, in this case:

• Parent model in Tabular format trains well, but
• Child model in Relational format, with the parent model, fails with error CUDA OOM (Out of Memory Error).

I have tried this on GCP with

• NVIDIA T4 x2 (16GB)
• NVIDIA L4 x2 (24GB)
• NVIDIA A100 (40GB)

I suspect the Relational format Child model fails because it requires both the Parent & Child tables to be loaded into GPU memory. But the dataset is tiny. How can I overcome the OOM error?

Do you have any suggestions?

[avsolatorio]

Hi @solitaryangler , you are correct; I think the issue is with the Relational model. In particular, it may be due to the large number of child rows a single parent row has. You can try setting the output_max_length to some number like 1024 or 2048 to see if it prevents the OOM. What this does, however, is it will limit the number of children that will be used in training the model.

Also, you can try experimenting with not using the trained parent model and instead training the full relational model from scratch. Just set parent_realtabformer_path=None. We got better results training the full relational model from scratch in one of our experiments. :)

So you can have something like:

child_model = REaLTabFormer(
    model_type="relational",
    parent_realtabformer_path=None,
    output_max_length=2048,
    train_size=0.8)

Hope this helps!

[solitaryangler]

Hi @avsolatorio

Thanks for your comments. You were absolutely right. Once I limited the child entries per parent row, then the model trained fine. I am closing this issue as resolved.

My apologies on the late reply, I was traveling.

[solitaryangler]

Closed

[avsolatorio]

Hello @solitaryangler, no worries! I am so glad that it's working now. 😀