On next pass: Can confirm that I get the same 28 pytest failures when running inside docker whether I'm on main or this branch, so I am not getting any signal from that.

For starting training and/or preparing data, this branch seems to work fine.

If someone feels like doing work for me, they should do the following:

1. Have a system with a gpu, up to date nvidia drivers, and nvidia container toolkit installed
2. Clone this branch
3. In the gpt-neox folder, run the following:

docker compose run gpt-neox -d # starts the container
docker compose exec gpt-neox bash # you will be in the container now
cd gpt-neox
python deepy.py train.py ./configs/pythia/14M.yml ./configs/docker/pythia-paths.yml

Tell me if it crashes and give me your errors if it does

If someone feels like doing work for me, they should do the following:

1. Have a system with a gpu, up to date nvidia drivers, and nvidia container toolkit installed
2. Clone this branch
3. In the gpt-neox folder, run the following:

docker compose run gpt-neox -d # starts the container
docker compose exec gpt-neox bash # you will be in the container now
cd gpt-neox
python deepy.py train.py ./configs/pythia/14M.yml ./configs/docker/pythia-paths.yml

Tell me if it crashes and give me your errors if it does

Got this, got an apex build problem, fixed it, the problem is that now apex takes forever to build, need to fork apex and only retain the fused adamw that we use and then it should be good.

Because the version bumps are sufficiently large, it probably makes sense to test it on a live cluster before merging the branch to make sure nothing very important is broken.