Working with newer Transformer versions

[vpgits]

Hi,
I have a keen interest on this research, and I wanted to experiment with it on my own without the heavy coupling the current code has with OpenAI Embeddings and Phi models. So, Ive attempted to rewrite the entire codebase trying to get it working with the newer transformer APIs. I've seen people having trouble trying to get it working with the new LLama 3.2 1B variant or LLama 3.1 Instruct variant ( I think the only variant that works right now is LLama 3 Instruct 8B).
I was initially planning to see what kind of behavior I would get with a SLM, but I keep running into issues where I dont really know what to do.
Apart from that I've noticed misaligned hyperparameters in the Makefile and the research paper, a subtle labeling bug where a part of the query is already there, KB having the correct answer always at the 0th index ect.
But my main pain point is in my implementation the loss plummets to zero within the first 100 steps from around 5.
I'd like to get your opinion on what a healthy training run would look like?

[xidulu]

Hi @vpgits , first author here, thank you for your interest!

I think the only variant that works right now is LLama 3 Instruct 8B

Correct.

I'd like to get your opinion on what a healthy training run would look like?

You should be able to see the loss function gradually goes down, and the models' outputs given a user query and a KB should be human-readable.

[vpgits]

Thank you for the prompt response.
Unfortunately my code converges within the first hundred steps. It cheats the loss function by generating part of the question and just repeating that again and again to make the loss go down.

[xidulu]

It cheats the loss function by generating part of the question and just repeating that again and again to make the loss go down.

This suggests that you should include the question part in the SFT / instruction tuning objective.

[vpgits]

I understand.
However when I attempt to train, with nearly identical hyperparameters with my rewrite on meta-llama/Llama-3.2-1B-Instruct I can see that the model just ignores the kb and attempts to just generate text as it pleases.
This is also accompanied by a lower training loss.

[vpgits]

Hi @xidulu ,
Do you have any insights on working with newer llama models that use RoPE embeddings. Should RoPE be applied to the kb queries as well?

[xidulu]

@vpgits

No positional embedding should not be applied on any KB related queries since KB items do not have order

I think applying RoPE on the query vectors output by the separate learned query head would not harm the performance, I don't think there would be a significant difference since the key vectors being looked up, i.e. the KB keys, don't have order or positional encoding.