├── .github ├── api-code-review.png ├── arcf-titan-ts-amp-curve.png ├── arcf-titan-ts-amp.png ├── arcf-titan-ts-curve.png ├── att_w_ts.png ├── cls_w_ts.png ├── dp_mp.png ├── in1k-titan-ts-amp.png ├── in1k-titan-ts-zero.png ├── in1k-titan-ts.png ├── logo.png ├── logo.svg ├── minGPT-titan-ts.png ├── minGPT-ts-result.png └── mlp_w_ts.png ├── .gitignore ├── INSTALL.md ├── LICENSE ├── README.md ├── dev ├── README.md ├── build_wheel.sh ├── run_test.sh └── upload.sh ├── docs ├── README.md ├── api │ ├── dist.md │ ├── nn.md │ └── ts.md ├── note │ ├── about.md │ └── dev.md └── tutorial │ ├── amp.md │ ├── face.md │ ├── in1k.md │ └── mgpt.md ├── projects ├── README.md ├── arcface │ ├── README.md │ ├── arcface_linear.py │ ├── arcface_loss.py │ ├── iresnet.py │ ├── main.py │ ├── utils.py │ └── valid.py ├── imagenet │ ├── README.md │ ├── main.py │ └── resnet.py ├── minGPT │ ├── README.md │ ├── main.py │ └── mingpt │ │ ├── __init__.py │ │ ├── model.py │ │ ├── trainer.py │ │ └── utils.py └── templates │ ├── README.md │ ├── ts_layer.py │ └── ts_loss.py ├── pyproject.toml ├── setup.cfg ├── setup.py ├── tests ├── README.md ├── test_cross_entropy.py ├── test_layers.py ├── test_modules.py ├── test_nn.py ├── test_serialization.py └── testing.py └── torchshard ├── __init__.py ├── _tensor.py ├── distributed ├── __init__.py ├── comm.py └── core.py ├── nn ├── __init__.py ├── _reduction.py ├── functional.py ├── init.py └── modules │ ├── __init__.py │ ├── linear.py │ ├── loss.py │ └── sparse.py ├── overrides.py ├── serialization.py └── utils.py /.github/api-code-review.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/api-code-review.png -------------------------------------------------------------------------------- /.github/arcf-titan-ts-amp-curve.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/arcf-titan-ts-amp-curve.png -------------------------------------------------------------------------------- /.github/arcf-titan-ts-amp.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/arcf-titan-ts-amp.png -------------------------------------------------------------------------------- /.github/arcf-titan-ts-curve.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/arcf-titan-ts-curve.png -------------------------------------------------------------------------------- /.github/att_w_ts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/att_w_ts.png -------------------------------------------------------------------------------- /.github/cls_w_ts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/cls_w_ts.png -------------------------------------------------------------------------------- /.github/dp_mp.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/dp_mp.png -------------------------------------------------------------------------------- /.github/in1k-titan-ts-amp.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/in1k-titan-ts-amp.png -------------------------------------------------------------------------------- /.github/in1k-titan-ts-zero.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/in1k-titan-ts-zero.png -------------------------------------------------------------------------------- /.github/in1k-titan-ts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/in1k-titan-ts.png -------------------------------------------------------------------------------- /.github/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/logo.png -------------------------------------------------------------------------------- /.github/logo.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/logo.svg -------------------------------------------------------------------------------- /.github/minGPT-titan-ts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/minGPT-titan-ts.png -------------------------------------------------------------------------------- /.github/minGPT-ts-result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/minGPT-ts-result.png -------------------------------------------------------------------------------- /.github/mlp_w_ts.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.github/mlp_w_ts.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/.gitignore -------------------------------------------------------------------------------- /INSTALL.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/INSTALL.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/README.md -------------------------------------------------------------------------------- /dev/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/dev/README.md -------------------------------------------------------------------------------- /dev/build_wheel.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/dev/build_wheel.sh -------------------------------------------------------------------------------- /dev/run_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CUDA_VISIBLE_DEVICES=0,1 python3 -m unittest discover -v -s tests -------------------------------------------------------------------------------- /dev/upload.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/dev/upload.sh -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/api/dist.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/api/dist.md -------------------------------------------------------------------------------- /docs/api/nn.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/api/nn.md -------------------------------------------------------------------------------- /docs/api/ts.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/api/ts.md -------------------------------------------------------------------------------- /docs/note/about.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/note/about.md -------------------------------------------------------------------------------- /docs/note/dev.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/note/dev.md -------------------------------------------------------------------------------- /docs/tutorial/amp.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/tutorial/amp.md -------------------------------------------------------------------------------- /docs/tutorial/face.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/tutorial/face.md -------------------------------------------------------------------------------- /docs/tutorial/in1k.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/tutorial/in1k.md -------------------------------------------------------------------------------- /docs/tutorial/mgpt.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/docs/tutorial/mgpt.md -------------------------------------------------------------------------------- /projects/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/README.md -------------------------------------------------------------------------------- /projects/arcface/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/README.md -------------------------------------------------------------------------------- /projects/arcface/arcface_linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/arcface_linear.py -------------------------------------------------------------------------------- /projects/arcface/arcface_loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/arcface_loss.py -------------------------------------------------------------------------------- /projects/arcface/iresnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/iresnet.py -------------------------------------------------------------------------------- /projects/arcface/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/main.py -------------------------------------------------------------------------------- /projects/arcface/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/utils.py -------------------------------------------------------------------------------- /projects/arcface/valid.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/arcface/valid.py -------------------------------------------------------------------------------- /projects/imagenet/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/imagenet/README.md -------------------------------------------------------------------------------- /projects/imagenet/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/imagenet/main.py -------------------------------------------------------------------------------- /projects/imagenet/resnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/imagenet/resnet.py -------------------------------------------------------------------------------- /projects/minGPT/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/minGPT/README.md -------------------------------------------------------------------------------- /projects/minGPT/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/minGPT/main.py -------------------------------------------------------------------------------- /projects/minGPT/mingpt/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /projects/minGPT/mingpt/model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/minGPT/mingpt/model.py -------------------------------------------------------------------------------- /projects/minGPT/mingpt/trainer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/minGPT/mingpt/trainer.py -------------------------------------------------------------------------------- /projects/minGPT/mingpt/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/minGPT/mingpt/utils.py -------------------------------------------------------------------------------- /projects/templates/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/templates/README.md -------------------------------------------------------------------------------- /projects/templates/ts_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/templates/ts_layer.py -------------------------------------------------------------------------------- /projects/templates/ts_loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/projects/templates/ts_loss.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/pyproject.toml -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/setup.cfg -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/setup.py -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/README.md -------------------------------------------------------------------------------- /tests/test_cross_entropy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/test_cross_entropy.py -------------------------------------------------------------------------------- /tests/test_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/test_layers.py -------------------------------------------------------------------------------- /tests/test_modules.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/test_modules.py -------------------------------------------------------------------------------- /tests/test_nn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/test_nn.py -------------------------------------------------------------------------------- /tests/test_serialization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/test_serialization.py -------------------------------------------------------------------------------- /tests/testing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/tests/testing.py -------------------------------------------------------------------------------- /torchshard/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/__init__.py -------------------------------------------------------------------------------- /torchshard/_tensor.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/_tensor.py -------------------------------------------------------------------------------- /torchshard/distributed/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/distributed/__init__.py -------------------------------------------------------------------------------- /torchshard/distributed/comm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/distributed/comm.py -------------------------------------------------------------------------------- /torchshard/distributed/core.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/distributed/core.py -------------------------------------------------------------------------------- /torchshard/nn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/__init__.py -------------------------------------------------------------------------------- /torchshard/nn/_reduction.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/_reduction.py -------------------------------------------------------------------------------- /torchshard/nn/functional.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/functional.py -------------------------------------------------------------------------------- /torchshard/nn/init.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/init.py -------------------------------------------------------------------------------- /torchshard/nn/modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/modules/__init__.py -------------------------------------------------------------------------------- /torchshard/nn/modules/linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/modules/linear.py -------------------------------------------------------------------------------- /torchshard/nn/modules/loss.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/modules/loss.py -------------------------------------------------------------------------------- /torchshard/nn/modules/sparse.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/nn/modules/sparse.py -------------------------------------------------------------------------------- /torchshard/overrides.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/overrides.py -------------------------------------------------------------------------------- /torchshard/serialization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/serialization.py -------------------------------------------------------------------------------- /torchshard/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/kaiyuyue/torchshard/HEAD/torchshard/utils.py --------------------------------------------------------------------------------