Running low on GPU memory on a cluster with ESM2 lowest config
I’m trying to train a model using esm2_t6_8M_UR50D, which is the smallest ESM2 model (6 layers, 8M parameters), and TrainingArguments withbatch_size=1,gradient_accumulation_steps=1andfp16=True. I’m training it on a cluster with four 32GB GPUs, and I’m seeing it using all four of them at their fullest. I also setPYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512beforehand. However, using this configuration, it runs out of memory at iteration 600 (out of 40k).What else can I do to make it work?
2023-12-04T23:58:35Z
@davguevwhat service are you using? are you using a cloud service or do you have your own cluster?
2023-12-05T00:39:32Z
Our lab has its own cluster with 4 GPUs and no one else was using them when I tried to train the model.