Major refactor (inc adding Pydantic)

[danbraunai]

Description

• Adds a models directory with model implementations and default configs
• Use pydantic for the base config (rather than argparse)
• Allow users to pass cli flags to update the config. These override the config file (and/or the pydantic defaults).

New usage (which is also written in the README and the train_llama.py script:

python train_llama.py [PATH/TO/CONFIG.yaml] [--key1 value1 --key2 value2 ...]

where

• PATH/TO/CONFIG.yaml contains the training config. If no path is provided, a default config will be used.
• --key1 value1 --key2 value2 ... override values in the config. Note that if you wish to update a
  nested value, you must use dotted notation (e.g. --train_dataset_config.name my_dataset).

To run on multiple GPUs, use

torchrun --standalone --nproc_per_node=N train_llama.py ...

where N is the number of GPUs to use.

Motivation and Context

It's currently quite hard to understand and work with the training scripts. This PR aims to set us towards a nicer library to use/contribute to.

How Has This Been Tested?

• unittest for one of the added functions utils.convert_dotted_args_to_nested_dict.
• No tests written for general config behaviour (it does look good after quick testing myself). These could be added.
• Single and multi-gpu training works according to the instructions in the README and the train_llama.py script. WANDB train loss goes down as desired.

Does this PR introduce a breaking change?

Yes. Complete override of config. Some config argument names have changed.

[lennart-finke]

Very cool, and ready for merge, with just one minor usability comment.

[lennart-finke]

Just one comment: Since compilation is on by default (which is probably good), we might want to include --compile 0 in the example prompt for local execution, since compilation on e.g. mac gives very ugly errors that might be hard for newcomers to trace back to this flag.

[danbraunai]

Good pickup, added.