search

danbraunai / simple_stories_train

Trains small LMs. Designed for training on SimpleStories
3 stars 1 forks source link

Add Typing to Llama Training #9

Closed lennart-finke closed 2 weeks ago

lennart-finke commented 2 months ago

Description

Added Jax and regular typing to llama_train.py

Related Issue

Should close #4.

How Has This Been Tested?

Executing the script in a Colab instance.

lennart-finke commented 2 months ago

Alright, will have to work on this some more as the checks indicate.

lennart-finke commented 2 months ago

Now it might go through. If someone has time, I recommend double checking the jaxtyping hints though, as I was not entirely sure everywhere.