├── .gitignore ├── LICENSE ├── README.md ├── docs ├── README.md └── sharding.md ├── examples ├── llama_fsdp.ipynb ├── llama_fsdp_tp.ipynb └── mnist_data_parallel.ipynb ├── scalax ├── __init__.py ├── sharding.py └── utils.py ├── setup.py └── tests ├── sharding_test.py └── utils_test.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/README.md -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/sharding.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/docs/sharding.md -------------------------------------------------------------------------------- /examples/llama_fsdp.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/examples/llama_fsdp.ipynb -------------------------------------------------------------------------------- /examples/llama_fsdp_tp.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/examples/llama_fsdp_tp.ipynb -------------------------------------------------------------------------------- /examples/mnist_data_parallel.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/examples/mnist_data_parallel.ipynb -------------------------------------------------------------------------------- /scalax/__init__.py: -------------------------------------------------------------------------------- 1 | from scalax import sharding, utils 2 | -------------------------------------------------------------------------------- /scalax/sharding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/scalax/sharding.py -------------------------------------------------------------------------------- /scalax/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/scalax/utils.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/setup.py -------------------------------------------------------------------------------- /tests/sharding_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/tests/sharding_test.py -------------------------------------------------------------------------------- /tests/utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/young-geng/scalax/HEAD/tests/utils_test.py --------------------------------------------------------------------------------