March 18, 2025

ikayaniaamirshahzad@gmail.com

PyTorch Transformer Stuck in Local Minima Occasionally



Hi, I am working on a project to pre-train a custom transformer model I developed and then fine-tune it for a downstream task. I am pre-training the model on an H100 cluster and this is working great. However, I am having some issues fine-tuning. I have been fine-tuning on two H100s using nn.DataParallel in a Jupyter Notebook. When I first spin up an instance to run this notebook (using PBS) my model fine-tunes great and the results are as I expect. However, several runs later, the model gets stuck in a local minima and my loss is stagnant. Between the model fine-tuning how I expect and getting stuck in a local minima I changed no code, just restarted my kernel. I also tried a new node and the first run there resulted in my training loss stuck again the local minima. I have tried several things:

  1. Only using one GPU (still gets stuck in a local minima)
  2. Setting seeds as well as CUDA based deterministics:
    1. torch.backends.cudnn.deterministic = True
    2. torch.backends.cudnn.benchmark = False

At first I thought my training loop was poorly set up, however, running the same seed twice, with a kernel reset in between, yielded the same exact results. I did this with two sets of seeds and the results from each seed matched its prior run. This leads me to be believe something is happening with CUDA in the H100. I am confident my training loop is set up properly and there is a problem with random weight initialization in the CUDA kernel.

I am not sure what is happening and am looking for some pointers. Should I try using a .py script instead of a Notebook? Is this a CUDA/GPU issue?

Any help would be greatly appreciated. Thanks!

submitted by /u/MartinW1255
[comments]



Source link

Leave a Comment