├── .github └── workflows │ ├── lint.yml │ ├── mypy.yml │ └── pytest.yml ├── .gitignore ├── .isort.cfg ├── LICENSE ├── README.md ├── code-of-conduct.md ├── docs ├── configurations.md ├── dataloader.md ├── evaluation.md ├── fine_tuning.md └── train_details.md ├── fms_fsdp ├── __init__.py ├── config │ ├── __init__.py │ └── training.py ├── policies │ ├── __init__.py │ ├── ac_handler.py │ ├── mixed_precision.py │ ├── param_init.py │ └── wrapping.py ├── readme.md ├── requirements.txt └── utils │ ├── __init__.py │ ├── checkpointing_utils.py │ ├── config_utils.py │ ├── dataloader_utils.py │ ├── dataset_utils.py │ └── train_utils.py ├── fms_to_hf_llama.py ├── fms_to_hf_mamba.py ├── images ├── loss_curve.png └── lr.png ├── main_training_llama.py ├── main_training_mamba.py ├── requirements-speculator.txt ├── requirements.txt ├── scripts ├── README_SPECULATOR.md ├── train.sh ├── train.slurm └── train_speculator.sh ├── setup.py ├── speculator ├── __init__.py ├── train_speculator.py └── train_speculator_utils.py ├── test-requirements.txt └── tests ├── conftest.py ├── test_datasets.py └── test_selective_ac.py /.github/workflows/lint.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/.github/workflows/lint.yml -------------------------------------------------------------------------------- /.github/workflows/mypy.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/.github/workflows/mypy.yml -------------------------------------------------------------------------------- /.github/workflows/pytest.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/.github/workflows/pytest.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/.gitignore -------------------------------------------------------------------------------- /.isort.cfg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/.isort.cfg -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/README.md -------------------------------------------------------------------------------- /code-of-conduct.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/code-of-conduct.md -------------------------------------------------------------------------------- /docs/configurations.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/docs/configurations.md -------------------------------------------------------------------------------- /docs/dataloader.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/docs/dataloader.md -------------------------------------------------------------------------------- /docs/evaluation.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/docs/evaluation.md -------------------------------------------------------------------------------- /docs/fine_tuning.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/docs/fine_tuning.md -------------------------------------------------------------------------------- /docs/train_details.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/docs/train_details.md -------------------------------------------------------------------------------- /fms_fsdp/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fms_fsdp/config/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/config/__init__.py -------------------------------------------------------------------------------- /fms_fsdp/config/training.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/config/training.py -------------------------------------------------------------------------------- /fms_fsdp/policies/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/policies/__init__.py -------------------------------------------------------------------------------- /fms_fsdp/policies/ac_handler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/policies/ac_handler.py -------------------------------------------------------------------------------- /fms_fsdp/policies/mixed_precision.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/policies/mixed_precision.py -------------------------------------------------------------------------------- /fms_fsdp/policies/param_init.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/policies/param_init.py -------------------------------------------------------------------------------- /fms_fsdp/policies/wrapping.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/policies/wrapping.py -------------------------------------------------------------------------------- /fms_fsdp/readme.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/readme.md -------------------------------------------------------------------------------- /fms_fsdp/requirements.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fms_fsdp/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /fms_fsdp/utils/checkpointing_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/utils/checkpointing_utils.py -------------------------------------------------------------------------------- /fms_fsdp/utils/config_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/utils/config_utils.py -------------------------------------------------------------------------------- /fms_fsdp/utils/dataloader_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/utils/dataloader_utils.py -------------------------------------------------------------------------------- /fms_fsdp/utils/dataset_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/utils/dataset_utils.py -------------------------------------------------------------------------------- /fms_fsdp/utils/train_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_fsdp/utils/train_utils.py -------------------------------------------------------------------------------- /fms_to_hf_llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_to_hf_llama.py -------------------------------------------------------------------------------- /fms_to_hf_mamba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/fms_to_hf_mamba.py -------------------------------------------------------------------------------- /images/loss_curve.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/images/loss_curve.png -------------------------------------------------------------------------------- /images/lr.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/images/lr.png -------------------------------------------------------------------------------- /main_training_llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/main_training_llama.py -------------------------------------------------------------------------------- /main_training_mamba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/main_training_mamba.py -------------------------------------------------------------------------------- /requirements-speculator.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/requirements-speculator.txt -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/requirements.txt -------------------------------------------------------------------------------- /scripts/README_SPECULATOR.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/scripts/README_SPECULATOR.md -------------------------------------------------------------------------------- /scripts/train.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/scripts/train.sh -------------------------------------------------------------------------------- /scripts/train.slurm: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/scripts/train.slurm -------------------------------------------------------------------------------- /scripts/train_speculator.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/scripts/train_speculator.sh -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/setup.py -------------------------------------------------------------------------------- /speculator/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /speculator/train_speculator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/speculator/train_speculator.py -------------------------------------------------------------------------------- /speculator/train_speculator_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/speculator/train_speculator_utils.py -------------------------------------------------------------------------------- /test-requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/test-requirements.txt -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/tests/conftest.py -------------------------------------------------------------------------------- /tests/test_datasets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/tests/test_datasets.py -------------------------------------------------------------------------------- /tests/test_selective_ac.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/foundation-model-stack/fms-fsdp/HEAD/tests/test_selective_ac.py --------------------------------------------------------------------------------