Frequent NaN Loss During Training

[Alvin-ne]

Hi HyperCore team,

I've been encountering frequent NaN loss issues while training models using the HyperCore library, which has hindered my training process. I've tried several mitigation strategies but haven't been able to resolve the problem. Here's a detailed breakdown:

During model training, the loss frequently becomes NaN. This occurs consistently across different hyperbolic models I've experimented with.

To address the NaN loss, I've tried the following approaches, but none have resolved the issue:
Using Riemannian Adam optimizer for optimizing manifold.c parameters.
Reducing learning rate to a very low value (lr = 5e-6) and weight decay (weight_decay = 1e-6).
Implementing learning rate schedulers (Linear warmup & Cosine decay).
Gradient clipping with max_grad = 1.0.

Could you provide guidance on potential causes of frequent NaN loss in HyperCore training and recommend additional strategies to mitigate this issue? For example, are there specific manifold parameter settings, stabilization techniques, or layer configurations that are more robust to numerical instability?

Let me know if you need more details about my model architecture, training setup, or specific modules used.

Thanks in advance for your help!