このリポジトリは、xLSTMを用いた言語モデルを学習するコードを実装したものです。
論文「xLSTM: Extended Long Short-Term Memory」に基づいて実装しています。
. ├── Makefile : Lint・Formatを実行するためのMakefile ├── README.md ├── cmd │ ├── docker.sh : Dockerコンテナのビルドとコンテナ起動するためのスクリプト │ └── train.sh : 学習を実行するためのスクリプト ├── docker │ └── Dockerfile ├── environment.yaml : conda環境の設定ファイル ├── pyproject.toml ├── requirements-dev.lock ├── requirements.lock ├── src │ ├── cfg │ │ ├── config_type.py │ │ ├── load_yaml_cfg.py │ │ └── yaml : 学習時の設定ファイル │ │ ├── 1.3b │ │ │ └── config.yaml │ │ ├── 125m │ │ │ └── config.yaml │ │ ├── 2.7b │ │ │ └── config.yaml │ │ ├── 350m │ │ │ └── config.yaml │ │ └── template │ │ └── config.yaml │ ├── dataset : データセットの読み込みを行うクラス │ │ ├── ja_cc_dataset.py │ │ ├── ja_wiki_dataset.py │ │ ├── nlp_dataset.py │ │ └── slim_pajama_dataset.py │ ├── experiment : 実験のためのクラス │ │ ├── setup │ │ │ ├── criterion.py │ │ │ ├── dataset.py │ │ │ ├── lr_scheduler.py │ │ │ ├── model.py │ │ │ └── optimizer.py │ │ ├── test │ │ │ └── generate.py │ │ └── train │ │ ├── train.py │ │ └── trainer.py │ ├── model : モデルの読み込みを行うクラス │ │ └── xlstm_model_wrapper.py │ └── utils.py └── tests
- xlstm : NX-AI社が公開しているxLSTMの公式実装
- transformers : tokenizerの読み込みに使用
- datasets : データセットの読み込みに使用
- omegaconf : 設定ファイルの読み込みに使用
- pyTorch : モデルの学習に使用
- wandb : ログの保存に使用
- ruff : Linter・Formatter
- mypy : 型チェック
- .envの作成
wandbのAPIキーを.envに記述してください。
$ cp .env.sample .env- Dockerコンテナのビルド
$ bash cmd/docker.sh build- Dockerコンテナの起動
$ bash cmd/docker.sh shell [GPU_ID](ex. 0,1,2,3)- 学習の実行
$ bash cmd/train.sh docker [CONFIG_PATH](ex. src/cfg/yaml/1.3b/config.yaml)- 依存関係の解決
- Ryeを用いる場合
$ rye sync- pipを用いる場合
$ pip install -r requirements.lock
$ source .venv/bin/activate- wandbへのログイン
$ (rye run) wandb login※ pipを利用している場合、
(rye run)を削除してください。
- 学習の実行
$ bash cmd/train.sh local [CONFIG_PATH](ex. src/cfg/yaml/1.3b/config.yaml)※ pipを利用している場合、cmd/train.shの
rye run pythonをpythonに変更してください。