Fix issues with Llama HF->NeoX conversion

[aurelion-source]

Resolves #1337 and #1342

• Following fix a GQA issue (#1314) #1315, the GQA code no longer splits heads based on num_q_heads. This PR updates tools/ckpts/convert_hf_llama_to_neox.py to concatenate tp-partitioned q, k, and v weights without per-head splitting.
• The current RMSNorm implementation incorrectly adds epsilon to the RMS instead of the variance. Fix is the same as RMSNorm epsilon implementation #1342, but ensures compatibility with partial RMSNorm.

[aflah02]

Hi @aurelion-source
I think there is a similar issue in the reverse direction (NeoX to HF) caused by the same changes as when I convert a GQA model to HF it starts generating gibberish until I modify the conversion file. The RMS Norm implementation also varies which causes discrepancies b/w Llama HF class and NeoX

[aurelion-source]

Hi @aurelion-source I think there is a similar issue in the reverse direction (NeoX to HF) caused by the same changes as when I convert a GQA model to HF it starts generating gibberish until I modify the conversion file. The RMS Norm implementation also varies which causes discrepancies b/w Llama HF class and NeoX

Hi @aflah02, thanks for pointing out the issue with the NeoX -> HF conversion. I've updated and tested the script to fix the problem.
This PR also addresses RMSNorm discrepancies to align with other implementations. Let me know if you run into any issues if you get a chance to test this.

[aflah02]

Thanks @aurelion-source
I'll try to test this next week and get back

One question I have is if I use the fused rms norm kernel will it be compatible with the HF version? or is it equivalent to the neox version

[Shetano]

Thanks @aurelion-source I'll try to test this next week and get back

One question I have is if I use the fused rms norm kernel will it be compatible with the HF version? or is it equivalent to the neox version

Yes, it should be compatible with HF. It uses apex's fused RMSNorm implementation.