TildeOpen30B distilation

[iPRET]

Hi!

I'd like to distill the TildeOpen model (LlamaForCausalLM class on huggingface) using the NeMo framework.
TildeOpen's basically Llama2/3, except it also uses YaRN as position embeddings.
Where do you propose I should start?

Thanks!
Ingus

[iPRET]

Update:
Found a prune and distill script on NVIDIA-NeMo/NeMo at commit 237e2c08ed (right before the "remove all the LLM stuff" commit)
Decided I'd work with that. Cause Megaton-Bridge doesn't have those scripts.

• It seems like megatron-LM already supports YaRN but NeMo at that commit didn't want to support it, so I added those parameters to NVIDIA-NeMo/NeMo.
• Implemented kbins packing, because otherwise I think only dense packing was supported.
• I suspect attention masks are ignored when using flash attention, so trying to currently fix that.

So far my sense is NeMo is much harder to adapt as compared to other open source frameworks.

[iPRET]

Update:
Overrode the attention mask and flash attention logic.
Fought with a bunch of bugs with data preparation, packing and whatever else.

But I'm finally getting close...
Currently actually managing to prune and start distillation.
Pruning seems to work fine afaict.
Distillation's a different story, the KL divergence seems to go down but the validation loss just goes up and up...

But I feel the result is coming within reach...
...In order to reach the goodness at the center of the pistachio, you have to break through the hard outer shell... Just gotta hope the inside is unlike the shell... still nutritious,.. not dried up, rotten and hard...

...I will distill TildeOpen with NeMo even if it's the last thing I do...

[iPRET]

Think I got it to work. With only ~3 months of smashing my head against the wall.

Overall vibe of NeMo - really hard to modify. Don't use if doing anything other than the beaten path.

Ingus