Has anyone had a good trained model

[retoriax]

Has anyone had got a good trained model and can share what dataset, model and training parameters/flags they used for training and eval?
I cant get good scores with kblam using Llama-3-8B-Instruct, text-embedding-3-large and the synthetic dataset.
Maybe also share what the final train loss and step count was.

[xidulu]

Can you provide some more details on the configurations you are using ?

Would be better if you can attach some training loss traces / plots.

Thanks!

[retoriax]

Sure, my latest attempt had the following configuration:
--dataset datasets
--train_dataset synthetic
--B 20
--gradient_accm_step 40
--total_steps 20001
--use_cached_embd
--encoder_spec BigOAI
--sep_query_head
--kb_token_layer_frequency 3
--use_data_aug
--key_embd_src key
--N 120000
--hf_model_spec "meta-llama/Llama-3.2-1B-Instruct"
--llm_type llama3
--save_step 400

I have trained it to step 4000+.
Heres the train_loss:

this is the eval settings:
python eval.py generation
--seed 1607
--dataset_dir ${TEST_DATASET_DIR}
--test_dataset synthetic.json
--llm_base_dir ${LLAMA_BASE_DIR}
--model_dir ${LLAMA_MODEL_CHKPT}
--encoder_dir ${ENCODER_FOR_LLAMA_CHKPT}
--save_dir ${ATTN_SAVE_DIR}
--kb_layer_frequency 3
--kb_size 100
--kb_scale_factor 100
--no-fancy_instruction
--encoder_spec BigOAI
--llm_type "llama3"
--precomputed_embed_keys_path ${TEST_DATASET_DIR}/synthetic_BigOAI_embd_key.npy
--precomputed_embed_values_path ${TEST_DATASET_DIR}/synthetic_BigOAI_embd_value.npy

for {LLAMA_MODEL_CHKPT} i used step 4000.

why are my scores so bad? i feel like it doesn't retrieve anything correctly

[xidulu]

The loss trajectory looks fine to me.

And the model is actually outputting readable text and sometime an answer pretty close to the ground truth, which is another good sign.

From my experience, the performance is usually better with

--kb_layer_frequency 1

Also I have never tried running the model with Llama 1B myself, which per memory during the rebuttal, has pretty concerning performance, the 8B model may have much better performance.

Maybe you could try adjusting the kb_layer_frequency first and then changing the base model if kb_layer_frequency does not help.

[retoriax]

Thank you very much @xidulu
I trained for 30h on a A100 and got good results for kb_size 100 and scale_factor 100.
{'rouge1': 0.7812977846708903, 'rouge2': 0.6080092027443471, 'rougeL': 0.7432360570209307, 'rougeLsum': 0.7433776824138043, 'bert_score_precision': 0.9007512772083283, 'bert_score_recall': 0.9034255331754685, 'bert_score_f1': 0.9011788907647132, 'mem_cost': 21168652288}
My train_loss was at 0.53 it took 5k steps.

I used kb_layer_frequency 2, with ada-002 and Llama 8B.
I created my own training dataset from scatch off of wikipedia entries.

Can anyone explain how --kb-scale-factor works exactly?
I can only guess for good values but is there a better approach?

[xidulu]

@retoriax
The metric number looks not bad at all!!!!

I created my own training dataset from scatch off of wikipedia entries.

So glad to hear KBLaM works on users' customized dataset !!!

Can anyone explain how --kb-scale-factor works exactly?

See this paragraph in the paper for a detailed introduction, kb-scale-factor fix the magnitude of the contribution from the KB's value embeddings regardless of the number of entries in the KB, it is introduced to make sure the model can work with arbitrary sizes of KB.

[vpgits]

Hi @retoriax ,
were you successful with training Llama 3.2 1B?

[retoriax]

Hi @vpgits,
unfortunately I wasn’t successful.

After xidulu mentioned “Also I have never tried running the model with Llama 1B myself, which per memory during the rebuttal, has pretty concerning performance, the 8B model may have much better performance.”
I immediately switched to the 8B models instead.

I hope you’ll have more success with the 1B models.

[vpgits]

I had to completely rewrite the modelling llama class to account for the new RoPE embeddings in llama models, , they have significantly restructured the api. And then had to rewrite the entire codebase to mimic the original repo. However i keep running into issues. I think my implementation might be wrong to be honest, but I dont see why a 1B param model should fail.