├── .deepsource.toml ├── .gitignore ├── LICENSE ├── README.md ├── config.yaml ├── images ├── axial.png ├── convit.png ├── mixer.png ├── normformer.png ├── paper_cogview_sandwich.png ├── reformer.png ├── revnet.png ├── rmsnorm_speed.png ├── sandwich.png └── scalenorm.png ├── img.png ├── inference.py ├── launch_on_tensorfork.py ├── main.py ├── requirements-dev.txt ├── requirements.txt ├── run.sh ├── script ├── inference.py ├── launch_multiple_runs.py ├── quantized_vid2tfrecord.py ├── start_on_multi_host_tpu.py └── text2tfrecord.py ├── setup.sh ├── setup_dev.sh ├── share_vm.py ├── src ├── backend.py ├── constants.py ├── context.py ├── data.py ├── main.py ├── model │ ├── activate.py │ ├── conv.py │ ├── loss.py │ ├── main.py │ ├── mixer.py │ ├── moe.py │ ├── norm.py │ └── reversible.py ├── optimizer.py └── utils │ ├── checkpoint.py │ └── wandblog.py ├── sweep.yaml ├── train_watcher.py └── unittests ├── consistency └── step.py └── grad ├── activation.py ├── backend.py ├── leak.py ├── loss.py └── norm.py /.deepsource.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/.deepsource.toml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/README.md -------------------------------------------------------------------------------- /config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/config.yaml -------------------------------------------------------------------------------- /images/axial.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/axial.png -------------------------------------------------------------------------------- /images/convit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/convit.png -------------------------------------------------------------------------------- /images/mixer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/mixer.png -------------------------------------------------------------------------------- /images/normformer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/normformer.png -------------------------------------------------------------------------------- /images/paper_cogview_sandwich.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/paper_cogview_sandwich.png -------------------------------------------------------------------------------- /images/reformer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/reformer.png -------------------------------------------------------------------------------- /images/revnet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/revnet.png -------------------------------------------------------------------------------- /images/rmsnorm_speed.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/rmsnorm_speed.png -------------------------------------------------------------------------------- /images/sandwich.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/sandwich.png -------------------------------------------------------------------------------- /images/scalenorm.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/images/scalenorm.png -------------------------------------------------------------------------------- /img.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/img.png -------------------------------------------------------------------------------- /inference.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/inference.py -------------------------------------------------------------------------------- /launch_on_tensorfork.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/launch_on_tensorfork.py -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/main.py -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/requirements-dev.txt -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/requirements.txt -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/run.sh -------------------------------------------------------------------------------- /script/inference.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/script/inference.py -------------------------------------------------------------------------------- /script/launch_multiple_runs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/script/launch_multiple_runs.py -------------------------------------------------------------------------------- /script/quantized_vid2tfrecord.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/script/quantized_vid2tfrecord.py -------------------------------------------------------------------------------- /script/start_on_multi_host_tpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/script/start_on_multi_host_tpu.py -------------------------------------------------------------------------------- /script/text2tfrecord.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/script/text2tfrecord.py -------------------------------------------------------------------------------- /setup.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/setup.sh -------------------------------------------------------------------------------- /setup_dev.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/setup_dev.sh -------------------------------------------------------------------------------- /share_vm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/share_vm.py -------------------------------------------------------------------------------- /src/backend.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/backend.py -------------------------------------------------------------------------------- /src/constants.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/constants.py -------------------------------------------------------------------------------- /src/context.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/context.py -------------------------------------------------------------------------------- /src/data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/data.py -------------------------------------------------------------------------------- /src/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/main.py -------------------------------------------------------------------------------- /src/model/activate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/activate.py -------------------------------------------------------------------------------- /src/model/conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/conv.py -------------------------------------------------------------------------------- /src/model/loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/loss.py -------------------------------------------------------------------------------- /src/model/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/main.py -------------------------------------------------------------------------------- /src/model/mixer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/mixer.py -------------------------------------------------------------------------------- /src/model/moe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/moe.py -------------------------------------------------------------------------------- /src/model/norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/norm.py -------------------------------------------------------------------------------- /src/model/reversible.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/model/reversible.py -------------------------------------------------------------------------------- /src/optimizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/optimizer.py -------------------------------------------------------------------------------- /src/utils/checkpoint.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/utils/checkpoint.py -------------------------------------------------------------------------------- /src/utils/wandblog.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/src/utils/wandblog.py -------------------------------------------------------------------------------- /sweep.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/sweep.yaml -------------------------------------------------------------------------------- /train_watcher.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/train_watcher.py -------------------------------------------------------------------------------- /unittests/consistency/step.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/unittests/consistency/step.py -------------------------------------------------------------------------------- /unittests/grad/activation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/unittests/grad/activation.py -------------------------------------------------------------------------------- /unittests/grad/backend.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/unittests/grad/backend.py -------------------------------------------------------------------------------- /unittests/grad/leak.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/unittests/grad/leak.py -------------------------------------------------------------------------------- /unittests/grad/loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/unittests/grad/loss.py -------------------------------------------------------------------------------- /unittests/grad/norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HomebrewML/Olmax/HEAD/unittests/grad/norm.py --------------------------------------------------------------------------------