Actor loss nan and Resizing model embedding

The model I use is GPT-2 124M. When resizing model embeddings during the training of STF and RW, I often encounter issues where the generated answers consist entirely of zeros. This causes both the log probabilities and actor loss to become NaN (Not a Number). I have noticed that resizing the embeddings can lead to the generation of token IDs that exceed the vocabulary size. I suspect this may be contributing to the problem. However, when I don't resize the model's embeddings and train STF and RW, I do not experience this issue during RLHF training. I don't know why.

Reproduce the GPT-2 124M training case with embedding resizing enabled for STF and RW, then inspect generated token IDs, vocabulary bounds, zero-only outputs, log probabilities, and actor loss. Compare the same training without resizing. Done means the resizing-related cause is identified and the training behavior is verified without invalid token IDs or NaN losses.