├── .github └── workflows │ └── action.yml ├── LICENSE ├── Makefile ├── README.md ├── benchmark_gpu.py ├── benchmark_tpu.py ├── conf ├── config.yaml └── model_config.yaml ├── data └── download_and_prepare_data.py ├── imgs └── 28b_TPU.png ├── main.py ├── prepareGPU.sh ├── prepareTPUVM.sh ├── requirements ├── tests ├── __init__.py ├── test_model_components.py ├── test_model_factory.py └── test_train_step.py └── transformer_tp ├── __init__ .py ├── models ├── __init__.py ├── replicated_utils.py └── transformer.py ├── partitioning ├── __init__.py └── partition.py ├── training ├── __init__.py └── train_functions.py └── utils ├── __init__.py ├── checkpoint.py └── helpers.py /.github/workflows/action.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/.github/workflows/action.yml -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/LICENSE -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/Makefile -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/README.md -------------------------------------------------------------------------------- /benchmark_gpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/benchmark_gpu.py -------------------------------------------------------------------------------- /benchmark_tpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/benchmark_tpu.py -------------------------------------------------------------------------------- /conf/config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/conf/config.yaml -------------------------------------------------------------------------------- /conf/model_config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/conf/model_config.yaml -------------------------------------------------------------------------------- /data/download_and_prepare_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/data/download_and_prepare_data.py -------------------------------------------------------------------------------- /imgs/28b_TPU.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/imgs/28b_TPU.png -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/main.py -------------------------------------------------------------------------------- /prepareGPU.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/prepareGPU.sh -------------------------------------------------------------------------------- /prepareTPUVM.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/prepareTPUVM.sh -------------------------------------------------------------------------------- /requirements: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/requirements -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_model_components.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/tests/test_model_components.py -------------------------------------------------------------------------------- /tests/test_model_factory.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/tests/test_model_factory.py -------------------------------------------------------------------------------- /tests/test_train_step.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/tests/test_train_step.py -------------------------------------------------------------------------------- /transformer_tp/__init__ .py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/__init__ .py -------------------------------------------------------------------------------- /transformer_tp/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/models/__init__.py -------------------------------------------------------------------------------- /transformer_tp/models/replicated_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/models/replicated_utils.py -------------------------------------------------------------------------------- /transformer_tp/models/transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/models/transformer.py -------------------------------------------------------------------------------- /transformer_tp/partitioning/__init__.py: -------------------------------------------------------------------------------- 1 | from .partition import create_opt_spec 2 | -------------------------------------------------------------------------------- /transformer_tp/partitioning/partition.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/partitioning/partition.py -------------------------------------------------------------------------------- /transformer_tp/training/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/training/__init__.py -------------------------------------------------------------------------------- /transformer_tp/training/train_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/training/train_functions.py -------------------------------------------------------------------------------- /transformer_tp/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/utils/__init__.py -------------------------------------------------------------------------------- /transformer_tp/utils/checkpoint.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/utils/checkpoint.py -------------------------------------------------------------------------------- /transformer_tp/utils/helpers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fattorib/transformer_shmap/HEAD/transformer_tp/utils/helpers.py --------------------------------------------------------------------------------