Training script

[jettjaniak]

No description provided.

[jettjaniak]

this should use argparse and take some arguments?

[jettjaniak]

float16 vs bfloat16 vs float32 - Can we just train everything on float32 to ensure reproducibility across all devices?

[SrGonao]

What can I do to finish this? Do you want me to take over @jannik-brinkmann

[jannik-brinkmann]

@SrGonao it would be great if someone could pickup the refactoring of the training - the functionality should be fine. @jettjaniak depending on your feedback, I can already start the trainings on the un-refactored version of the code, to provide new models to the evals team. If you want to wait for the refactoring to be done, I could work on it on Thursday, if no one else could do it before then :)

[jettjaniak]

All the comments need to be resolved, but below are some of top priorities:

• training independent of tokenizer, PretokDataset etc. - just use a tokenized HF dataset
• revamp TrainingConfig (waiting for try hydra for training configs #48)
• remove things related to model compilation, ddp and all things that we don't need
• remove amp, everything should train on float32
• define a function performing a single training step (wandb should be outside of this)

[jettjaniak]

I don't think we need this script. We're hosting the dataset on HF, so we should host it in the llama2c github repo. How we generated it is perhaps more interesting and could be it's own script, but whatever

[jaidhyani]

I added a bunch of comments to try to make this obnoxiously-huge PR slightly less of a pain to review.

[jaidhyani]

We added this when we added the llama2c submodule. Ironically, we also removed llama2c in the course of developing this PR. Technically we don't need this anymore, but it's not a bad idea to have it around for any submodules we add in the future.

[jaidhyani]

mamba really wants to run on CUDA, and there are optional-but-preferred dependencies for doing so. We want to use that when we can, because otherwise training mamba models would take foreverrrrrrrrr - but only on platforms that support it. We split out requirements into those that don't require cuda (-nocuda) and those that do (still in requirements.txt, which automatically includes -nocuda requirements). Github CI doesn't run in a CUDA env because why would Microsoft give away that much GPU compute? So we use the non-cuda requirements.txt for CI.

[jaidhyani]

--black config is now implicit in pyproject.toml

[jaidhyani]

I think these are llama2c artifacts? Also no technically longer needed, but on the other hand these are generally things we'd want to exclude from git if anything with these names ever showed up.

[jaidhyani]

Debugging wandb integration involved a lot of wandb artifacts being created and I was too lazy to change directories.

[jaidhyani]

This is a weird one. Iteration parameters are always derivable from config and dataset size, but it's convenient to store them anyway, if only to cut down on how many parameters we're passing around. Originally implemented while trying to break the original script into functional pieces, and then never revisited because that seemed like too much effort.

[jaidhyani]

The main training function! Given a GigaConfig, does setup and runs the training loop. Most of the actual logic lives in train_step.

[jaidhyani]