I'm training a BERT sequence classifier on a custom dataset. When the training starts, the loss is at around ~0.4 in a few steps. I print the absolute sum of gradients for each layer/item in the model and the values are high. The model converges initially but when left to be trained for a few hours and sometimes even early as well it gets stuck. I am calculating gradients with the below code. Also the logs are at - https://pastecode.io/s/v2s3mr3e (initial convergence) . I'm printing gradients, loss, metrics and logits.
for name, param in model.named_parameters(): print(name, param.grad.abs().sum())
While the model is stuck, the loss value is around ~0.69 with worse performance metrics (precision/recall) on the training set but the gradients are very small compared to the initial training phase. Also it seems that predictions are swinging with most of the values predicted as either 0 or 1. Following are the logs for the stuck phase - https://pastecode.io/s/cjuxog44
Training code Link - Training Code
It seems that the model is stuck at a local minima where the gradient values are relatively smaller even though the loss is high. How can I mitigate it ? One option I see is using a higher learning rate or a cyclic learning rate but not sure if that's the right approach since the the learning rate is 5e-5 with LR scheduler disabled. Below is the plot for Loss, Bert pooler and classifier gradients sum over steps.
Also the data is 50-50 balanced. Batch size is 32. I'm using AdamW. I have also tried SGD but the convergence is very slow.
Or there might be some error/reason which I am not able to identify. Please help