├── .github └── workflows │ └── build.yml ├── .gitignore ├── AUTHORS ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── examples ├── mnist.py ├── mnist_dataset.py └── toy_model_tpu.py ├── mesh_tensorflow ├── __init__.py ├── auto_mtf │ ├── README.md │ ├── __init__.py │ ├── api.py │ ├── api_test.py │ ├── graph_interface.py │ ├── graph_interface_test.py │ ├── layout_optimizer.py │ ├── layout_optimizer_test.py │ ├── memory_estimator.py │ ├── memory_estimator_test.py │ ├── print_cp_model_solution.py │ ├── scheduler.py │ ├── scheduler_test.py │ ├── valid_layouts.py │ └── valid_layouts_test.py ├── beam_search.py ├── bert │ ├── __init__.py │ ├── bert.py │ ├── config │ │ ├── base.json │ │ ├── base_renamed.json │ │ ├── base_v2.json │ │ ├── base_v2_relu.json │ │ ├── default_init.json │ │ ├── direct.json │ │ ├── intermediate_dropout.json │ │ ├── large.json │ │ ├── large_v2.json │ │ ├── moe.json │ │ ├── no_bias.json │ │ ├── relative.json │ │ ├── xl_v2.json │ │ └── xxl_v2.json │ ├── optimization.py │ ├── run_classifier.py │ ├── run_pretraining.py │ ├── run_squad.py │ └── tokenization.py ├── experimental │ ├── __init__.py │ ├── data_aug_lib.py │ ├── data_aug_lib_test.py │ ├── input_reader.py │ ├── input_reader_test.py │ ├── model_executor.py │ ├── offline_data_aug.py │ └── unet.py ├── import_test.py ├── layers.py ├── layers_test.py ├── ops.py ├── ops_test.py ├── ops_with_redefined_builtins.py ├── optimize.py ├── placement_mesh_impl.py ├── simd_mesh_impl.py ├── simd_mesh_impl_test.py ├── test_utils.py ├── test_utils_test.py ├── tpu_variables.py ├── transformer │ ├── __init__.py │ ├── adaptive_softmax.py │ ├── adaptive_softmax_test.py │ ├── attention.py │ ├── dataset.py │ ├── dataset_test.py │ ├── evolved_transformer.py │ ├── fixup_layers.py │ ├── funnel_transformer.py │ ├── funnel_transformer_test.py │ ├── gin │ │ ├── 8x8_dp_r4.gin │ │ ├── __init__.py │ │ ├── defaults.gin │ │ ├── layouts │ │ │ ├── 16x16_m8.gin │ │ │ ├── 2x2_dp.gin │ │ │ ├── 2x2_dp_mp.gin │ │ │ ├── 2x2_mp.gin │ │ │ ├── 4x4_dp.gin │ │ │ ├── 8x16_m2.gin │ │ │ ├── 8x16_m8.gin │ │ │ ├── 8x8_dp.gin │ │ │ ├── 8x8_m4.gin │ │ │ ├── 8x8_mp.gin │ │ │ └── single.gin │ │ ├── learning_rate_schedules │ │ │ ├── constant.gin │ │ │ └── slanted_triangular.gin │ │ ├── models │ │ │ ├── aiayn_base.gin │ │ │ ├── aiayn_large.gin │ │ │ ├── bi_base.gin │ │ │ ├── bi_l12.gin │ │ │ ├── bi_l12_w1.gin │ │ │ ├── bi_l12_w2.gin │ │ │ ├── bi_l12_w3.gin │ │ │ ├── bi_moe.gin │ │ │ ├── bi_w1.gin │ │ │ ├── lm_base.gin │ │ │ ├── lm_moe.gin │ │ │ ├── lm_w1.gin │ │ │ └── tiny.gin │ │ ├── problems │ │ │ ├── lm1b.gin │ │ │ ├── lm1b_untok.gin │ │ │ ├── t2t_ende.gin │ │ │ ├── t2t_enfr.gin │ │ │ └── t2t_lm1b.gin │ │ └── transparent_attention.gin │ ├── heterogeneous_moe.py │ ├── learning_rate_schedules.py │ ├── learning_rate_schedules_test.py │ ├── main.py │ ├── memory_layers.py │ ├── memory_layers_test.py │ ├── moe.py │ ├── t2t_vocabulary.py │ ├── transformer.py │ ├── transformer_layers.py │ ├── transformer_layers_test.py │ ├── universal_transformer.py │ ├── utils.py │ ├── utils_test.py │ ├── vocab_embeddings.py │ ├── vocab_embeddings_test.py │ └── vocabulary.py ├── utils.py └── utils_test.py ├── oss_scripts ├── oss_pip_install.sh ├── oss_release.sh └── oss_tests.sh ├── setup.cfg ├── setup.py ├── transformer_data_splitting.png └── transformer_model_splitting.png /.github/workflows/build.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/.github/workflows/build.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/.gitignore -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/AUTHORS -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/README.md -------------------------------------------------------------------------------- /examples/mnist.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/examples/mnist.py -------------------------------------------------------------------------------- /examples/mnist_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/examples/mnist_dataset.py -------------------------------------------------------------------------------- /examples/toy_model_tpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/examples/toy_model_tpu.py -------------------------------------------------------------------------------- /mesh_tensorflow/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/__init__.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/README.md -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/__init__.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/api.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/api.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/api_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/api_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/graph_interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/graph_interface.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/graph_interface_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/graph_interface_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/layout_optimizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/layout_optimizer.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/layout_optimizer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/layout_optimizer_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/memory_estimator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/memory_estimator.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/memory_estimator_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/memory_estimator_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/print_cp_model_solution.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/print_cp_model_solution.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/scheduler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/scheduler.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/scheduler_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/scheduler_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/valid_layouts.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/valid_layouts.py -------------------------------------------------------------------------------- /mesh_tensorflow/auto_mtf/valid_layouts_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/auto_mtf/valid_layouts_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/beam_search.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/beam_search.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/__init__.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/bert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/bert.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/base.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/base.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/base_renamed.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/base_renamed.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/base_v2.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/base_v2.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/base_v2_relu.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/base_v2_relu.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/default_init.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/default_init.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/direct.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/direct.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/intermediate_dropout.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/intermediate_dropout.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/large.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/large.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/large_v2.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/large_v2.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/moe.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/moe.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/no_bias.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/no_bias.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/relative.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/relative.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/xl_v2.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/xl_v2.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/config/xxl_v2.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/config/xxl_v2.json -------------------------------------------------------------------------------- /mesh_tensorflow/bert/optimization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/optimization.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/run_classifier.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/run_classifier.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/run_pretraining.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/run_pretraining.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/run_squad.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/run_squad.py -------------------------------------------------------------------------------- /mesh_tensorflow/bert/tokenization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/bert/tokenization.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/__init__.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/data_aug_lib.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/data_aug_lib.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/data_aug_lib_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/data_aug_lib_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/input_reader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/input_reader.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/input_reader_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/input_reader_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/model_executor.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/model_executor.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/offline_data_aug.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/offline_data_aug.py -------------------------------------------------------------------------------- /mesh_tensorflow/experimental/unet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/experimental/unet.py -------------------------------------------------------------------------------- /mesh_tensorflow/import_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/import_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/layers.py -------------------------------------------------------------------------------- /mesh_tensorflow/layers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/layers_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/ops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/ops.py -------------------------------------------------------------------------------- /mesh_tensorflow/ops_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/ops_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/ops_with_redefined_builtins.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/ops_with_redefined_builtins.py -------------------------------------------------------------------------------- /mesh_tensorflow/optimize.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/optimize.py -------------------------------------------------------------------------------- /mesh_tensorflow/placement_mesh_impl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/placement_mesh_impl.py -------------------------------------------------------------------------------- /mesh_tensorflow/simd_mesh_impl.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/simd_mesh_impl.py -------------------------------------------------------------------------------- /mesh_tensorflow/simd_mesh_impl_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/simd_mesh_impl_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/test_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/test_utils.py -------------------------------------------------------------------------------- /mesh_tensorflow/test_utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/test_utils_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/tpu_variables.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/tpu_variables.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/__init__.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/adaptive_softmax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/adaptive_softmax.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/adaptive_softmax_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/adaptive_softmax_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/attention.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/dataset.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/dataset_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/dataset_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/evolved_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/evolved_transformer.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/fixup_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/fixup_layers.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/funnel_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/funnel_transformer.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/funnel_transformer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/funnel_transformer_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/8x8_dp_r4.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/8x8_dp_r4.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/__init__.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/defaults.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/defaults.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/16x16_m8.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/16x16_m8.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/2x2_dp.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/2x2_dp.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/2x2_dp_mp.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/2x2_dp_mp.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/2x2_mp.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/2x2_mp.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/4x4_dp.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/4x4_dp.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/8x16_m2.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/8x16_m2.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/8x16_m8.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/8x16_m8.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/8x8_dp.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/8x8_dp.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/8x8_m4.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/8x8_m4.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/8x8_mp.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/8x8_mp.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/layouts/single.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/layouts/single.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/learning_rate_schedules/constant.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/learning_rate_schedules/constant.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/learning_rate_schedules/slanted_triangular.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/learning_rate_schedules/slanted_triangular.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/aiayn_base.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/aiayn_base.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/aiayn_large.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/aiayn_large.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_base.gin: -------------------------------------------------------------------------------- 1 | # -*-Python-*- 2 | 3 | utils.run.model_type = "bitransformer" 4 | -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_l12.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/bi_l12.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_l12_w1.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/bi_l12_w1.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_l12_w2.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/bi_l12_w2.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_l12_w3.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/bi_l12_w3.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_moe.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/bi_moe.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/bi_w1.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/bi_w1.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/lm_base.gin: -------------------------------------------------------------------------------- 1 | # -*-Python-*- 2 | 3 | utils.run.model_type = "lm" 4 | -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/lm_moe.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/lm_moe.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/lm_w1.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/lm_w1.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/models/tiny.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/models/tiny.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/problems/lm1b.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/problems/lm1b.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/problems/lm1b_untok.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/problems/lm1b_untok.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/problems/t2t_ende.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/problems/t2t_ende.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/problems/t2t_enfr.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/problems/t2t_enfr.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/problems/t2t_lm1b.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/problems/t2t_lm1b.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/gin/transparent_attention.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/gin/transparent_attention.gin -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/heterogeneous_moe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/heterogeneous_moe.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/learning_rate_schedules.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/learning_rate_schedules.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/learning_rate_schedules_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/learning_rate_schedules_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/main.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/memory_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/memory_layers.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/memory_layers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/memory_layers_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/moe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/moe.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/t2t_vocabulary.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/t2t_vocabulary.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/transformer.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/transformer_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/transformer_layers.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/transformer_layers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/transformer_layers_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/universal_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/universal_transformer.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/utils.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/utils_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/vocab_embeddings.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/vocab_embeddings.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/vocab_embeddings_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/vocab_embeddings_test.py -------------------------------------------------------------------------------- /mesh_tensorflow/transformer/vocabulary.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/transformer/vocabulary.py -------------------------------------------------------------------------------- /mesh_tensorflow/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/utils.py -------------------------------------------------------------------------------- /mesh_tensorflow/utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/mesh_tensorflow/utils_test.py -------------------------------------------------------------------------------- /oss_scripts/oss_pip_install.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/oss_scripts/oss_pip_install.sh -------------------------------------------------------------------------------- /oss_scripts/oss_release.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/oss_scripts/oss_release.sh -------------------------------------------------------------------------------- /oss_scripts/oss_tests.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/oss_scripts/oss_tests.sh -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/setup.cfg -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/setup.py -------------------------------------------------------------------------------- /transformer_data_splitting.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/transformer_data_splitting.png -------------------------------------------------------------------------------- /transformer_model_splitting.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/mesh/HEAD/transformer_model_splitting.png --------------------------------------------------------------------------------