#1400·gpt-neox

Add round trip conversion based tests for import and export from HF

Author: aflah02Created Jun 27, 2026Updated Jun 27, 2026
Labelsfeature request

I think it would be nice to have some tests to catch regression between import and export of HF checkpoints into the library so that changes elsewhere in the codebase don't break the import/export scripts without raising alarms. This has happened in the past (see: https://github.com/EleutherAI/gpt-neox/pull/1345)

My proposal - Round trip conversion based testing: Take supported HF models say Llama 2 7b, bring it into NeoX and then export back to HF.

Things to check -

  • Match the original and final model's tensors.
  • Match logprobs etc. b/w the original HF and the converted NeoX model
  • For the above can be tested by running the same set of evals on both models
  • Repeat export multiple times with different MP/PP configs to cover all cases.

I'd like to keep this test as a locally running test but one which users have modifying the model files/export scripts have to run and submit the generated report in their PR to ensure everything is in order. I can also imagine this being a test in the CI but I'm not sure if we can then test multiple MP/PP configs though that might also be hard for users to do.