├── .flake8 ├── .github ├── ISSUE_TEMPLATE │ ├── bug-report.yml │ ├── config.yml │ ├── feature-request.yml │ └── rfc.yml └── workflows │ ├── issue.yml │ ├── lint.yaml │ └── release.yml ├── .gitignore ├── .gitmodules ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── configs ├── delta_net_1B.json ├── delta_net_340M.json ├── gated_deltanet_1B.json ├── gated_deltanet_340M.json ├── gla_340M.json ├── gla_7B.json ├── gsa_1B.json ├── gsa_340M.json ├── hgrn2_340M.json ├── mamba2_1B.json ├── mamba2_340M.json ├── mamba_1B.json ├── mamba_340M.json ├── samba_1B.json ├── sba_340m.json ├── transformer_1B.json ├── transformer_340M.json └── transformer_7B.json ├── custom_models └── sba │ ├── __init__.py │ ├── config_sba.py │ ├── modeling_sba.py │ ├── sba_layer.py │ └── stickbreaking_attention │ ├── __init__.py │ ├── sb_attn │ ├── __init__.py │ ├── sb_bwd.py │ └── sb_fwd.py │ ├── sb_ref.py │ ├── sb_varlen │ ├── __init__.py │ ├── sb_varlen_bwd.py │ ├── sb_varlen_fwd.py │ └── softplus.py │ └── utils.py ├── flame ├── __init__.py ├── components │ ├── __init__.py │ └── checkpoint.py ├── config_manager.py ├── data.py ├── models │ ├── __init__.py │ ├── activation_offloading.py │ ├── fla.toml │ ├── parallelize_fla.py │ └── pipeline_fla.py ├── tools │ ├── __init__.py │ └── utils.py ├── train.py └── utils │ ├── __init__.py │ ├── convert_dcp_to_hf.py │ ├── convert_hf_to_dcp.py │ └── preprocess.py ├── pyproject.toml ├── setup.py ├── tests └── torchtitan ├── train.sh └── utils └── reshard.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 127 3 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug-report.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/ISSUE_TEMPLATE/bug-report.yml -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/ISSUE_TEMPLATE/config.yml -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature-request.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/ISSUE_TEMPLATE/feature-request.yml -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/rfc.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/ISSUE_TEMPLATE/rfc.yml -------------------------------------------------------------------------------- /.github/workflows/issue.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/workflows/issue.yml -------------------------------------------------------------------------------- /.github/workflows/lint.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/workflows/lint.yaml -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.github/workflows/release.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.gitmodules -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/README.md -------------------------------------------------------------------------------- /configs/delta_net_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/delta_net_1B.json -------------------------------------------------------------------------------- /configs/delta_net_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/delta_net_340M.json -------------------------------------------------------------------------------- /configs/gated_deltanet_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/gated_deltanet_1B.json -------------------------------------------------------------------------------- /configs/gated_deltanet_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/gated_deltanet_340M.json -------------------------------------------------------------------------------- /configs/gla_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/gla_340M.json -------------------------------------------------------------------------------- /configs/gla_7B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/gla_7B.json -------------------------------------------------------------------------------- /configs/gsa_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/gsa_1B.json -------------------------------------------------------------------------------- /configs/gsa_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/gsa_340M.json -------------------------------------------------------------------------------- /configs/hgrn2_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/hgrn2_340M.json -------------------------------------------------------------------------------- /configs/mamba2_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/mamba2_1B.json -------------------------------------------------------------------------------- /configs/mamba2_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/mamba2_340M.json -------------------------------------------------------------------------------- /configs/mamba_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/mamba_1B.json -------------------------------------------------------------------------------- /configs/mamba_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/mamba_340M.json -------------------------------------------------------------------------------- /configs/samba_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/samba_1B.json -------------------------------------------------------------------------------- /configs/sba_340m.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/sba_340m.json -------------------------------------------------------------------------------- /configs/transformer_1B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/transformer_1B.json -------------------------------------------------------------------------------- /configs/transformer_340M.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/transformer_340M.json -------------------------------------------------------------------------------- /configs/transformer_7B.json: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/configs/transformer_7B.json -------------------------------------------------------------------------------- /custom_models/sba/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/__init__.py -------------------------------------------------------------------------------- /custom_models/sba/config_sba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/config_sba.py -------------------------------------------------------------------------------- /custom_models/sba/modeling_sba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/modeling_sba.py -------------------------------------------------------------------------------- /custom_models/sba/sba_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/sba_layer.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/__init__.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_attn/__init__.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_attn/sb_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_attn/sb_bwd.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_attn/sb_fwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_attn/sb_fwd.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_ref.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_ref.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_varlen/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_varlen/__init__.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_varlen/sb_varlen_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_varlen/sb_varlen_bwd.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_varlen/sb_varlen_fwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_varlen/sb_varlen_fwd.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/sb_varlen/softplus.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/sb_varlen/softplus.py -------------------------------------------------------------------------------- /custom_models/sba/stickbreaking_attention/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/custom_models/sba/stickbreaking_attention/utils.py -------------------------------------------------------------------------------- /flame/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.1.0" 2 | -------------------------------------------------------------------------------- /flame/components/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /flame/components/checkpoint.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/components/checkpoint.py -------------------------------------------------------------------------------- /flame/config_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/config_manager.py -------------------------------------------------------------------------------- /flame/data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/data.py -------------------------------------------------------------------------------- /flame/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /flame/models/activation_offloading.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/models/activation_offloading.py -------------------------------------------------------------------------------- /flame/models/fla.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/models/fla.toml -------------------------------------------------------------------------------- /flame/models/parallelize_fla.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/models/parallelize_fla.py -------------------------------------------------------------------------------- /flame/models/pipeline_fla.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/models/pipeline_fla.py -------------------------------------------------------------------------------- /flame/tools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /flame/tools/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/tools/utils.py -------------------------------------------------------------------------------- /flame/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/train.py -------------------------------------------------------------------------------- /flame/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /flame/utils/convert_dcp_to_hf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/utils/convert_dcp_to_hf.py -------------------------------------------------------------------------------- /flame/utils/convert_hf_to_dcp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/utils/convert_hf_to_dcp.py -------------------------------------------------------------------------------- /flame/utils/preprocess.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/flame/utils/preprocess.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/pyproject.toml -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/setup.py -------------------------------------------------------------------------------- /tests/torchtitan: -------------------------------------------------------------------------------- 1 | ../torchtitan -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/train.sh -------------------------------------------------------------------------------- /utils/reshard.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fla-org/flame/HEAD/utils/reshard.py --------------------------------------------------------------------------------