diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..66cb585b26f73322c31cc00a5a567377a3f98eca 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,35 +1,52 @@ *.7z filter=lfs diff=lfs merge=lfs -text *.arrow filter=lfs diff=lfs merge=lfs -text *.bin filter=lfs diff=lfs merge=lfs -text +*.bin.* filter=lfs diff=lfs merge=lfs -text *.bz2 filter=lfs diff=lfs merge=lfs -text -*.ckpt filter=lfs diff=lfs merge=lfs -text *.ftz filter=lfs diff=lfs merge=lfs -text *.gz filter=lfs diff=lfs merge=lfs -text *.h5 filter=lfs diff=lfs merge=lfs -text *.joblib filter=lfs diff=lfs merge=lfs -text *.lfs.* filter=lfs diff=lfs merge=lfs -text -*.mlmodel filter=lfs diff=lfs merge=lfs -text *.model filter=lfs diff=lfs merge=lfs -text *.msgpack filter=lfs diff=lfs merge=lfs -text -*.npy filter=lfs diff=lfs merge=lfs -text -*.npz filter=lfs diff=lfs merge=lfs -text *.onnx filter=lfs diff=lfs merge=lfs -text *.ot filter=lfs diff=lfs merge=lfs -text *.parquet filter=lfs diff=lfs merge=lfs -text *.pb filter=lfs diff=lfs merge=lfs -text -*.pickle filter=lfs diff=lfs merge=lfs -text -*.pkl filter=lfs diff=lfs merge=lfs -text -*.pt filter=lfs diff=lfs merge=lfs -text *.pth filter=lfs diff=lfs merge=lfs -text *.rar filter=lfs diff=lfs merge=lfs -text -*.safetensors filter=lfs diff=lfs merge=lfs -text saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.tar.* filter=lfs diff=lfs merge=lfs -text -*.tar filter=lfs diff=lfs merge=lfs -text *.tflite filter=lfs diff=lfs merge=lfs -text *.tgz filter=lfs diff=lfs merge=lfs -text -*.wasm filter=lfs diff=lfs merge=lfs -text *.xz filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text +*.zstandard filter=lfs diff=lfs merge=lfs -text +*.tfevents* filter=lfs diff=lfs merge=lfs -text +*.db* filter=lfs diff=lfs merge=lfs -text +*.ark* filter=lfs diff=lfs merge=lfs -text +**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text +**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text +**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.gguf* filter=lfs diff=lfs merge=lfs -text +*.ggml filter=lfs diff=lfs merge=lfs -text +*.llamafile* filter=lfs diff=lfs merge=lfs -text +*.pt2 filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +checkpoints/evo2_nemo_7b/weights/__0_1.distcp filter=lfs diff=lfs merge=lfs -text +checkpoints/evo2_nemo_7b/weights/__0_0.distcp filter=lfs diff=lfs merge=lfs -text +checkpoints/evo2_nemo_7b/weights/common.pt filter=lfs diff=lfs merge=lfs -text +*.fa filter=lfs diff=lfs merge=lfs -text +*.fa.gz filter=lfs diff=lfs merge=lfs -text +checkpoints/evo2_nemo_7b/weights/.metadata filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..28ba4c0e5571f523f3cae1e518261c6deb2a0486 --- /dev/null +++ b/LICENSE @@ -0,0 +1,467 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +-- + +This repository also contains code from Hugging Face Inc., Google Research, +and Facebook (from their Fairseq project). Files from these +organizations have notices at the top of each file. Below are licenses +used in those files, as indicated. + + +------------- LICENSE FOR NVIDIA code -------------- + + +# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +------------- LICENSE FOR huggingface and Google Research code -------------- + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +------------- LICENSE FOR Facebook Fairseq code -------------- + +MIT License + +Copyright (c) Facebook, Inc. and its affiliates. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..70b05ba9980728862ff80149922dc72fbfd5af1e --- /dev/null +++ b/README.md @@ -0,0 +1,321 @@ +--- +frameworks: +- "" +language: +- en +- zh +license: apache-2.0 +tags: +- OneScience +- life-science +- genomic-sequence-modeling +- DNA/RNA +tasks: [] +--- + +

+ Evo2 +

+ +# Model Introduction + +Evo2 is a large-scale foundation model for genomic sequences. It can generate, complete, and predict functions for DNA/RNA sequences. The model supports pretraining on large-scale genomic datasets and can be rapidly adapted to downstream tasks such as regulatory element prediction and variant effect assessment, providing a unified sequence modeling tool for genomics research. + +# Model Description + +Evo2 is built on a hybrid architecture that combines Hyena operators with convolutional gating, and supports context windows at the million-token scale. The model is pretrained on multi-species genomic data containing trillions of base pairs. It takes raw nucleotide sequences as input and supports autoregressive generation, sequence classification, and embedding extraction. Evo2 is compatible with multi-GPU distributed training and flexible data mixing, achieves leading performance on benchmarks such as GenBench, and is suitable for pretraining, fine-tuning, and downstream functional validation. + +# Applicable Scenarios + +| Scenario | Description | +| --- | --- | +| DNA/RNA sequence generation | Generate subsequent genomic sequence fragments from an input prompt | +| FASTA sequence prediction | Read FASTA input and output sequence prediction results | +| Local offline inference validation | Use the example inputs in this repository and a local checkpoint to verify the inference workflow end to end | +| Mini training workflow validation | Use `data/data_mini/genome_data/` to validate the training scripts, data loading, and checkpoint workflow | +| Cluster training adaptation | Use `scripts/train_slurm.py` to run training in a Slurm/cluster environment | +| Data preprocessing | Generate the `.bin/.idx` data required by the Byte-Level tokenizer from FASTA/JSON data | + + +# Usage + +## 1. Using OneCode + +You can try intelligent one-click AI4S programming through the OneCode online environment: + +[Try intelligent one-click AI4S programming](https://gitee.com/link?target=https%3A%2F%2Fweb-2069360198568017922-iaaj.ksai.scnet.cn%3A58043%2Fhome) + +## 2. Manual Installation and Usage + +**Hardware Requirements** + +- Running on a GPU or DCU is recommended. +- A CPU can be used for import checks and small-configuration connectivity validation, but full training and inference will be slow. +- DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster. + +### Download the Model Package + +```bash +modelscope download --model OneScience/evo2 --local_dir ./evo2 +cd evo2 +``` + +### Install the Runtime Environment +#### DCU Environment + +```bash +# Activate DTK and CONDA first +conda create -n onescience311 python=3.11 -y +conda activate onescience311 +# uv installation is supported +pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai +``` + +If the scripts need to access the OneScience source code, set: + +```bash +export ONESCIENCE_ROOT=/path/to/onescience +``` + +When run from this directory, the scripts automatically add the project root and `model/` to `sys.path`. You can also manually verify that the local package can be imported: + +```bash +export PYTHONPATH=$(pwd)/model:${PYTHONPATH:-} +python -c "import evo2; print('evo2 import ok')" +``` + +## 3. Quick Start + +### Quick Checks + +View the entry-point arguments: + +```bash +python scripts/infer.py --help +python scripts/predict.py --help +python scripts/train.py --help +python scripts/train_slurm.py --help +``` + +## 4. Data and Weights + +The example training data has been placed at: + +```text +data/data_mini/genome_data/ +``` + +It contains: + +* `chr20.fa`, `chr21.fa`, `chr22.fa`, and their compressed versions +* Merged FASTA: `chr20_21_22.fa` +* Preprocessed binary data for training, validation, and testing: `preprocessed_data/` + +The `dataset_prefix` in `config/genome_data_config.yaml` is a relative path. During actual runs, it should be used together with: + +```bash +--dataset-dir data/data_mini/genome_data +``` + +By default, the 7B NeMo checkpoint is placed inside the repository: + +```text +checkpoints/evo2_nemo_7b/ +``` + +The inference and prediction entry points read this directory by default; training does not load a checkpoint by default. To fine-tune from existing weights or resume training, explicitly pass `--ckpt-dir checkpoints/evo2_nemo_7b` in the training command. You can also use `EVO2_CKPT_DIR` or `--ckpt-dir` to override it with another checkpoint path: + +```bash +export EVO2_CKPT_DIR=/path/to/evo2_nemo_7b +``` + +The weight files are located under `/path/to/evo2/checkpoints/evo2_nemo_7b/weights`. + +## 5. Prompt-Based Generation Inference + +`infer.py` writes results to `--output-file`, but it does not automatically create the parent directory. Create the output directory before the first run: + +```bash +mkdir -p outputs +``` + +Minimal example: + +```bash +python scripts/infer.py \ + --prompt "ATGCGT" \ + --output-file outputs/evo2_generation.txt +``` + +Common parameters: + +| Parameter | Default | Description | +| --- | --- | --- | +| `--ckpt-dir` | `checkpoints/evo2_nemo_7b` | NeMo checkpoint directory; can be overridden with `EVO2_CKPT_DIR` | +| `--prompt` | Built-in E. coli classification-label prompt | Input sequence or text prompt | +| `--max-new-tokens` | `1024` | Maximum number of tokens to generate | +| `--temperature` | `1.0` | Sampling temperature | +| `--top-k` | `0` | top-k sampling | +| `--top-p` | `0.0` | top-p sampling | +| `--tensor-parallel-size` | `1` | Tensor parallel size | +| `--output-file` | Empty | Write to a file when specified; otherwise print logs | + +## 6. FASTA Prediction + +The directory includes an example FASTA: + +```text +data/predict_example.fa +``` + +Run prediction: + +```bash +python scripts/predict.py \ + --fasta data/predict_example.fa \ + --output-dir outputs/predict +``` + +`predict.py` automatically creates `--output-dir`. If `--fasta` is not passed, it reads `data/predict_example.fa` by default; if `--output-dir` is not passed, it writes output to `outputs/predict` by default. + +Common parameters: + +| Parameter | Default | Description | +| --- | --- | --- | +| `--fasta` | `data/predict_example.fa` | Input FASTA file | +| `--ckpt-dir` | `checkpoints/evo2_nemo_7b` | NeMo checkpoint directory; can be overridden with `EVO2_CKPT_DIR` | +| `--output-dir` | `outputs/predict` | Prediction result directory | +| `--batch-size` | `1` | Prediction batch size | +| `--model-size` | `7b` | Model specification; available values are determined by NeMo `HYENA_MODEL_OPTIONS` | +| `--output-log-prob-seqs` | Disabled | Output sequence log probability | + +`data/data_mini/genome_data/chr20.fa`, `chr21.fa`, and `chr22.fa` are relatively large chromosome-level FASTA files. Direct prediction may consume substantial GPU memory. It is recommended to split them into shorter fragments first, or use a multi-GPU parallel configuration. + +## 7. Training + +The training entry point requires one of two options: pass `-d/--dataset-config`, or use `--mock-data`. For the built-in example data, `config/genome_data_config.yaml` and `data/data_mini/genome_data` are recommended. + +Example for training the 1B architecture from scratch or validating the training workflow: + +```bash +python scripts/train.py \ + -d config/genome_data_config.yaml \ + --dataset-dir data/data_mini/genome_data \ + --model-size 1b \ + --result-dir results_1b \ + --devices 8 \ + --num-nodes 1 \ + --seq-length 8192 \ + --micro-batch-size 2 \ + --lr 0.0001 \ + --warmup-steps 5 \ + --max-steps 1000 \ + --clip-grad 1 \ + --wd 0.01 \ + --activation-checkpoint-recompute-num-layers 1 \ + --val-check-interval 50 \ + --limit-val-batches 2 +``` + +This example does not load a checkpoint and is suitable for validating the training workflow when 1B weights are unavailable. To fine-tune from a 1B NeMo checkpoint, additionally pass `--ckpt-dir /path/to/evo2_nemo_1b`. + +Example for training the 7B long-context architecture from scratch or validating the training workflow: + +```bash +python scripts/train.py \ + -d config/genome_data_config.yaml \ + --dataset-dir data/data_mini/genome_data \ + --model-size 7b_arc_longcontext \ + --result-dir results_7b \ + --devices 8 \ + --num-nodes 1 \ + --seq-length 1024 \ + --micro-batch-size 1 \ + --lr 0.0001 \ + --warmup-steps 5 \ + --max-steps 1000 \ + --clip-grad 1 \ + --wd 0.01 \ + --activation-checkpoint-recompute-num-layers 1 \ + --val-check-interval 50 \ + --limit-val-batches 2 +``` + +This example does not load a checkpoint. If you need to fine-tune from an existing 7B NeMo checkpoint or resume training, additionally pass `--ckpt-dir checkpoints/evo2_nemo_7b` or another checkpoint path. + +In a Slurm/cluster environment, replace the entry point with: + +```bash +python scripts/train_slurm.py ... +``` + +Common training parameters: + +The table below covers the main parameters used in the training examples above and in `scripts/train_evo2_1b.sh` and `scripts/train_evo2_7b.sh`. For more advanced parameters, run `python scripts/train.py --help`. + +| Parameter | Default | Description | +| --- | --- | --- | +| `-d, --dataset-config` | Required, or use `--mock-data` | Training data configuration YAML | +| `--mock-data` | Disabled | Do not read real data; use mock data for connectivity testing. Mutually exclusive with `-d/--dataset-config` | +| `--dataset-dir` | `EVO2_DATASET_DIR` or `data/data_mini/genome_data` | Data root directory | +| `--ckpt-dir` | Empty | Initial checkpoint directory; pass explicitly when fine-tuning from existing weights or resuming training | +| `--model-size` | `7b` | Model specification. Common values include `1b`, `7b`, and `7b_arc_longcontext`; use `test` for testing | +| `--devices` | `1` | Number of devices on a single node | +| `--num-nodes` | `1` | Number of nodes | +| `--tensor-parallel-size` | `1` | Tensor parallel size | +| `--pipeline-model-parallel-size` | `1` | Pipeline model parallel size | +| `--context-parallel-size` | `1` | Context parallel size | +| `--sequence-parallel` | Disabled | Enable sequence parallel, usually used together with tensor parallelism | +| `--seq-length` | `8192` | Training sequence length | +| `--micro-batch-size` | `1` | Micro batch size | +| `--global-batch-size` | Automatically inferred | Global batch size | +| `--grad-acc-batches` | `1` | Number of gradient accumulation batches, used to infer the global batch size | +| `--lr` | `3e-4` | Learning rate, explicitly set to `0.0001` in the examples | +| `--min-lr` | `3e-5` | Minimum learning rate of the cosine scheduler | +| `--warmup-steps` | `2500` | Number of warmup steps, explicitly set to `5` in the examples | +| `--constant-steps` | `80000` | Number of steps for which the learning rate stays constant | +| `--max-steps` | `500000` | Number of training steps | +| `--wd` | `0.01` | Optimizer weight decay | +| `--clip-grad` | `1.0` | Gradient clipping threshold | +| `--workers` | `8` | Number of DataLoader workers | +| `--activation-checkpoint-recompute-num-layers` | Empty | Override the number of activation recomputation layers; explicitly set to `1` in the examples | +| `--val-check-interval` | Empty | Validation and checkpoint check interval; explicitly set to `50` in the examples | +| `--limit-val-batches` | `20` | Maximum number of batches to run per validation; explicitly set to `2` in the examples | +| `--result-dir` | `./results` | Log and result directory | +| `--experiment-name` | `evo2` | Experiment name, which affects the log directory | +| `--disable-checkpointing` | Not passed by default; checkpoint callback enabled | Pass this parameter during connectivity tests to disable the checkpoint callback | +| `--no-save-last-checkpoint` | Save the last checkpoint by default | Disable saving the last checkpoint; often used together with `--disable-checkpointing` for quick tests | +| `--save-top-k` | `5` | Number of best checkpoints to save | +| `--ckpt-async-save` | Disabled | Enable asynchronous checkpoint saving | + +## 8. Data Preprocessing + +The current example training data has already been preprocessed. To process new FASTA or JSON data, use: + +```bash +bash scripts/tools/data_process/preprocess_data_fasta.sh +bash scripts/tools/data_process/preprocess_data_json.sh +``` + +Related configurations: + +```text +config/genome_preprocess_config.yaml +config/genome_data_config.yaml +``` + +# Official OneScience Information + +| Platform | OneScience Main Repository | Skills Repository | +| --- | --- | --- | +| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills | +| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills | + +# Citations and License + +- Original Evo2 paper: [Genome modelling and design across all domains of life with Evo 2](https://doi.org/10.1038/s41586-026-10176-5). +- The original Evo2 code uses the Apache License 2.0. See `LICENSE` in this repository. Original license source: [ArcInstitute/evo2 LICENSE](https://github.com/ArcInstitute/evo2/blob/main/LICENSE). +- The Evo2 7B model weights are sourced from Hugging Face [`arcinstitute/evo2_7b`](https://huggingface.co/arcinstitute/evo2_7b). In this repository, `checkpoints/evo2_nemo_7b` is the converted NeMo checkpoint format. The original weight license is Apache-2.0. +- The example dataset `data/data_mini/genome_data/` was downloaded and preprocessed from the hg38 `chr20`, `chr21`, and `chr22` chromosome FASTA files provided by the UCSC Genome Browser. It is used only to validate the training and inference workflows. See the data download directory: [UCSC hg38 chromosomes](https://hgdownload.soe.ucsc.edu/goldenpath/hg38/chromosomes/). +- Use of UCSC hg38 data must comply with the UCSC Genome Browser data usage terms. See: [UCSC Terms of Use](https://genome.ucsc.edu/license/). +- If you use this repository, the model weights, or generated results in research, it is recommended to cite the original Evo2 paper, the model weight source, the UCSC hg38 data source, and relevant OneScience project information, and to add citations for downstream analysis tools or datasets according to the actual task. diff --git a/checkpoints/evo2_nemo_7b/context/io.json b/checkpoints/evo2_nemo_7b/context/io.json new file mode 100644 index 0000000000000000000000000000000000000000..a46deac2f6ffc2fda8d06f7492f3e5d592eecb3e --- /dev/null +++ b/checkpoints/evo2_nemo_7b/context/io.json @@ -0,0 +1 @@ +{"root": {"type": "ref", "key": "trainer_context_1"}, "objects": {"tuple_1": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "tensor_model_parallel_size"], ["Index(index=1)", "pipeline_model_parallel_comm_backend"], ["Index(index=2)", "pipeline_model_parallel_size"], ["Index(index=3)", "virtual_pipeline_model_parallel_size"], ["Index(index=4)", "sequence_parallel"], ["Index(index=5)", "context_parallel_size"], ["Index(index=6)", "expert_model_parallel_size"], ["Index(index=7)", "expert_tensor_parallel_size"], ["Index(index=8)", "moe_extended_tp"], ["Index(index=9)", "use_te_rng_tracker"], ["Index(index=10)", "pipeline_dtype"], ["Index(index=11)", "microbatch_group_size_per_vp_stage"], ["Index(index=12)", "num_layers_in_first_pipeline_stage"], ["Index(index=13)", "num_layers_in_last_pipeline_stage"], ["Index(index=14)", "pipeline_model_parallel_layout"], ["Index(index=15)", "account_for_embedding_in_pipeline_split"], ["Index(index=16)", "account_for_loss_in_pipeline_split"]], "metadata": null}, "dict_1": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_2": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.collections.llm.gpt.model.hyena", "name": "Hyena7bConfig"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_1"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_1"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_2"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "hyena7b_config_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='tensor_model_parallel_size')", {"type": "leaf", "value": 1, "paths": [".model.config.tensor_model_parallel_size"]}], ["Attr(name='pipeline_model_parallel_comm_backend')", {"type": "leaf", "value": null, "paths": [".model.config.pipeline_model_parallel_comm_backend"]}], ["Attr(name='pipeline_model_parallel_size')", {"type": "leaf", "value": 1, "paths": [".model.config.pipeline_model_parallel_size"]}], ["Attr(name='virtual_pipeline_model_parallel_size')", {"type": "leaf", "value": null, "paths": [".model.config.virtual_pipeline_model_parallel_size"]}], ["Attr(name='sequence_parallel')", {"type": "leaf", "value": false, "paths": [".model.config.sequence_parallel"]}], ["Attr(name='context_parallel_size')", {"type": "leaf", "value": 1, "paths": [".model.config.context_parallel_size"]}], ["Attr(name='expert_model_parallel_size')", {"type": "leaf", "value": 1, "paths": [".model.config.expert_model_parallel_size"]}], ["Attr(name='expert_tensor_parallel_size')", {"type": "leaf", "value": null, "paths": [".model.config.expert_tensor_parallel_size"]}], ["Attr(name='moe_extended_tp')", {"type": "leaf", "value": false, "paths": [".model.config.moe_extended_tp"]}], ["Attr(name='use_te_rng_tracker')", {"type": "leaf", "value": false, "paths": [".model.config.use_te_rng_tracker"]}], ["Attr(name='pipeline_dtype')", {"type": "leaf", "value": null, "paths": [".model.config.pipeline_dtype"]}], ["Attr(name='microbatch_group_size_per_vp_stage')", {"type": "leaf", "value": 1, "paths": [".model.config.microbatch_group_size_per_vp_stage"]}], ["Attr(name='num_layers_in_first_pipeline_stage')", {"type": "leaf", "value": null, "paths": [".model.config.num_layers_in_first_pipeline_stage"]}], ["Attr(name='num_layers_in_last_pipeline_stage')", {"type": "leaf", "value": null, "paths": [".model.config.num_layers_in_last_pipeline_stage"]}], ["Attr(name='pipeline_model_parallel_layout')", {"type": "leaf", "value": null, "paths": [".model.config.pipeline_model_parallel_layout"]}], ["Attr(name='account_for_embedding_in_pipeline_split')", {"type": "leaf", "value": false, "paths": [".model.config.account_for_embedding_in_pipeline_split"]}], ["Attr(name='account_for_loss_in_pipeline_split')", {"type": "leaf", "value": false, "paths": [".model.config.account_for_loss_in_pipeline_split"]}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_1"}, "paths": [".model.config"]}, "tuple_2": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "optimizer"], ["Index(index=1)", "lr"], ["Index(index=2)", "min_lr"], ["Index(index=3)", "decoupled_lr"], ["Index(index=4)", "decoupled_min_lr"], ["Index(index=5)", "weight_decay"], ["Index(index=6)", "fp8_recipe"], ["Index(index=7)", "fp16"], ["Index(index=8)", "bf16"], ["Index(index=9)", "reuse_grad_buf_for_mxfp8_param_ag"], ["Index(index=10)", "params_dtype"], ["Index(index=11)", "use_precision_aware_optimizer"], ["Index(index=12)", "store_param_remainders"], ["Index(index=13)", "main_grads_dtype"], ["Index(index=14)", "main_params_dtype"], ["Index(index=15)", "exp_avg_dtype"], ["Index(index=16)", "exp_avg_sq_dtype"], ["Index(index=17)", "loss_scale"], ["Index(index=18)", "initial_loss_scale"], ["Index(index=19)", "min_loss_scale"], ["Index(index=20)", "loss_scale_window"], ["Index(index=21)", "hysteresis"], ["Index(index=22)", "adam_beta1"], ["Index(index=23)", "adam_beta2"], ["Index(index=24)", "adam_eps"], ["Index(index=25)", "sgd_momentum"], ["Index(index=26)", "use_distributed_optimizer"], ["Index(index=27)", "overlap_param_gather_with_optimizer_step"], ["Index(index=28)", "optimizer_cpu_offload"], ["Index(index=29)", "optimizer_offload_fraction"], ["Index(index=30)", "use_torch_optimizer_for_cpu_offload"], ["Index(index=31)", "overlap_cpu_optimizer_d2h_h2d"], ["Index(index=32)", "pin_cpu_grads"], ["Index(index=33)", "pin_cpu_params"], ["Index(index=34)", "clip_grad"], ["Index(index=35)", "log_num_zeros_in_grad"], ["Index(index=36)", "barrier_with_L1_time"], ["Index(index=37)", "timers"], ["Index(index=38)", "config_logger_dir"]], "metadata": null}, "dict_3": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_4": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_2": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "megatron.core.optimizer.optimizer_config", "name": "OptimizerConfig"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_2"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_3"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_4"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "optimizer_config_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='optimizer')", {"type": "leaf", "value": "adam", "paths": [".model.optim.config.optimizer"]}], ["Attr(name='lr')", {"type": "leaf", "value": 0.0001, "paths": [".model.optim.config.lr"]}], ["Attr(name='min_lr')", {"type": "leaf", "value": null, "paths": [".model.optim.config.min_lr"]}], ["Attr(name='decoupled_lr')", {"type": "leaf", "value": null, "paths": [".model.optim.config.decoupled_lr"]}], ["Attr(name='decoupled_min_lr')", {"type": "leaf", "value": null, "paths": [".model.optim.config.decoupled_min_lr"]}], ["Attr(name='weight_decay')", {"type": "leaf", "value": 0.01, "paths": [".model.optim.config.weight_decay"]}], ["Attr(name='fp8_recipe')", {"type": "leaf", "value": null, "paths": [".model.optim.config.fp8_recipe"]}], ["Attr(name='fp16')", {"type": "leaf", "value": false, "paths": [".model.optim.config.fp16"]}], ["Attr(name='bf16')", {"type": "leaf", "value": false, "paths": [".model.optim.config.bf16"]}], ["Attr(name='reuse_grad_buf_for_mxfp8_param_ag')", {"type": "leaf", "value": false, "paths": [".model.optim.config.reuse_grad_buf_for_mxfp8_param_ag"]}], ["Attr(name='params_dtype')", {"type": "pyref", "module": "torch", "name": "float32", "paths": [".model.optim.config.params_dtype", ".model.optim.config.main_grads_dtype", ".model.optim.config.main_params_dtype", ".model.optim.config.exp_avg_dtype", ".model.optim.config.exp_avg_sq_dtype"]}], ["Attr(name='use_precision_aware_optimizer')", {"type": "leaf", "value": false, "paths": [".model.optim.config.use_precision_aware_optimizer"]}], ["Attr(name='store_param_remainders')", {"type": "leaf", "value": true, "paths": [".model.optim.config.store_param_remainders"]}], ["Attr(name='main_grads_dtype')", {"type": "pyref", "module": "torch", "name": "float32", "paths": [".model.optim.config.params_dtype", ".model.optim.config.main_grads_dtype", ".model.optim.config.main_params_dtype", ".model.optim.config.exp_avg_dtype", ".model.optim.config.exp_avg_sq_dtype"]}], ["Attr(name='main_params_dtype')", {"type": "pyref", "module": "torch", "name": "float32", "paths": [".model.optim.config.params_dtype", ".model.optim.config.main_grads_dtype", ".model.optim.config.main_params_dtype", ".model.optim.config.exp_avg_dtype", ".model.optim.config.exp_avg_sq_dtype"]}], ["Attr(name='exp_avg_dtype')", {"type": "pyref", "module": "torch", "name": "float32", "paths": [".model.optim.config.params_dtype", ".model.optim.config.main_grads_dtype", ".model.optim.config.main_params_dtype", ".model.optim.config.exp_avg_dtype", ".model.optim.config.exp_avg_sq_dtype"]}], ["Attr(name='exp_avg_sq_dtype')", {"type": "pyref", "module": "torch", "name": "float32", "paths": [".model.optim.config.params_dtype", ".model.optim.config.main_grads_dtype", ".model.optim.config.main_params_dtype", ".model.optim.config.exp_avg_dtype", ".model.optim.config.exp_avg_sq_dtype"]}], ["Attr(name='loss_scale')", {"type": "leaf", "value": null, "paths": [".model.optim.config.loss_scale"]}], ["Attr(name='initial_loss_scale')", {"type": "leaf", "value": 4294967296, "paths": [".model.optim.config.initial_loss_scale"]}], ["Attr(name='min_loss_scale')", {"type": "leaf", "value": 1.0, "paths": [".model.optim.config.min_loss_scale"]}], ["Attr(name='loss_scale_window')", {"type": "leaf", "value": 1000, "paths": [".model.optim.config.loss_scale_window"]}], ["Attr(name='hysteresis')", {"type": "leaf", "value": 2, "paths": [".model.optim.config.hysteresis"]}], ["Attr(name='adam_beta1')", {"type": "leaf", "value": 0.9, "paths": [".model.optim.config.adam_beta1"]}], ["Attr(name='adam_beta2')", {"type": "leaf", "value": 0.999, "paths": [".model.optim.config.adam_beta2"]}], ["Attr(name='adam_eps')", {"type": "leaf", "value": 1e-08, "paths": [".model.optim.config.adam_eps"]}], ["Attr(name='sgd_momentum')", {"type": "leaf", "value": 0.9, "paths": [".model.optim.config.sgd_momentum"]}], ["Attr(name='use_distributed_optimizer')", {"type": "leaf", "value": true, "paths": [".model.optim.config.use_distributed_optimizer"]}], ["Attr(name='overlap_param_gather_with_optimizer_step')", {"type": "leaf", "value": false, "paths": [".model.optim.config.overlap_param_gather_with_optimizer_step"]}], ["Attr(name='optimizer_cpu_offload')", {"type": "leaf", "value": false, "paths": [".model.optim.config.optimizer_cpu_offload"]}], ["Attr(name='optimizer_offload_fraction')", {"type": "leaf", "value": 0.0, "paths": [".model.optim.config.optimizer_offload_fraction"]}], ["Attr(name='use_torch_optimizer_for_cpu_offload')", {"type": "leaf", "value": false, "paths": [".model.optim.config.use_torch_optimizer_for_cpu_offload"]}], ["Attr(name='overlap_cpu_optimizer_d2h_h2d')", {"type": "leaf", "value": false, "paths": [".model.optim.config.overlap_cpu_optimizer_d2h_h2d"]}], ["Attr(name='pin_cpu_grads')", {"type": "leaf", "value": true, "paths": [".model.optim.config.pin_cpu_grads"]}], ["Attr(name='pin_cpu_params')", {"type": "leaf", "value": true, "paths": [".model.optim.config.pin_cpu_params"]}], ["Attr(name='clip_grad')", {"type": "leaf", "value": 1.0, "paths": [".model.optim.config.clip_grad"]}], ["Attr(name='log_num_zeros_in_grad')", {"type": "leaf", "value": false, "paths": [".model.optim.config.log_num_zeros_in_grad"]}], ["Attr(name='barrier_with_L1_time')", {"type": "leaf", "value": false, "paths": [".model.optim.config.barrier_with_L1_time"]}], ["Attr(name='timers')", {"type": "leaf", "value": null, "paths": [".model.optim.config.timers"]}], ["Attr(name='config_logger_dir')", {"type": "leaf", "value": "", "paths": [".model.optim.config.config_logger_dir"]}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_2"}, "paths": [".model.optim.config"]}, "tuple_3": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "config"]], "metadata": null}, "dict_5": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_6": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_3": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.lightning.pytorch.optim.megatron", "name": "MegatronOptimizerModule"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_3"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_5"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_6"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "megatron_optimizer_module_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='config')", {"type": "ref", "key": "optimizer_config_1"}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_3"}, "paths": [".model.optim"]}, "dict_7": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}, "paths": [".model.tokenizer.special_tokens"]}, "tuple_4": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "special_tokens"]], "metadata": null}, "dict_8": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_9": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_4": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.collections.common.tokenizers.bytelevel_tokenizers", "name": "ByteLevelTokenizer"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_4"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_8"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_9"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "byte_level_tokenizer_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='special_tokens')", {"type": "ref", "key": "dict_7"}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_4"}, "paths": [".model.tokenizer"]}, "tuple_5": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "config"], ["Index(index=1)", "optim"], ["Index(index=2)", "tokenizer"]], "metadata": null}, "dict_10": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_11": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_5": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.collections.llm.gpt.model.hyena", "name": "HyenaModel"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_5"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_10"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_11"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "hyena_model_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='config')", {"type": "ref", "key": "hyena7b_config_1"}], ["Attr(name='optim')", {"type": "ref", "key": "megatron_optimizer_module_1"}], ["Attr(name='tokenizer')", {"type": "ref", "key": "byte_level_tokenizer_1"}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_5"}, "paths": [".model"]}, "tuple_6": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "always_save_context"]], "metadata": null}, "dict_12": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [["Key(key='always_save_context')", {"type": "leaf", "value": true, "paths": [".trainer.strategy.kwargs['always_save_context']"]}]], "metadata": {"type": "ref", "key": "tuple_6"}, "paths": [".trainer.strategy.kwargs"]}, "tuple_7": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "ckpt_save_optimizer"], ["Index(index=1)", "save_ckpt_format"], ["Index(index=2)", "ckpt_async_save"], ["Index(index=3)", "kwargs"]], "metadata": null}, "dict_13": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_14": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_6": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.lightning.pytorch.strategies.megatron_strategy", "name": "MegatronStrategy"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_7"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_13"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_14"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "megatron_strategy_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='ckpt_save_optimizer')", {"type": "leaf", "value": false, "paths": [".trainer.strategy.ckpt_save_optimizer"]}], ["Attr(name='save_ckpt_format')", {"type": "leaf", "value": "torch_dist", "paths": [".trainer.strategy.save_ckpt_format"]}], ["Attr(name='ckpt_async_save')", {"type": "leaf", "value": false, "paths": [".trainer.strategy.ckpt_async_save"]}], ["Attr(name='kwargs')", {"type": "ref", "key": "dict_12"}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_6"}, "paths": [".trainer.strategy"]}, "tuple_8": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "accelerator"], ["Index(index=1)", "strategy"], ["Index(index=2)", "devices"]], "metadata": null}, "dict_15": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_16": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_7": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.lightning.pytorch.trainer", "name": "Trainer"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_8"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_15"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_16"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "trainer_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='accelerator')", {"type": "leaf", "value": "cpu", "paths": [".trainer.accelerator"]}], ["Attr(name='strategy')", {"type": "ref", "key": "megatron_strategy_1"}], ["Attr(name='devices')", {"type": "leaf", "value": 1, "paths": [".trainer.devices"]}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_7"}, "paths": [".trainer"]}, "dict_17": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}, "paths": [".extra"]}, "tuple_9": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [["Index(index=0)", "model"], ["Index(index=1)", "trainer"], ["Index(index=2)", "extra"]], "metadata": null}, "dict_18": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "dict_19": {"type": {"type": "pyref", "module": "builtins", "name": "dict"}, "items": [], "metadata": {"type": {"type": "pyref", "module": "builtins", "name": "tuple"}, "items": [], "metadata": null}}, "buildable_traverser_metadata_8": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}, "items": [["Attr(name='fn_or_cls')", {"type": "pyref", "module": "nemo.lightning.io.pl", "name": "TrainerContext"}], ["Attr(name='argument_names')", {"type": "ref", "key": "tuple_9"}], ["Attr(name='argument_tags')", {"type": "ref", "key": "dict_18"}], ["Attr(name='argument_history')", {"type": "ref", "key": "dict_19"}]], "metadata": {"type": "pyref", "module": "fiddle._src.config", "name": "BuildableTraverserMetadata"}}, "trainer_context_1": {"type": {"type": "pyref", "module": "fiddle._src.config", "name": "Config"}, "items": [["Attr(name='model')", {"type": "ref", "key": "hyena_model_1"}], ["Attr(name='trainer')", {"type": "ref", "key": "trainer_1"}], ["Attr(name='extra')", {"type": "ref", "key": "dict_17"}]], "metadata": {"type": "ref", "key": "buildable_traverser_metadata_8"}, "paths": [""]}}, "refcounts": {"tuple_1": 1, "dict_1": 1, "dict_2": 1, "buildable_traverser_metadata_1": 1, "hyena7b_config_1": 1, "tuple_2": 1, "dict_3": 1, "dict_4": 1, "buildable_traverser_metadata_2": 1, "optimizer_config_1": 1, "tuple_3": 1, "dict_5": 1, "dict_6": 1, "buildable_traverser_metadata_3": 1, "megatron_optimizer_module_1": 1, "dict_7": 1, "tuple_4": 1, "dict_8": 1, "dict_9": 1, "buildable_traverser_metadata_4": 1, "byte_level_tokenizer_1": 1, "tuple_5": 1, "dict_10": 1, "dict_11": 1, "buildable_traverser_metadata_5": 1, "hyena_model_1": 1, "tuple_6": 1, "dict_12": 1, "tuple_7": 1, "dict_13": 1, "dict_14": 1, "buildable_traverser_metadata_6": 1, "megatron_strategy_1": 1, "tuple_8": 1, "dict_15": 1, "dict_16": 1, "buildable_traverser_metadata_7": 1, "trainer_1": 1, "dict_17": 1, "tuple_9": 1, "dict_18": 1, "dict_19": 1, "buildable_traverser_metadata_8": 1, "trainer_context_1": 1}, "version": "0.0.1"} \ No newline at end of file diff --git a/checkpoints/evo2_nemo_7b/context/model.yaml b/checkpoints/evo2_nemo_7b/context/model.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a4e204458086658b97d64467372d80d29f5147e9 --- /dev/null +++ b/checkpoints/evo2_nemo_7b/context/model.yaml @@ -0,0 +1,302 @@ +_target_: nemo.collections.llm.gpt.model.hyena.HyenaModel +config: + _cpu_offloading_context: null + _target_: nemo.collections.llm.gpt.model.hyena.Hyena7bConfig + account_for_embedding_in_pipeline_split: false + account_for_loss_in_pipeline_split: false + activation_func: + _call_: false + _target_: torch._C._nn.gelu + activation_func_fp8_input_store: false + add_attn_proj_bias: true + add_bias_linear: false + add_bias_output: false + add_qkv_bias: false + apply_query_key_layer_scaling: false + apply_residual_connection_post_layernorm: false + apply_rope_fusion: true + async_tensor_model_parallel_allreduce: false + attention_backend: + _call_: true + _target_: megatron.core.transformer.enums.AttnBackend + attention_dropout: 0.0 + attention_softmax_in_fp32: true + autocast_dtype: null + barrier_with_L1_time: true + batch_p2p_comm: true + batch_p2p_sync: true + bf16: true + bias_activation_fusion: true + bias_dropout_add_fusion: true + bias_dropout_fusion: false + calculate_per_token_loss: false + clone_scatter_output_in_embedding: true + config_logger_dir: '' + context_parallel_size: 1 + cp_comm_type: null + cpu_offloading: false + cpu_offloading_activations: true + cpu_offloading_num_layers: 0 + cpu_offloading_weights: true + cross_entropy_fusion_impl: native + cross_entropy_loss_fusion: false + cuda_graph_retain_backward_graph: false + cuda_graph_scope: full + cuda_graph_use_single_mempool: false + cuda_graph_warmup_steps: 3 + data_step_fn: + _call_: false + _target_: nemo.collections.llm.gpt.model.base.gpt_data_step + deallocate_pipeline_outputs: false + defer_embedding_wgrad_compute: false + delay_wgrad_compute: false + deterministic_mode: false + disable_bf16_reduced_precision_matmul: false + disable_parameter_transpose_cache: false + distribute_saved_activations: null + embedding_init_method: null + embedding_init_method_std: null + enable_autocast: false + enable_cuda_graph: false + expert_model_parallel_size: 1 + expert_tensor_parallel_size: null + external_cuda_graph: false + ffn_hidden_size: 11008 + finalize_model_grads_func: null + first_last_layers_bf16: false + flash_decode: false + forward_step_fn: + _call_: false + _target_: nemo.collections.llm.gpt.model.hyena.hyena_forward_step + fp16: false + fp16_lm_cross_entropy: false + fp32_residual_connection: true + fp8: null + fp8_amax_compute_algo: most_recent + fp8_amax_history_len: 1 + fp8_dot_product_attention: false + fp8_interval: 1 + fp8_margin: 0 + fp8_multi_head_attention: false + fp8_param: false + fp8_recipe: delayed + fp8_wgrad: true + gated_linear_unit: true + get_attention_mask_from_fusion: false + grad_scale_func: null + grad_sync_func: null + gradient_accumulation_fusion: false + hetereogenous_dist_checkpoint: false + heterogeneous_block_specs: false + hidden_dropout: 0.0 + hidden_size: 4096 + hierarchical_context_parallel_sizes: null + hybrid_attention_ratio: 0.0 + hybrid_mlp_ratio: 0.0 + hybrid_override_pattern: SDH*SDHSDH*SDHSDH*SDHSDH*SDHSDH* + hyena_filter_no_wd: true + hyena_init_method: small_init + hyena_output_layer_init_method: wang_init + inference_rng_tracker: false + init_method: null + init_method_std: 0.02 + init_model_with_meta_device: false + is_hybrid_model: false + kv_channels: null + layernorm_epsilon: 1.0e-06 + layernorm_zero_centered_gamma: false + make_vocab_size_divisible_by: 8 + mamba_head_dim: 64 + mamba_num_groups: 8 + mamba_num_heads: null + mamba_state_dim: 128 + mapping_type: base + masked_softmax_fusion: false + memory_efficient_layer_norm: false + microbatch_group_size_per_vp_stage: 1 + mlp_chunks_for_prefill: 1 + moe_apply_probs_on_input: false + moe_aux_loss_coeff: 0 + moe_deepep_num_sms: 20 + moe_enable_deepep: false + moe_expert_capacity_factor: null + moe_extended_tp: false + moe_ffn_hidden_size: null + moe_grouped_gemm: false + moe_input_jitter_eps: null + moe_layer_freq: 1 + moe_layer_recompute: false + moe_pad_expert_input_to_capacity: false + moe_per_layer_logging: false + moe_permute_fusion: false + moe_router_bias_update_rate: 0.001 + moe_router_dtype: null + moe_router_enable_expert_bias: false + moe_router_force_load_balancing: false + moe_router_group_topk: null + moe_router_load_balancing_type: aux_loss + moe_router_num_groups: null + moe_router_padding_for_fp8: false + moe_router_pre_softmax: false + moe_router_score_function: softmax + moe_router_topk: 2 + moe_router_topk_limited_devices: null + moe_router_topk_scaling_factor: null + moe_shared_expert_intermediate_size: null + moe_shared_expert_overlap: false + moe_token_dispatcher_type: allgather + moe_token_drop_policy: probs + moe_token_dropping: false + moe_use_legacy_grouped_gemm: false + moe_z_loss_coeff: null + mrope_section: null + mtp_loss_scaling_factor: null + mtp_num_layers: null + multi_latent_attention: false + no_rope_freq: null + no_sync_func: null + normalization: RMSNorm + num_attention_heads: 32 + num_groups_hyena: 4096 + num_groups_hyena_medium: 256 + num_groups_hyena_short: 256 + num_layers: 32 + num_layers_at_end_in_bf16: 1 + num_layers_at_start_in_bf16: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + num_microbatches_with_partial_activation_checkpoints: null + num_moe_experts: null + num_query_groups: null + output_layer_init_method: null + overlap_p2p_comm: false + overlap_p2p_comm_warmup_flush: false + parallel_output: true + param_sync_func: null + params_dtype: + _call_: false + _target_: torch.bfloat16 + perform_initialization: true + persist_layer_norm: false + pipeline_dtype: null + pipeline_model_parallel_comm_backend: null + pipeline_model_parallel_layout: null + pipeline_model_parallel_size: 1 + pipeline_model_parallel_split_rank: null + position_embedding_type: rope + post_process: true + pre_process: true + qk_layernorm: false + quant_recipe: null + recompute_granularity: full + recompute_method: uniform + recompute_modules: null + recompute_num_layers: 4 + remove_activation_post_first_layer: true + rotary_base: 10000 + rotary_interleaved: false + rotary_percent: 1.0 + seq_len_interpolation_factor: null + seq_length: 8192 + sequence_parallel: false + softmax_scale: null + symmetric_ar_type: null + tensor_model_parallel_size: 1 + test_mode: false + timers: null + to_upper: normalized_weighted + tokenizer_library: byte-level + tokenizer_model_path: null + tp_comm_atomic_ag: false + tp_comm_atomic_rs: false + tp_comm_bootstrap_backend: nccl + tp_comm_bulk_dgrad: true + tp_comm_bulk_wgrad: true + tp_comm_overlap: false + tp_comm_overlap_ag: true + tp_comm_overlap_disable_fc1: false + tp_comm_overlap_disable_qkv: false + tp_comm_overlap_rs: true + tp_comm_overlap_rs_dgrad: false + tp_comm_split_ag: true + tp_comm_split_rs: true + tp_only_amax_red: false + use_b2b_causal_conv1d: false + use_cpu_initialization: false + use_kitchen: false + use_mamba_mem_eff_path: true + use_ring_exchange_p2p: false + use_short_conv_bias: false + use_te: true + use_te_rng_tracker: false + variable_seq_lengths: false + virtual_pipeline_model_parallel_size: null + vortex_style_fp8: false + wgrad_deferral_limit: 0 + window_size: null +model_context_managers: [] +model_transform: null +optim: + _target_: nemo.lightning.pytorch.optim.megatron.MegatronOptimizerModule + config: + _target_: megatron.core.optimizer.optimizer_config.OptimizerConfig + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-08 + barrier_with_L1_time: false + bf16: false + clip_grad: 1.0 + config_logger_dir: '' + decoupled_lr: null + decoupled_min_lr: null + exp_avg_dtype: + _call_: false + _target_: torch.float32 + exp_avg_sq_dtype: + _call_: false + _target_: torch.float32 + fp16: false + fp8_recipe: null + hysteresis: 2 + initial_loss_scale: 4294967296 + log_num_zeros_in_grad: false + loss_scale: null + loss_scale_window: 1000 + lr: 0.0001 + main_grads_dtype: + _call_: false + _target_: torch.float32 + main_params_dtype: + _call_: false + _target_: torch.float32 + min_loss_scale: 1.0 + min_lr: null + optimizer: adam + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + overlap_cpu_optimizer_d2h_h2d: false + overlap_param_gather_with_optimizer_step: false + params_dtype: + _call_: false + _target_: torch.float32 + pin_cpu_grads: true + pin_cpu_params: true + reuse_grad_buf_for_mxfp8_param_ag: false + sgd_momentum: 0.9 + store_param_remainders: true + timers: null + use_distributed_optimizer: true + use_precision_aware_optimizer: false + use_torch_optimizer_for_cpu_offload: false + weight_decay: 0.01 + lr_mult: 1.0 + lr_scheduler: null + no_weight_decay_cond: null + scale_lr_cond: null +tokenizer: + _bos_id: null + _eos_id: 0 + _pad_id: 1 + _target_: nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer + special_tokens: {} + vocab_size: 512 diff --git a/checkpoints/evo2_nemo_7b/weights/.metadata b/checkpoints/evo2_nemo_7b/weights/.metadata new file mode 100644 index 0000000000000000000000000000000000000000..e843b600e38362168dacb9625c091ee88b93a5b4 --- /dev/null +++ b/checkpoints/evo2_nemo_7b/weights/.metadata @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9136c78b4ec769bcbcaae649bbb9fd2b8f9d563b1547279a77c5eb09fd27bd36 +size 180057 diff --git a/checkpoints/evo2_nemo_7b/weights/__0_0.distcp b/checkpoints/evo2_nemo_7b/weights/__0_0.distcp new file mode 100644 index 0000000000000000000000000000000000000000..adbc54481cf13e524dcfcdf376b81e6b4b00a32d --- /dev/null +++ b/checkpoints/evo2_nemo_7b/weights/__0_0.distcp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d0adb229a170471e48f3526b7c52e31d01f41b617f7deeabc06bfbf37d15928d +size 6482259480 diff --git a/checkpoints/evo2_nemo_7b/weights/__0_1.distcp b/checkpoints/evo2_nemo_7b/weights/__0_1.distcp new file mode 100644 index 0000000000000000000000000000000000000000..516e697e3f1356f0dcba4bf715fe1e10d99c491f --- /dev/null +++ b/checkpoints/evo2_nemo_7b/weights/__0_1.distcp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:35c492de6332f4e423f26186c76115ee43ecdc45eb1289b006d40226c520c60c +size 6482293876 diff --git a/checkpoints/evo2_nemo_7b/weights/common.pt b/checkpoints/evo2_nemo_7b/weights/common.pt new file mode 100644 index 0000000000000000000000000000000000000000..bdf44783c8e9e7e28476275899beb00c4332ab75 --- /dev/null +++ b/checkpoints/evo2_nemo_7b/weights/common.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:74760670a728a4aed7567efcb15c1572ce7388cf389dbf5102c55569ea31498a +size 2839 diff --git a/checkpoints/evo2_nemo_7b/weights/metadata.json b/checkpoints/evo2_nemo_7b/weights/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c9dbde0bcf4d2d993122dcc7d6bcb1eef8b6fb77 --- /dev/null +++ b/checkpoints/evo2_nemo_7b/weights/metadata.json @@ -0,0 +1 @@ +{"sharded_backend": "torch_dist", "sharded_backend_version": 1, "common_backend": "torch", "common_backend_version": 1} \ No newline at end of file diff --git a/config/config.yaml b/config/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..89f810e270a47f418895e4849f2bdc3bccdf64dc --- /dev/null +++ b/config/config.yaml @@ -0,0 +1,13 @@ +model: + name: Evo2 + package_path: ../model/evo2 + base_dependency: OneScience +runtime: + pythonpath: ../model + onescience_root_env: ONESCIENCE_ROOT +train: + entry: ../scripts/train.py + config: ./opengenome2.yml +inference: + entry: ../scripts/infer.py + checkpoint_dir: ../checkpoints \ No newline at end of file diff --git a/config/genome_data_config.yaml b/config/genome_data_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ac4b9fa76137312591ab45ae382ff1d657ae27e2 --- /dev/null +++ b/config/genome_data_config.yaml @@ -0,0 +1,9 @@ +- dataset_prefix: ./preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train + dataset_split: train + dataset_weight: 1.0 +- dataset_prefix: ./preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val + dataset_split: validation + dataset_weight: 1.0 +- dataset_prefix: ./preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test + dataset_split: test + dataset_weight: 1.0 \ No newline at end of file diff --git a/config/genome_preprocess_config.yaml b/config/genome_preprocess_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8f91c7c2c6f20412d1b5d41f02ba019c14aadf43 --- /dev/null +++ b/config/genome_preprocess_config.yaml @@ -0,0 +1,31 @@ + +- datapaths: ["chr20_21_22.fa"] + output_dir: "./preprocessed_data" + output_prefix: chr20_21_22_uint8_distinct + train_split: 0.9 + valid_split: 0.05 + test_split: 0.05 + overwrite: True + embed_reverse_complement: true + random_reverse_complement: 0.0 + random_lineage_dropout: 0.0 + include_sequence_id: false + transcribe: "back_transcribe" + force_uppercase: false + indexed_dataset_dtype: "uint8" + tokenizer_type: "Byte-Level" + vocab_file: null + vocab_size: null + merges_file: null + pretrained_tokenizer_model: null + special_tokens: null + fast_hf_tokenizer: true + append_eod: true + enforce_sample_length: null + ftfy: false + workers: 1 + preproc_concurrency: 100000 + chunksize: 25 + drop_empty_sequences: true + nnn_filter: false # If you split your fasta on NNN (in human these are contigs), then you should set this to true. + seed: 12342 # Not relevant because we are not using random reverse complement or lineage dropout. diff --git a/config/opengenome2.yml b/config/opengenome2.yml new file mode 100644 index 0000000000000000000000000000000000000000..dfb319546f37a55055b196cfecd89dd02e501f18 --- /dev/null +++ b/config/opengenome2.yml @@ -0,0 +1,177 @@ +{ + "train-data-paths": [ + "/datasets/evo2/metagenomics/pretraining_data_metagenomics/data_metagenomics_train_text_CharLevelTokenizer_document", + "/datasets/evo2/gtdb_imgpr/pretraining_data_gtdb_imgpr/data_gtdb_imgpr_train_text_CharLevelTokenizer_document", + "/datasets/evo2/imgvr/pretraining_data_imgvr/data_imgvr_train_text_CharLevelTokenizer_document", + "/datasets/evo2/hg38/pretraining_data_hg38/data_hg38_all_text_CharLevelTokenizer_document", + "/datasets/evo2/ncrna/pretraining_data_ncrna/data_ncrna_train_text_CharLevelTokenizer_document", + "/datasets/evo2/mrna_splice/pretraining_data_mrna_splice/data_mrna_splice_train_text_CharLevelTokenizer_document", + "/datasets/evo2/promoters/pretraining_data_promoters/data_promoters_train_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch1_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch2_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch3_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch4_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch5_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch6_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch7_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch8_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch1_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch2_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch3_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch4_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch5_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch6_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch7_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch8_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch1_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch2_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch3_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch4_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch5_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch6_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch7_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch8_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch1_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch2_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch3_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch4_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch5_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch6_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch7_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch8_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch1_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch2_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch3_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch4_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch5_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch6_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch7_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_train_batch8_protista_text_CharLevelTokenizer_document", + + ], + + # note no hg38 + "valid-data-paths": [ + "/datasets/evo2/metagenomics/pretraining_data_metagenomics/data_metagenomics_valid_text_CharLevelTokenizer_document", + "/datasets/evo2/gtdb_imgpr/pretraining_data_gtdb_imgpr/data_gtdb_imgpr_valid_text_CharLevelTokenizer_document", + "/datasets/evo2/imgvr/pretraining_data_imgvr/data_imgvr_valid_text_CharLevelTokenizer_document", + "/datasets/evo2/ncrna/pretraining_data_ncrna/data_ncrna_valid_text_CharLevelTokenizer_document", + "/datasets/evo2/mrna_splice/pretraining_data_mrna_splice/data_mrna_splice_valid_text_CharLevelTokenizer_document", + "/datasets/evo2/promoters/pretraining_data_promoters/data_promoters_valid_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch1_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch2_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch3_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch4_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch5_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch6_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch7_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch8_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch1_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch2_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch3_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch4_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch5_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch6_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch7_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch8_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch1_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch2_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch3_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch4_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch5_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch6_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch7_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch8_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch1_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch2_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch3_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch4_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch5_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch6_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch7_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch8_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch1_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch2_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch3_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch4_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch5_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch6_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch7_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_valid_batch8_protista_text_CharLevelTokenizer_document", + ], + + # note no hg38 + "test-data-paths": [ + "/datasets/evo2/metagenomics/pretraining_data_metagenomics/data_metagenomics_test_text_CharLevelTokenizer_document", + "/datasets/evo2/gtdb_imgpr/pretraining_data_gtdb_imgpr/data_gtdb_imgpr_test_text_CharLevelTokenizer_document", + "/datasets/evo2/imgvr/pretraining_data_imgvr/data_imgvr_test_text_CharLevelTokenizer_document", + "/datasets/evo2/ncrna/pretraining_data_ncrna/data_ncrna_test_text_CharLevelTokenizer_document", + "/datasets/evo2/mrna_splice/pretraining_data_mrna_splice/data_mrna_splice_test_text_CharLevelTokenizer_document", + "/datasets/evo2/promoters/pretraining_data_promoters/data_promoters_test_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch1_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch2_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch3_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch4_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch5_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch6_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch7_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch8_animalia_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch1_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch2_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch3_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch4_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch5_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch6_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch7_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch8_fungi_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch1_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch2_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch3_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch4_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch5_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch6_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch7_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch8_plantae_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch1_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch2_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch3_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch4_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch5_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch6_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch7_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch8_chromista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch1_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch2_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch3_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch4_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch5_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch6_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch7_protista_text_CharLevelTokenizer_document", + "/datasets/evo2/eukaryote_ncbi/pretraining_data/data_test_batch8_protista_text_CharLevelTokenizer_document", + ], +"train-data-weights": [ + 0.09, 0.15, 0.017, 0.003, 0.02, 0.09, 0.0002, 0.05625, 0.05625, + 0.05625, 0.05625, 0.05625, 0.05625, 0.05625, 0.05625, 0.01, 0.01, 0.01, 0.01, + 0.01, 0.01, 0.01, 0.01, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, + 0.010725, 0.010725, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, + 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001 + ], + "test-data-weights": [ + 0.09, 0.15, 0.017, 0.02, 0.09, 0.0002, 0.05625, 0.05625, 0.05625, 0.05625, + 0.05625, 0.05625, 0.05625, 0.05625, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, + 0.01, 0.01, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, + 0.010725, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, + 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001 + ], + "valid-data-weights": [ + 0.09, 0.15, 0.017, 0.02, 0.09, 0.0002, 0.05625, 0.05625, 0.05625, + 0.05625, 0.05625, 0.05625, 0.05625, 0.05625, 0.01, 0.01, 0.01, 0.01, 0.01, + 0.01, 0.01, 0.01, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, 0.010725, + 0.010725, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, 0.00075, + 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001 + ], + "checkpoint_validation_with_forward_pass": False, + # "use_wandb": false, + "wandb_host": "https://api.wandb.ai", + "wandb_project": "evo2" +} \ No newline at end of file diff --git a/configuration.json b/configuration.json new file mode 100644 index 0000000000000000000000000000000000000000..afa1075f14489a2a3799df734fd32c127be42b10 --- /dev/null +++ b/configuration.json @@ -0,0 +1,15 @@ +{ + "model_name": "Evo2", + "framework": "PyTorch/NeMo", + "task": "genomic-sequence-modeling", + "entry_points": { + "train": "scripts/train.py", + "inference": "scripts/infer.py", + "predict": "scripts/predict.py" + }, + "source_package": "model/evo2", + "base_dependency": "onescience", + "config": "config/config.yaml", + "weight_dir": "checkpoints", + "license": "Apache-2.0" +} \ No newline at end of file diff --git a/data/data_mini/genome_data/chr20.fa b/data/data_mini/genome_data/chr20.fa new file mode 100644 index 0000000000000000000000000000000000000000..440894fc51f5179f51e95c77e21dfffe5bef1d11 --- /dev/null +++ b/data/data_mini/genome_data/chr20.fa @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a704b3a25c698bb5219b750d0ba6a333ddb0043904b4921992d5f709771c2666 +size 65733058 diff --git a/data/data_mini/genome_data/chr20.fa.gz b/data/data_mini/genome_data/chr20.fa.gz new file mode 100644 index 0000000000000000000000000000000000000000..a7e7aad20b5f9e6a46569276070026689a5c1612 --- /dev/null +++ b/data/data_mini/genome_data/chr20.fa.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:271ca0ee4247c2fa30c8491764ad8acd4bdde8eb3cf314f0720235010ae85c9f +size 20350590 diff --git a/data/data_mini/genome_data/chr20_21_22.fa b/data/data_mini/genome_data/chr20_21_22.fa new file mode 100644 index 0000000000000000000000000000000000000000..6e738329efff6dd05bd8ea9593791c441cedfe8f --- /dev/null +++ b/data/data_mini/genome_data/chr20_21_22.fa @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5c65fe1f46bb7949ab4e670a004e01f4815501f5f949c5b75cd06cd6df229059 +size 165212093 diff --git a/data/data_mini/genome_data/chr21.fa b/data/data_mini/genome_data/chr21.fa new file mode 100644 index 0000000000000000000000000000000000000000..2cccd38875fb83924516e55878841007b6d8e21e --- /dev/null +++ b/data/data_mini/genome_data/chr21.fa @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:35c71b68436d1a278ecb6a1e875af3ba4020738a028a7feac769a6d62790ae1f +size 47644190 diff --git a/data/data_mini/genome_data/chr21.fa.gz b/data/data_mini/genome_data/chr21.fa.gz new file mode 100644 index 0000000000000000000000000000000000000000..c1b4178c158f157be36ce581ccb54a6a85cf7c52 --- /dev/null +++ b/data/data_mini/genome_data/chr21.fa.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c979ca1e5065c2521a50773473e0d0cc018fd6f3e9bb3aa90493fe7b45d57d1b +size 12709705 diff --git a/data/data_mini/genome_data/chr22.fa b/data/data_mini/genome_data/chr22.fa new file mode 100644 index 0000000000000000000000000000000000000000..2eed7edf12b6396521a68fbe08435dee0df3c79a --- /dev/null +++ b/data/data_mini/genome_data/chr22.fa @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ce3ee1ca39356238f7aee438a40a88b4f1b9d80b316b263e16fb12402212d10f +size 51834845 diff --git a/data/data_mini/genome_data/chr22.fa.gz b/data/data_mini/genome_data/chr22.fa.gz new file mode 100644 index 0000000000000000000000000000000000000000..2dcd29557935fcacafbb681c45f394944a88f77c --- /dev/null +++ b/data/data_mini/genome_data/chr22.fa.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:05f9d97d6fbfd08a44ca45b50837ca2ae9c471f35ba79dffec04d2cb5eaaf695 +size 12255678 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test.bin b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test.bin new file mode 100644 index 0000000000000000000000000000000000000000..9122b488cc134ac9ca8460eb1e6124287604cc19 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c03f70b9b67e85e9dd37277ec7d18ec69ff8bf18e822dfb8300d944c18c2e5a2 +size 101636938 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test.idx b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test.idx new file mode 100644 index 0000000000000000000000000000000000000000..7a399d2c9ad0d2ae612a7cb4dbef5e354c81f5cc Binary files /dev/null and b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test.idx differ diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-description.txt b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-description.txt new file mode 100644 index 0000000000000000000000000000000000000000..57c5d9faa876838ff7d38cd2db1db88ff094a5a0 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-description.txt @@ -0,0 +1,13 @@ +{ + "class": "GPTDataset", + "dataset_path": "data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test", + "num_samples": 16, + "index_split": "test", + "random_seed": 1234, + "sequence_length": 8192, + "split": null, + "split_matrix": null, + "tokenizer": { + "class": "nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer" + } +} \ No newline at end of file diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-document_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-document_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..773d280ae20ab69994f7e5855ce0cb77a95600a6 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-document_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28d4b63fdadd72ed19f8fad1638746569959d59ee50d3e10d67810fced9dc78f +size 136 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-sample_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-sample_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..44be2cac9798a926c98d6705643f7a04f5c16f82 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-sample_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b4b37f6c8f8a94fd66ed63489802ee9b633c786b3a3c94ac900fb0e5ef4f617d +size 99384 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-shuffle_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-shuffle_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..3f6659410d0a6df051db296118ab6c3ae7a97c76 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/342ca946fce706bd1bebace29d20c889-GPTDataset-test-shuffle_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a68f8b87962dfc3d24757b2869592ab15741d2326a29cdb30b4df150912cfe7a +size 49752 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-description.txt b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-description.txt new file mode 100644 index 0000000000000000000000000000000000000000..29d92c904e5f6464c0aef4b4a78d8f777248be90 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-description.txt @@ -0,0 +1,13 @@ +{ + "class": "GPTDataset", + "dataset_path": "data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test", + "num_samples": 8, + "index_split": "test", + "random_seed": 1234, + "sequence_length": 1024, + "split": null, + "split_matrix": null, + "tokenizer": { + "class": "nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer" + } +} \ No newline at end of file diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-document_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-document_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..773d280ae20ab69994f7e5855ce0cb77a95600a6 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-document_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28d4b63fdadd72ed19f8fad1638746569959d59ee50d3e10d67810fced9dc78f +size 136 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-sample_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-sample_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..f7c5ea617013c6c88ee70b6893fb8700ab33dbbe --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-sample_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:53f2655dac75ebf3e3bb0454ffe82e61bcda93a47b3c8df90401d123e0be8bb5 +size 794168 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-shuffle_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-shuffle_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..4e6cd5601379b85d14f1bb158ae6d9fdc8f71cfd --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_test/cache/GPTDataset_indices/bdd3bfe88a1201b3e8f83953f7072ca6-GPTDataset-test-shuffle_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c33ad031b614ec8043d296f2ecc4ebd4a21d1cf636f40ebfa0af56bc7fa0eb8c +size 397144 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train.bin b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train.bin new file mode 100644 index 0000000000000000000000000000000000000000..9a8f1639d7391b818f12c8f31e576f5923e0aa93 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3f3d2d68af61f62bd11cbc0c13648029637baaf6a3a0d2ce228db6724ac60b7e +size 128888336 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train.idx b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train.idx new file mode 100644 index 0000000000000000000000000000000000000000..863d1f8ef7d3e2d8f86362f8e8d0986478e5bc0b Binary files /dev/null and b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train.idx differ diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-description.txt b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-description.txt new file mode 100644 index 0000000000000000000000000000000000000000..e716f86591255ac83c6ad30d73b5efd01930a8ad --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-description.txt @@ -0,0 +1,13 @@ +{ + "class": "GPTDataset", + "dataset_path": "data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train", + "num_samples": 8000, + "index_split": "train", + "random_seed": 1234, + "sequence_length": 1024, + "split": null, + "split_matrix": null, + "tokenizer": { + "class": "nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer" + } +} \ No newline at end of file diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-document_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-document_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..773d280ae20ab69994f7e5855ce0cb77a95600a6 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-document_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28d4b63fdadd72ed19f8fad1638746569959d59ee50d3e10d67810fced9dc78f +size 136 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-sample_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-sample_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..f4c48d39d3c52a4f46678028889166d9e958d0d3 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-sample_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34f9e9c57b454b1b4888a5492aed2014952633e9b6c93639bca9c4217532d2d9 +size 1007072 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-shuffle_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-shuffle_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..fe422b2332e6e3c65047b58d46d86a61ec11ba84 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/746855b0ca028bf7efd5622b8169f84e-GPTDataset-train-shuffle_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d0e161e823144594d9eef6ffd50ca778f8e1ae1f0831cf1d040c02a9f056771 +size 503596 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-description.txt b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-description.txt new file mode 100644 index 0000000000000000000000000000000000000000..f624d4299ab82529037367311d0a30ea987a5dc2 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-description.txt @@ -0,0 +1,13 @@ +{ + "class": "GPTDataset", + "dataset_path": "data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train", + "num_samples": 16000, + "index_split": "train", + "random_seed": 1234, + "sequence_length": 8192, + "split": null, + "split_matrix": null, + "tokenizer": { + "class": "nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer" + } +} \ No newline at end of file diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-document_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-document_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..0330a481fef199b01f1d4bac6063b3edf85d479c --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-document_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bc6782cd85e561058197445aef6cf1e89bb2897b69dacb26af801ba986f2322e +size 144 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-sample_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-sample_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..f18006e8dd7bcfeaa7c385567e70b3dcac435c62 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-sample_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ee4ef251a29baa3cdcbbefbd7a4876eccf2b69d1e1d326513a342b29382e167f +size 251864 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-shuffle_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-shuffle_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..163dc024822ea49f8901a98a82759f7248daf74e --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_train/cache/GPTDataset_indices/9f44c7ed4a3ebd5a2cfaed3f9be3b954-GPTDataset-train-shuffle_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cdfdd5f02d57271a14de63780596ea8db0044311e6990ba05cdeaf5c0181255a +size 125992 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val.bin b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val.bin new file mode 100644 index 0000000000000000000000000000000000000000..e59b97c1eb453ed3edbfee179c22435232fad7dd --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:596393be374e2a0aa08e7f4f8ba3c42a21c76f825c1b3f49c76c0f5dd44e7cd5 +size 93419968 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val.idx b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val.idx new file mode 100644 index 0000000000000000000000000000000000000000..6528f9207edb4cb806298f496ff597b89741b9fc Binary files /dev/null and b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val.idx differ diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-description.txt b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-description.txt new file mode 100644 index 0000000000000000000000000000000000000000..9658a8e0b365da0d8d9b42047b90f66459356c80 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-description.txt @@ -0,0 +1,13 @@ +{ + "class": "GPTDataset", + "dataset_path": "data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val", + "num_samples": 672, + "index_split": "valid", + "random_seed": 1234, + "sequence_length": 8192, + "split": null, + "split_matrix": null, + "tokenizer": { + "class": "nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer" + } +} \ No newline at end of file diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-document_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-document_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..773d280ae20ab69994f7e5855ce0cb77a95600a6 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-document_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28d4b63fdadd72ed19f8fad1638746569959d59ee50d3e10d67810fced9dc78f +size 136 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-sample_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-sample_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..383d5562b25d7bed9f0070509ebe3683bf1cf563 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-sample_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8ab709b3b7a8600968430e26d46a77c8d7924914d52ec076847841aa73edd81 +size 91360 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-shuffle_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-shuffle_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..7b6267a75c3d5198d3b0bb3c183ca6d83c4ffcbb --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/a8235fab7adf408b89a5c5806a781046-GPTDataset-valid-shuffle_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2e7754043b6dd3baae0db20e3d0b8a7f04743c973a5f856f78f720c144ba8763 +size 45740 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-description.txt b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-description.txt new file mode 100644 index 0000000000000000000000000000000000000000..af6e38e4c5976cec60588ad5e726d19c2acb4a80 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-description.txt @@ -0,0 +1,13 @@ +{ + "class": "GPTDataset", + "dataset_path": "data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val", + "num_samples": 336, + "index_split": "valid", + "random_seed": 1234, + "sequence_length": 1024, + "split": null, + "split_matrix": null, + "tokenizer": { + "class": "nemo.collections.common.tokenizers.bytelevel_tokenizers.ByteLevelTokenizer" + } +} \ No newline at end of file diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-document_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-document_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..773d280ae20ab69994f7e5855ce0cb77a95600a6 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-document_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:28d4b63fdadd72ed19f8fad1638746569959d59ee50d3e10d67810fced9dc78f +size 136 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-sample_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-sample_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..9dc64457ff925d642ab9453d6c61b65b422ae279 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-sample_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2ca3ac536738dbf74d06d9bfa4d62571aa432e54dae531100a622a59ae697a19 +size 729976 diff --git a/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-shuffle_index.npy b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-shuffle_index.npy new file mode 100644 index 0000000000000000000000000000000000000000..20f26a3f5d6197ae855b9d47f79947fa4d1b7565 --- /dev/null +++ b/data/data_mini/genome_data/preprocessed_data/chr20_21_22_uint8_distinct_byte-level_val/cache/GPTDataset_indices/cccf4722f9a5df3869a82454d596a5e6-GPTDataset-valid-shuffle_index.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fe3d86dcad9cbb2128b4d2995a2dd9dd2755387da0be2459d9f6ead8e910801e +size 365048 diff --git a/data/predict_example.fa b/data/predict_example.fa new file mode 100644 index 0000000000000000000000000000000000000000..27ea0f02517c66b578d826c3d14fe0d7a7486faf --- /dev/null +++ b/data/predict_example.fa @@ -0,0 +1,2 @@ +>example_sequence_1 +ACGTGCTAGCTAGGATCCGATCGTACGATCGATCGTACGTAGCTAGCTGACTGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGCATCGATCGATCGTACGATCGATCGTAGCTAGCTAGC diff --git a/data/prompts.csv b/data/prompts.csv new file mode 100644 index 0000000000000000000000000000000000000000..134a9a39a1a67160fffc0cab98a3ee3ba407c76d --- /dev/null +++ b/data/prompts.csv @@ -0,0 +1,5 @@ +Sequence,Name,Percent +GAATAGGAACAGCTCCGGTCTACAGCTCCCAGCGTGAGCGACGCAGAAGACGGTGATTTCTGCATTTCCATCTGAGGTACCGGGTTCATCTCACTAGGGAGTGCCAGACAGTGGGCGCAGGCCAGTGTGTGTGCGCACCGTGCGCGAGCCGAAGCAGGGCGAGGCATTGCCTCACCTGGGAAGCGCAAGGGGTCAGGGAGTTCCCTTTCCGAGTCAAAGAAAGGGGTGATGGACGCACCTGGAAAATCGGGTCACTCCCACCCGAATATTGCGCTTTTCAGACCGGCTTAAGAAACGGCGCACCACGAGACTATATCCCACACCTGGCTCAGAGGGTCCTACGCCCACGGAATCTCGCTGATTGCTAGCACAGCAGTCTGAGATCAAACTGCAAGGCGGCAACGAGGCTGGGGGAGGGGCGCCCGCCATTGCCCAGGCTTGCTTAGGTAAACAAAGCAGCCGGGAAGCTCGAACTGGGTGGAGCCCACCACAGCTCAAGGAGGCCTGCCTGCCTCTGTAGGCTCCACCTCTGGGGGCAGGGCACAGACAAACAAAAAGGCAGCAGTAACCTCTGCAGACTTAAGTGTCCCTGTCTGACAGCTTTGAAGAGAGCAGTGGTTCTCCCAGCACGCAGCTGGAGATCTGAGAACGGGCAGACTGCCTCCTCAAGTGGGTCCCTGACCCCTGACCCCCGAGCAGCCTAACTGGGAGGCACCCCCCAGCAGGGGCACACTGACACCTCACACGGCAGGGTATTCCAACAGACCTGCAGCTGAGGGTCCTGTCTGTTAGAAGGAAAACTAACAACCAGAAAGGACATCTACACCGAAAACCCATCTGTACATCACCATCATCAAAGACCAAAAGTAGATAAAACCACAAAGATGGGGAAAAAACAGAACAGAAAAACTGGAAACTCTAAAACGCAGAGCGCCTCTCCTCCTCCAAAGGAACGCAGTTCCTCACCAGCAACAGAACAAAGCTGGATGGAGAATGATTTTGACGAGCTGAGAGAAGAAGGCTTCAGACGATCAAATTACTCTGAGCTACGGGAGGACATTCAAACCAAAGGCAAAGAAGTTGAAAACTTTGAAAAAAATTTAGAAGAATGTATAACTAGAATAACCAATACAGAGAAGTGCTTAAAGGAGCTGATGGAGCTGAAAACCAAGGCTCGAGAACTACGTGAAGAATGCAGAAGCCTCAGGAGCCGATGCGATCAACTGGAAGAAAGGGTATCAGCAATGGAAGATGAAATGAATGAAATGAAGCGAGAAGGGAAGTTTAGAGAAAAAAGAATAAAAAGAAATGAGCAAAGCCTCCAAGAAATATGGGACTATGTGAAAAGACCAAATCTACGTCTGATTGGTGTACCTGAAAGTGATGTGGAGAATGGAACCAAGTTGGAAAACACTCTGCAGGATATTATCCAGGAGAACTTCCCCAATCTAGCAAGGCAGGCCAACGTTCAGATTCAGGAAATACAGAGAACGCCACAAAGATACTCCTCGAGAAGAGCAACTCCAAGACACATAATTGTCAGATTCACCAAAGTTGAAATGAAGGAAAAAATGTTAAGGGCAGCCAGAGAGAAAGGTCGGGTTACCCTCAAAGGGAAGCCTATCAGACTAACAGCAGATCTCTCGGCAGAAACCCTACAAGCCAGAAGAGAGTGGGGGCCAATATTCAACATTCTTAAAGAAAAGAATTTTCAACCCAGAATTTCATTTCCAGCCAAACTAAGCTTCATAAGTGAAGGAGAAAGAAAATACTTTACAGACAAGCAAATGCTGAGAGATTTTGTCACCACCAGGCCTACCCTAAAAGAGCTCCTGAAGGAAGCACTAAACATGGAAAGGAACAACCGGTACCAGCCGCTGCAAAATCATGCCAAAATGTAAAGACCATCGAGACTAGGAAGAAACTGCATCAACTAATGAGCAAAATCACCAGCTAACATCATAATGACAGGATCAAATTCACACATAACAATATTAACTTTAAATATAAATGGACTAAATTCTGCAATTAAAAGACACAGACTGGCAAGTTGGATAAAGAGTCAAGACCCATCAGTGTGCTGTATTCAGGAAACCCATCTCATGTGCAGAGACACACATAGGCTCAAAATAAAAGGATGGAGGAAGATCTACCAAGCAAATGGAAAACAAAAAAGGCAGGGGTTGCAATCCTAGTCTCTGATAAAACAGACTTTAAACCAACAAAGATCAAAAGAGACAAAGAAGGCCATTACATAATGGTAAAGGGATCAATTCAACAAGAGGAGCTAACTATCCTAAATATTTATGCACCCAATACAGGAGCACCCAGATTCATAAAGCAAGTCCTGAGTGACCTACAAAGAGACTTAGACTCCCACACATTAATAATGGGAGACTTTAACACCCCACTGTCAATATTAGACAGATCAACGAGACAGAAAGTCAACAAGGATACCCAGGAATTGAACTCAGCTCTGCACCAAGCAGACCTAATAGACATCTACAGAACTCTCCACCCCAAATCAACAGAATATACATTTTTTTCAGCACCACACCACACCTATTCCAAAATCGACCACATAGTTGGAAGTAAAGCTCTCCTCAGCAAATGTAAAAGAACAGAAATTATAACAAACTATCTCTCAGACCACAGTGCAATCAAACTAGAACTCAGGATTAAGAATCTCACTCAAAGCCGCTCAACTACATGGAAACTGAACAACCTGCTCCTGAATGACTACTGGGTACATAACGAAATGAAGGCAGAAATAAAGATGTTCTTTGAAACCAACGAGAACAAAGACACCACATACCAGAATCTCTGGGACGCATTCAAAGCAGTGTGTAGAGGGAAATTTATAGCACTAAATGCCTACAAGAGAAAGCAGGAAAGATCCAAAATTGACACCCTAACATCACAATTAAAAGAACTAGAAAAGCAAGAGCAAACACATTCAAAAGCTAGCAGAAGGCAAGAAATAACTAAAATCAGAGCAGAACTGAAGGAAATAGAGACACAAAAAACCCTTCAAAAAATCAATGAATCCAGGAGCTGGTTTTTTGAAAGGATCAACAAAATTGATAGACCGCTAGCAAGACTAATAAAGAAAAAAAGAGAGAAGAATCAAATAGACACAATAAAAAATGATAAAGGGGATATCACCACCGATCCCACAGAAATACAAACTACCATCAGAGAATACTACAAACACCTCTACGCAAATAAACTAGAAAATCTAGAAGAAATGGATACATTCCTCGACACATACACTCTCCCAAGACTAAAACAGGAAGAAGTTGAATCTCTGAATGGACCAATAACAGGCTCTGAAATTGTGGCAATAATCAATAGTTTACCAACCAAAAAGAGTCCAGGACCAGATGGATTCACAGCCGAATTCTACCAGAGGTACAAGGAGGAACTGGTACCATTCCTTCTGAAACTATTCCAATCAATAGAAAAAGAGGGAATCCTCCCTAACTCATTTTATGAGGCCAGCATCATTCTGATACCAAAGCCGGGCAGAGACACAACCAAAAAAGAGAATTTTAGACCAATATCCTTGATGAACATTGATGCAAAAATCCTCAATAAAATACTGGCAAACCGAATCCAGCAGCACATCAAAAAGCTTATCCACCATGATCAAGTGGGCTTCATCCCTGGGATGCAAGGCTGGTTCAATATACGCAAATCAATAAATGTAATCCAGCATATAAACAGAGCCAAAGACAAAAACCACATGATTATCTCAATAGATGCAGAAAAAGCCTTTGACAAAATTCAACAACCCTTCATGCTAAAAACTCTCAATAAATTAGGTATTGATGGGACGTATTTCAAAATAATAAGAGCTATCTATGACAAACCCACAGCCAATATCATACTGAATGGGCAAAAACTGGAAGCATTCCCTTTGAAAACTGGCACAAGACAGGGATGCCCTCTCTCACCGCTCCTATTCAACATAGTGTTGGAAGTTCTGGCCAGGGCAATCAGGCAGGAGAAGGAAATAAAGGGTATTCAATTAGGAAAAGAGGAAGTCAAATTGTCCCTGTTTGCAGACGACATGATTGTTTATCTAGAAAACCCCATTGTCTCAGCCCAAAATCTCCTTAAGCTGATAAGCAACTTCAGCAAAGTCTCAGGATACAAAATCAATGTACAAAAATCACAAGCATTCTTATACACCAACAACAGACAAACAGAGAGCCAAATCATGGGTGAACTCCCATTCACAATTGCTTCAAAGAGGATAAAATACCTAAGAATCCAACTTACAAGGGATGTGAAGGACCTCTTCAAGGAGAACTACAAACCACTGCTCAAGGAAATAAAAGAGGACACAAACAAATGGAAGAACATTCCATGCTCATGGGTAGGAAGAATCAATATCGTGAAAATGGCCATACTGCCCAAGGTAATTTACAGATTCAATGCCATCCCCATCAAGCTACCAATGACTTTCTTCACAGAATTGGAAAAAACTACTTTAAAGTTCATATGGAACCAAAAAAGAGCCCGCATTGCCAAGTCAATCCTAAGCCAAAAGAACAAAGCTGGAGGCATCACACTACCTTACTTCAAACTATACTACAAGGCTACAGTAACCAAAACAGCATGGTACTGGTACCAAAACAGAGATATAGATCAATGGAACAGAACAGAGCCCTCAGAAATAATGCCACATATCTACAACTATCTGATCTTTGACAAACCTGAGAAAAACAAGCAATGGGGAAAGGATTCCCTATTTAATAAATGGTGCTGGGAAAACTGGCTAGCCATATGTAGAAAGCTGAAACTGGATCTCTTCCTTACACCTTTATACAAAAATCAATTCAAGATGGATTAAAGATTTAAACGTTAAACCTAAAACCATAAAAACCCTAGAAGAAAACCTAGGCATTACCATTCAGGACATAGGCGTGGGCAAGGACTTCATGTCCAAAACACCAAAAGCAATGGCAACAAAAGACAAAATTGACAAATGGGATCTAATTAAACTAAAGAGCTTCTGCACAGCAAAAGAAACTACCATCAGAGTGAACAGGCAACCTACAACATGGGAGAAAATTTTCGCAACCTACTCATCTGACAAAGGGCTAATATCCAGAATCTACAATGAACTCAAACAAATTTACAAGAAAAAAACAAACAACCCCATCAAAAAGTGGGCGAAGGACATGAACAGACACTTCTCAAAAGAAGACATTTATGCAGCCAAAAAACACATGAAGAAATGCTCATCATCACTGGCCATCAGAGAAATGCAAATCAAAACCACTATGAGATATCATCTCACACCAGTTAGAATGGCAATCATTAAAAAGTCAGGAAACAACAGGTGCTGGAGAGGATGCGGAGAAATAGGAACACTTTTACACTGTTGGTGGGACTGTAAACTAGTTCAACCATTGTGGAAGTCAGTGTGGCGATTCCTCAGGGATCTAGAACTAGAAATACCATTTGACCCAGCCATCCCATTACTGGGTATATACCCAGAGGACTATAAATCATGCTGCTATAAAGACACATGCACTCGTATGTTTATTGCGGCACTATTCACAATAGCAAAAACTTGGAACCAACCCAAATGTCCAACAATGATAGACTGGATTAAGAAAATGTGGCACATATACACCATGGAATATTATGCAGCCATAAAAAATGATGAGTTCATATCCTTTGTAGGGACATGGATGAAATTGGAAACCATCATTCTCAGTAAACTATCGCAAGAACAAAAAACCAAACACCGCATATTCTCACTCATAGGTGGGAATTGAACAATGAGATCACATGGACACAGGAAGGGGAATATCACACTCTGGGGACTGTGGTGGGGTCGGGGGAGGGGGGAGGGGTAGCATTGGGAGATATACCTAATGCTAGATGACACATTAGTGGGTGCAGCGCACCAGCATGGCACATGTATACATATGTAACTAACCTGCACAATGTGCACATGTACCCTAAAACTTAGAGTATAATTAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAGATCACACCACTGCACTCCAGCCTGGGTGTCAAAGCGAGACCCTGTCTCAGGAAAAAAAAAAAAAAAAAAAAAAAAAGGCTTAATTGATTGAACCAGATTCGAGAAAACAGTGCTAAATTATAATTTTCTCAATACTGTAAATATTTTTCAATCTTCAGCTTCATTAACTTCTATAATTGAAATTATCCCAATTATTACCTGACATGTACTAAAATTCCCTAAAATGGATCTTGAGTAACATTTTCACAGTACGATAATTTTTCTCTCTGTATATATTTATATAGTCACATATATGCACATACATTATACAAGCATTACTTTTCTATAACTGTAAGGTCAGAATTTGAAGTTGTGTTTTCTTTATCTTTTTATTTCCAATACTTGGCATCAAGTTGATATTCATTAGAAGTAAAGGAGGAAGGAAATGAATAATCTTCAGATACTAAGAACATTACACTTAAATTATTATTAAATCTAATTTGCATTCTCATATATGGCTTAGCT,L1RE2,100% +GACACCATCGAATGGCGCAAAACCTTTCGCGGTATGGCATGATAGCGCCCGGAAGAGAGTCAATTCAGGGTGGTGAATGTGAAACCAGTAACGTTATACGATGTCGCAGAGTATGCCGGTGTCTCTTATCAGACCGTTTCCCGCGTGGTGAACCAGGCCAGCCACGTTTCTGCGAAAACGCGGGAAAAAGTGGAAGCGGCGATGGCGGAGCTGAATTACATTCCCAACCGCGTGGCACAACAACTGGCGGGCAAACAGTCGTTGCTGATTGGCGTTGCCACCTCCAGTCTGGCCCTGCACGCGCCGTCGCAAATTGTCGCGGCGATTAAATCTCGCGCCGATCAACTGGGTGCCAGCGTGGTGGTGTCGATGGTAGAACGAAGCGGCGTCGAAGCCTGTAAAGCGGCGGTGCACAATCTTCTCGCGCAACGCGTCAGTGGGCTGATCATTAACTATCCGCTGGATGACCAGGATGCCATTGCTGTGGAAGCTGCCTGCACTAATGTTCCGGCGTTATTTCTTGATGTCTCTGACCAGACACCCATCAACAGTATTATTTTCTCCCATGAAGACGGTACGCGACTGGGCGTGGAGCATCTGGTCGCATTGGGTCACCAGCAAATCGCGCTGTTAGCGGGCCCATTAAGTTCTGTCTCGGCGCGTCTGCGTCTGGCTGGCTGGCATAAATATCTCACTCGCAATCAAATTCAGCCGATAGCGGAACGGGAAGGCGACTGGAGTGCCATGTCCGGTTTTCAACAAACCATGCAAATGCTGAATGAGGGCATCGTTCCCACTGCGATGCTGGTTGCCAACGATCAGATGGCGCTGGGCGCAATGCGCGCCATTACCGAGTCCGGGCTGCGCGTTGGTGCGGATATCTCGGTAGTGGGATACGACGATACCGAAGACAGCTCATGTTATATCCCGCCGTCAACCACCATCAAACAGGATTTTCGCCTGCTGGGGCAAACCAGCGTGGACCGCTTGCTGCAACTCTCTCAGGGCCAGGCGGTGAAGGGCAATCAGCTGTTGCCCGTCTCACTGGTGAAAAGAAAAACCACCCTGGCGCCCAATACGCAAACCGCCTCTCCCCGCGCGTTGGCCGATTCATTAATGCAGCTGGCACGACAGGTTTCCCGACTGGAAAGCGGGCAGTGAGCGCAACGCAATTAATGTGAGTTAGCTCACTCATTAGGCACCCCAGGCTTTACACTTTATGCTTCCGGCTCGTATGTTGTGTGGAATTGTGAGCGGATAACAATTTCACACAGGAAACAGCTATGACCATGATTACGGATTCACTGGCCGTCGTTTTACAACGTCGTGACTGGGAAAACCCTGGCGTTACCCAACTTAATCGCCTTGCAGCACATCCCCCTTTCGCCAGCTGGCGCTGGAGTGCGATCTTCCTGAGGCCGATACTGTCGTCGTCCCCTCAAACTGGCAGATGCACGGTTACGATGCGCCCATCTACACCAACGTAACCTATCCCATTACGGTCAATCCGCCGTTTGTTCCCACGGAGAATCCGACGGGTTGTTACTCGCTCACATTTAATGTTGATGAAAGCTGGCTACAGGAAGGCCAGACGCGAATTATTTTTGATGGCGTTAACTCGGCGTTTCATCTGTGGTGCAACGGGCGCTGGGTCGGTTACGGCCAGGACAGTCGTTTGCCGTCTGAATTTGACCTGAGCGCATTTTTACGCGCCGGAGAAAACCGCCTCGCGGTGATGGTGCTGCGTTGGAGTGACGGCAGTTATCTGGAAGATCAGGATATGTGGCGGATGAGCGGCATTTTCCGTGACGTCTCGTTGCTGCATAAACCGACTACACAAATCAGCGATTTCCATGTTGCCACTCGCTTTAATGATGATTTCAGCCGCGCTGTACTGGAGGCTGAAGTTCAGATGTGCGGCGAGTTGCGTGACTACCTACGGGTAACAGTTTCTTTATGGCAGGGTGAAACGCAGGTCGCCAGCGGCACCGCGCCTTTCGGCGGTGAAATTATCGATGAGCGTGGTGGTTATGCCGATCGCGTCACACTACGTCTGAACGTCGAAAACCCGAAACTGTGGAGCGCCGAAATCCCGAATCTCTATCGTGCGGTGGTTGAACTGCACACCGCCGACGGCACGCTGATTGAAGCAGAAGCCTGCGATGTCGGTTTCCGCGAGGTGCGGATTGAAAATGGTCTGCTGCTGCTGAACGGCAAGCCGTTGCTGATTCGAGGCGTTAACCGTCACGAGCATCATCCTCTGCATGGTCAGGTCATGGATGAGCAGACGATGGTGCAGGATATCCTGCTGATGAAGCAGAACAACTTTAACGCCGTGCGCTGTTCGCATTATCCGAACCATCCGCTGTGGTACACGCTGTGCGACCGCTACGGCCTGTATGTGGTGGATGAAGCCAATATTGAAACCCACGGCATGGTGCCAATGAATCGTCTGACCGATGATCCGCGCTGGCTACCGGCGATGAGCGAACGCGTAACGCGAATGGTGCAGCGCGATCGTAATCACCCGAGTGTGATCATCTGGTCGCTGGGGAATGAATCAGGCCACGGCGCTAATCACGACGCGCTGTATCGCTGGATCAAATCTGTCGATCCTTCCCGCCCGGTGCAGTATGAAGGCGGCGGAGCCGACACCACGGCCACCGATATTATTTGCCCGATGTACGCGCGCGTGGATGAAGACCAGCCCTTCCCGGCTGTGCCGAAATGGTCCATCAAAAAATGGCTTTCGCTACCTGGAGAGACGCGCCCGCTGATCCTTTGCGAATACGCCCACGCGATGGGTAACAGTCTTGGCGGTTTCGCTAAATACTGGCAGGCGTTTCGTCAGTATCCCCGTTTACAGGGCGGCTTCGTCTGGGACTGGGTGGATCAGTCGCTGATTAAATATGATGAAAACGGCAACCCGTGGTCGGCTTACGGCGGTGATTTTGGCGATACGCCGAACGATCGCCAGTTCTGTATGAACGGTCTGGTCTTTGCCGACCGCACGCCGCATCCAGCGCTGACGGAAGCAAAACACCAGCAGCAGTTTTTCCAGTTCCGTTTATCCGGGCAAACCATCGAAGTGACCAGCGAATACCTGTTCCGTCATAGCGATAACGAGCTCCTGCACTGGATGGTGGCGCTGGATGGTAAGCCGCTGGCAAGCGGTGAAGTGCCTCTGGATGTCGCTCCACAAGGTAAACAGTTGATTGAACTGCCTGAACTACCGCAGCCGGAGAGCGCCGGGCAACTCTGGCTCACAGTACGCGTAGTGCAACCGAACGCGACCGCATGGTCAGAAGCCGGGCACATCAGCGCCTGGCAGCAGTGGCGTCTGGCGGAAAACCTCAGTGTGACGCTCCCCGCCGCGTCCCACGCCATCCCGCATCTGACCACCAGCGAAATGGATTTTTGCATCGAGCTGGGTAATAAGCGTTGGCAATTTAACCGCCAGTCAGGCTTTCTTTCACAGATGTGGATTGGCGATAAAAAACAACTGCTGACGCCGCTGCGCGATCAGTTCACCCGTGCACCGCTGGATAACGACATTGGCGTAAGTGAAGCGACCCGCATTGACCCTAACGCCTGGGTCGAACGCTGGAAGGCGGCGGGCCATTACCAGGCCGAAGCAGCGTTGTTGCAGTGCACGGCAGATACACTTGCTGATGCGGTGCTGATTACGACCGCTCACGCGTGGCAGCATCAGGGGAAAACCTTATTTATCAGCCGGAAAACCTACCGGATTGATGGTAGTGGTCAAATGGCGATTACCGTTGATGTTGAAGTGGCGAGCGATACACCGCATCCGGCGCGGATTGGCCTGAACTGCCAGCTGGCGCAGGTAGCAGAGCGGGTAAACTGGCTCGGATTAGGGCCGCAAGAAAACTATCCCGACCGCCTTACTGCCGCCTGTTTTGACCGCTGGGATCTGCCATTGTCAGACATGTATACCCCGTACGTCTTCCCGAGCGAAAACGGTCTGCGCTGCGGGACGCGCGAATTGAATTATGGCCCACACCAGTGGCGCGGCGACTTCCAGTTCAACATCAGCCGCTACAGTCAACAGCAACTGATGGAAACCAGCCATCGCCATCTGCTGCACGCGGAAGAAGGCACATGGCTGAATATCGACGGTTTCCATATGGGGATTGGTGGCGACGACTCCTGGAGCCCGTCAGTATCGGCGGAATTCCAGCTGAGCGCCGGTCGCTACCATTACCAGTTGGTCTGGTGTCAAAAATAATAATAACCGGGCAGGCCATGTCTGCCCGTATTTCGCGTAAGGAAATCCATTATGTACTATTTAAAAAACACAAACTTTTGGATGTTCGGTTTATTCTTTTTCTTTTACTTTTTTATCATGGGAGCCTACTTCCCGTTTTTCCCGATTTGGCTACATGACATCAACCATATCAGCAAAAGTGATACGGGTATTATTTTGCCGCTATTTCTCTGTTCTCGCTATTATTCCAACCGCTGTTTGGTCTGCTTTCTGACAAACTCGGGCTGCGCAAATACCTGCTGTGGATTATTACCGGCATGTTAGTGATGTTTGCGCCGTTCTTTATTTTTATCTTCGGGCCACTGTTACAATACAACATTTTAGTAGGATCGATTGTTGGTGGTATTTATCTAGGCTTTTGTTTTAACGCCGGTGCGCCAGCAGTAGAGGCATTTATTGAGAAAGTCAGCCGTCGCAGTAATTTCGAATTTGGTCGCGCGCGGATGTTTGGCTGTGTTGGCTGGGCGCTGTGTGCCTCGATTGTCGGCATCATGTTCACCATCAATAATCAGTTTGTTTTCTGGCTGGGCTCTGGCTGTGCACTCATCCTCGCCGTTTTACTCTTTTTCGCCAAAACGGATGCGCCCTCTTCTGCCACGGTTGCCAATGCGGTAGGTGCCAACCATTCGGCATTTAGCCTTAAGCTGGCACTGGAACTGTTCAGACAGCCAAAACTGTGGTTTTTGTCACTGTATGTTATTGGCGTTTCCTGCACCTACGATGTTTTTGACCAACAGTTTGCTAATTTCTTTACTTCGTTCTTTGCTACCGGTGAACAGGGTACGCGGGTATTTGGCTACGTAACGACAATGGGCGAATTACTTAACGCCTCGATTATGTTCTTTGCGCCACTGATCATTAATCGCATCGGTGGGAAAAACGCCCTGCTGCTGGCTGGCACTATTATGTCTGTACGTATTATTGGCTCATCGTTCGCCACCTCAGCGCTGGAAGTGGTTATTCTGAAAACGCTGCATATGTTTGAAGTACCGTTCCTGCTGGTGGGCTGCTTTAAATATATTACCAGCCAGTTTGAAGTGCGTTTTTCAGCGACGATTTATCTGGTCTGTTTCTGCTTCTTTAAGCAACTGGCGATGATTTTTATGTCTGTACTGGCGGGCAATATGTATGAAAGCATCGGTTTCCAGGGCGCTTATCTGGTGCTGGGTCTGGTGGCGCTGGGCTTCACCTTAATTTCCGTGTTCACGCTTAGCGGCCCCGGCCCGCTTTCCCTGCTGCGTCGTCAGGTGAATGAAGTCGCTTAAGCAATCAATGTCGGATGCGGCGCGACGCTTATCCGACCAACATATCATAACGGAGTGATCGCATTGAACATGCCAATGACCGAAAGAATAAGAGCAGGCAAGCTATTTACCGATATGTGCGAAGGCTTACCGGAAAAAAGACTTCGTGGGAAAACGTTAATGTATGAGTTTAATCACTCGCATCCATCAGAAGTTGAAAAAAGAGAAAGCCTGATTAAAGAAATGTTTGCCACGGTAGGGGAAAACGCCTGGGTAGAACCGCCTGTCTATTTCTCTTACGGTTCCAACATCCATATAGGCCGCAATTTTTATGCAAATTTCAATTTAACCATTGTCGATGACTACACGGTAACAATCGGTGATAACGTACTGATTGCACCCAACGTTACTCTTTCCGTTACGGGACACCCTGTACACCATGAATTGAGAAAAAACGGCGAGATGTACTCTTTTCCGATAACGATTGGCAATAACGTCTGGATCGGAAGTCATGTGGTTATTAATCCAGGCGTCACCATCGGGGATAATTCTGTTATTGGCGCGGGTAGTATCGTCACAAAAGACATTCCACCAAACGTCGTGGCGGCTGGCGTTCCTTGTCGGGTTATTCGCGAAATAAACGACCGGGATAAGCACTATTATTTCAAAGATTATAAAGTTGAATCGTCAGTTTAAATTATAAAAATTGCCTGATACGCTGCGCTTATCAGGCCTACAAGTTCAGCGATCTACATTAGCCGCATCCGGCATGAACAAAGCGCAGGAACAAGCGTCGCATCATGCCTCTTTGACCCACAGCTGCGGAAAACGTACTGGTGCAAAACGCAGGGTTATGATCATCAGCCCAACGACGCACAGCGCATGAAATGCCCAGTCCATCAGGTAATTGCCGCTGATACTACGCAGCACGCCAGAAAACCACGGGGCAAGCCCGGCGATGATAAAACCGATTCCCTGCATAAACGCCACCAGCTTGCCAGCAATAGCCGGTTGCACAGAGTGATCGAGCGCCAGCAGCAAACAGAGCGGAAACGCGCCGCCCAGACCTAACCCACACCATCGCCCACAATACCGGCAATTGCATCGGCAGCCAGATAAAGCCGCAGAACCCCACCAGTTGTAACACCAGCGCCAGCATTAACAGTTTGCGCCGATCCTGATGGCGAGCCATAGCAGGCATCAGCAAAGCTCCTGCGGCTTGCCCAAGCGTCATCAATGCCAGTAAGGAACCGCTGTACTGCGCGCTGGCACCAATCTCAATATAGAAAGCGGGTAACCAGGCAATCAGGCTGGCGTAACCGCCGTTAATCAGACCGAAGTAAACACCCAGCGTCCACGCGCGGGGAGTGAATACCACGCGAACCGGAGTGGTTGTTGTCTTGTGGGAAGAGGCGACCTCGCGGGCGCTTTGCCACCACCAGGCAAAGAGCGCAACAACGGCAGGCAGCGCCACCAGGCGAGTGTTTGATACCAGGTTTCGCTATGTTGAACTAACCAGGGCGT,ECOLAC,100% +GTTAATGTAGCTTAAAACAAAAGCAAGGTACTGAAAATACCTAGACGAGTATATCCAACTCCATAAACAACAAAGGTTTGGTCCCGGCCTTCTTATTGGTTACTAGGAAACTTATACATGCAAGTATCCGCCCGCCAGTGAATACGCCTTCTAAATCATCACTGATCAAAGAGAGCTGGCATCAAGCACACACCCCAAGTGTAGCTCATGACGTCTCGCCTAGCCACACCCCCACGGGAAACAGCAGTAGTAAATATTTAGCAATTAACAAAAGTTAGACTAAGTTATCCTAATAAAGGACTGGTCAATTTCGTGCCAGCAACCGCGGCCATACGATTAGTCCAAATTAATAAGCATACGGCGTAAAGCGTATTAGAAGAATTAAAAAAATAAAGTTAAATCTTATACTAGCTGTTTAAAGCTCAAGATAAGACATAAATAGCCTACGAAAGTGACTTTAATAATCCTAAACATACGATAGCTAGGGTACAAACTGAGATTAGATACCTCACTATGCCTAGCCCTAAACTTTGATAGCTACCTTTACAAAGCTATCCGCCAGAGAACTACTAGCCAGAGCTTAAAACTTAAAGGACTTGGCGGTGCTTTATATCCACCTAGGGGAGCCTGTCTCGTAACCGATGAACCCCGATACACCTTACCGTCACTTGCTAATTCAGTCCATATACCACCATCTTCAGCAAACCCCTATAGGGCACAAAAGTGAGCTTAATCATAACCCATGAAAAAGTTAGGCCGAGGTGTCGCCTACGTGACGGTCAAAGATGGGCTACATTTTCTATTATAGAATAGACAAACGGATACCACTCTGAAATGGGTGGTTGAAGGCGGATTTAGTAGTAAACTAAGAATAGAGAGCTTAATTGAACAAGGCCATGAAGCGCGTACACACCGCCCGTCACTCTCCTCAAGTACCTCCACATCAAACAATCATATTACAGATTTAAACAAATACAAGAGGAGACAAGTCGTAACAAGGTAAGCGTACTGGAAAGTGTGCTTGGGTAACTCAAAGTGTAGCTTAACAAAAAGCATCTGGCTTACACCTAGAAGACCTCATTCACAATGATCACTTTGAACTAAATCTAGCCCTACCAACCTTACACCCAACTCTCACACTACATTAAATTAAAACATTCATTTATCAAAAAGTATAGGAGATAGAAATTTCACTAAGGCGCAATAGAGATAGTACCGCAAGGGAATGATGAAAGATAATTTAATAGTAAAAAATAGCAAGGATTAACCCCTTTACCTTTTGCATAATGAATTAACTAGAAAAATCTGACAAAGAGAACTACAGCCAGAAACCCCGAAATCAGACGAGCTATCTGATAGTAATCCCCAGGATCAATTCATCTATGTGGCAAAATAGTGAAAAAACTTACAGATAGAGGTGAAATACCAATCGAGCCTGATGATAGCTGGTTGTCCAGAAATAGAATTTCAGTTCTACCTAAAACTTACCACAAAAACAAAATAATTCCAATGTAAGTTTTAGAGATATTCAAAAGGGGTACAGCTCTTTTGACCAAGGATACAACCTTGATTAGCGAGTAAATTCACCATTAATTTCATAGTTGGCTTGGAAGCAGCCATCAATTAAGAAAGCGTTAAAGCTCAACAACCAACCAAACTAAAAAATCCCAAGAATTAATTAATGATCTCCTAAACATAATACTGGACTAATCTATATAAATAGAAGAAATAATGTTAGTATAAGTAATAAGAAGTATTTCTCCCTGCATAAGCTTATATCAGATCGGATGCCCACTGATAGTTAACAATCAAATAATTAAATACAAAAATAAAACCTTTATTACACCAATTGTTAACCCAACACAGGCATGCTTAAGGGAAAGATTAAAAGAAGGAAAAGGAACTCGGCAAACATAAACCCCGCCTGTTTACCAAAAACATCACCTCGAGCATTACTAGTATTCGAGGCACTGCCTGCCCAGTGACCAAGTGTTAAACGGCCGCGGTACTCTGACCGTGCAAAGGTAGCATAATCATTTGTTCCTTAATTAGGGACTTGTATGAACGGCCACACGAGGGTTTAACTGTCTCTTTCCTCTAATCAATGAAATTGACCTTCTCGTGAAGAGGCGAGAATAAACATATAAGACGAGAAGACCCTATGGAGCTTAAATTAACTAATTTAATTGCTATCCTATAAATCTACAAGATACAACTAAACAGCATAATAAATTAACAATTTTGGTTGGGGTGACCTCGGAGAAGAAAAAAACCTCCGAACGATATTATAATTCAGACTTTACAAGTCAAGATTCACTAATCGCTTATTGACCCAATACTTGATCAACGGAACAAGTTACCCTAGGGATAACAGCGCAATCCTACTCTAGAGTCCCTATCGACAGCAGGGTTTACGACCTCGATGTTGGATCAGGACATCCTAATGGTGCAGCCGCTATTAAGGGTTCGTTTGTTCAACGATTAAAGTCCTACGTGATCTGAGTTCAGACCGGAGCAATCCAGGTCGGTTTCTATCTATAGTTTATTTATTCCAGTACGAAAGGACAGAAAAAATGAGGCCAATCTTACCAAGACGCCTTCAGCTAAATTTATGAATAAATCTCAATCTAGATAAGCTAAACCACCCAATCCAAGAACAGGATTTGTTAAGATAGCAAAAATTGGTTACTGCATAAAACTTAAGCTTTTACTTACGGAGGTTCAACTCCTCTTCTTAACAATGTTCTTGATTAATGTCCTAACAGTAACCTTGCCTATCCTTCTAGCAGTAGCCTTCCTCACCTTAGTTGAACGAAAGGCCTTAGGCTACATACAACTTCGTAAAGGCCCCAATGTAGTAGGACCCTACGGTCTTCTTCAACCTATCGCAGATGCAATCAAGCTATTTACCAAAGAACCCGTCTATCCACAAACCTCCTCAAAATTCCTATTTACCATTGCCCCAATTCTAGCCCTAACCTTAGCCCTAACTGTATGAGCTCCTCTTCCAATACCATATCCCCTAATTAACTTAAATCTAAGCCTATTATTTATTCTCGCAATATCAAGTCTGATAGTTTACTCCATCCTATGATCAGGCTGAGCATCAAATTCAAAATACGCCCTCATAGGAGCCCTACGAGCAGTAGCCCAAACCATCTCCTATGAAGTCTCTATAACAACTATTATCTTATCAATAGTACTAATAAATGGGTCCTTTACACTAACCGCCTTCGCTACAACACAAGAACATCTATGACTAATCTTTCCTATATGACCCCTAATAATAATATGATTTACATCAACTTTAGCAGAAACTAACCGAGCTCCATTTGATTTAACCGAAGGAGAATCAGAACTAGTTTCCGGCTTCAACGTCGAATATTCAGCTGGCCCTTTCGCCCTATTTTTTATAGCCGAATACGCTAACATTATCATAATAAATGCCCTTACTGTAATTCTATTTATAGGAACCTCTTGTAACCCCCAAATACCAGAAATTAGCACCATCAACTTTGTCGTAAAAACTATAATCTTAACTATCTGCTTTTTATGAGTACGAGCATCTTACCCACGATTCCGATACGACCAACTAATATATCTCCTCTGAAAAAATTTTCTTCCACTAACTCTAGCCCTATGCATGTGACACATCTCAATCTTGATTTCACTAGCATGCATCCCACCACAAGCATAGAAATATGTCTGACAAAAGAATTACTTTGATAGAGTAAATTATAGAGGTCTAAACCCTCTTATTTCTAGAATTGCAGGAATCGAACCTAAACTCGAGAATTCAAAAATCTCAGTGCTACCAATTACACCATATCCTACTAGTAAGGTCAGCTAAATTAAGCTATCGGGCCCATACCCCGAAAATGTCGGATTACACCCCTCCCATACTAATAAACCCACTAGCCCTTAGCCTAATCCTAACAACACTACTCGCAGGAACACTAATTACTATAATAAGCTCCCATTGACTAACAGCCTGAATAGGACTAGAAATGAACATACTTACTATAATTCCTATTCTAATAAAGACAACCAATCCACGATCCACAGAAGCCGCCACGAAATACTTTATAACCCAAGCCACAGCATCCATAATACTTATAATAGCTTTAACAATTAATCTAATATACTCAGGACAATGATCAATTATAAAAATAACCAACCCTGTAGCATCAAATGTAGCATTAATAGCTCTAATAACCAAACTAGGCTCAGCCCCATTCCACTTCTGAGTCCCAGAAGTAACGCAAGGAGTCGAACTCACATCAGGAATGATCTTATTAACCTGACAAAAATTAGCACCATTATCTCTATTATATCAAATAGCCACCTACACCAACACCAACCTAATCTATCTTTCTGGCCTACTTTCAATCCTAATTGGAGGATGAGGAGGCCTAAACCAAACACAACTACGAAAAATCTTAGCCTACTCTTCAATCTCCCATATAGGCTGAATACTCATTATTTTACCCTTTAACCCTACCCTTACTCTCCTGAACTTAGCCATTTACATTTTACTAACACTATCCATCTTTATAATCCTAGCAAATACCCTCACAACTTCAATATCATCCCTAACTCTAATATGAAACAAAACACCTGCAATAACCATTATACTTATAACTACCTTGCTATCCTTAGGAGGACTACCTCCACTTTCAGGATTCACACCTAAATGACTTATAATCCATGAACTGACTAAAAACAACAGTATTATCATACCACTAACCATAGCCATTATAACACTACTAAACATGTACTTCTACATACGACTAATCTACTACTCATCACTCACAATTCTTCCATCTACAAATAATATGAAAATAACCTGACAATTTACTAGTACTAAACACACAATAATATTACCGACTCTAATCACCCTATCTAACATACTACTTCCCCTAACCCCAATAATTTCAATGCTAGAATAGGAATTTAGGTTAAAACAGACCAAGAGCCTTCAAAGCCCTAAGTAAGTGCATTATACTTAACTCCTGAAATAAGGACTGCAAGATATCACCTTACATCAACTGAATGCAAATCAGACGCTTTAACTAAACTAAGCCCTTCTAGATTGGAGGGCTTCAATCCCACGAAAATCTTAGTTAACAACTAAACACCCTAGCTAACTGGCTTCAATCTACTTCTCCCGCCTTGAAGGAGGGGGAAAAAGGCGGGAGAAGCCCCGGCAGAATTGAAGCTGCTTTTTTCGAATTTGCAGTTCGACATGTTTACACTTTCAAGGCCTGGTAAAAAGAGATTACTCTCTGTGGTTAGATTTACAGTCTAATACTTACTCAGCCATTTTACCTATGTTTGCTAACCGCTGACTATATTCAACAAACCACAAAGACATTGGGACACTGTATCTATTATTTGGTGCTTGAGCTGGTATAGTAGGGACTGCTTTTAGTATCCTAATTCGGGCAGAACTAGGTCAACCAGGCTCTCTTCTTGGAGACGACCAAATCTATAATGTTATTGTCACAGCACACGCCTTTGTAATAATCTTCTTTATAGTTATGCCAATTATAATTGGAGGCTTTGGAAACTGATTAATTCCACTTATAATCGGAGCACCTGATATAGCTTTTCCTCGAATAAACAATATGAGTTTTTGACTACTACCTCCATCTTTCCTACTACTCTTGGCATCTTCCATAGTAGAAGCTGGGACAGGCACTGGTTGGACCGTATACCCTCCCCTGGCAGGAAACCTAGCCCATGCAGGAGCTTCTGTAGATTTAACTATTTTTTCACTTCACCTTGCAGGAGTATCCTCTATTTTAAGTGCAATTAATTTTATCACTACCATCATTAACATAAAACCTCCAGCTATGTCTCAATATCACATACCCTTATTTGTATGGTCCATTTTAGTTACAGCCGTCCTTCTTCTTCTATCCCTCCCAGTTCTAGCAGCAGGTATTACAATATTATTAACGGACCGCAACCTCAATACTACTTTCTTTGACCCTGCAGGAGGAGGAGACCCAATTCTATACCAACACCTATTCTGGTTTTTTGGACACCCTGAAGTCTATATTCTAATTCTCCCAGGATTTGGAATAGTTTCTCATATCGTTACGTACTACTCAGGGAAAAAAGAACCCTTCGGTTATATAGGAATAGTATGGGCTAT,NC_007596.2Mammuthusprimigeniusmitochondrion,30% +GATCACAGGTCTATCACCCTATTAACCACTCACGGGAGCTCTCCATGCATTTGGTATTTTCGTCTGGGGGGTATGCACGCGATAGCATTGCGAGACGCTGGAGCCGGAGCACCCTATGTCGCAGTATCTGTCTTTGATTCCTGCCTCATCCTATTATTTATCGCACCTACGTTCAATATTACAGGCGAACATACTTACTAAAGTGTGTTAATTAATTAATGCTTGTAGGACATAATAATAACAATTGAATGTCTGCACAGCCACTTTCCACACAGACATCATAACAAAAAATTTCCACCAAACCCCCCCTCCCCCGCTTCTGGCCACAGCACTTAAACACATCTCTGCCAAACCCCAAAAACAAAGAACCCTAACACCAGCCTAACCAGATTTCAAATTTTATCTTTTGGCGGTATGCACTTTTAACAGTCACCCCCCAACTAACACATTATTTTCCCCTCCCACTCCCATACTACTAATCTCATCAATACAACCCCCGCCCATCCTACCCAGCACACACACACCGCTGCTAACCCCATACCCCGAACCAACCAAACCCCAAAGACACCCCCCACAGTTTATGTAGCTTACCTCCTCAAAGCAATACACTGAAAATGTTTAGACGGGCTCACATCACCCCATAAACAAATAGGTTTGGTCCTAGCCTTTCTATTAGCTCTTAGTAAGATTACACATGCAAGCATCCCCGTTCCAGTGAGTTCACCCTCTAAATCACCACGATCAAAAGGAACAAGCATCAAGCACGCAGCAATGCAGCTCAAAACGCTTAGCCTAGCCACACCCCCACGGGAAACAGCAGTGATTAACCTTTAGCAATAAACGAAAGTTTAACTAAGCTATACTAACCCCAGGGTTGGTCAATTTCGTGCCAGCCACCGCGGTCACACGATTAACCCAAGTCAATAGAAGCCGGCGTAAAGAGTGTTTTAGATCACCCCCTCCCCAATAAAGCTAAAACTCACCTGAGTTGTAAAAAACTCCAGTTGACACAAAATAGACTACGAAAGTGGCTTTAACATATCTGAACACACAATAGCTAAGACCCAAACTGGGATTAGATACCCCACTATGCTTAGCCCTAAACCTCAACAGTTAAATCAACAAAACTGCTCGCCAGAACACTACGAGCCACAGCTTAAAACTCAAAGGACCTGGCGGTGCTTCATATCCCTCTAGAGGAGCCTGTTCTGTAATCGATAAACCCCGATCAACCTCACCACCTCTTGCTCAGCCTATATACCGCCATCTTCAGCAAACCCTGATGAAGGCTACAAAGTAAGCGCAAGTACCCACGTAAAGACGTTAGGTCAAGGTGTAGCCCATGAGGTGGCAAGAAATGGGCTACATTTTCTACCCCAGAAAACTACGATAGCCCTTATGAAACTTAAGGGTCGAAGGTGGATTTAGCAGTAAACTAAGAGTAGAGTGCTTAGTTGAACAGGGCCCTGAAGCGCGTACACACCGCCCGTCACCCTCCTCAAGTATACTTCAAAGGACATTTAACTAAAACCCCTACGCATTTATATAGAGGAGACAAGTCGTAACATGGTAAGTGTACTGGAAAGTGCACTTGGACGAACCAGAGTGTAGCTTAACACAAAGCACCCAACTTACACTTAGGAGATTTCAACTTAACTTGACCGCTCTGAGCTAAACCTAGCCCCAAACCCACTCCACCTTACTACCAGACAACCTTAGCCAAACCATTTACCCAAATAAAGTATAGGCGATAGAAATTGAAACCTGGCGCAATAGATATAGTACCGCAAGGGAAAGATGAAAAATTATAACCAAGCATAATATAGCAAGGACTAACCCCTATACCTTCTGCATAATGAATTAACTAGAAATAACTTTGCAAGGAGAGCCAAAGCTAAGACCCCCGAAACCAGACGAGCTACCTAAGAACAGCTAAAAGAGCACACCCGTCTATGTAGCAAAATAGTGGGAAGATTTATAGGTAGAGGCGACAAACCTACCGAGCCTGGTGATAGCTGGTTGTCCAAGATAGAATCTTAGTTCAACTTTAAATTTGCCCACAGAACCCTCTAAATCCCCTTGTAAATTTAACTGTTAGTCCAAAGAGGAACAGCTCTTTGGACACTAGGAAAAAACCTTGTAGAGAGAGTAAAAAATTTAACACCCATAGTAGGCCTAAAAGCAGCCACCAATTAAGAAAGCGTTCAAGCTCAACACCCACTACCTAAAAAATCCCAAACATATAACTGAACTCCTCACACCCAATTGGACCAATCTATCACCCTATAGAAGAACTAATGTTAGTATAAGTAACATGAAAACATTCTCCTCCGCATAAGCCTGCGTCAGATTAAAACACTGAACTGACAATTAACAGCCCAATATCTACAATCAACCAACAAGTCATTATTACCCTCACTGTCAACCCAACACAGGCATGCTCATAAGGAAAGGTTAAAAAAAGTAAAAGGAACTCGGCAAATCTTACCCCGCCTGTTTACCAAAAACATCACCTCTAGCATCACCAGTATTAGAGGCACCGCCTGCCCAGTGACACATGTTTAACGGCCGCGGTACCCTAACCGTGCAAAGGTAGCATAATCACTTGTTCCTTAAATAGGGACCTGTATGAATGGCTCCACGAGGGTTCAGCTGTCTCTTACTTTTAACCAGTGAAATTGACCTGCCCGTGAAGAGGCGGGCATAACACAGCAAGACGAGAAGACCCTATGGAGCTTTAATTTATTAATGCAAACAGTACCTAACAAACCCACAGGTCCTAAACTACCAAACCTGCATTAAAAATTTCGGTTGGGGCGACCTCGGAGCAGAACCCAACCTCCGAGCAGTACATGCTAAGACTTCACCAGTCAAAGCGAACTACTATACTCAATTGATCCAATAACTTGACCAACGGAACAAGTTACCCTAGGGATAACAGCGCAATCCTATTCTAGAGTCCATATCAACAATAGGGTTTACGACCTCGATGTTGGATCAGGACATCCCGATGGTGCAGCCGCTATTAAAGGTTCGTTTGTTCAACGATTAAAGTCCTACGTGATCTGAGTTCAGACCGGAGTAATCCAGGTCGGTTTCTATCTACNTTCAAATTCCTCCCTGTACGAAAGGACAAGAGAAATAAGGCCTACTTCACAAAGCGCCTTCCCCCGTAAATGATATCATCTCAACTTAGTATTATACCCACACCCACCCAAGAACAGGGTTTGTTAAGATGGCAGAGCCCGGTAATCGCATAAAACTTAAAACTTTACAGTCAGAGGTTCAATTCCTCTTCTTAACAACATACCCATGGCCAACCTCCTACTCCTCATTGTACCCATTCTAATCGCAATGGCATTCCTAATGCTTACCGAACGAAAAATTCTAGGCTATATACAACTACGCAAAGGCCCCAACGTTGTAGGCCCCTACGGGCTACTACAACCCTTCGCTGACGCCATAAAACTCTTCACCAAAGAGCCCCTAAAACCCGCCACATCTACCATCACCCTCTACATCACCGCCCCGACCTTAGCTCTCACCATCGCTCTTCTACTATGAACCCCCCTCCCCATACCCAACCCCCTGGTCAACCTCAACCTAGGCCTCCTATTTATTCTAGCCACCTCTAGCCTAGCCGTTTACTCAATCCTCTGATCAGGGTGAGCATCAAACTCAAACTACGCCCTGATCGGCGCACTGCGAGCAGTAGCCCAAACAATCTCATATGAAGTCACCCTAGCCATCATTCTACTATCAACATTACTAATAAGTGGCTCCTTTAACCTCTCCACCCTTATCACAACACAAGAACACCTCTGATTACTCCTGCCATCATGACCCTTGGCCATAATATGATTTATCTCCACACTAGCAGAGACCAACCGAACCCCCTTCGACCTTGCCGAAGGGGAGTCCGAACTAGTCTCAGGCTTCAACATCGAATACGCCGCAGGCCCCTTCGCCCTATTCTTCATAGCCGAATACACAAACATTATTATAATAAACACCCTCACCACTACAATCTTCCTAGGAACAACATATGACGCACTCTCCCCTGAACTCTACACAACATATTTTGTCACCAAGACCCTACTTCTAACCTCCCTGTTCTTATGAATTCGAACAGCATACCCCCGATTCCGCTACGACCAACTCATACACCTCCTATGAAAAAACTTCCTACCACTCACCCTAGCATTACTTATATGATATGTCTCCATACCCATTACAATCTCCAGCATTCCCCCTCAAACCTAAGAAATATGTCTGATAAAAGAGTTACTTTGATAGAGTAAATAATAGGAGCTTAAACCCCCTTATTTCTAGGACTATGAGAATCGAACCCATCCCTGAGAATCCAAAATTCTCCGTGCCACCTATCACACCCCATCCTAAAGTAAGGTCAGCTAAATAAGCTATCGGGCCCATACCCCGAAAATGTTGGTTATACCCTTCCCGTACTAATTAATCCCCTGGCCCAACCCGTCATCTACTCTACCATCTTTGCAGGCACACTCATCACAGCGCTAAGCTCGCACTGATTTTTTACCTGAGTAGGCCTAGAAATAAACATGCTAGCTTTTATTCCAGTTCTAACCAAAAAAATAAACCCTCGTTCCACAGAAGCTGCCATCAAGTATTTCCTCACGCAAGCAACCGCATCCATAATCCTTCTAATAGCTATCCTCTTCAACAATATACTCTCCGGACAATGAACCATAACCAATACTACCAATCAATACTCATCATTAATAATCATAATAGCTATAGCAATAAAACTAGGAATAGCCCCCTTTCACTTCTGAGTCCCAGAGGTTACCCAAGGCACCCCTCTGACATCCGGCCTGCTTCTTCTCACATGACAAAAACTAGCCCCCATCTCAATCATATACCAAATCTCTCCCTCACTAAACGTAAGCCTTCTCCTCACTCTCTCAATCTTATCCATCATAGCAGGCAGTTGAGGTGGATTAAACCAAACCCAGCTACGCAAAATCTTAGCATACTCCTCAATTACCCACATAGGATGAATAATAGCAGTTCTACCGTACAACCCTAACATAACCATTCTTAATTTAACTATTTATATTATCCTAACTACTACCGCATTCCTACTACTCAACTTAAACTCCAGCACCACGACCCTACTACTATCTCGCACCTGAAACAAGCTAACATGACTAACACCCTTAATTCCATCCACCCTCCTCTCCCTAGGAGGCCTGCCCCCGCTAACCGGCTTTTTGCCCAAATGGGCCATTATCGAAGAATTCACAAAAAACAATAGCCTCATCATCCCCACCATCATAGCCACCATCACCCTCCTTAACCTCTACTTCTACCTACGCCTAATCTACTCCACCTCAATCACACTACTCCCCATATCTAACAACGTAAAAATAAAATGACAGTTTGAACATACAAAACCCACCCCATTCCTCCCCACACTCATCGCCCTTACCACGCTACTCCTACCTATCTCCCCTTTTATACTAATAATCTTATAGAAATTTAGGTTAAATACAGACCAAGAGCCTTCAAAGCCCTCAGTAAGTTGCAATACTTAATTTCTGTAACAGCTAAGGACTGCAAAACCCCACTCTGCATCAACTGAACGCAAATCAGCCACTTTAATTAAGCTAAGCCCTTACTAGACCAATGGGACTTAAACCCACAAACACTTAGTTAACAGCTAAGCACCCTAATCAACTGGCTTCAATCTACTTCTCCCGCCGCCGGGAAAAAAGGCGGGAGAAGCCCCGGCAGGTTTGAAGCTGCTTCTTCGAATTTGCAATTCAATATGAAAATCACCTCGGAGCTGGTAAAAAGAGGCCTAACCCCTGTCTTTAGATTTACAGTCCAATGCTTCACTCAGCCATTTTACCTCACCCCCACTGATGTTCGCCGACCGTTGACTATTCTCTACAAACCACAAAGACATTGGAACACTATACCTATTATTCGGCGCATGAGCTGGAGTCCTAGGCACAGCTCTAAGCCTCCTTATTCGAGCCGAGCTGGGCCAGCCAGGCAACCTTCTAGGTAACGACCACATCTACAACGTTATCGTCACAGCCCATGCATTTGTAATAATCTTCTTCATAGTAATACCCATCATAATCGGAGGCTTTGGCAACTGACTAGTTCCCCTAATAATCGGTGCCCCCGATATGGCGTTTCCCCGCATAAACAACATAAGCTTCTGACTCTTACCTCCCTCTCTCCTACTCCTGCTCGCATCTGCTATAGTGGAGGCCGGAGCAGGAACAGGTTGAACAGTCTACCCTCCCTTAGCAGGGAACTACTCCCACCCTGGAGCCTCCGTAGACCTAACCATCTTCTCCTTACACCTAGCAGGTGTCTCCTCTATCTTAGGGGCCATCAATTTCATCACAACAATTATCAATATAAAACCCCCTGCCATAACCCAATACCAAACGCCCCTCTTCGTCTGATCCGTCCTAATCACAGCAGTCCTACTTCTCCTATCTCTCCCAGTCCTAGCTGCTGGCATCACTATACTACTAACAGACCGCAACCTCAACACCACCTTCTTCGACCCCGCCGGAGGAGGAGACCCCATTCTATACCAACACCTATTCTGATTTTTCGGTCACCCTGAAGTTTATATTCTTATCCTACCAGGCTTCGGAATAATCTCCCATATTGTAACTTACTACTCCGGAAAAAAAGAACCATTTGGATACATAGGTATGGTCTGAGCTATGATATCAATTGGCTTCCTAGGGTTTATCGTGTGAGCACACCATATATTTACAGTAGGAATAGACGTAGACACACGAGCATATTTCACCTCCGCTACCATAATCATCGCTATCCCCACCGGCGTCAAAGTATTTAGCTGACTCGCCACACTCCACGGAAGCAATATGAAATGATCTGCTGCAGTGCTCTGAGCCCTAGGATTCATCTTTCTTTTCACCGTAGGTGGCCTGACTGGCATTGTATTAGCAAACTCATCACTAGACATCGTACTACACGACACGTACTACGTTGTAGCCCACTTCCACTATGTCCTATCAATAGGAGCTGTATTTGCCATCATAGGAGGCTTCATTCACTGATTTCCCCTATTCTCAGGCTACACCCTAGACCAAACCTACGCCAAAATCCATTTCACTATCATATTCATCGGCGTAAATCTAACTTTCTTCCCACAACACTTTCTCGGCCTATCCGGAATGCCCCGACGTTACTCGGACTACCCCGATGCATACACCACATGAAACATCCTATCATCTGTAGGCTCATTCATTTCTCTAACAGCAGTAATATTAATAATTTTCATGATTTGAGAAGCCTTCGCTTCGAAGCGAAAAGTCCTAATAGTAGAAGAACCCTCCATAAACCTGGAGTGACTATATGGATGCCCCCCACCCTACCACACATTCGAAGAACCCGTATACATAAAATCTAGACAAAAAAGGAAGGAATCGAACCCCCCAAAGCTGGTTTCAAGCCAACCCCATGGCCTCCATGACTTTTTCAAAAAGGTATTAGAAAAACCATTTCATAACTTTGTCAAAGTTAAATTATAGGCTAAATCCTATATATCTTAATGGCACATGCAGCGCAAGTAGGTCTACAAG,NC_012920.1_homosapiens_mitochondrion,50% diff --git a/model/evo2/__init__.py b/model/evo2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f2f43ba713c267a9baaa1007696b25de8646216 --- /dev/null +++ b/model/evo2/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/model/evo2/api.py b/model/evo2/api.py new file mode 100644 index 0000000000000000000000000000000000000000..cd18b285b2815194566d9d745a82063dfb025750 --- /dev/null +++ b/model/evo2/api.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from abc import ABC, abstractmethod +from typing import Generic, Sequence, TypeVar + +from megatron.core.transformer.module import MegatronModule +from nemo.lightning.megatron_parallel import DataT, MegatronLossReduction + + +__all__: Sequence[str] = ( + "BionemoMegatronModel", + "MegatronLossReduction", # re-export Megatron's loss definition as it's a core part of the bionemo-llm API + "MegatronLossType", + "MegatronModelType", +) + + +class BionemoMegatronModel(MegatronModule, Generic[DataT], ABC): + """Models that use Megatron must be a MegatronModule type. + + The only major difference is the explicit `forward` pass method signature that makes this class compatible + with bionemo-core's `Model` structural type. + """ + + @abstractmethod + def forward(self, *args, **kwargs) -> DataT: # noqa: D102 + raise NotImplementedError() + + +# Typechecking: ensure that the bionemo megatron model abstraction is compliant with bionemo-core's Model +# _: type[Model] = BionemoMegatronModel + + +MegatronModelType = TypeVar("MegatronModelType", bound=MegatronModule) +# bound=BionemoMegatronModel) + +MegatronLossType = TypeVar("MegatronLossType", bound=MegatronLossReduction) diff --git a/model/evo2/data/__init__.py b/model/evo2/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f2f43ba713c267a9baaa1007696b25de8646216 --- /dev/null +++ b/model/evo2/data/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/model/evo2/data/fasta_dataset.py b/model/evo2/data/fasta_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..be6c2474c5396f3f38f132dc577b9295f8c016a6 --- /dev/null +++ b/model/evo2/data/fasta_dataset.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import json +from pathlib import Path + +import torch + +from bionemo.noodles.nvfaidx import NvFaidx + + +class SimpleFastaDataset(torch.utils.data.Dataset): + """A simple dataset for Evo2 prediction. + + Currently, this will not work for pre-training or fine-tuning, as that would require: + 1) including "labels" in the input and 2) offsetting/rolling either the labels or + input_ids to handle the off-by-one token prediction alignment. + """ + + def __init__(self, fasta_path: Path, tokenizer, prepend_bos: bool = True): + """Initialize the dataset.""" + super().__init__() + self.fasta = NvFaidx(fasta_path) + self.seqids = sorted(self.fasta.keys()) + self.tokenizer = tokenizer + self.prepend_bos = prepend_bos # needed for getting predictions for the requested set of tokens. + + def write_idx_map(self, output_dir: Path): + """Write the index map to the output directory.""" + with open(output_dir / "seq_idx_map.json", "w") as f: + json.dump({seqid: idx for idx, seqid in enumerate(self.seqids)}, f) + + def __len__(self): + """Get the length of the dataset.""" + return len(self.seqids) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + """Get an item from the dataset.""" + sequence = self.fasta[self.seqids[idx]].sequence().upper() + tokenized_seq = self.tokenizer.text_to_ids(sequence) + if self.prepend_bos: # in pretraining we use EOS to start new sequences. + tokens: list[int] = [self.tokenizer.eod] + tokenized_seq + else: + tokens: list[int] = tokenized_seq + loss_mask = torch.ones_like(torch.tensor(tokens, dtype=torch.long), dtype=torch.long) + if self.prepend_bos: + loss_mask[0] = ( + 0 # mask the eos token which we use for causal offsetting. Later in predict we take the output + ) + # for the first [:-1] tokens which align with the sequence starting after the EOS. + return { + "tokens": torch.tensor(tokens, dtype=torch.long), + "position_ids": torch.arange(len(tokens), dtype=torch.long), + "seq_idx": torch.tensor(idx, dtype=torch.long), + "loss_mask": loss_mask, + } diff --git a/model/evo2/data/indexed_dataset.py b/model/evo2/data/indexed_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..f7ec34c892e881d45a558e0d7846e71855858326 --- /dev/null +++ b/model/evo2/data/indexed_dataset.py @@ -0,0 +1,73 @@ +"""Compatibility helpers for building Megatron indexed datasets. + +The legacy JSON preprocessing script expects ``indexed_dataset.make_builder``. +In the standalone package we keep that small API surface and delegate the +actual dataset writing to Megatron Core. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def _dtype_from_vocab_size(vocab_size: int | None) -> type[np.number]: + """Choose a compact dtype that can hold token ids for the vocabulary.""" + if vocab_size is None: + return np.int32 + if vocab_size <= np.iinfo(np.uint8).max + 1: + return np.uint8 + if vocab_size <= np.iinfo(np.uint16).max + 1: + return np.uint16 + return np.int32 + + +class _IndexedDatasetBuilderAdapter: + """Adapter matching the small builder API used by preprocess_data_json.""" + + def __init__(self, bin_path: str, dtype: type[np.number]) -> None: + from megatron.core.datasets.indexed_dataset import IndexedDatasetBuilder + + self.dtype = dtype + self._builder = IndexedDatasetBuilder(bin_path=bin_path, dtype=dtype) + + def add_item(self, item: Any) -> None: + import torch + + if isinstance(item, torch.Tensor): + tensor = item + else: + tensor = torch.as_tensor(np.asarray(item, dtype=self.dtype)) + self._builder.add_item(tensor) + + def end_document(self) -> None: + self._builder.end_document() + + def finalize(self, idx_path: str) -> None: + self._builder.finalize(idx_path) + + +def make_builder( + out_file: str, + impl: str = "mmap", + vocab_size: int | None = None, + dtype: type[np.number] | None = None, +) -> _IndexedDatasetBuilderAdapter: + """Create a Megatron indexed dataset builder. + + Args: + out_file: Output ``.bin`` file path. + impl: Dataset implementation. Megatron Core writes mmap datasets here. + vocab_size: Optional vocabulary size used to infer storage dtype. + dtype: Optional explicit NumPy dtype. + """ + if impl != "mmap": + raise ValueError( + "Only the 'mmap' dataset implementation is supported by this " + "standalone Evo2 compatibility layer." + ) + return _IndexedDatasetBuilderAdapter( + bin_path=out_file, + dtype=dtype or _dtype_from_vocab_size(vocab_size), + ) diff --git a/model/evo2/data/tokenizer.py b/model/evo2/data/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..7b86467dedd537382a0f92a6c5d3ecd3f6de1595 --- /dev/null +++ b/model/evo2/data/tokenizer.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import ftfy +from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec +from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer + +from evo2.utils.config import Evo2PreprocessingConfig + + +class Evo2Tokenizer: + """Tokenizer for Evo2.""" + + def __init__(self, params: Evo2PreprocessingConfig | None = None): + """Initialize the Evo2Tokenizer.""" + # Pass all NeMo2/Megatron-compliant parameters associated with config.Evo2PreprocessingConfig. + self.params: Evo2PreprocessingConfig = params if params is not None else Evo2PreprocessingConfig() + self.tokenizer: TokenizerSpec = get_nmt_tokenizer( + library=self.params.tokenizer_type.lower(), + vocab_file=str(self.params.vocab_file) if self.params.vocab_file is not None else None, + merges_file=str(self.params.merges_file) if self.params.merges_file is not None else None, + model_name=self.params.tokenizer_model_name, + tokenizer_model=self.params.pretrained_tokenizer_model, + special_tokens=self.params.special_tokens, + use_fast=self.params.fast_hf_tokenizer, + ) + + def tokenize( + self, + text: str | list[str], + use_ftfy: bool = False, + enforce_sample_length: None | int = None, + append_eod: bool = False, + drop_empty_sequences: bool = False, + ): + """Tokenize the input text data for Evo2.""" + if isinstance(text, str): + text = [text] + # Tokenize a document or batch of strings. + doc_ids = [] + for l, t in enumerate(text): + if use_ftfy: + t = ftfy.fix_text(t) + # Tokenize the string. + text_ids: list = self.tokenizer.text_to_ids(t) + if drop_empty_sequences and len(text_ids) == 0: + continue + # Append EOD token (EOD ID: 0) if appropriate. + eod_length = int(append_eod and l == len(text) - 1) + token_length = len(text_ids) + eod_length + text_ids += [0] * eod_length + if enforce_sample_length is not None: + # Pad shorter sequences (Pad ID: 1) and except excessive sequences. + if token_length > enforce_sample_length: + raise ValueError( + "Detected input text with a length greater than the maximum " + f"possible sample length of {enforce_sample_length}.)" + ) + else: + text_ids += [1] * (enforce_sample_length - token_length) + # Append to document. + doc_ids.append(text_ids) + return doc_ids diff --git a/model/evo2/data/transcript_extraction.py b/model/evo2/data/transcript_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..9790b4184f1fef6aafb3ec08730575f9d4d89381 --- /dev/null +++ b/model/evo2/data/transcript_extraction.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import argparse +import math +import re +import sys +from collections import defaultdict + +from nemo.utils import logging + +from bionemo.noodles import complement_sequence, reverse_sequence +from bionemo.noodles.nvfaidx import NvFaidx + + +def parse_gtf_attributes(attributes: str): + """Parses the attributes field of a GTF file line into a dictionary. + + Args: + attributes (str): The attributes field from a GTF file line. + + Returns: + dict: A dictionary of attribute key-value pairs. + """ + # Split on all semicolons that are not inside quotes + attributes = re.split(r';(?=(?:[^"]*"[^"]*")*[^"]*$)', attributes) + out = {} + for a in attributes: + if len(a) == 0: + continue + key = a.split()[0] + value = a.split('"')[1] + out[key] = value + return out + + +def extract_transcript_exons(gtf_path: str, only_longest_transcript: bool): + """Extracts transcript exons from a GTF file and optionally keeps only the longest transcript per gene. + + Args: + gtf_path (str): Path to the GTF file. + only_longest_transcript (bool): Whether to keep only the longest transcript per gene. + + Returns: + dict: A dictionary containing transcript and exon information. + """ + genes = defaultdict(set) + gene2transcripts = defaultdict(set) + transcripts = {} + exons = {} + exon2transcript = {} + transcript2gene = {} + transcript2exon = defaultdict(set) + skip_transcripts = set() + + gtf_fields = ["seqname", "source", "feature", "start", "end", "score", "strand", "frame", "attribute"] + with open(gtf_path) as infile: + for line in infile: + # skip header lines + if line.startswith("#"): + continue + line = line.strip().split("\t") + if len(line) < 9: + continue + + # parse the attributes into a dictionary + line = dict(zip(gtf_fields, line)) + attribs = parse_gtf_attributes(line["attribute"]) + + if line["feature"] == "gene": + contig, start, end, strand = line["seqname"], line["start"], line["end"], line["strand"] + start, end = int(line["start"]) - 1, int(line["end"]) + gene_id = attribs.get("gene_id", None) + if not gene_id: + continue + genes[gene_id].add((contig, start, end, strand)) + + elif line["feature"] == "exon": + contig, start, end, strand = line["seqname"], line["start"], line["end"], line["strand"] + start, end = int(line["start"]) - 1, int(line["end"]) + gene_id = attribs.get("gene_id", None) + if not gene_id: + continue + transcript_id = attribs["transcript_id"] + gene2transcripts[gene_id].add(transcript_id) + + # Skip exons that have already been handled and are likely errors + if transcript_id in skip_transcripts: + continue + exon_number = int(attribs["exon_number"]) + + exon_id = (gene_id, transcript_id, exon_number) + if exon_id in exons: + del exons[exon_id] + if transcript_id in transcripts: + del transcripts[transcript_id] + if transcript_id in transcript2exon: + del transcript2exon[transcript_id] + skip_transcripts.add(transcript_id) + continue + + exons[exon_id] = {"seqname": contig, "start": start, "end": end, "strand": strand} + if exon_id in exon2transcript: + raise Exception("Exon Already Exists in exon2transcript") + exon2transcript[exon_id] = transcript_id + transcript2exon[transcript_id].add(exon_id) + + elif line["feature"] == "transcript": + contig, start, end, strand = line["seqname"], line["start"], line["end"], line["strand"] + start, end = int(line["start"]) - 1, int(line["end"]) + gene_id = attribs.get("gene_id", None) + if not gene_id: + continue + gbkey = attribs["gbkey"] + transcript_biotype = attribs["transcript_biotype"] + transcript_id = attribs["transcript_id"] + if transcript_id in skip_transcripts: + continue + + transcripts[transcript_id] = { + "seqname": contig, + "start": start, + "end": end, + "strand": strand, + "gbkey": gbkey, + "transcript_biotype": transcript_biotype, + } + transcript2gene[transcript_id] = gene_id + gene2transcripts[gene_id].add(transcript_id) + + if only_longest_transcript: + transcript_lengths = defaultdict(int) + for exon in exons: + transcript_lengths[exon[1]] += exons[exon]["end"] - exons[exon]["start"] + + keep_transcripts = {} + keep_exons = {} + keep_exon2transcript = {} + keep_transcript2gene = {} + keep_transcript2exon = defaultdict(set) + keep_skip_transcripts = set() + + for gene in gene2transcripts: + this_transcripts = gene2transcripts[gene] + this_transcript_lengths = [(transcript, transcript_lengths[transcript]) for transcript in this_transcripts] + longest_transcript = max(this_transcript_lengths, key=lambda x: x[1])[0] + keep_transcripts[longest_transcript] = dict(transcripts[longest_transcript]) + for exon in transcript2exon[longest_transcript]: + keep_exons[exon] = dict(exons[exon]) + keep_exon2transcript[exon] = longest_transcript + keep_transcript2exon[longest_transcript].add(exon) + keep_transcript2gene[longest_transcript] = gene + + transcripts = keep_transcripts + exons = keep_exons + exon2transcript = keep_exon2transcript + transcript2gene = keep_transcript2gene + transcript2exon = keep_transcript2exon + skip_transcripts = keep_skip_transcripts + + return { + "transcripts": transcripts, + "exons": exons, + "exon2transcript": exon2transcript, + "transcript2gene": transcript2gene, + "transcript2exon": transcript2exon, + } + + +def extract_default_transcript_sequences(transcript_info, fasta_records, output_file): + """Extracts default transcript sequences from the provided transcript information and writes them to an output file. + + Args: + transcript_info (dict): Dictionary containing transcript and exon information. + fasta_records (NvFaidx): Indexed FASTA records. + output_file (TextIO): File object to write the output sequences. + """ + for transcript_id in transcript_info["transcripts"]: + gene_id = transcript_info["transcript2gene"][transcript_id] + this_exons = sorted(transcript_info["transcript2exon"][transcript_id], key=lambda x: x[-1]) + + seqname = None + exon_qc_failed = False + if len(this_exons) > 1: + for i in range(1, len(this_exons)): + this_exon = this_exons[i] + prev_exon = this_exons[i - 1] + this_coords = transcript_info["exons"][this_exon] + prev_coords = transcript_info["exons"][prev_exon] + if this_coords["strand"] != prev_coords["strand"]: + exon_qc_failed = True + if this_coords["strand"] == "+" and this_coords["start"] < prev_coords["start"]: + exon_qc_failed = True + if this_coords["strand"] == "-" and this_coords["start"] > prev_coords["start"]: + exon_qc_failed = True + if this_coords["seqname"] != prev_coords["seqname"]: + exon_qc_failed = True + + if exon_qc_failed: + continue + + transcript_seq = "" + for exon in this_exons: + coords = transcript_info["exons"][exon] + if seqname is None: + seqname = coords["seqname"] + exon_seq = str(fasta_records[coords["seqname"]][coords["start"] : coords["end"]]) + if coords["strand"] == "-": + exon_seq = reverse_sequence(complement_sequence(exon_seq)) + transcript_seq += exon_seq + + print(f">{seqname}|{gene_id}|{transcript_id}\n{transcript_seq}", file=output_file) + + +def extract_stitched_transcript_sequences( + transcript_info, fasta_records, output_file, stitch_token="@", promoter_size=1024, intron_window=32, overlap=False +): + """Extracts stitched transcript sequences from the provided transcript information and writes them to an output file. + + The "stitched" word refers to the process of combining sequences from different regions of the genome to form a single, + continuous transcript sequence. + This includes: + Promoter Region: A specified number of base pairs (bp) upstream of the transcript start site. + Exons: The coding regions of the transcript. + Intron Windows: A specified number of bp from the neighboring introns around each exon. + + The stitch_token is used to denote the boundaries between + these regions in the stitched transcript sequences. + + Args: + transcript_info (dict): Dictionary containing transcript and exon information. + fasta_records (NvFaidx): Indexed FASTA records. + output_file (TextIO): File object to write the output sequences. + stitch_token (str, optional): Token to use for stitching sequences. Defaults to "@". + promoter_size (int, optional): Number of bp to include in the promoter region. Defaults to 1024. + intron_window (int, optional): Number of bp to include from neighboring introns. Defaults to 32. + overlap (bool, optional): Whether to allow overlap of neighboring intron windows. Defaults to False. + """ + for transcript_id in transcript_info["transcripts"]: + gene_id = transcript_info["transcript2gene"][transcript_id] + this_exons = sorted(transcript_info["transcript2exon"][transcript_id], key=lambda x: x[-1]) + + exon_qc_failed = False + if len(this_exons) > 1: + for i in range(1, len(this_exons)): + this_exon = this_exons[i] + prev_exon = this_exons[i - 1] + this_coords = transcript_info["exons"][this_exon] + prev_coords = transcript_info["exons"][prev_exon] + if this_coords["strand"] != prev_coords["strand"]: + exon_qc_failed = True + if this_coords["strand"] == "+" and this_coords["start"] < prev_coords["start"]: + exon_qc_failed = True + if this_coords["strand"] == "-" and this_coords["start"] > prev_coords["start"]: + exon_qc_failed = True + if this_coords["seqname"] != prev_coords["seqname"]: + exon_qc_failed = True + + if exon_qc_failed: + continue + + transcript_seq = "" + seqname = None + for i in range(len(this_exons)): + # Previous Exon + prev_exon = this_exons[i - 1] if i > 0 else None + prev_coords = transcript_info["exons"].get(prev_exon, None) + # Current Exon + cur_exon = this_exons[i] + cur_coords = transcript_info["exons"].get(cur_exon, None) + exon_number = cur_exon[-1] + if seqname is None: + seqname = cur_coords["seqname"] + # Next Exon + next_exon = this_exons[i + 1] if i < len(this_exons) - 1 else None + next_coords = transcript_info["exons"].get(next_exon, None) + # Extract the stitched spliced sequence without overlapping intron windows. + intron_window_left = ( + min(intron_window, math.floor(abs(cur_coords["start"] - prev_coords["end"]) / 2)) + if not overlap and prev_coords is not None + else intron_window + ) + intron_window_right = ( + min(intron_window, math.ceil(abs(next_coords["start"] - cur_coords["end"]) / 2)) + if not overlap and next_coords is not None + else intron_window + ) + if cur_coords["strand"] == "+" and exon_number == 1: + exon_start = cur_coords["start"] - promoter_size + exon_end = cur_coords["end"] + intron_window_right + elif cur_coords["strand"] == "-" and exon_number == 1: + exon_start = cur_coords["start"] - intron_window_left + exon_end = cur_coords["end"] + promoter_size + else: + exon_start = cur_coords["start"] - intron_window_left + exon_end = cur_coords["end"] + intron_window_right + exon_seq = str(fasta_records[cur_coords["seqname"]][exon_start:exon_end]) + if cur_coords["strand"] == "-": + exon_seq = stitch_token + reverse_sequence(complement_sequence(exon_seq)) + transcript_seq += exon_seq + + if stitch_token and len(stitch_token) > 0: + transcript_seq = transcript_seq[len(stitch_token) :] + + print(f">{seqname}|{gene_id}|{transcript_id}\n{transcript_seq}", file=output_file) + + +def run(args): + """Main function to run the transcript extraction process based on command line arguments. + + Args: + args (argparse.Namespace): Parsed command line arguments. + """ + with open(args.output_path, "w") if args.output_path is not None else sys.stdout as output_file: + if args.verbose: + logging.info("Indexing FASTA file...") + + fasta_index = NvFaidx(args.fasta_path) + + if args.transcript_type == "default": + if args.verbose: + logging.info("Extracting default transcripts...") + if args.only_longest_transcript: + logging.info("Only extracting the longest transcript per gene.") + else: + logging.info("Extracting all transcripts regardless of length.") + + elif args.transcript_type == "stitched": + if args.verbose: + logging.info("Extracting stitched transcripts...") + if args.only_longest_transcript: + logging.info("Only extracting the longest transcript per gene.") + else: + logging.info("Extracting all transcripts regardless of length.") + + transcript_info = extract_transcript_exons(args.gtf_path, args.only_longest_transcript) + + if args.transcript_type == "default": + extract_default_transcript_sequences(transcript_info, fasta_index, output_file) + elif args.transcript_type == "stitched": + extract_stitched_transcript_sequences( + transcript_info, + fasta_index, + output_file, + promoter_size=args.stitched_promoter, + intron_window=args.stitched_intron, + overlap=args.stitched_overlap, + ) + + +def parse_args(): + """Parses command line arguments for the transcript extraction script. + + Returns: + argparse.Namespace: Parsed command line arguments. + """ + ap = argparse.ArgumentParser(description="Extract spliced transcripts from a FASTA and GTF.") + ap.add_argument("--fasta-path", type=str, required=True, help="Path to FASTA file to extract transcripts from.") + ap.add_argument( + "--gtf-path", + type=str, + required=True, + help="Path to gene transfer format (GTF) file associated with the FASTA.", + ) + ap.add_argument("--output-path", type=str, default=None, help="Path to output FASTA file.") + ap.add_argument( + "--transcript-type", + type=str, + default="default", + choices=["default", "stitched"], + help="Type of transcript to extract from the GTF and FASTA files for splicing. 'Stitched' transcripts include 1024 bp of sequence from the promoter and 32 bp around each exon.", + ) + ap.add_argument( + "--stitched-promoter", + type=int, + default=1024, + help="Number of bp to include in the promoter region when --transcript-type=stitched is used. Defaults to 1024.", + ) + ap.add_argument( + "--stitched-intron", + type=int, + default=32, + help="Number of bp to include from neighboring introns when --transcript-type=stitched is used. Defaults to 32.", + ) + ap.add_argument( + "--stitched-overlap", + action="store_true", + help="Allow overlap of neighboring intron windows when --transcript-type=stitched is used. Defaults to False, i.e. prevents overlap by shortening the intron windows for a contiguous splice.", + ) + ap.add_argument( + "--only-longest-transcript", action="store_true", help="Only extract the longest transcript per gene." + ) + ap.add_argument("-v", "--verbose", action="store_true", help="Turn on verbose log messages.") + return ap.parse_args() + + +def main(): + """Entry point for the script. Parses arguments and runs the extraction process.""" + args = parse_args() + if args.verbose: + logging.info(args) + run(args) + + +if __name__ == "__main__": + main() diff --git a/model/evo2/lightning.py b/model/evo2/lightning.py new file mode 100644 index 0000000000000000000000000000000000000000..7cb4a7778b8e9833d2986e09fe5e47164db2ebe7 --- /dev/null +++ b/model/evo2/lightning.py @@ -0,0 +1,495 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Callable, Generic, Iterable, Iterator, List, Literal, Optional, Sequence, Tuple, TypeVar, Union, Protocol, Type +from abc import ABC, abstractmethod + +import lightning.pytorch as pl +import torch.distributed +from megatron.core import parallel_state +from megatron.core.optimizer.optimizer_config import OptimizerConfig +from nemo.lightning import io as nlio +from nemo.lightning.megatron_parallel import ( + DataT, + MegatronLossReduction, + ReductionT, +) +from nemo.lightning.pytorch.optim import MegatronOptimizerModule +from torch import Tensor + +# from bionemo.core.model.config import BionemoTrainableModelConfig +from evo2.api import MegatronLossType, MegatronModelType +# from bionemo.llm.api import MegatronLossType, MegatronModelType + +__all__: Sequence[str] = ( + "BionemoLightningModule", + "LightningPassthroughPredictionMixin", + "PassthroughLossReduction", + "batch_collator", + "default_megatron_optimizer", + "get_dtype_device", +) + + +T = TypeVar("T") +BatchT = TypeVar("BatchT") + + +ModelOutput = TypeVar("ModelOutput", Tensor, list[Tensor], tuple[Tensor], dict[str, Tensor], covariant=True) +"""A Model's forward pass may produce a tensor, multiple tensors, or named tensors. +""" + +LossType = TypeVar("LossType") +"""Stand-in for a loss function; no constraints. +""" + + +class Model(Protocol[ModelOutput]): + """Lightweight interface for a model: must have a forward method.""" + + def forward(self, *args, **kwargs) -> ModelOutput: + """Prediction / forward-step for a model.""" + ... + + +ModelType = TypeVar("ModelType", bound=Model) +"""Generic type for things that have a forward pass. +""" + + +class BionemoModelConfig(Generic[ModelType], ABC): + """An abstract class for model configuration.""" + + @abstractmethod + def configure_model(self, *args, **kwargs) -> Model: + """Configures the model.""" + raise NotImplementedError() + + +class BionemoTrainableModelConfig(Generic[ModelType, LossType], BionemoModelConfig[ModelType]): + """An abstract class for trainable model configuration.""" + + @abstractmethod + def get_loss_reduction_class(self) -> Type[LossType]: + """Returns the loss reduction class.""" + raise NotImplementedError() + + + + +def some_first(seq: Iterable[Optional[T]]) -> T: + """Returns the first non-None value from the sequence or fails""" # noqa: D415 + for s in seq: + if s is not None: + return s + raise ValueError("non-None value not found") + + +def get_dtype_device(torch_object) -> Tuple[torch.dtype, torch.device]: # noqa: D103 + match torch_object: + case []: + raise ValueError("Looking up dtype on an empty list") + case {**data} if not data: + raise ValueError("Looking up dtype on an empty dict") + case Tensor(dtype=dtype, device=device): + return dtype, device + case torch.nn.Module() as m: + try: + p = next(m.parameters()) + except StopIteration as e: + raise ValueError("Cannot get dtype on a torch module with no parameters.") from e + return p.dtype, p.device + case dict(keys=_, values=values): + val = some_first(values()) + return get_dtype_device(val) + case list() as l: + val = some_first(l) + return get_dtype_device(val) + case _: + raise TypeError("Got something we didnt expect") + + +# NOTE(SKH): These types are all wrong, but are close. The inner type must always be a Tensor, but the outer container should be generic. +def batch_collator( + batches: Optional[Union[Tuple[ReductionT], List[ReductionT]]], + batch_dim: int = 0, + seq_dim: int = 1, + batch_dim_key_defaults: dict[str, int] = {"token_logits": 1}, + seq_dim_key_defaults: dict[str, int] = {"token_logits": 0}, + preferred_gpu: int = 0, +) -> Optional[ReductionT]: + """Takes a sequence of batches and collates them into a single batch. + + This is distinct from the standard pytorch default_collator since it does + not add the batch dimension, it's assumed the batch + dimension is already present in the input, as would be the case when + parallelizing across minibatches. + + IMPORTANT: The underlying data primitive _must_ be a torch Tensor. The input to this function is a recurisve type, + there can be any amount of nesting between dictionaries, tuples, and lists, as long as the inner type is a n-d Tensor. + + Examples: + Outer container = Dict: + [{'a': Tensor([1]), 'b': Tensor([2])}, {'a': Tensor([2]), 'b': Tensor([3])}] -> {'a': Tensor([1, 2]), 'b': Tensor([2, 3])} + Outer container = List: + [[Tensor([1]), Tensor([2])], [Tensor([2]), Tensor([3])]] -> [Tensor([1, 2]), Tensor([2, 3])] + Outer container = Tuple: + ([Tensor([1]), Tensor([2])], [Tensor([2]), Tensor([3])]) -> (Tensor([1, 2]), Tensor([2, 3])) + + Args: + batches (Optional[Sequence[ReductionT]]): sequence of batches to collate into a single batch. + batch_dim: If you know that the batch dim for the batch you are concatenating is not the 0th dimension (for + example it is sequence first) then supply that dimension. + seq_dim: If you know that the sequence dim for the batch you are concatenating is not the 1st dimension (for + example it is sequence first) then supply that dimension. This is used for padding to the max length. + batch_dim_key_defaults (dictionary of keys to integers): If your batch is a dictionary and you know that some + keys have non-standard (0) batch dimensions, supply those here. By default "token_logits" has batch dim 1 + and otherwise all keys are assumed to have batch dim 0. + seq_dim_key_defaults (dictionary of keys to integers): If your batch is a dictionary and you know that some + keys have non-standard (1) sequence dimensions, supply those here. By default "token_logits" has seq dim 0 + and otherwise all keys are assumed to have seq dim 1. + preferred_gpu: If any of the tensors are on any GPU, all of them will be moved to this GPU. 0 by default. + + Returns: + A single batch of the same type as the elements of your input sequence. + """ + match batches: + # Handle base-cases for batch concatenation, either a list of None or a list of tensors + case [None, *_]: + return None + case [Tensor(), *_]: + # If any tensor is on a GPU, move all to preferred GPU + if any(t.is_cuda for t in batches): + device = torch.device(f"cuda:{preferred_gpu}") + batches = [t.to(device) for t in batches] + # First shortcut if all tensors are 1D (they have at least one batch dim, and it must be at 0) + if len(batches) > 0 and isinstance(batches[0], Tensor) and batches[0].ndim == 1: + return torch.cat(batches, dim=0) + # Find max sequence length across all tensors + max_seq_len = max(batch.size(seq_dim) for batch in batches) + # Pad each tensor to max length along seq_dim + padded_batches = [] + for batch in batches: + # Initialize padding tuple - needs 2 values per dim, starting from last dim + # e.g. for 3D tensor: [left_pad_dim2, right_pad_dim2, left_pad_dim1, right_pad_dim1, left_pad_dim0, right_pad_dim0] + pad_size = [0] * (2 * batch.ndim) + # Calculate padding needed at end of sequence dimension + pad_amount = max_seq_len - batch.size(seq_dim) + # Pad end of sequence dimension by putting padding amount in correct position + # For seq_dim=1 in 3D tensor: [0, 0, 0, pad_amount, 0, 0] + pad_size[2 * (batch.ndim - 1 - seq_dim) + 1] = pad_amount + padded_batch = torch.nn.functional.pad(batch, tuple(pad_size)) + padded_batches.append(padded_batch) + padded_batch = torch.cat(padded_batches, dim=batch_dim) + assert padded_batch.size(seq_dim) == max_seq_len + return padded_batch + # Next 3 calls are the recursive calls into the sub-structures of the batch. We handle dictionaries, tuples, and lists + case [dict(), *_]: + return { + key: batch_collator( + [batch[key] for batch in batches], + batch_dim=batch_dim_key_defaults.get(key, batch_dim), + seq_dim=seq_dim_key_defaults.get(key, seq_dim), + batch_dim_key_defaults=batch_dim_key_defaults, + seq_dim_key_defaults=seq_dim_key_defaults, + preferred_gpu=preferred_gpu, + ) + for key in batches[0] + } + case [tuple(), *_]: + return tuple( + batch_collator( + [batch[i] for batch in batches], + batch_dim=batch_dim, + seq_dim=seq_dim, + batch_dim_key_defaults=batch_dim_key_defaults, + seq_dim_key_defaults=seq_dim_key_defaults, + preferred_gpu=preferred_gpu, + ) + for i in range(len(batches[0])) + ) + case [list(), *_]: + return [ + batch_collator( + [batch[i] for batch in batches], + batch_dim=batch_dim, + seq_dim=seq_dim, + batch_dim_key_defaults=batch_dim_key_defaults, + seq_dim_key_defaults=seq_dim_key_defaults, + preferred_gpu=preferred_gpu, + ) + for i in range(len(batches[0])) + ] + # Final cases shouldn't happen, an empty sequence (no batches), or "other". + case []: + raise ValueError("Cannot process an empty sequence") + case _: + raise ValueError("Unsupported input structure in batch_collator") + + +# TODO(@jstjohn): Properly use the Generic for DataT and ReductionT usage. Define our own batch/output types. +# TODO(@skothenhill): Re-think the generics here- the way that `batch_collator` is expressed, `batches` should be a recursive generic type. +class PassthroughLossReduction(MegatronLossReduction, Generic[DataT]): + """A workaround for nemo/megatron to perform inference. + + Internally in NeMo2.0 the forward step is always expected to return a loss reduction class, and forward is + expected to return a loss. This class hijacks that mechanism to instead pass through the forward output unperturbed + as the loss (to enable inference in the predict step), and then the reduce method is used to collate the batch of + forward outputs into a single batch. This supports the model forward output being a tensor, dict, tuple, or list of + tensors. The inner type _must always be a Tensor_. + """ + + def forward(self, batch: DataT, forward_out: DataT) -> Tuple[Tensor, DataT]: + """Passes through the `forward_out` value as the 2nd tuple element. + + Args: + batch: The batch of data that was passed through the model to generate output. NOTE: this value is ignored. + forward_out: The output from your model's forward pass. + + Returns: + A tuple containing the loss tensor (dummy in this case) and the forward output (unmodified). + """ + return torch.zeros((1, 1)), forward_out + + def reduce(self, forward_out: List[DataT]) -> DataT: + """Collates list of model's outputs into a single output.""" + return batch_collator(forward_out) + + +class LightningPassthroughPredictionMixin: + """A mixin that allows your model to do inference on the predict step by hijacking nemo's loss reduction mechanism.""" + + def predict_loss_reduction(self) -> PassthroughLossReduction: + """For the predict step, pass through the forward pass output.""" + return PassthroughLossReduction() + + +ForwardStep = Callable[[MegatronModelType, DataT], DataT] +"""Megatron-compatible forward pass function. +""" + +DataStep = Callable[[Iterator[DataT]], DataT] +"""Batches together an iterator of individual examples. + +Necessary for compatability with Megatron. This function type is similiar to the collate function of PyTorch. + +A `DataStep` function takes an iterator over individual examples. Each example may be a tensor, sequence of tensors, +or a set of named tensors (provided as a `dict` mapping `str` names to each `Tensor`). Each iteration must +yield the same type. + +The output of this function will mirror the same structure of each yielded example. It will be a concatenation of all +of the examples in the iterator. +""" + + +class BionemoLightningModule( + Generic[MegatronModelType, MegatronLossType], + pl.LightningModule, + nlio.IOMixin, + nlio.ConnectorMixin, + LightningPassthroughPredictionMixin, +): + """Reusable PyTorch Lightning module for Megatron models that is compatible with NeMo's conventions.""" + + def __init__( + self, + config: BionemoTrainableModelConfig[MegatronModelType, MegatronLossType], + forward_step: ForwardStep, + data_step: DataStep, + optimizer: MegatronOptimizerModule, + model_transform: Optional[Callable[[MegatronModelType], MegatronModelType]] = None, + configure_init_model_parallel: bool = False, + **model_construct_args, + ) -> None: + """Constructor. + + Args: + config: Serializable configuration object that allows one to construct a new model instance and loss + function. Necessary for Megatron-based training as the model itself cannot be serialized and + distributed to nodes. Instead, we serialize the procedure for making the model and distribute that. + forward_step: Performs forward pass using the model and a batch of data. + data_step: Custom batch-creating function for the model. + optimizer: Megatron-compatible distributed optimizer instance. Defaults to using ADAM with a 1e-4 learning + rate. + model_construct_args: Optional. Any arguments necessary to construct the model in the `config`'s + `configure_model` method. + model_transform: Optional. The model transform function. + configure_init_model_parallel: Optional. Whether to initialize the model parallel at configuration time. + **model_construct_args: Optional. Arguments necessary for the supplied model configuration's + `configure_model` method, which will make an instance of the model. + """ + super().__init__() + self.config = config + self.module_construct_args: Optional[dict[str, Any]] = model_construct_args + # ***must** be set up in configure_model() -- megatron constraint + # also, must be called `module`: nemo expects the actual model to be stored this way + self.module: Optional[MegatronModelType] = None + self.loss_reduction_class: type[MegatronLossType] = config.get_loss_reduction_class() + self.optim = optimizer + self.optim.connect(self) # This will bind the `configure_optimizers` method + self._data_step = data_step + self._forward_step = forward_step + self.model_transform = model_transform + self.configure_init_model_parallel = configure_init_model_parallel + # configure metrics + self.train_metric = self.config.train_metric.get_instance() if self.config.train_metric else None + self.valid_metric = self.config.valid_metric.get_instance() if self.config.valid_metric else None + + def configure_model(self) -> None: + """Updates internal state: instantiates the model from the object's config, assigns to `model` attribute. + + NOTE: this method is idempotent; successive calls have no effect. The model is only initialized once. + + Raises: + ValueError iff the internal config's configure_model method returns None. + """ + if self.configure_init_model_parallel: + self.trainer.strategy._init_model_parallel = True + if self.module is None: + if self.module_construct_args is None: + module_construct_args = {} + elif "model_construct_args" in self.module_construct_args: + # Not sure why this is needed, but it seems "model_construct_args" ends up as a key inside this dict. + module_construct_args = self.module_construct_args["model_construct_args"] + else: + module_construct_args = self.module_construct_args + + model: MegatronModelType = self.config.configure_model(**module_construct_args) + self.module = model + if self.module is None: + raise ValueError("Invalid semantics: configure_model method **MUST** initialize the model.") + + def is_on_logging_device(self): + """Return True if last stage of pipeline parallel and first tensor parallel rank.""" + return parallel_state.is_pipeline_last_stage() and parallel_state.get_tensor_model_parallel_rank() == 0 + + def forward(self, *args, **kwargs) -> DataT: + """Call the forward method of the underlying model, and return whatever it outputs.""" + # safe to do because configure_model is idempotent + self.configure_model() + assert self.module is not None, "ERROR: configure_model() method has been incorrectly overridden!" + prediction = self.module(*args, **kwargs) # for now just pass through to the underlying model + return prediction + + def data_step(self, dataloader_iter: Iterator[DataT]) -> DataT: # noqa: D102 + return self._data_step(dataloader_iter) + + def forward_step(self, batch) -> Tensor: + """Megatron-required: the training forward step for the model, which is required to produce the loss. + + Normally, the forward pass of a model means its inference. Loss is computed using the predictions + from the forward pass against labels. Megatron unfortunately conflates these two different concepts + and instead has models "forward" method produce the loss. See the Megatron docs for details: + https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/pipeline_parallel/schedules.py#L170 + + To get actual predictions, use the :func:`forward` method instead. + """ + # safe to do because configure_model is idempotent + self.configure_model() + assert self.module is not None + return self._forward_step(self.module, batch) + + def update_metric( + self, batch, outputs, metric, task: Literal["pretraining", "classification", "regression"] + ) -> None: + """Update metric for logging.""" + match task: + case "pretraining": + logits = outputs["token_logits"].detach().transpose(0, 1) # [s, b, v] -> [b, s, v] + metric(logits, batch["labels"]) + case "classification": + classification_output = outputs["classification_output"] + num_classes = classification_output.shape[-1] + labels = batch["labels"] + if classification_output.ndim == 3: # token-level classification + classification_output = classification_output.reshape(-1, num_classes)[ + batch["loss_mask"].view(-1) + ] # shape [-1, num_classes] + assert classification_output.ndim == 2 + + labels = batch["labels"].reshape(-1)[batch["loss_mask"].view(-1)] + metric( + classification_output.reshape(-1, num_classes), + labels.reshape(-1), + ) + case "regression": + regression_output = outputs["regression_output"] + metric(regression_output, batch["labels"]) + case _: + raise NotImplementedError(f"unrecognized task {task}") + + def training_step(self, batch, batch_idx: Optional[int] = None) -> Tensor: + """In mcore the loss-function is part of the forward-pass when labels are provided.""" + outputs = self.forward_step(batch) + if self.train_metric is not None: + if self.is_on_logging_device(): + self.update_metric(batch, outputs, self.train_metric, self.config.train_metric.task) + + self.log( + self.config.train_metric.metric_name, + self.train_metric, + on_step=True, + on_epoch=False, + prog_bar=True, + ) + + return outputs + + def validation_step(self, batch, batch_idx: Optional[int] = None) -> Tensor: + """In mcore the loss-function is part of the forward-pass when labels are provided.""" + outputs = self.forward_step(batch) + if self.valid_metric is not None and self.is_on_logging_device(): + self.update_metric(batch, outputs, self.valid_metric, self.config.valid_metric.task) + + return outputs + + def predict_step(self, batch, batch_idx: Optional[int] = None) -> Tensor: + """Alias for forward_step.""" + if len(batch) == 0: + return + return self.forward_step(batch) + + def training_loss_reduction(self) -> MegatronLossType: + """This is the function that takes batch['loss_mask'] and the logits output by the model and reduces the loss.""" + return self.loss_reduction_class() + + def validation_loss_reduction(self) -> MegatronLossType: # noqa: D102 + return self.loss_reduction_class(validation_step=True) + + def test_loss_reduction(self) -> MegatronLossType: # noqa: D102 + return self.loss_reduction_class(validation_step=True) + + def on_validation_epoch_end(self): # noqa: D102 + if self.valid_metric is None: + return + + if self.trainer.sanity_checking: + self.valid_metric.reset() # clean up sanity runs + return + + self.log( + self.config.valid_metric.metric_name, + self.valid_metric, + on_step=False, + on_epoch=True, + prog_bar=True, + ) + + +def default_megatron_optimizer() -> MegatronOptimizerModule: + """Default distributed optimizer uses Adam with a 1e-4 learning rate.""" + return MegatronOptimizerModule( + config=OptimizerConfig(lr=1e-4, optimizer="adam", use_distributed_optimizer=True), + ) diff --git a/model/evo2/models/__init__.py b/model/evo2/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8a80a47c85a0e874f79507f7db35153165f2f7fb --- /dev/null +++ b/model/evo2/models/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from evo2.models.mamba import MAMBA_MODEL_OPTIONS, HybridMambaConfig8BEvo2Loss, MambaModel + + +__all__ = ["MAMBA_MODEL_OPTIONS", "HybridMambaConfig8BEvo2Loss", "MambaModel"] diff --git a/model/evo2/models/mamba.py b/model/evo2/models/mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..872ce767201813d93a8de86a265311a537bfd5bd --- /dev/null +++ b/model/evo2/models/mamba.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from dataclasses import dataclass +from typing import Callable + +import megatron.core.models.mamba.mamba_model +import torch +import torch.nn.functional as F +from megatron.core import parallel_state +from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import GPTInferenceWrapper +from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import InferenceWrapperConfig +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.utils import WrappedTensor, deprecate_inference_params +from nemo.collections.llm.gpt.model.base import GPTModel, gpt_data_step +from nemo.collections.llm.gpt.model.megatron.hyena.hyena_utils import make_upper_case, reweighted_cross_entropy +from nemo.collections.llm.gpt.model.ssm import ( + NemotronHConfigBase, +) +from nemo.lightning import get_vocab_size +from typing_extensions import override + +from evo2.utils.loss.embedding_variance import SquaredErrorTargetedVarianceLoss + + +logger = logging.getLogger(__name__) + + +def mamba_forward_step(model, batch) -> torch.Tensor: + """Forward step function for Mamba models, similar to hyena_forward_step. + + Args: + model: The Mamba model + batch: Dictionary containing input batch data + + Returns: + torch.Tensor: Output from the model forward pass + """ + forward_args = { + "input_ids": batch["tokens"], + "position_ids": batch["position_ids"], + "labels": batch["labels"], + "loss_mask": batch["loss_mask"], + } + forward_args["attention_mask"] = None + return model(**forward_args) + + +class MambaModel(GPTModel): + """Mamba model that extends GPTModel for integration with NeMo. + + Note that the loss calculation is handled by CustomMCoreMambaModel instead. + """ + + @override + def get_inference_wrapper( + self, params_dtype, inference_batch_times_seqlen_threshold, inference_max_seq_length=8192 + ) -> GPTInferenceWrapper: + """Gets the inference wrapper for the Mamba model.""" + from megatron.core.models.mamba import MambaModel as MCoreMambaModel + + # Find MCoreMambaModel instance + mcore_model = self.module + while mcore_model: + if isinstance(mcore_model, (MCoreMambaModel, Evo2StyleMCoreMambaModel)): + break + mcore_model = getattr(mcore_model, "module", None) + if mcore_model is None or not isinstance(mcore_model, (MCoreMambaModel, Evo2StyleMCoreMambaModel)): + raise ValueError("Mamba model instance not found in the model structure.") + + vocab_size = None + if self.tokenizer is not None: + vocab_size = self.tokenizer.vocab_size + elif hasattr(self.config, "vocab_size"): + vocab_size = self.config.vocab_size + else: + raise ValueError("Unable to find vocab size.") + + inference_wrapper_config = InferenceWrapperConfig( + hidden_size=mcore_model.config.hidden_size, + params_dtype=params_dtype, + inference_batch_times_seqlen_threshold=inference_batch_times_seqlen_threshold, + padded_vocab_size=vocab_size, + inference_max_seq_length=inference_max_seq_length, + ) + + model_inference_wrapper = GPTInferenceWrapper(mcore_model, inference_wrapper_config) + return model_inference_wrapper + + @override + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + decoder_input: torch.Tensor | None = None, + inference_context=None, + packed_seq_params=None, + inference_params=None, + runtime_gather_output: bool | None = None, + loss_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward pass that delegates to CustomMCoreMambaModel, which handles loss calculation.""" + extra_kwargs = {"packed_seq_params": packed_seq_params} if packed_seq_params is not None else {} + output_tensor = self.module( + input_ids, + position_ids, + attention_mask, + decoder_input=decoder_input, + labels=labels, # Pass labels to the Megatron module + inference_params=inference_params, + inference_context=inference_context, + runtime_gather_output=runtime_gather_output, + loss_mask=loss_mask, # Pass loss_mask to the Megatron module + **extra_kwargs, + ) + + # Return whatever CustomMCoreMambaModel.forward returns + # (logits during inference, loss during training) + return output_tensor + + +# Custom MCoreMambaModel with reweighted loss calculation +class Evo2StyleMCoreMambaModel(megatron.core.models.mamba.mamba_model.MambaModel): + """Custom version of MCoreMambaModel that implements reweighted loss calculation. + + Note that this is similar to the HyenaModel for uppercase/lowercase handling. + """ + + def __init__(self, *args, **kwargs): + """Initializes `Evo2StyleMCoreMambaModel` with unique parameters for the Evo2 variant of `MCoreMambaModel`.""" + super().__init__(*args, **kwargs) + if self.config.use_targeted_variance_loss: + if not hasattr(self.config, "embedding_init_method_std"): + logger.warning("embedding_init_method_std is not supported in this config, please upgrade Megatron-LM") + # 1.0 is the suggested value for embedding_init_method_std from the + # [Spike No More](https://arxiv.org/abs/2312.16903) paper. + embedding_init_method_std: float = getattr(self.config, "embedding_init_method_std", 1.0) + self.targeted_variance_loss = SquaredErrorTargetedVarianceLoss( + loss_coeff=self.config.targeted_variance_loss_loss_coeff, + var_target=embedding_init_method_std**2, + ) + + @override + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + decoder_input: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + inference_context=None, + runtime_gather_output: bool | None = None, + *, + inference_params=None, + loss_mask: torch.Tensor | None = None, + ): + """Forward pass with custom loss calculation for uppercase/lowercase reweighting. + + Note that this mimics the behavior in hyena_model.py lines 273-292. + + Forward function of the Mamba model. This function passes the input tensors + through the embedding layer, and then the decoder and finally into the post + processing layer (optional). + + It either returns the Loss values if labels are given or the final hidden units + """ + # If decoder_input is provided (not None), then input_ids and position_ids are ignored. + # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input. + + inference_context = deprecate_inference_params(inference_context, inference_params) + + # Decoder embedding. + if decoder_input is not None: + pass + elif self.pre_process: + decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) + else: + # intermediate stage of pipeline + # decoder will get hidden_states from encoder.input_tensor + decoder_input = None + + rotary_pos_emb = None + if self.position_embedding_type == "rope": + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, self.decoder, decoder_input, self.config + ) + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len) + + # Wrap decoder_input to allow the decoder (MambaBlock) to delete the + # reference held by this caller function, enabling early garbage collection + # for inference. + if inference_context is not None and not self.training: + decoder_input = WrappedTensor(decoder_input) + + # The following assert will currently fail when running inference. + # Commented out for now. + # TODO (duncan/rwaleffe): (1) confirm that the externally-generated + # attention mask is not needed and is ignored by the model in + # inference mode, (2) reduce the size of the externally-generated + # attention mask to prevent CPU OOM (as we did for training), (3) + # force the attention mask passed to the model in inference mode to + # be None, so this assert will succeed. + # assert attention_mask is None, "The attention mask is ignored and should be set to None" + + # Run decoder. + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + ) + + if not self.post_process: + return hidden_states + + # logits and loss + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() + + if ( + not self.training + and inference_context is not None + and inference_context.materialize_only_last_token_logits + ): + hidden_states = hidden_states[-1, :, :].unsqueeze(0) + + logits, _ = self.output_layer(hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output) + + if labels is None: + # [s b h] => [b s h] + return logits.transpose(0, 1).contiguous() + + # Apply reweighted loss calculation for uppercase/lowercase handling + labels, lowercase_mask = make_upper_case(labels) + loss = self.compute_language_model_loss(labels, logits) + normalize_per_batch = True if self.config.to_upper == "normalized_weighted" else False + loss = reweighted_cross_entropy( + loss, + (labels, loss_mask, lowercase_mask), + lowercase_weight=self.config.lowercase_loss_reweighting, + normalize_per_batch=normalize_per_batch, + ) + if self.training and self.config.use_targeted_variance_loss: + # Only use this in training, not validation etc. + var_loss = self.targeted_variance_loss(self.embedding.word_embeddings.weight) + loss += var_loss + return loss + + +def mamba_no_weight_decay_cond(name, param, exclude_embeddings: bool = False): + """Condition for no weight decay for Mamba parameters. + + Note that this follows the same pattern as in the original Mamba implementation. + """ + # Mamba-specific parameters that should not have weight decay + if ( + name.endswith("dt_bias") + or name.endswith("A_log") + or name.endswith("D") + or ("embedding" in name and exclude_embeddings) + or getattr(param, "_no_weight_decay", False) + ): + no_wd = True + # All other parameters - use default MCore behavior: + # Do not regularize biases and norm parameters + # (See megatron.core.optimizer._get_pram_groups) + # TODO exclude embeddings + else: + no_wd = name.endswith(".bias") or len(param.shape) == 1 + return no_wd + + +def mamba_no_weight_decay_cond_with_embeddings(name, param): + """Condition for no weight decay for Mamba parameters with embeddings. + + Note that this follows the same pattern as in the original Mamba implementation but also skips WD on embeddings. + """ + return mamba_no_weight_decay_cond(name, param, exclude_embeddings=True) + + +@dataclass +class HybridMambaConfig8BEvo2Loss(NemotronHConfigBase): + """Config for 8B hybrid Mamba model.""" + + hybrid_override_pattern: str = "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-" + num_layers: int = 52 + seq_length: int = 8192 + hidden_size: int = 4096 + mamba_ssm_ngroups: int = 8 + mamba_state_dim: int = 128 + mamba_head_dim: int = 64 + ffn_hidden_size: int = 21504 + num_attention_heads: int = 32 + init_method_std: float = 0.014 + num_query_groups: int = 8 + make_vocab_size_divisible_by: int = 128 + tokenizer_library: str = "byte-level" # Use Evo2 tokenizer + tokenizer_name: str = None + masked_softmax_fusion: bool = True + apply_query_key_layer_scaling: bool = False + persist_layer_norm: bool = True + attention_softmax_in_fp32: bool = False + vocab_size: int = 512 + first_last_layers_bf16: bool = True + is_hybrid_model: bool = True + forward_step_fn: Callable = mamba_forward_step + data_step_fn: Callable = gpt_data_step + # Set a reasonable default for to_upper to match HyenaModel behavior + to_upper: str = "normalized_weighted" + # Set lowercase loss reweighting factor + lowercase_loss_reweighting: float = 1.0 + activation_func: Callable = lambda x: torch.square(F.relu(x)) # lambda x: torch.pow(F.relu(x), 2) + # The trainer is responsible for using this when initializing the optimizer state: + # opt = MegatronOptimizerModule(opt_config, sched, no_weight_decay_cond=model_config.hyena_no_weight_decay_cond_fn) + hyena_no_weight_decay_cond_fn: Callable = mamba_no_weight_decay_cond + spike_no_more_embedding_init: bool = False # TODO: remove this. + layernorm_embeddings: bool = False + # If set to true, use targeted variance loss which encourages the word embedding weight variances + # to be close to a target value (1.0). + use_targeted_variance_loss: bool = False + targeted_variance_loss_loss_coeff: float = 0.1 + + def __post_init__(self): + """Post-init logic for Evo2 to enable backwards compatibility with old configs.""" + # Specific post_init logic for Evo2 to enable backwards compatibility with old configs. + if not hasattr(self, "embedding_init_method_std"): + raise ValueError("embedding_init_method_std is not supported in this config, please upgrade Megatron-LM") + if self.spike_no_more_embedding_init and self.embedding_init_method_std is None: + logger.warning( + "spike_no_more_embedding_init is deprecated, please set " + "embedding_init_method_std=[desired_stdev] in the future. To get the old behavior set to 1.0. " + "For now setting to 1.0." + ) + self.embedding_init_method_std = 1.0 + # Continue with the remaining post-init logic defined in NemotronHConfigBase and/or TransformerConfig. + super().__post_init__() + + @override + def configure_model( + self, tokenizer, pre_process=None, post_process=None, vp_stage: int | None = None + ) -> Evo2StyleMCoreMambaModel: + """Configures the model for training or inference.""" + mamba_stack_spec = self.mamba_stack_spec + if not isinstance(mamba_stack_spec, ModuleSpec): + mamba_stack_spec = mamba_stack_spec() + + assert getattr(self, "virtual_pipeline_model_parallel_size", None) is None and vp_stage is None, ( + "Virtual pipeline model parallelism is temporarily unsupported in SSM/Mamaba " + "models due to upstream MCore MambaModel API dependency" + ) + return Evo2StyleMCoreMambaModel( + self, + mamba_stack_spec=mamba_stack_spec, + vocab_size=get_vocab_size(self, tokenizer.vocab_size, self.make_vocab_size_divisible_by), + max_sequence_length=self.seq_length, + hybrid_attention_ratio=self.hybrid_attention_ratio, + hybrid_mlp_ratio=self.hybrid_mlp_ratio, + hybrid_override_pattern=self.hybrid_override_pattern, + position_embedding_type=self.position_embedding_type, + rotary_percent=self.rotary_percent, + rotary_base=self.rotary_base, + seq_len_interpolation_factor=self.seq_len_interpolation_factor, + pre_process=pre_process or parallel_state.is_pipeline_first_stage(), + post_process=post_process or parallel_state.is_pipeline_last_stage(), + ) + + +# Dictionary mapping model size names to config classes +MAMBA_MODEL_OPTIONS: dict[str, type[NemotronHConfigBase]] = { + "hybrid_mamba_8b": HybridMambaConfig8BEvo2Loss, +} diff --git a/model/evo2/utils/__init__.py b/model/evo2/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..337a2b1397075a443fc417e0aff8c7dd81a71477 --- /dev/null +++ b/model/evo2/utils/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/model/evo2/utils/callbacks.py b/model/evo2/utils/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..19d93916cdce61d24f3801d3b330a6296ea6ce9f --- /dev/null +++ b/model/evo2/utils/callbacks.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import logging +import os +from typing import Any, Literal, Sequence + +import lightning.pytorch as pl +import torch +from lightning.pytorch.callbacks import BasePredictionWriter + +# from bionemo.llm.lightning import batch_collator +from evo2.lightning import batch_collator + + +IntervalT = Literal["epoch", "batch"] + + +class PredictionWriter(BasePredictionWriter, pl.Callback): + """A callback that writes predictions to disk at specified intervals during training. + + Logits, Embeddings, Hiddens, Input IDs, and Labels may all be saved to the disk depending on trainer configuration. + Batch Idxs are provided for each prediction in the same dictionary. These must be used to maintain order between + multi device predictions and single device predictions. + """ + + def __init__( + self, + output_dir: str | os.PathLike, + write_interval: IntervalT, + batch_dim_key_defaults: dict[str, int] | None = None, + seq_dim_key_defaults: dict[str, int] | None = None, + ): + """Initializes the callback. + + Args: + output_dir: The directory where predictions will be written. + write_interval: The interval at which predictions will be written (batch, epoch). Epoch may not be used with multi-device trainers. + batch_dim_key_defaults: The default batch dimension for each key, if different from the standard 0. + seq_dim_key_defaults: The default sequence dimension for each key, if different from the standard 1. + """ + super().__init__(write_interval) + self.write_interval = write_interval + self.output_dir = str(output_dir) + self.batch_dim_key_defaults = batch_dim_key_defaults + self.seq_dim_key_defaults = seq_dim_key_defaults + + def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, *args, **kwargs) -> None: # noqa: D417 + """Invoked with Trainer.fit, validate, test, and predict are called. Will immediately fail when 'write_interval' is 'epoch' and 'trainer.num_devices' > 1. + + Args: + trainer: The Trainer instance. + pl_module: The LightningModule instance. + """ + if trainer.num_devices > 1 and self.write_interval == "epoch": + raise ValueError( + "Multi-GPU predictions are not permitted as outputs are not ordered and batch indices are lost." + ) + + def write_on_batch_end( + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + prediction: Any, + batch_indices: Sequence[int], + batch: Any, + batch_idx: int, + dataloader_idx: int, + ) -> None: + """Writes predictions to disk at the end of each batch. + + Predictions files follow the naming pattern, where rank is the active GPU in which the predictions were made. + predictions__rank_{rank}__batch_{batch_idx}.pt + + Args: + trainer: The Trainer instance. + pl_module: The LightningModule instance. + prediction: The prediction made by the model. + batch_indices: The indices of the batch. + batch: The batch data. + batch_idx: The index of the batch. + dataloader_idx: The index of the dataloader. + """ + # this will create N (num processes) files in `output_dir` each containing + # the predictions of it's respective rank + result_path = os.path.join(self.output_dir, f"predictions__rank_{trainer.global_rank}__batch_{batch_idx}.pt") + + # batch_indices is not captured due to a lightning bug when return_predictions = False + # we use input IDs in the prediction to map the result to input. + + # NOTE store the batch_idx so we do not need to rely on filenames for reconstruction of inputs. This is wrapped + # in a tensor and list container to ensure compatibility with batch_collator. + prediction["batch_idx"] = torch.tensor([batch_idx], dtype=torch.int64) + + torch.save(prediction, result_path) + logging.info(f"Inference predictions are stored in {result_path}\n{prediction.keys()}") + + def write_on_epoch_end( + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + predictions: Any, + batch_indices: Sequence[int], + ) -> None: + """Writes predictions to disk at the end of each epoch. + + Writing all predictions on epoch end is memory intensive. It is recommended to use the batch writer instead for + large predictions. + + Multi-device predictions will likely yield predictions in an order that is inconsistent with single device predictions and the input data. + + Args: + trainer: The Trainer instance. + pl_module: The LightningModule instance. + predictions: The predictions made by the model. + batch_indices: The indices of the batch. + + Raises: + Multi-GPU predictions are output in an inconsistent order with multiple devices. + """ + # this will create N (num processes) files in `output_dir` each containing + # the predictions of it's respective rank + + result_path = os.path.join(self.output_dir, f"predictions__rank_{trainer.global_rank}.pt") + + # collate multiple batches / ignore empty ones + collate_kwargs = {} + if self.batch_dim_key_defaults is not None: + collate_kwargs["batch_dim_key_defaults"] = self.batch_dim_key_defaults + if self.seq_dim_key_defaults is not None: + collate_kwargs["seq_dim_key_defaults"] = self.seq_dim_key_defaults + + prediction = batch_collator([item for item in predictions if item is not None], **collate_kwargs) + + # batch_indices is not captured due to a lightning bug when return_predictions = False + # we use input IDs in the prediction to map the result to input + if isinstance(prediction, dict): + keys = prediction.keys() + else: + keys = "tensor" + torch.save(prediction, result_path) + logging.info(f"Inference predictions are stored in {result_path}\n{keys}") diff --git a/model/evo2/utils/checkpoint/README.md b/model/evo2/utils/checkpoint/README.md new file mode 100644 index 0000000000000000000000000000000000000000..af57551de78a6da6b4ac6e4b820e791da380c8e3 --- /dev/null +++ b/model/evo2/utils/checkpoint/README.md @@ -0,0 +1,65 @@ +# Evo2 Checkpoint Conversion Library + +This library contains helper scripts for converting checkpoint formats for Evo2. + +## Converting ZeRO-1 / PyTorch Checkpoints to NeMo2 Checkpoints + +To convert a single PyTorch or ZeRO-1 checkpoints (`.pt`) into NeMo2 format, run the following command: +``` +python sub-packages/bionemo-evo2/src/bionemo/evo2/utils/checkpoint/convert_to_nemo.py --model-path --output-dir --model-size --ckpt-format +``` +where `--model-size` can be set to `7b` or `40b` (or their `_arc_1m` variants with modified GLU dimensions) and `--ckpt-format` can be set to `torch_dist` or `zarr`. + +The NeMo2 checkpoint should have the following structure for `torch_dist`: +``` +default--val_loss=2.3738-epoch=0-consumed_samples=800.0-last +├── context +│ ├── io.json +│ └── model.yaml +└── weights + ├── __*_*.distcp + ├── common.pt + └── metadata.json +``` +and the following structure for `zarr`: +``` +interleaved_hyena_7b_fix_shape +├── context +│ ├── io.json +│ └── model.yaml +└── weights + ├── common.pt + ├── metadata.json + └── # Example: module.decoder.layers.0.mixer.dense + └── shard_*_*.pt +``` + +## Converting ZeRO-1 MP{N} to ZeRO-1 MP1 + +To convert sharded (MP>1) ZeRO-1 checkpoints to un-sharded (MP1) checkpoints (or any order of model parallelism) compatible with the `convert_to_nemo.py` conversion script, you can run the following command: +``` +python sub-packages/bionemo-evo2/src/bionemo/evo2/utils/checkpoint/convert_checkpoint_model_parallel_evo2.py --source_dir --output_dir --mp_size +``` + +ZeRO-1 checkpoints should have the following structure: +``` +arc_7b_tp8_pretrained_ckpt/global_step199400 +└── mp_rank_*_model_states.pt +``` + +## Converting ZeRO-3 to ZeRO-1 + +To convert ZeRO-3 checkpoints into ZeRO-1 checkpoints, run the following command: +``` +python sub-packages/bionemo-evo2/src/bionemo/evo2/utils/checkpoint/convert_zero3_to_zero1.py --overwrite --mp_size +``` + +ZeRO-3 checkpoints should have the following structure: +``` +arc_40b_zero3_w32_mp8_test_notfinal_ckpt/global_step1 +├── bf16_zero_pp_rank_*_mp_rank_*_optim_states.pt +├── configs +│ ├── 40b_test_chkpt.yml +│ └── opengenome.yml +└── zero_pp_rank_*_mp_rank_*_model_states.pt +``` diff --git a/model/evo2/utils/checkpoint/__init__.py b/model/evo2/utils/checkpoint/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..337a2b1397075a443fc417e0aff8c7dd81a71477 --- /dev/null +++ b/model/evo2/utils/checkpoint/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/model/evo2/utils/checkpoint/convert_checkpoint_model_parallel_evo2.py b/model/evo2/utils/checkpoint/convert_checkpoint_model_parallel_evo2.py new file mode 100644 index 0000000000000000000000000000000000000000..9938e6bc8e052ef5c51ac495e40577b4ed33c87a --- /dev/null +++ b/model/evo2/utils/checkpoint/convert_checkpoint_model_parallel_evo2.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""This script converts (potentially sharded) ZeRo1 checkpoint parameters to the desired level of model tensor parallelism for the Evo 2 architecture. + +It only supports Zero-1 checkpoints and does not convert any optimizer state, +only the parameters. + +Usage: + python convert_checkpoint_model_parallel_evo2.py \ + --input-checkpoint-dir /path/to/input/checkpoint/global_step1000 \ + --output-checkpoint-dir /path/to/output/checkpoint_mp2/global_step1000 \ + --output-model-parallelism 2 +""" + +import argparse +import os +import re +from collections import OrderedDict +from glob import glob +from pathlib import Path +from typing import List, Optional, Set, Union + +import torch +from nemo.utils import logging +from params import EVO2_PARAMS, Param + + +DEVICE = "cpu" +DEFAULT_PARAM_PATTERN = r"sequential\.\d+\.(.+)" + + +def get_args(): + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Convert checkpoint parameters to desired model parallelism.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--source_dir", + type=str, + required=True, + help="Path to the input checkpoint directory containing ZeRo1 checkpoint shards, i.e. mp_rank_*_model_states.pt.", + ) + parser.add_argument( + "--glob-pattern", + type=str, + default="mp_rank_*_model_states.pt", + required=False, + help="Filename pattern to glob for ZeRo1 checkpoint shards.", + ) + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Path to the output checkpoint directory to dump the --mp_size converted model checkpoint (ZeRo1).", + ) + parser.add_argument("--mp_size", type=int, required=True, help="Desired output model parallelism to convert to.") + parser.add_argument( + "--exclude-extra", + action="store_true", + help="Exclude extra states in the conversion. Default to False, i.e. include extra states.", + ) + parser.add_argument("--verbose", action="store_true", help="Print more information about the conversion.") + args = parser.parse_args() + return args + + +def concatenate_tensors_across_shards( + tensor_name: str, + data_shards: List[OrderedDict[str, torch.Tensor]], + partition_dim: int, + hidden_dim: Optional[int] = None, + verbose: bool = False, +) -> torch.Tensor: + """Concatenate tensor shards across multiple shards. + + Args: + tensor_name (str): Name of the tensor to concatenate. + data_shards (List[OrderedDict[str, torch.Tensor]]): List of data shards containing tensors. + partition_dim (int): Dimension along which to partition the tensor. + hidden_dim (int, optional): Hidden dimension of the tensor. Defaults to None. + verbose (bool, optional): Whether to print detailed information. Defaults to False. + + Returns: + torch.Tensor: Concatenated tensor. + """ + # Retrieve tensor shards. + tensors = [shard["module"][tensor_name] for shard in data_shards] + + # Check shape of tensors without tensor parallelism, i.e. stored in all shards of the checkpoint. + if partition_dim is None: + for i, tensor in enumerate(tensors): + if not torch.allclose(tensors[0], tensor): + logging.info( + f"WARNING: Synchronized params differ for param {tensor_name}: abs max diff = {(tensors[0] - tensor).abs().max()}." + ) + # Get the distribution of tensors[0] and tensor. + if verbose: + ref_tensor = tensors[0].flatten().to(torch.float32) + ref_min, ref_max = ref_tensor.min(), ref_tensor.max() + + q = torch.tensor([0.25, 0.5, 0.75], device=ref_tensor.device) + ref_quantiles = ref_tensor.quantile(q) + logging.info(f"rank0 tensor: min={ref_min}, max={ref_max} quantiles={ref_quantiles}") + + target_tensor = tensor.flatten().to(torch.float32) + target_min, target_max = target_tensor.min(), target_tensor.max() + target_quantiles = target_tensor.quantile(q) + logging.info(f"rank{i} tensor: min={target_min}, max={target_max} quantiles={target_quantiles}") + + logging.info(f"rank0 tensor distribution:\n {ref_tensor.histc(100, min=ref_min, max=ref_max)}") + logging.info(f"rank{i} distribution:\n {target_tensor.histc(100, min=ref_min, max=ref_max)}") + + logging.info(f"tensor {tensor_name} not partitioned, returning rank0 tensor {tensors[0].shape}") + return tensors[0] + # Check for sharding across the hidden dimension. + elif partition_dim == hidden_dim: + raise ValueError(f"Detected sharding for {tensor_name} across hidden dimension at index {hidden_dim}.") + + # Check that the tensors have a consistent hidden dimension. + expected_dim = None + if hidden_dim is not None: + for tensor in tensors: + if expected_dim is None: + # Store expected hidden dimension for all tensors. + expected_dim = tensor.shape[hidden_dim] + if not tensor.shape[hidden_dim] == expected_dim: + raise ValueError(f"Tensor {tensor_name} has invalid hidden shape {tensor.shape}.") + + # Concatenate shards. + return torch.cat(tensors, dim=partition_dim) + + +def split_tensor_across_shards( + data_shards: List[OrderedDict], + tensor: torch.Tensor, + tensor_name: str, + partition_dim: int, +) -> None: + """Split a tensor across multiple shards. + + Args: + data_shards (List[OrderedDict]): List of data shards to store the split tensors. + tensor (torch.Tensor): Tensor to split. + tensor_name (str): Name of the tensor. + partition_dim (int): Dimension along which to partition the tensor. + """ + if partition_dim is None: + # No sharding. Synchronize weights across all shards. + for data_shard in data_shards: + data_shard["module"][tensor_name] = tensor + data_shard["param_shapes"][tensor_name] = tensor.shape + else: + # Split the tensor along the partition dimension across shards. + n_shards = len(data_shards) + if tensor.shape[partition_dim] % n_shards != 0: + raise ValueError( + f"Cannot shard {tensor_name} of dimension {tensor.shape[partition_dim]} across {n_shards} evenly." + ) + for chunk, data_shard in zip( + torch.chunk(tensor, chunks=n_shards, dim=partition_dim), + data_shards, + ): + data_shard["module"][tensor_name] = chunk.clone() + data_shard["param_shapes"][tensor_name] = chunk.shape + + +def format_output_filename(shard: int) -> str: + """Format the output filename for a given shard index. + + Args: + shard (int): Shard index. + + Returns: + str: Formatted output filename. + """ + return f"mp_rank_{str(shard).zfill(2)}_model_states.pt" + + +def check_params( + detected: List[str], + expected: Union[Set[str], List[str]], + buffers: Set[str], + param_pattern: str = DEFAULT_PARAM_PATTERN, + verbose: bool = False, +): + """Check that all model parameters are expected. + + Args: + detected (List[str]): Detected model parameters names. + expected (Set[str]): Expected model parameters names. + buffers (Set[str]): Set of buffer names. + param_pattern (str, optional): Regex pattern to match parameter names. Defaults to DEFAULT_PARAM_PATTERN. + verbose (bool, optional): Whether to print detailed information. Defaults to False. + """ + # Expected model parameters. + expected = set(expected) if not isinstance(expected, set) else expected + # Detected model parameters. + model_param_names = [] + for k in detected: + match = re.search(param_pattern, k) + if match is not None: + model_param_names.append(match.group(1)) + else: + logging.info(f"Could not match {k}") + detected_param_set = set(model_param_names) + if verbose: + logging.info("Detected Params:\n {detected_params}".format(detected_params="\n ".join(detected_param_set))) + + # Log unexpected model parameters. + missing_params = expected - detected_param_set + extra_params = detected_param_set - expected + extra_params = [param for param in extra_params if param not in buffers] + extra_params = [param for param in extra_params if not param.endswith("._extra_state")] + if len(extra_params) > 0: + logging.info(f"WARNING: detected extra params: {extra_params}") + if len(missing_params) > 0: + logging.info(f"WARNING: missing params: {missing_params}") + if not (extra_params or missing_params): + logging.info("No missing or extra params detected!") + + +def convert_model_weights( + input_data_shards: List[OrderedDict], + output_data_shards: List[OrderedDict], + model_parameter_names: List[str], + param_list: List[Param], + verbose: bool = False, + exclude_extra: bool = False, +): + """Convert model weights from input model parallelism to output model parallelism. + + Args: + input_data_shards (List[OrderedDict]): List of input data shards. + output_data_shards (List[OrderedDict]): List of output data shards. + model_parameter_names (List[str]): List of model parameter names. + param_list (List[Param]): List of parameter information. + verbose (bool, optional): Whether to print detailed information. Defaults to False. + exclude_extra (bool, optional): Whether to exclude extra states in the conversion. Defaults to False. + """ + logging.info( + f"Converting {len(model_parameter_names)} parameters from {len(input_data_shards)} input shards to {len(output_data_shards)} output shards..." + ) + converted = 0 + skipped = 0 + for model_parameter in model_parameter_names: + if args.verbose: + logging.info(f"Processing {model_parameter}...") + + # Ignore FP8 extra state. + if model_parameter.endswith("._extra_state"): + if "extra_state" in model_parameter: + logging.info(f"Ignoring {model_parameter} -> contains extra state.") + skipped += 1 + continue + + # Get the partition dimension and hidden dimension of each parameter. + param_info = None + for param in param_list: + if ".".join(model_parameter.split(".")[2:]) == param.name: + if param_info is None: + param_info = param + else: + raise ValueError( + f"Found more than one matching model parallelism parameter for {model_parameter}: {param_info}, {param}" + ) + if param_info is None: + raise ValueError(f"Could not find {model_parameter} among known parameters.") + + # Concatenate shards. + concatenated_tensor = concatenate_tensors_across_shards( + model_parameter, input_data_shards, param_info.partition_dim, param_info.hidden_dim, verbose=verbose + ) + # Split into shards. + split_tensor_across_shards( + output_data_shards, + concatenated_tensor, + model_parameter, + param_info.partition_dim, + ) + converted += 1 + logging.info(f"Converted {converted} of {len(model_parameter_names)} parameters (skipped {skipped} params).") + num_params = len(output_data_shards[0]["module"]) + logging.info(f"Total Params: {num_params}") + if not all(num_params == len(shard["module"]) for shard in output_data_shards): + raise ValueError("Shards have different number of parameters, which is not permitted in model parallelism.") + + if not exclude_extra: + logging.info("Adding extra states from rank0 input shard...") + rank0_model = input_data_shards[0]["module"] + for k in rank0_model.keys(): + for i, output_shard in enumerate(output_data_shards): + if k not in output_shard["module"]: + if i == 0: + logging.info(f"Adding {k} to output shards.") + output_shard["module"][k] = rank0_model[k] + new_params = len(output_data_shards[0]["module"]) - num_params + logging.info(f"Added {new_params} extra states, total params: {num_params + new_params}") + if not all(num_params + new_params == len(shard["module"]) for shard in output_data_shards): + raise ValueError("Shards have different number of parameters after adding extra states.") + + for shard_idx, output_data_shard in enumerate(output_data_shards): + output_path = Path(output_data_shard["output_dir"]) / format_output_filename(shard_idx) + torch.save( + output_data_shard, + output_path, + ) + logging.info(f"Converted checkpoint saved to: {output_path}") + + +def convert_zero1_model_parallel_checkpoint( + source_dir: str, + output_dir: str, + glob_pattern: str = "mp_rank_*_model_states.pt", + model_parallel: int = 8, + param_list: List[Param] = EVO2_PARAMS, + exclude_extra_params: bool = False, + verbose: bool = False, +): + """Convert sharded ZeRo1 checkpoint to desired model parallelism. + + Args: + source_dir (str): Path to the input checkpoint directory. + output_dir (str): Path to the output checkpoint directory. + glob_pattern (str): Filename pattern to glob for ZeRo1 checkpoint shards. Defaults to "mp_rank_*_model_states.pt". + model_parallel (int): Desired output model parallelism. Defaults to 8. + param_list (List[Param]): List of parameter information. Defaults to EVO2_PARAMS. + exclude_extra_params (bool): Whether to exclude extra states in the conversion. Defaults to False. + verbose (bool): Whether to print detailed information. Defaults to False. + """ + # Argument validation. + if not os.path.exists(source_dir): + raise ValueError(f"Input checkpoint dir ({source_dir}) not found.") + os.makedirs(output_dir, exist_ok=True) + logging.info(f"Converting checkpoint from {source_dir} to {output_dir}") + + # Identify all checkpoint model path files. + parameter_paths = sorted(glob(f"{source_dir}/{glob_pattern}")) + if len(parameter_paths) == 0: + raise ValueError(f"No parameter files found in {source_dir}") + + # Load all shards from the ZeRo1 checkpoint. + input_data_shards = [torch.load(path, map_location=DEVICE) for path in parameter_paths] + buffers = {buf for x in input_data_shards for buf in x.get("buffer_names", [])} + + # Initialize output MP shards. + output_data_shards = [ + { + "module": OrderedDict(), + "param_shapes": OrderedDict(), + "dp_world_size": input_data_shards[0]["dp_world_size"], + "output_dir": output_dir, + } + for _ in range(model_parallel) + ] + model_parameter_names = input_data_shards[0]["module"].keys() + + # Check no missing or extra params + check_params( + detected=list(model_parameter_names), + expected={param.name for param in param_list}, + buffers=buffers, + verbose=verbose, + ) + # Convert the checkpoint + convert_model_weights( + input_data_shards, + output_data_shards, + model_parameter_names, + param_list, + verbose=verbose, + exclude_extra=exclude_extra_params, + ) + logging.info("Done!") + + +if __name__ == "__main__": + args = get_args() + convert_zero1_model_parallel_checkpoint( + args.source_dir, + args.output_dir, + args.glob_pattern, + args.mp_size, + EVO2_PARAMS, + args.exclude_extra, + args.verbose, + ) diff --git a/model/evo2/utils/checkpoint/convert_to_nemo.py b/model/evo2/utils/checkpoint/convert_to_nemo.py new file mode 100644 index 0000000000000000000000000000000000000000..8a75d957ca47354dcca8138f69540ec95e9e66e5 --- /dev/null +++ b/model/evo2/utils/checkpoint/convert_to_nemo.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import argparse +import logging +from pathlib import Path + +from nemo.collections.llm.gpt.model.hyena import ( + HYENA_MODEL_OPTIONS, + HuggingFaceSavannaHyenaImporter, + HyenaConfig, + HyenaModel, + PyTorchHyenaImporter, +) +from nemo.lightning import io, teardown + + +def parse_args(): + """Parse command-line arguments.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--model-path", + type=str, + required=True, + help="Path to the Evo2 un-sharded (MP1) model checkpoint file, or a Hugging Face model name. Any model " + "from the Savanna Evo2 family is supported such as 'hf://arcinstitute/savanna_evo2_1b_base'.", + ) + parser.add_argument("--output-dir", type=str, required=True, help="Output directory path for the converted model.") + parser.add_argument( + "--model-size", + type=str, + choices=sorted(HYENA_MODEL_OPTIONS.keys()), + default="1b", + help="Model architecture to use, choose between 1b, 7b, 40b, or test (a sub-model of 4 layers, " + "less than 1B parameters). '*_arc_longcontext' models have GLU / FFN dimensions that support 1M " + "context length when trained with TP>>8.", + ) + parser.add_argument( + "--strip-optimizer", + action="store_true", + help="Strip the optimizer state from the model checkpoint, this works on nemo2 format checkpoints.", + ) + return parser.parse_args() + + +@io.model_importer(HyenaModel, "pytorch") +class HyenaOptimizerRemover(io.ModelConnector["HyenaModel", HyenaModel]): + """Removes the optimizer state from a nemo2 format model checkpoint.""" + + def __new__(cls, path: str, model_config=None): + """Creates a new importer instance. + + Args: + path: Path to the PyTorch model + model_config: Optional model configuration + + Returns: + PyTorchHyenaImporter instance + """ + instance = super().__new__(cls, path) + instance.model_config = model_config + return instance + + def init(self) -> HyenaModel: + """Initializes a new HyenaModel instance. + + Returns: + HyenaModel: Initialized model + """ + return HyenaModel(self.config, tokenizer=self.tokenizer) + + def get_source_model(self): + """Returns the source model.""" + model, _ = self.nemo_load(self) + return model + + def apply(self, output_path: Path, checkpoint_format: str = "torch_dist") -> Path: + """Applies the model conversion from PyTorch to NeMo format. + + Args: + output_path: Path to save the converted model + checkpoint_format: Format for saving checkpoints + + Returns: + Path: Path to the saved NeMo model + """ + source = self.get_source_model() + + target = self.init() + trainer = self.nemo_setup(target, ckpt_async_save=False, save_ckpt_format=checkpoint_format) + source.to(self.config.params_dtype) + target.to(self.config.params_dtype) + self.convert_state(source, target) + self.nemo_save(output_path, trainer) + + logging.info(f"Converted Hyena model to Nemo, model saved to {output_path}") + + teardown(trainer, target) + del trainer, target + + return output_path + + def convert_state(self, source, target): + """Converts the state dictionary from source format to target format. + + Args: + source: Source model state + target: Target model + + Returns: + Result of applying state transforms + """ + mapping = {k: k for k in source.module.state_dict().keys()} + return io.apply_transforms( + source, + target, + mapping=mapping, + ) + + @property + def tokenizer(self): + """Gets the tokenizer for the model. + + Returns: + Tokenizer instance + """ + from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer + + tokenizer = get_nmt_tokenizer( + library=self.model_config.tokenizer_library, + ) + + return tokenizer + + @property + def config(self) -> HyenaConfig: + """Gets the model configuration. + + Returns: + HyenaConfig: Model configuration + """ + return self.model_config + + +def main(): + """Convert a PyTorch Evo2 model checkpoint to a NeMo model checkpoint.""" + args = parse_args() + + evo2_config = HYENA_MODEL_OPTIONS[args.model_size]() + + if args.strip_optimizer: + importer = HyenaOptimizerRemover(args.model_path, model_config=evo2_config) + assert not args.model_path.startswith("hf://"), "Strip optimizer only works on local nemo2 format checkpoints." + elif args.model_path.startswith("hf://"): + importer = HuggingFaceSavannaHyenaImporter(args.model_path.lstrip("hf://"), model_config=evo2_config) + else: + importer = PyTorchHyenaImporter(args.model_path, model_config=evo2_config) + importer.apply(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/model/evo2/utils/checkpoint/convert_zero3_to_zero1.py b/model/evo2/utils/checkpoint/convert_zero3_to_zero1.py new file mode 100644 index 0000000000000000000000000000000000000000..28296f90791eca0ce04620c6d67cb76685081bf6 --- /dev/null +++ b/model/evo2/utils/checkpoint/convert_zero3_to_zero1.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +#!/usr/bin/env python + +import argparse +import os +import time +from multiprocessing import Pool +from typing import List, Optional + +import zero3_conversion_lib +from zero3_conversion_lib import get_elapsed, process_single_rank + + +def convert_zero_checkpoint_to_fp32_state_dict( + checkpoint_dir: str, + output_dir: str, + tag: Optional[str] = None, + exclude_frozen_parameters: bool = False, + mp_size: int = 8, + overwrite: bool = False, + num_workers: int = 1, + ranks_to_process: Optional[List[int]] = None, +): + """Converts a DeepSpeed Zero-3 checkpoint to a PyTorch FP32 state_dict. + + Args: + checkpoint_dir (str): Path to the desired checkpoint folder. + output_dir (str): Directory to save the PyTorch FP32 state_dict output files. + tag (Optional[str]): Checkpoint tag used as a unique identifier or sub-directory that contains the checkpoint. + exclude_frozen_parameters (bool): Whether to exclude frozen parameters. + mp_size (int): Model parallel size of the source checkpoint. + overwrite (bool): Whether to overwrite existing MP shards. + num_workers (int): Number of workers to use for processing. + ranks_to_process (Optional[List[int]]): List of ranks to process. + + Raises: + FileNotFoundError: If the checkpoint directory does not exist. + """ + ds_checkpoint_dir = os.path.join(checkpoint_dir, tag) if tag is not None else checkpoint_dir + + if not os.path.isdir(ds_checkpoint_dir): + raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist") + + output_dir = os.path.join(output_dir, tag) if tag is not None else output_dir + if not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + + num_workers = min(num_workers, mp_size) + + if ranks_to_process is not None: + ranks_to_process = list(ranks_to_process) + assert len(ranks_to_process) <= mp_size, f"Expected {mp_size} ranks to process, got {len(ranks_to_process)}" + assert all(0 <= r < mp_size for r in ranks_to_process), ( + f"Expected ranks to be in range [0, {mp_size}), got {ranks_to_process}" + ) + else: + ranks_to_process = list(range(mp_size)) + + print(f"Processing ranks: {ranks_to_process}", flush=True) + + start = time.time() + if num_workers > 1: + with Pool(num_workers) as p: + p.starmap( + process_single_rank, + [(i, ds_checkpoint_dir, output_dir, overwrite, exclude_frozen_parameters) for i in ranks_to_process], + ) + else: + for i in ranks_to_process: + process_single_rank(i, ds_checkpoint_dir, output_dir, overwrite, exclude_frozen_parameters) + + total_time = get_elapsed(time.time() - start) + print(f"All done!\n-> Total time: {total_time}\n-> All outputs written to {os.path.abspath(output_dir)}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "checkpoint_dir", type=str, help="path to the desired checkpoint folder, e.g., path/checkpoint-12" + ) + parser.add_argument( + "output_dir", + type=str, + help="directory to the pytorch fp32 state_dict output files(e.g. path/checkpoint-12-output/)", + ) + parser.add_argument("--overwrite", action="store_true", help="Overwrite existing MP shards") + parser.add_argument( + "-t", + "--tag", + type=str, + default=None, + help="Checkpoint tag used as a unique identifier or sub-directory that contains the checkpoint, e.g. 'global_step1' or 'latest'.", + ) + parser.add_argument("--exclude_frozen_parameters", action="store_true", help="exclude frozen parameters") + parser.add_argument("-d", "--debug", action="store_true", help="enable debug") + parser.add_argument("--mp_size", required=True, type=int, help="Model parallel size of source checkpoint") + parser.add_argument("--rank_start", default=None, type=int, help="Start rank to process") + parser.add_argument("--rank_end", default=None, type=int, help="End rank to process") + parser.add_argument("--num_workers", default=1, type=int, help="Number of workers to use for processing") + args = parser.parse_args() + + if args.rank_start is not None: + if args.rank_end is None: + args.rank_end = args.mp_size - 1 + else: + assert args.rank_end < args.mp_size, "Expected end rank to be less than mp_size" + + assert args.rank_start < args.rank_end, "Expected start rank to be less than end rank" + assert args.rank_start >= 0, "Expected start rank to be greater than 0" + args.ranks_to_process = list(range(args.rank_start, args.rank_end + 1)) + else: + args.ranks_to_process = list(range(args.mp_size)) + + print("Args:") + for k, v in args.__dict__.items(): + print(f" {k}: {v}", flush=True) + print("") + zero3_conversion_lib.debug = args.debug + + convert_zero_checkpoint_to_fp32_state_dict( + args.checkpoint_dir, + args.output_dir, + tag=args.tag, + exclude_frozen_parameters=args.exclude_frozen_parameters, + mp_size=args.mp_size, + overwrite=args.overwrite, + num_workers=args.num_workers, + ranks_to_process=args.ranks_to_process, + ) diff --git a/model/evo2/utils/checkpoint/params.py b/model/evo2/utils/checkpoint/params.py new file mode 100644 index 0000000000000000000000000000000000000000..45b2cf3012d149d2d1b82ad41dadd0625510c7da --- /dev/null +++ b/model/evo2/utils/checkpoint/params.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from dataclasses import dataclass + + +@dataclass +class Param: + """A dataclass representing a parameter in a checkpoint. + + Attributes: + name (str): The name of the parameter in the checkpoint. + partition_dim (int): The dimension index that gets sharded. `None` for no sharding. + hidden_dim (int): The hidden dimension index. `None` for no hidden dimension. + """ + + name: str # Name of the parameter in the checkpoint. + partition_dim: int # The dimension index that gets sharded. `None` for no sharding. + hidden_dim: int # The hidden dimension index. `None` for no hidden dimension. + + +EVO2_PARAMS = [ + # Only layer_00. + Param("word_embeddings.weight", 0, 1), # torch.Size([64, 8192]) + Param("input_layernorm.weight", None, 0), # torch.Size([8192]) + Param("post_attention_layernorm.weight", None, 0), # torch.Size([8192]) + Param("pre_mlp_layernorm.weight", None, 0), # torch.Size([8192]) + Param("outer_mlp_layernorm.weight", None, 0), # torch.Size([8192]) + Param("mixer.dense_projection.weight", 0, 1), # torch.Size([3072, 8192]), + Param("mixer.hyena_proj_conv.short_conv_weight", 0, None), # torch.Size([3072, 3]), + Param("mixer.mixer.conv_bias", 0, None), # torch.Size([1024]), + Param("mixer.mixer.filter.decay", 0, None), # torch.Size([64, 8192]), + Param("mixer.mixer.filter.gamma", 0, None), # torch.Size([1024, 16]), + Param("mixer.mixer.filter.h", 0, None), # torch.Size([64, 8192]), + Param("mixer.mixer.filter.p", 0, None), # torch.Size([1024, 16]), + Param("mixer.mixer.filter.R", 0, None), # torch.Size([1024, 16]), + Param("mixer.mixer.filter.t", None, 0), # torch.Size([1, 1, seqlen]), + Param("mixer.mixer.short_conv.short_conv_weight", 0, None), # torch.Size([64, 1, 7]), + Param("mixer.rotary_emb.inv_freq", None, None), # torch.Size([64]) + Param("mixer.dense.weight", 1, 0), # torch.Size([8192, 2048]), + Param("mixer.dense.bias", None, 0), # torch.Size([8192]) + Param("mlp.w1.weight", 0, 1), # torch.Size([2736, 8192]), + Param("mlp.w2.weight", 0, 1), # torch.Size([2736, 8192]), + Param("mlp.w3.weight", 1, 0), # torch.Size([8192, 2736]), + # Only last layer. + Param("norm.weight", None, 0), # torch.Size([8192]), +] diff --git a/model/evo2/utils/checkpoint/zero3_conversion_lib.py b/model/evo2/utils/checkpoint/zero3_conversion_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..9b34959c9f7e56b7d0cc1e16060500ff2aacbe76 --- /dev/null +++ b/model/evo2/utils/checkpoint/zero3_conversion_lib.py @@ -0,0 +1,699 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Helper utility for converting ZeRO3 and ZeRO2 checkpoints to PyTorch.""" + +import glob +import math +import os +import re +import time +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Set + +import psutil +import torch +from tqdm import tqdm + + +BUFFER_NAMES = "buffer_names" +DS_VERSION = "ds_version" +FP32_FLAT_GROUPS = "fp32_flat_groups" +FROZEN_PARAM_FRAGMENTS = "frozen_param_fragments" +FROZEN_PARAM_SHAPES = "frozen_param_shapes" +OPTIMIZER_STATE_DICT = "optimizer_state_dict" +PARAM_SHAPES = "param_shapes" +PARTITION_COUNT = "partition_count" +SINGLE_PARTITION_OF_FP32_GROUPS = "single_partition_of_fp32_groups" +ZERO_STAGE = "zero_stage" +EXTRA_STATE = "._extra_state" + + +@dataclass +class ZeroModelState: + """A dataclass representing the state of a ZeRO model. + + Attributes: + buffers (Dict): Buffers in the model state. + extra_states (Dict): Extra states in the model state. + param_shapes (List): Shapes of the parameters. + shared_params (List): Shared parameters in the model state. + ds_version (int): Version of the DeepSpeed checkpoint. + frozen_param_shapes (Dict): Shapes of the frozen parameters. + frozen_param_fragments (Dict): Fragments of the frozen parameters. + """ + + buffers: Dict + extra_states: Dict + param_shapes: List + shared_params: List + ds_version: int + frozen_param_shapes: Dict + frozen_param_fragments: Dict + + +debug = 0 +device = torch.device("cpu") + + +def profile_memory_decorator(func: Iterable): + """A decorator to profile memory usage of a function. + + Args: + func (Iterable): The function to be decorated. + + Returns: + wrapper: The decorated function with memory profiling. + """ + + def profile_memory(): + pid = os.getpid() + process = psutil.Process(pid) + memory_info = process.memory_info() + print_pid(f"{pid}: RSS = {memory_info.rss / 1024**2:.2f} MB") + + def wrapper(*args, **kwargs): + profile_memory() + func(*args, **kwargs) + profile_memory() + + return wrapper + + +def print_pid(msg: str): + """Prints the process ID along with a message. + + Args: + msg (str): The message to be printed. + """ + pid = os.getpid() + print(f"{pid=}:{msg}") + + +def atoi(text: str): + """Converts a string to an integer if it is a digit, otherwise returns the string. + + Args: + text (str): The text to be converted. + + Returns: + int or str: The converted integer or the original string. + """ + return int(text) if text.isdigit() else text + + +def natural_keys(text: str): + """Sorts a list in human order. + + Args: + text (str): The text to be sorted. + + Returns: + list: The sorted list. + + Note: + alist.sort(key=natural_keys) sorts in human order. + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + """ + return [atoi(c) for c in re.split(r"(\d+)", text)] + + +def get_checkpoint_files(checkpoint_dir: str, glob_pattern: str): + """Retrieves checkpoint files from a directory based on a glob pattern. + + Args: + checkpoint_dir (str): The directory to search for checkpoint files. + glob_pattern (str): The glob pattern to match files. + + Returns: + list: A sorted list of checkpoint files. + + Raises: + FileNotFoundError: If no files matching the glob pattern are found. + """ + # XXX: need to test that this simple glob rule works for multi-node setup too + ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys) + + if len(ckpt_files) == 0: + raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'") + + return ckpt_files + + +def get_model_files_by_rank(checkpoint_dir: str, rank: int): + """Retrieves model files for a specific rank from a checkpoint directory. + + Args: + checkpoint_dir (str): The directory to search for model files. + rank (int): The rank to search for. + + Returns: + list: A list of model files for the specified rank. + """ + return get_checkpoint_files(checkpoint_dir, f"*mp_rank_{rank:02}_model_states.pt") + + +def get_optim_files_by_rank(checkpoint_dir: str, rank: int): + """Retrieves optimizer files for a specific rank from a checkpoint directory. + + Args: + checkpoint_dir (str): The directory to search for optimizer files. + rank (int): The rank to search for. + + Returns: + list: A list of optimizer files for the specified rank. + """ + return get_checkpoint_files(checkpoint_dir, f"*mp_rank_{rank:02}_optim_states.pt") + + +def create_ds_output_path(rank: int): + """Creates the output path for a DeepSpeed checkpoint. + + Args: + rank (int): The rank to create the output path for. + + Returns: + str: The output path for the DeepSpeed checkpoint. + """ + return f"mp_rank_{rank:02}_model_states.pt" + + +def create_zero3_model_state_path(dp_rank: int, mp_rank: int): + """Creates the path for a ZeRO3 model state file. + + Args: + dp_rank (int): The data parallel rank. + mp_rank (int): The model parallel rank. + + Returns: + str: The path for the ZeRO3 model state file. + """ + return f"zero_pp_rank_{dp_rank}_mp_rank_{mp_rank:02}_model_states.pt" + + +def create_zero3_optim_state_path(dp_rank: int, mp_rank: int): + """Creates the path for a ZeRO3 optimizer state file. + + Args: + dp_rank (int): The data parallel rank. + mp_rank (int): The model parallel rank. + + Returns: + str: The path for the ZeRO3 optimizer state file. + """ + return f"bf16_zero_pp_rank_{dp_rank}_mp_rank_{mp_rank:02}_optim_states.pt" + + +def get_model_state_file(checkpoint_dir: str, zero_stage: int): + """Retrieves the model state file from a checkpoint directory based on the ZeRO stage. + + Args: + checkpoint_dir (str): The directory to search for the model state file. + zero_stage (int): The ZeRO stage to search for. + + Returns: + str: The path to the model state file. + + Raises: + FileNotFoundError: If the directory or model state file is not found. + ValueError: If the ZeRO stage is not supported. + """ + if not os.path.isdir(checkpoint_dir): + raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") + + # there should be only one file + if zero_stage <= 2: + file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt") + elif zero_stage == 3: + file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") + else: + raise ValueError(f"Unsupported zero stage {zero_stage}. Expected 1, 2, or 3") + + if not os.path.exists(file): + raise FileNotFoundError(f"can't find model states file at '{file}'") + + return file + + +def parse_model_states(files: Set[str]): + """Parses model state files and returns a list of ZeroModelState objects. + + Args: + files (Set[str]): A set of file paths to parse. + + Returns: + List[ZeroModelState]: A list of parsed ZeroModelState objects. + + Raises: + ValueError: If a file is not a model state checkpoint. + """ + zero_model_states = [] + for file in files: + state_dict = torch.load(file, map_location=device) + + if BUFFER_NAMES not in state_dict: + raise ValueError(f"{file} is not a model state checkpoint") + buffer_names = state_dict[BUFFER_NAMES] + if debug: + print_pid("Found buffers:", buffer_names) + + # recover just the buffers while restoring them to fp32 if they were saved in fp16 + buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names} + + extra_states = {k: v for k, v in state_dict["module"].items() if k.endswith(EXTRA_STATE)} + + # collect parameters that are included in param_shapes + param_shapes = state_dict[PARAM_SHAPES] + param_names = [] + for s in param_shapes: + for name in s.keys(): + param_names.append(name) + + # update with frozen parameters + frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None) + if frozen_param_shapes is not None: + if debug: + print_pid(f"Found frozen_param_shapes: {frozen_param_shapes}") + param_names += list(frozen_param_shapes.keys()) + + # handle shared params + shared_params = [[k, v] for k, v in state_dict["shared_params"].items()] + + ds_version = state_dict.get(DS_VERSION, None) + + frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None) + + z_model_state = ZeroModelState( + buffers=buffers, + extra_states=extra_states, + param_shapes=param_shapes, + shared_params=shared_params, + ds_version=ds_version, + frozen_param_shapes=frozen_param_shapes, + frozen_param_fragments=frozen_param_fragments, + ) + zero_model_states.append(z_model_state) + + return zero_model_states + + +def parse_optim_states(files: Set[str], ds_checkpoint_dir: str): + """Parses optimizer state files and returns the ZeRO stage, world size, and fp32 flat groups. + + Args: + files (Set[str]): A set of file paths to parse. + ds_checkpoint_dir (str): The directory containing the DeepSpeed checkpoint. + + Returns: + tuple: A tuple containing the ZeRO stage, world size, and fp32 flat groups. + + Raises: + ValueError: If a file is not a ZeRO checkpoint or if the number of files does not match the expected world size. + """ + total_files = len(files) + state_dicts = [] + for f in files: + state_dict = torch.load(f, map_location=device) + # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights + # and also handle the case where it was already removed by another helper script + state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) + state_dict[OPTIMIZER_STATE_DICT] = { + FP32_FLAT_GROUPS: state_dict[OPTIMIZER_STATE_DICT][FP32_FLAT_GROUPS], + ZERO_STAGE: state_dict[OPTIMIZER_STATE_DICT][ZERO_STAGE], + PARTITION_COUNT: state_dict[OPTIMIZER_STATE_DICT][PARTITION_COUNT], + } + state_dicts.append(state_dict) + + if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]: + raise ValueError(f"{files[0]} is not a zero checkpoint") + zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE] + world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT] + + # For ZeRO-2 each param group can have different partition_count as data parallelism for expert + # parameters can be different from data parallelism for non-expert parameters. So we can just + # use the max of the partition_count to get the dp world_size. + + if type(world_size) is list: + world_size = max(world_size) + + if world_size != total_files: + raise ValueError( + f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. " + "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes." + ) + + # the groups are named differently in each stage + if zero_stage <= 2: + fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS + elif zero_stage == 3: + fp32_groups_key = FP32_FLAT_GROUPS + else: + raise ValueError(f"unknown zero stage {zero_stage}") + + if zero_stage <= 2: + fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))] + elif zero_stage == 3: + # if there is more than one param group, there will be multiple flattened tensors - one + # flattened tensor per group - for simplicity merge them into a single tensor + # + # XXX: could make the script more memory efficient for when there are multiple groups - it + # will require matching the sub-lists of param_shapes for each param group flattened tensor + + fp32_flat_groups = [ + torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts)) + ] + + return zero_stage, world_size, fp32_flat_groups + + +def _get_fp32_state_dict_from_zero_checkpoint( + ds_checkpoint_dir: str, rank: int, exclude_frozen_parameters: bool = False +): + """Returns the fp32 state dictionary reconstructed from a ZeRO checkpoint. + + Args: + ds_checkpoint_dir (str): Path to the DeepSpeed checkpoint folder. + rank (int): The rank to process. + exclude_frozen_parameters (bool): Whether to exclude frozen parameters. + + Returns: + OrderedDict: The reconstructed fp32 state dictionary. + """ + print_pid(f"Processing zero checkpoint '{ds_checkpoint_dir}'") + + # optim_files = get_optim_files(ds_checkpoint_dir) + # zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir) + + optim_files = get_optim_files_by_rank(ds_checkpoint_dir, rank=rank) + optim_files_check = get_checkpoint_files(ds_checkpoint_dir, f"bf16*_{rank:02d}_optim_states.pt") + assert set(optim_files) == set(optim_files_check), f"Expected {optim_files_check}, got {optim_files}" + # check ordering as well + for f1, f2 in zip(optim_files, optim_files_check): + assert os.path.basename(f1) == os.path.basename(f2), ( + f"Found mismatching optim files for rank {rank}: {os.path.basename(f1)} != {os.path.basename(f2)}" + ) + print_pid(f" -> Optim files for rank {rank}: {len(optim_files)}") + + if debug: + print_pid(f"{optim_files=}") + + if os.environ.get("ZERO3_CONVERSION_DEBUG", "0") == "1": + breakpoint() + + zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir) + assert len(optim_files) == world_size, f"Expected {world_size} optim files, got {len(optim_files)}" + if debug: + print_pid( + f" -> rank{rank} stage: {zero_stage} {world_size=} {len(fp32_flat_groups)=} {fp32_flat_groups.shape=}" + ) + + model_files = get_model_files_by_rank(ds_checkpoint_dir, rank=rank) + model_files_check = get_checkpoint_files(ds_checkpoint_dir, f"zero_*_mp_rank_{rank:02d}_model_states.pt") + assert set(model_files) == set(model_files_check), f"Expected {model_files_check}, got {model_files}" + + for f1, f2 in zip(model_files, model_files_check): + assert os.path.basename(f1) == os.path.basename(f2), ( + f"Found mismatching optim files for rank {rank}: {os.path.basename(f1)} != {os.path.basename(f2)}" + ) + print_pid(f" -> Model files for rank {rank}: {len(model_files)}") + + assert len(optim_files) == len(model_files), ( + f"Expected same number of optim and model files: {len(optim_files)} != {len(model_files)}" + ) + assert len(optim_files) > 0, f"Expected at least one optim file, got {len(optim_files)}" + + zero_model_states = parse_model_states(model_files) + print_pid(f"Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}") + + return _get_fp32_state_dict_from_zero3_checkpoint( + world_size, fp32_flat_groups, zero_model_states, exclude_frozen_parameters + ) + + +def zero3_partitioned_param_info(unpartitioned_numel: int, world_size: int): + """Returns the partitioned and padding number of elements for a parameter. + + Args: + unpartitioned_numel (int): The number of elements in the unpartitioned parameter. + world_size (int): The world size. + + Returns: + tuple: A tuple containing the partitioned number of elements and the padding number of elements. + """ + remainder = unpartitioned_numel % world_size + padding_numel = (world_size - remainder) if remainder else 0 + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + return partitioned_numel, padding_numel + + +def _zero3_merge_frozen_params(state_dict: Dict[str, Any], world_size: int, zero_model_states: List[ZeroModelState]): + """Merges frozen parameters into the state dictionary. + + Args: + state_dict (Dict[str, Any]): The state dictionary to update. + world_size (int): The world size. + zero_model_states (List[ZeroModelState]): The list of ZeroModelState objects. + + Returns: + None + """ + if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: + return + + if debug: + for i in range(world_size): + num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values()) + print_pid(f"rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}") + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size + print_pid(f"Frozen params: Have {avail_numel} numels to process.") + print_pid(f"Frozen params: Need {wanted_numel} numels in {wanted_params} params") + + total_params = 0 + total_numel = 0 + for name, shape in zero_model_states[0].frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states) + state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape) + + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) + + if debug: + print_pid( + f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + print_pid(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +# @profile_memory_decorator +def _zero3_merge_trainable_params( + state_dict: Dict[str, Any], + world_size: int, + fp32_flat_groups: List[torch.Tensor], + zero_model_states: List[ZeroModelState], +): + """Merges trainable parameters into the state dictionary. + + Args: + state_dict (Dict[str, Any]): The state dictionary to update. + world_size (int): The world size. + fp32_flat_groups (List[torch.Tensor]): The list of fp32 flat groups. + zero_model_states (List[ZeroModelState]): The list of ZeroModelState objects. + + Returns: + None + """ + if os.environ.get("ZERO3_CONVERSION_DEBUG", "0") == "1": + breakpoint() + + param_shapes = zero_model_states[0].param_shapes + avail_numel = fp32_flat_groups[0].numel() * world_size + # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each + # param, re-consolidating each param, while dealing with padding if any + + # merge list of dicts, preserving order + param_shapes = {k: v for d in param_shapes for k, v in d.items()} + + if debug: + for i in range(world_size): + print_pid(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}") + + wanted_params = len(param_shapes) + wanted_numel = sum(shape.numel() for shape in param_shapes.values()) + # not asserting if there is a mismatch due to possible padding + avail_numel = fp32_flat_groups[0].numel() * world_size + print_pid(f"Trainable params: Have {avail_numel} numels to process.") + print_pid(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + offset = 0 + total_numel = 0 + total_params = 0 + pid = os.getpid() + for name, shape in tqdm(param_shapes.items(), desc=f"{pid=}: Gathering Sharded Weights"): + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + total_params += 1 + # NOTE: partitioned_numel includes padding, padding applies if unpartitioned_numel is not divisible by world_size + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) + + if debug: + print_pid( + f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + # XXX: memory usage doubles here + state_dict[name] = ( + torch.cat(tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)), 0) + .narrow(0, 0, unpartitioned_numel) + .view(shape) + ) + offset += partitioned_numel + + offset *= world_size + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print_pid(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements") + + +def _get_fp32_state_dict_from_zero3_checkpoint( + world_size: int, + fp32_flat_groups: List[torch.Tensor], + zero_model_states: List[ZeroModelState], + exclude_frozen_parameters: bool, +): + """Returns the fp32 state dictionary reconstructed from a ZeRO3 checkpoint. + + Args: + world_size (int): The world size. + fp32_flat_groups (List[torch.Tensor]): The list of fp32 flat groups. + zero_model_states (List[ZeroModelState]): The list of ZeroModelState objects. + exclude_frozen_parameters (bool): Whether to exclude frozen parameters. + + Returns: + OrderedDict: The reconstructed fp32 state dictionary. + """ + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print_pid(f"added {len(buffers)} buffers") + + # extra state (e.g., fp8) + extra_states = zero_model_states[0].extra_states + state_dict.update(extra_states) + if debug: + print_pid(f"added {len(extra_states)} extra_states") + + if not exclude_frozen_parameters: + _zero3_merge_frozen_params(state_dict, world_size, zero_model_states) + + _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def get_elapsed(t: float): + """Converts elapsed time in seconds to a formatted string. + + Args: + t (float): The elapsed time in seconds. + + Returns: + str: The formatted elapsed time as a string. + """ + minutes = t // 60 + seconds = t % 60 + if minutes > 0: + total_time = f"{minutes:.0f}min{seconds:.0f}s" + else: + total_time = f"{seconds:.1f}s" + return total_time + + +def process_single_rank( + rank: int, + ds_checkpoint_dir: str, + output_dir: str, + overwrite: bool = False, + exclude_frozen_parameters: bool = False, +): + """Processes a single rank to gather and save the state dictionary. + + Args: + rank (int): The rank to process. + ds_checkpoint_dir (str): Path to the DeepSpeed checkpoint folder. + output_dir (str): Directory to save the output. + overwrite (bool): Whether to overwrite existing files. Default is False. + exclude_frozen_parameters (bool): Whether to exclude frozen parameters. Default is False. + """ + print_pid(f"Gathering rank {rank} state_dict...") + + start = time.time() + output_path = os.path.join(output_dir, create_ds_output_path(rank)) + if os.path.exists(output_path) and not overwrite: + print_pid(f"Output path {output_path} exists, skipping") + return + + print_pid(f" -> Gathering data parallel partitions for mp rank {rank}...") + + if os.environ.get("ZERO3_CONVERSION_DEBUG", "0") == "1": + breakpoint() + + state_dict = _get_fp32_state_dict_from_zero_checkpoint( + ds_checkpoint_dir=ds_checkpoint_dir, rank=rank, exclude_frozen_parameters=exclude_frozen_parameters + ) + print_pid(f" -> Done processing rank {rank} state_dict, gathered {len(state_dict)} params") + + checkpoint = { + "module": state_dict, + "param_shapes": OrderedDict(), + "dp_world_size": 1, + } + + for param, value in state_dict.items(): + if isinstance(value, torch.Tensor): + checkpoint["param_shapes"][param] = value.shape + + print_pid(f" -> Saving mp rank {rank} checkpoint to {output_path}") + torch.save(checkpoint, f"{output_path}") + + total_time = get_elapsed(time.time() - start) + print_pid(f" -> rank {rank} took {total_time}") diff --git a/model/evo2/utils/config.py b/model/evo2/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..a66ea521021e9a4c6751aaf57aa00d91246ee817 --- /dev/null +++ b/model/evo2/utils/config.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from pathlib import Path +from typing import Literal + +from pydantic import BaseModel + + +class Evo2TaxonomyLineage(BaseModel): + """Pydantic model class that defines the source lineage of a DNA sequence.""" + + domain: None | str = None + phylum: None | str = None + clazz: None | str = None + order: None | str = None + family: None | str = None + genus: None | str = None + species: None | str = None + + +class Evo2PreprocessingConfig(BaseModel): + """Pydantic model class specifying the configuration schema for a preprocessed IndexedDataset (.bin, .idx).""" + + # Paths + datapaths: list[Path] = [] + output_dir: None | Path = None + output_prefix: None | str = None + # Random Datasplit + train_split: float = 0.7 + valid_split: float = 0.2 + test_split: float = 0.1 + # Overwrite existing binaries. Otherwise, skip already preprocessed datasets. + overwrite: bool = False + # Raw Preprocessing Transforms + embed_reverse_complement: bool = False + random_reverse_complement: float = 0.0 + random_lineage_dropout: float = 0.0 + transcribe: None | Literal["transcribe", "back_transcribe"] = None + force_uppercase: bool = False + indexed_dataset_dtype: str = "uint8" + # Tokenization Transforms + append_eod: bool = True + enforce_sample_length: None | int = None + ftfy: bool = False + # NeMo Tokenizer Configuration + tokenizer_type: Literal[ + "Byte-Level", + "HuggingFace", + "SentencePiece", + "Regex", + "Megatron", + "Tiktoken", + ] = "Byte-Level" + vocab_file: None | Path = None + vocab_size: None | int = 512 + merges_file: None | Path = None + tokenizer_model_name: None | str = None + pretrained_tokenizer_model: None | str = None + special_tokens: None | dict[str, str] = {} + fast_hf_tokenizer: bool = False + # Compute Configuration + # NOTE: If preprocessing a large amount of short individual sequences (< 1000 bp), do NOT use + # multiprocessing (workers > 1) because sequence-level parallel IPC will dominate the preprocessing time! + workers: int = 1 + preproc_concurrency: int = 100000 + chunksize: int = 1 + # Filters + drop_empty_sequences: bool = False + nnn_filter: bool = False + # RNG + seed: None | int = None + # Evo2 Taxonomic Lineage Tags + # SeqID Sub-String Indexing: "ABC" will have taxonomy data from "A". + taxonomy_data: dict[str, Evo2TaxonomyLineage] = {} + # Periodicity of injecting phylogenetic lineage tags in the sequence prior to tokenization. + prompt_spacer_length: int = 131072 diff --git a/model/evo2/utils/datamodule_utils.py b/model/evo2/utils/datamodule_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..029ff7d12ba1d821b3c4e38a47cca6d40185b4a0 --- /dev/null +++ b/model/evo2/utils/datamodule_utils.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import Any, Dict, List, Union + + +def float_or_int_or_none(value: Union[str, float, int, None]) -> Union[float, int, None]: + """Converts a given value into a float, int, or None. + + Args: + value (Union[str, float, int, None]): A value that can be either a string, float, int, or None. + + Returns: + Union[float, int, None]: A float, int, or None based on the input value. + + If the input value is None or "None", it returns None. + If the input value is an int or float, it returns the same value. + If the input value is a string, it tries to convert it into an int if possible, otherwise into a float. + """ + if value is None or value == "None": + return + if isinstance(value, (int, float)): + return value + if value.isdigit(): + return int(value) + return float(value) + + +def parse_kwargs_to_arglist(kwargs: Dict[str, Any]) -> List[str]: + """Converts a dictionary of keyword arguments into a list of command-line arguments. + + Args: + kwargs (Dict[str, Any]): A dictionary where keys are argument names and values are argument values. + + Returns: + A list of strings, where each string is a command-line argument in the format '--argument-name value'. + """ + arglist = [] + for k, v in kwargs.items(): + arglist.extend([f"--{k.replace('_', '-')}", str(v)]) + return arglist + + +def infer_global_batch_size( + micro_batch_size: int, + num_nodes: int, + devices: int, + accumulate_grad_batches: int = 1, + tensor_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + context_model_parallel_size: int = 1, +) -> int: + """Infers the global batch size based on the micro batch size, number of nodes, devices, accumulation of gradient batches, and model parallel sizes. + + Args: + micro_batch_size (int): The micro batch size. + num_nodes (int): The number of nodes. + devices (int): The number of devices. + accumulate_grad_batches (int): The accumulation of gradient batches. Defaults to 1. + tensor_model_parallel_size (int): The tensor model parallel size. Defaults to 1. + pipeline_model_parallel_size (int): The pipeline model parallel size. Defaults to 1. + context_model_parallel_size (int): The context model parallel size. Defaults to 1. + + Returns: + int: The global batch size. + """ + if not all( + isinstance(arg, int) + for arg in [ + micro_batch_size, + num_nodes, + devices, + accumulate_grad_batches, + tensor_model_parallel_size, + pipeline_model_parallel_size, + context_model_parallel_size, + ] + ): + raise ValueError( + f"All arguments must be of type int, got {type(micro_batch_size)}, {type(num_nodes)}, {type(devices)}, " + f"{type(accumulate_grad_batches)}, {type(tensor_model_parallel_size)}, {type(pipeline_model_parallel_size)}, and {type(context_model_parallel_size)}" + ) + if micro_batch_size <= 0: + raise ValueError(f"micro_batch_size must be greater than 0, got {micro_batch_size}") + if num_nodes <= 0: + raise ValueError(f"num_nodes must be greater than 0, got {num_nodes}") + if devices <= 0: + raise ValueError(f"devices must be greater than 0, got {devices}") + if accumulate_grad_batches <= 0: + raise ValueError(f"accumulate_grad_batches must be greater than 0, got {accumulate_grad_batches}") + if tensor_model_parallel_size <= 0: + raise ValueError(f"tensor_model_parallel_size must be greater than 0, got {tensor_model_parallel_size}") + if pipeline_model_parallel_size <= 0: + raise ValueError(f"pipeline_model_parallel_size must be greater than 0, got {pipeline_model_parallel_size}") + if context_model_parallel_size <= 0: + raise ValueError(f"context_model_parallel_size must be greater than 0, got {context_model_parallel_size}") + + world_size = num_nodes * devices + if world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * context_model_parallel_size) != 0: + raise ValueError( + f"world_size must be divisible by tensor_model_parallel_size * pipeline_model_parallel_size * context_model_parallel_size, " + f"got {world_size} and TP{tensor_model_parallel_size} * PP{pipeline_model_parallel_size} * CP{context_model_parallel_size}" + ) + + model_parallel_size = tensor_model_parallel_size * pipeline_model_parallel_size * context_model_parallel_size + data_parallel_size = world_size // model_parallel_size + global_batch_size = micro_batch_size * data_parallel_size * accumulate_grad_batches + return global_batch_size + + +def infer_num_samples( + limit_batches: Union[float, int, str, None], num_samples_in_dataset: int, global_batch_size: int, stage: str +): + """Infers the number of samples based on the limit_batches parameter, the length of the dataset, and the global batch size. + + Args: + limit_batches (Union[float, int, str, None]): The limit on the number of batches. Can be a float + between 0 and 1, an integer, a string, or None. If None, defaults to 1.0. + num_samples_in_dataset (int): The number of samples in the dataset. + global_batch_size (int): The global batch size. + stage (str): The stage of the training. + + Returns: + int: The number of samples from the limit. + + Raises: + ValueError: If the limited number of samples is less than the global batch size, or if the + limit_batches parameter is invalid. + + If limit_batches is a float between 0 and 1, the number of samples is inferred as a fraction of the number of samples + in the dataset. If limit_batches is an integer greater than or equal to 1, the number of limited samples is inferred + as the product of limit_batches and global batch size. If limit_batches is None, it defaults to 1.0, indicating that + all dataset samples should be used. + """ + limit_batches = 1.0 if limit_batches is None else limit_batches # validation data does not require upsampling + if 0 < limit_batches <= 1.0 and isinstance(limit_batches, float): + num_limited_samples = int(num_samples_in_dataset * limit_batches) + if num_limited_samples < global_batch_size: + raise ValueError( + "The limited number of %s samples %s is less than the global batch size %s" + % (stage, num_limited_samples, global_batch_size) + ) + elif limit_batches >= 1 and isinstance(limit_batches, int): + num_limited_samples = int(limit_batches * global_batch_size) + else: + raise ValueError("Invalid choice of limit_%s_batches size: %s" % (stage, limit_batches)) + + return num_limited_samples diff --git a/model/evo2/utils/logger_utils.py b/model/evo2/utils/logger_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f3278d8fa3030158b6a62196e57dbd1e1fc2546e --- /dev/null +++ b/model/evo2/utils/logger_utils.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pathlib +from typing import Any, Dict, List, Optional, Sequence + +from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger +from nemo.lightning.nemo_logger import NeMoLogger +from nemo.lightning.pytorch import callbacks as nemo_callbacks +from nemo.utils import logging +from pydantic import BaseModel + + +__all__: Sequence[str] = ( + "WandbConfig", + "setup_nemo_lightning_logger", +) + + +class WandbConfig(BaseModel): + """Note: `name` controls the exp name is handled by the NeMoLogger so it is ommitted here. + `directory` is also omitted since it is set by the NeMoLogger. + + Args: + entity: The team posting this run (default: your username or your default team) + project: The name of the project to which this run will belong. + name: Display name for the run. By default it is set by NeMoLogger to experiment name + tags: Tags associated with this run. + group: A unique string shared by all runs in a given group + job_type: Type of run, which is useful when you're grouping runs together into larger experiments. + offline: Run offline (data can be streamed later to wandb servers). + id: Sets the version, mainly used to resume a previous run. + anonymous: Enables or explicitly disables anonymous logging. + """ # noqa: D205 + + entity: str | None # The team posting this run (default: your username or your default team) + project: str # The name of the project to which this run will belong. + name: str | None = None # Display name for the run. By default, it is set by NeMoLogger to experiment name + # save_dir: #Path where data is saved. "This is handled by NeMoLogger" + tags: List[str] | None # Tags associated with this run. + group: str | None # A unique string shared by all runs in a given group. + job_type: str | None = ( + None # Type of run, which is useful when you're grouping runs together into larger experiments. + ) + offline: bool # Run offline (data can be streamed later to wandb servers). + id: str | None # Sets the version, mainly used to resume a previous run. + anonymous: bool # Enables or explicitly disables anonymous logging. + log_model: bool # Save checkpoints in wandb dir to upload on W&B servers. + + +def setup_nemo_lightning_logger( + name: str | None = None, + root_dir: str | pathlib.Path = "./results", + initialize_tensorboard_logger: bool = False, + wandb_config: Optional[WandbConfig] = None, + ckpt_callback: Optional[nemo_callbacks.ModelCheckpoint] = None, + **kwargs: Dict[str, Any], +) -> NeMoLogger: + """Setup the logger for the experiment. + + Arguments: + name: The name of the experiment. Results go into `root_dir`/`name` + root_dir: The root directory to create the `name` directory in for saving run results. + initialize_tensorboard_logger: Whether to initialize the tensorboard logger. + wandb_config: The remaining configuration options for the wandb logger. + ckpt_callback: The checkpoint callback to use, must be a child of the pytorch lightning ModelCheckpoint callback. + NOTE the type annotation in the underlying NeMoCheckpoint constructor is incorrect. + **kwargs: The kwargs for the NeMoLogger. + + Returns: + NeMoLogger: NeMo logger instance. + """ + # The directory that the logger will save to + save_dir = pathlib.Path(root_dir) / name + save_dir.mkdir(parents=True, exist_ok=True) + + version = "dev" + if wandb_config is not None: + if wandb_config.name is None: + wandb_config.name = name + wandb_logger = WandbLogger(save_dir=save_dir, **wandb_config.model_dump()) + else: + wandb_logger = None + logging.warning("WandB is currently turned off.") + if initialize_tensorboard_logger: + tb_logger = TensorBoardLogger(save_dir=root_dir, name=name, version=version) + else: + tb_logger = None + logging.warning("User-set tensorboard is currently turned off. Internally one may still be set by NeMo2.") + + logger: NeMoLogger = NeMoLogger( + name=name, + log_dir=str(root_dir), + tensorboard=tb_logger, + wandb=wandb_logger, + ckpt=ckpt_callback, + version=version, + update_logger_directory=False, + **kwargs, + ) + # Needed so that the trainer can find an output directory for the profiler + logger.save_dir = save_dir + return logger diff --git a/model/evo2/utils/logging/__init__.py b/model/evo2/utils/logging/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..83e4f91e1a8871ad0f359d81c3c8f36c31fbdfe9 --- /dev/null +++ b/model/evo2/utils/logging/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/model/evo2/utils/logging/callbacks.py b/model/evo2/utils/logging/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..42808d9e79f8510462fc0877eb778acc5df71faa --- /dev/null +++ b/model/evo2/utils/logging/callbacks.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright The Lightning AI team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from lightning.pytorch import Callback +from megatron.core import parallel_state +from megatron.core.tensor_parallel.mappings import _gather_along_last_dim + + +class TEVCallback(Callback): + """Callback for logging TEV statistics before each optimizer step. + + This callback handles different parallelism strategies: + - Pipeline Parallelism: Only computes on first pipeline stage + - Tensor Parallelism: Gathers embedding shards across TP ranks + - Context Parallelism: Gathers across CP ranks + - Data Parallelism: Only logs on rank 0 of each model parallel group + """ + + @torch.no_grad() + def on_before_optimizer_step(self, trainer, pl_module, optimizer) -> None: + """Called before each optimizer step during training. + + This method calculates and logs Token Embedding Variance (TEV) statistics: + 1. Gets embedding parameter only on pipeline rank 0 (where embeddings live) + 2. Gathers embedding shards across tensor and context parallel ranks + 3. Calculates the token embedding variance (TEV) + 4. Logs the mean and standard deviation of TEV values only on data parallel rank 0 + + Args: + trainer: The Lightning trainer instance + pl_module: The current Lightning module being trained + optimizer: The optimizer being used + + Note: + The callback assumes embeddings live on pipeline rank 0, which is the standard + configuration in Megatron-LM. + """ + # Only compute on pipeline rank 0 where embeddings live + if not parallel_state.is_pipeline_first_stage(): + return + + # Get all named parameters from the model + named_params = dict(pl_module.named_parameters()) + + # Find all parameter keys containing 'embed' + embed_keys = [key for key in named_params.keys() if "embed" in key] + + # Validate we have exactly one embedding layer + if len(embed_keys) == 0: + raise ValueError("No embed keys found.") + if len(embed_keys) > 1: + raise ValueError("Multiple embed keys found.") + + # Get the embedding parameter + embed = named_params[embed_keys[0]] + + # If using tensor parallelism, gather embedding shards + if parallel_state.get_tensor_model_parallel_world_size() > 1: + embed = _gather_along_last_dim(embed, group=parallel_state.get_tensor_model_parallel_group()) + + # If using context parallelism, gather across context parallel ranks + if parallel_state.get_context_parallel_world_size() > 1: + world_size = parallel_state.get_context_parallel_world_size() + dim_size = list(embed.size()) + dim_size[0] = dim_size[0] * world_size + + output = torch.empty(dim_size, dtype=embed.dtype, device=torch.cuda.current_device()) + torch.distributed.all_gather_into_tensor( + output, embed.contiguous(), group=parallel_state.get_context_parallel_group() + ) + embed = output + + # Calculate token embedding variance (TEV) + # First center the embeddings by subtracting the mean + # Then calculate the mean squared deviation (variance) + # Finally take the square root to get standard deviation + tev = torch.sqrt(torch.mean(torch.pow(embed - embed.mean(dim=0), 2), dim=0)) + + # Calculate statistics of the TEV values + tev_mean = torch.mean(tev).item() + tev_sd = torch.std(tev).item() + + # Only log on data parallel rank 0 to avoid duplicate logging + if parallel_state.get_data_parallel_rank() == 0: + # Log the TEV statistics + pl_module.log("tev_mean", tev_mean, on_step=True, on_epoch=False, sync_dist=False) + pl_module.log("tev_sd", tev_sd, on_step=True, on_epoch=False, sync_dist=False) diff --git a/model/evo2/utils/loss/__init__.py b/model/evo2/utils/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..83e4f91e1a8871ad0f359d81c3c8f36c31fbdfe9 --- /dev/null +++ b/model/evo2/utils/loss/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/model/evo2/utils/loss/embedding_variance.py b/model/evo2/utils/loss/embedding_variance.py new file mode 100644 index 0000000000000000000000000000000000000000..e73fdc4a6d4f6ea227234b608b82ea1f0da834ac --- /dev/null +++ b/model/evo2/utils/loss/embedding_variance.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# This assumes you have a parallel_state module similar to Megatron-LM's. +# If not, you'll need to pass tp_world_size and tp_group directly if they +# are obtained differently. +# from megatron.core import parallel_state # Example import + + +import torch +import torch.distributed +from megatron.core import parallel_state +from torch.autograd import Function + + +class SquaredErrorTargetedVarianceLossFunction(Function): + """This loss function is used to calculate the loss based on the squared difference between the global mean of per-word variances and target.""" + + @staticmethod + def forward(ctx, we_weight: torch.Tensor, loss_coeff: float, var_target: float) -> torch.Tensor: + """Calculates a loss based on the squared difference between the global mean of per-word variances and target. + + Assumes vocab-parallel sharding for we_weight (dim 0 is sharded). + + Args: + ctx (torch.autograd.FunctionContext): Context object for backward pass. + we_weight (torch.Tensor): Local shard of embedding weights (V_local, H). + loss_coeff (float): Loss coefficient. + var_target (float): Targeted variance for the embedding weights. + + Returns: + torch.Tensor: Scalar loss value. + + weights + """ + if not we_weight.is_floating_point(): + we_weight = we_weight.float() + + V_local, H = we_weight.shape # V_local: words on this rank, H: embedding dim + + # Save dimensions for backward pass + ctx.H_embedding_dim = H + ctx.V_local_word_count = V_local + ctx.loss_coeff = loss_coeff + ctx.var_target = var_target + + # Handle H=0 edge case (embedding dimension is zero) + if H == 0: + ctx.is_H_dim_zero = True + # Mean variance is 0 if H=0. Loss is based on (0 - VAR_TARGET)^2. + loss_value = loss_coeff * (0.0 - var_target) ** 2 + final_loss_tensor = torch.tensor(loss_value, device=we_weight.device, dtype=we_weight.dtype) + # Save we_weight for shape, None for we_mean_per_word and V_final (as they are not well-defined or zero) + ctx.save_for_backward(we_weight, None, None) + return final_loss_tensor + ctx.is_H_dim_zero = False + + # Get TP info (assuming parallel_state is globally accessible) + # Ensure parallel_state is imported and available in the execution scope. + # from some_module import parallel_state # Make sure this is accessible + tp_world_size = parallel_state.get_tensor_model_parallel_world_size() or 1 + tp_group = parallel_state.get_tensor_model_parallel_group() # Can be None + ctx.tp_world_size_val = tp_world_size + + # 1. Per-word mean (across embedding dimension H) + # Shape: (V_local, 1) + we_mean_per_word = we_weight.mean(dim=1, keepdim=True) + + # 2. Per-word variance (across embedding dimension H) + # we_sq_diffs_per_word shape: (V_local, H) + we_sq_diffs_per_word = (we_weight - we_mean_per_word) ** 2 + # we_var_per_word_local shape: (V_local,) (biased variance) + we_var_per_word_local = we_sq_diffs_per_word.mean(dim=1, keepdim=False) + + # 3. Mean of these per-word variances *on this local rank* + # v_local_mean_of_vars shape: scalar tensor + v_local_mean_of_vars = torch.tensor(0.0, device=we_weight.device, dtype=we_weight.dtype) + if V_local > 0: # Avoid NaN from mean of empty tensor if V_local is 0 + v_local_mean_of_vars = we_var_per_word_local.mean(dim=0, keepdim=False) + + # 4. Globally average these local mean variances + # V_final_globally_avg_var is the V in the loss formula L = alpha*(V-T)^2 + V_final_globally_avg_var = v_local_mean_of_vars.clone() + if tp_world_size > 1: + # Computes V_final = (1/tp_world_size) * sum(v_local_mean_of_vars from each rank) + V_final_globally_avg_var /= tp_world_size + torch.distributed.all_reduce(V_final_globally_avg_var, group=tp_group, op=torch.distributed.ReduceOp.SUM) + + # 5. Calculate final loss: LOSS_COEFF * (V_final - VAR_TARGET)^2 + final_loss = loss_coeff * (V_final_globally_avg_var - var_target) ** 2 + + # Save tensors needed for gradient computation in backward + ctx.save_for_backward(we_weight, we_mean_per_word, V_final_globally_avg_var) + # Other necessary scalars (H, V_local, tp_world_size) are already on ctx. + + return final_loss + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]: + """Backward pass for the SquaredErrorTargetedVarianceLossFunction.""" + we_weight, we_mean_per_word, V_final_saved = ctx.saved_tensors + + # Handle H=0 edge case (gradient is zero) + if getattr(ctx, "is_H_dim_zero", False): + return torch.zeros_like(we_weight), None, None # Grad for we_weight only + + H = ctx.H_embedding_dim + V_local = ctx.V_local_word_count + tp_world_size = ctx.tp_world_size_val + loss_coeff = ctx.loss_coeff + var_target = ctx.var_target + + # Handle V_local=0 edge case (no words on this rank, so no gradient) + if V_local == 0: + return torch.zeros_like(we_weight), None, None # Grad for we_weight only + + # Chain rule: d(TotalLoss)/dw = d(TotalLoss)/d(final_loss) * d(final_loss)/dw + # grad_output is d(TotalLoss)/d(final_loss) + + # 1. Calculate d(final_loss) / d(V_final_saved) + # final_loss = LOSS_COEFF * (V_final_saved - VAR_TARGET)**2 + # dL_dV_final is d(final_loss) / d(V_final_saved) + dL_dV_final = loss_coeff * 2.0 * (V_final_saved - var_target) + + # grad_V_final is d(TotalLoss) / d(V_final_saved) + grad_V_final = grad_output * dL_dV_final # Scalar + + # 2. Propagate gradient from V_final_saved to v_local_mean_of_vars (on current rank) + # V_final_saved = (1/tp_world_size) * sum_k(v_local_mean_of_vars_k) + # So, d(V_final_saved) / d(v_local_mean_of_vars_current_rank) = 1 / tp_world_size + # grad_v_local_mean is d(TotalLoss) / d(v_local_mean_of_vars_current_rank) + grad_v_local_mean = grad_V_final * (1.0 / tp_world_size) # Scalar + + # 3. Propagate gradient from v_local_mean_of_vars to we_var_per_word_local_i + # v_local_mean_of_vars = mean(we_var_per_word_local) = (1/V_local) * sum_i(we_var_per_word_local_i) + # So, d(v_local_mean_of_vars) / d(we_var_per_word_local_i) = 1 / V_local + # The coefficient to apply for the next step of chain rule: + # This is grad_v_local_mean scaled by (1/V_local) + # This represents d(TotalLoss)/d(we_var_per_word_local_i), assuming it's uniform. + coeff_for_per_word_var_grad = grad_v_local_mean * (1.0 / V_local) # Scalar + + # 4. Propagate gradient from we_var_per_word_local_i to we_weight_ik + # we_var_per_word_local_i = (1/H) * sum_k (we_weight_ik - we_mean_per_word_i[0])^2 + # d(we_var_per_word_local_i) / d(we_weight_ik) = (2/H) * (we_weight_ik - we_mean_per_word_i[0]) + # The term (we_weight_ik - we_mean_per_word_i[0]) is (we_weight - we_mean_per_word) + + # Combine coefficients for the (we_weight - we_mean_per_word) term: + # This is coeff_for_per_word_var_grad * (2/H) + final_scalar_coefficient = coeff_for_per_word_var_grad * (2.0 / H) + + grad_we_weight = final_scalar_coefficient * (we_weight - we_mean_per_word) + + # The forward function only takes we_weight as a tensor input requiring grad, the other two inputs + # are floats and do not get gradients. + return grad_we_weight, None, None + + +class SquaredErrorTargetedVarianceLoss(torch.nn.Module): + """Applies a loss that will encourage variance of some parameter to be close to var_target.""" + + def __init__(self, loss_coeff: float = 0.1, var_target: float = 1.0): + """Applies a loss that will encourage variance of some parameter to be close to var_target. + + Args: + loss_coeff: Loss coefficient. Defaults to 0.1. + var_target: targetted variance for the embedding weights. Defaults to 1.0. + """ + super().__init__() + self.loss_coeff = loss_coeff + self.var_target = var_target + + def forward(self, we_weight: torch.Tensor) -> torch.Tensor: + """Applies the loss to the embedding weights with the user requested loss coefficient and targeted variance. + + Args: + we_weight: Embedding weights. + + Returns: + torch.Tensor: Loss value. + """ + return SquaredErrorTargetedVarianceLossFunction.apply(we_weight, self.loss_coeff, self.var_target) diff --git a/scripts/infer.py b/scripts/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..2a1e18742d796588c9851d2270ce7c2e954084ef --- /dev/null +++ b/scripts/infer.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +import sys +import time +from typing import Literal, Optional + +import nemo.lightning as nl +import torch +from megatron.core.inference.common_inference_params import CommonInferenceParams +from megatron.core.inference.inference_request import InferenceRequest +from nemo.collections.llm import inference +from nemo.utils import logging + + +CheckpointFormats = Literal["torch_dist", "zarr"] +DEFAULT_CKPT_DIR = os.environ.get( + "EVO2_CKPT_DIR", + os.path.join(_PROJECT_ROOT, "checkpoints", "evo2_nemo_7b"), +) + + +def parse_args(): + """Parse arguments for Evo2 inference.""" + ap = argparse.ArgumentParser() + + # generation args: + default_prompt = ( + "|d__Bacteria;" + + "p__Pseudomonadota;" + + "c__Gammaproteobacteria;" + + "o__Enterobacterales;" + + "f__Enterobacteriaceae;" + + "g__Escherichia;" + + "s__Escherichia|" + ) + ap.add_argument( + "--prompt", + type=str, + default=default_prompt, + help="Prompt to generate text from Evo2. Defaults to a phylogenetic lineage tag for E coli.", + ) + ap.add_argument( + "--ckpt-dir", + type=str, + default=DEFAULT_CKPT_DIR, + help="Path to checkpoint directory containing pre-trained Evo2 model.", + ) + ap.add_argument("--temperature", type=float, default=1.0, help="Temperature during sampling for generation.") + ap.add_argument("--top-k", type=int, default=0, help="Top K during sampling for generation.") + ap.add_argument("--top-p", type=float, default=0.0, help="Top P during sampling for generation.") + ap.add_argument("--max-new-tokens", type=int, default=1024, help="Maximum number of tokens to generate.") + ap.add_argument("--seed", type=int, default=None, help="Random seed for generation.") + # compute args: + ap.add_argument("--tensor-parallel-size", type=int, default=1, help="Order of tensor parallelism. Defaults to 1.") + ap.add_argument( + "--pipeline-model-parallel-size", type=int, default=1, help="Order of pipeline parallelism. Defaults to 1." + ) + ap.add_argument( + "--context-parallel-size", type=int, default=1, help="Order of context parallelism. Defaults to 1." + ) + # output args: + ap.add_argument( + "--output-file", + type=str, + default=None, + help="Output file containing the generated text produced by the Evo2 model. If not provided, the output will be logged.", + ) + # extra: + ap.add_argument( + "--ckpt-format", + type=str, + choices=["torch_dist", "zarr"], + default="torch_dist", + help="Specify checkpoint format to use. Defaults to 'torch_dist', as 'zarr' is deprecated.", + ) + ap.add_argument( + "--fp8", + action="store_true", + default=False, + help="Whether to use vortex style FP8. Defaults to False.", + ) + ap.add_argument( + "--flash-decode", + action="store_true", + default=False, + help="Whether to use flash decode. Defaults to True.", + ) + return ap.parse_args() + + +def infer( + prompt: str, + ckpt_dir: str, + temperature: float, + top_k: int, + top_p: float, + max_new_tokens: int, + tensor_parallel_size: int, + pipeline_model_parallel_size: int, + context_parallel_size: int, + output_file: Optional[str] = None, + ckpt_format: CheckpointFormats = "torch_dist", + seed: Optional[int] = None, + vortex_style_fp8: bool = False, + flash_decode: bool = False, + return_log_probs: bool = False, +) -> list[InferenceRequest]: + """Inference workflow for Evo2. + + Args: + prompt (str): Prompt to generate text from Evo2. + ckpt_dir (str): Path to checkpoint directory containing pre-trained Evo2 model. + temperature (float): Temperature during sampling for generation. + top_k (int): Top K during sampling for generation. + top_p (float): Top P during sampling for generation. + max_new_tokens (int): Maximum number of tokens to generate. + tensor_parallel_size (int): Order of tensor parallelism. + pipeline_model_parallel_size (int): Order of pipeline parallelism. + context_parallel_size (int): Order of context parallelism. + output_file (str): Output file containing the generated text produced by the Evo2 model. + ckpt_format (CheckpointFormats): Checkpoint format to use. + seed (int): Random seed for generation. + vortex_style_fp8 (bool): Whether to use vortex style FP8. + flash_decode (bool): Whether to use flash decode. + return_log_probs (bool): Whether to return log probabilities. + + Returns: + None + """ + model_parallel_size = tensor_parallel_size * pipeline_model_parallel_size * context_parallel_size + if model_parallel_size > torch.cuda.device_count(): + raise ValueError( + f"Requested model parallel size {model_parallel_size} is greater than the " + f"number of available CUDA devices {torch.cuda.device_count()}" + ) + # Create PTL trainer. + trainer = nl.Trainer( + accelerator="gpu", + devices=model_parallel_size, + strategy=nl.MegatronStrategy( + tensor_model_parallel_size=tensor_parallel_size, + pipeline_model_parallel_size=pipeline_model_parallel_size, + context_parallel_size=context_parallel_size, + pipeline_dtype=torch.bfloat16, + ckpt_load_optimizer=False, # Needs to be false for a normal model checkpoint. + ckpt_save_optimizer=False, + ckpt_async_save=False, + save_ckpt_format=ckpt_format, + ckpt_load_strictness="log_all", + ), + log_every_n_steps=1, + limit_val_batches=10, + num_sanity_val_steps=0, + plugins=nl.MegatronMixedPrecision( + precision="bf16-mixed", + params_dtype=torch.bfloat16, + ), + ) + inference_wrapped_model, mcore_tokenizer = inference.setup_model_and_tokenizer( + path=ckpt_dir, + trainer=trainer, + params_dtype=torch.bfloat16, + inference_batch_times_seqlen_threshold=8192, # TODO + inference_max_seq_length=8192, # TODO + recompute_granularity=None, + recompute_num_layers=None, + recompute_method=None, + vortex_style_fp8=vortex_style_fp8, + flash_decode=flash_decode, + enable_flash_decode=flash_decode, + ) + t0 = time.perf_counter_ns() + # TODO: fix return type in NeMo inference.generate (it is a list[InferenceRequest] not a dict) + results: list[InferenceRequest] = inference.generate( + model=inference_wrapped_model, + max_batch_size=1, # vortex only supports batch size 1 + tokenizer=mcore_tokenizer, + prompts=[prompt], + random_seed=seed, + inference_params=CommonInferenceParams( + temperature=temperature, + top_k=top_k, + top_p=top_p, + return_log_probs=return_log_probs, + num_tokens_to_generate=max_new_tokens, + ), + ) + dt = (time.perf_counter_ns() - t0) / 1e9 # seconds + tokens_per_sec = (len(results[0].generated_text) + 1) / dt # +1 for the prompt + + print(f"Inference time: {dt} seconds, {tokens_per_sec} tokens/sec", file=sys.stderr) + if torch.distributed.get_rank() == 0: + if output_file is None: + logging.info(results) + else: + with open(output_file, "w") as f: + f.write(f"{results[0]}\n") + + return results + + +def main(): + """Main function for Evo2 inference.""" + # Parse args. + args = parse_args() + if not args.ckpt_dir: + raise SystemExit( + "ERROR: Evo2 checkpoint path is required. Put the checkpoint under " + "checkpoints/evo2_nemo_7b, set EVO2_CKPT_DIR, or pass --ckpt-dir." + ) + if not os.path.isdir(args.ckpt_dir): + raise SystemExit(f"ERROR: Evo2 checkpoint directory does not exist: {args.ckpt_dir}") + infer( + prompt=args.prompt, + ckpt_dir=args.ckpt_dir, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + max_new_tokens=args.max_new_tokens, + tensor_parallel_size=args.tensor_parallel_size, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + context_parallel_size=args.context_parallel_size, + output_file=args.output_file, + ckpt_format=args.ckpt_format, + seed=args.seed, + vortex_style_fp8=args.fp8, # Vortex only applied FP8 to some layers. + flash_decode=args.flash_decode, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/inference.sh b/scripts/inference.sh new file mode 100644 index 0000000000000000000000000000000000000000..393746893286e9861c87c136f4e28e0850bf8c02 --- /dev/null +++ b/scripts/inference.sh @@ -0,0 +1,41 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +cd "${PROJECT_ROOT}" + +CKPT_DIR="${CKPT_DIR:-${EVO2_CKPT_DIR:-${PROJECT_ROOT}/checkpoints/evo2_nemo_7b}}" +PROMPT="${PROMPT:-ATGCGT}" +OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/outputs}" +OUTPUT_FILE="${OUTPUT_FILE:-${OUTPUT_DIR}/evo2_generation.txt}" +RUN_WITH_SRUN="${RUN_WITH_SRUN:-0}" + +if [[ -z "${CKPT_DIR}" ]]; then + echo "ERROR: Evo2 checkpoint path is required. Put it under checkpoints/evo2_nemo_7b or set EVO2_CKPT_DIR/CKPT_DIR." >&2 + exit 1 +fi +if [[ ! -d "${CKPT_DIR}" ]]; then + echo "ERROR: Evo2 checkpoint directory does not exist: ${CKPT_DIR}" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" + +echo "EVO2_PROJECT_ROOT: ${PROJECT_ROOT}" +echo "EVO2_CKPT_DIR: ${CKPT_DIR}" +echo "EVO2_OUTPUT_FILE: ${OUTPUT_FILE}" + +if [ "${RUN_WITH_SRUN}" = "1" ]; then + srun python scripts/infer.py \ + --ckpt-dir "${CKPT_DIR}" \ + --prompt "${PROMPT}" \ + --output-file "${OUTPUT_FILE}" +else + python scripts/infer.py \ + --ckpt-dir "${CKPT_DIR}" \ + --prompt "${PROMPT}" \ + --output-file "${OUTPUT_FILE}" +fi + +echo "Inference completed. Result file: ${OUTPUT_FILE}" diff --git a/scripts/predict.py b/scripts/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..6fc28b0c4632979014e79a30659d822fbf17e3a8 --- /dev/null +++ b/scripts/predict.py @@ -0,0 +1,535 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +import tempfile +from pathlib import Path +from typing import Literal, Optional + +import nemo.lightning as nl +import torch +from lightning.pytorch import LightningDataModule +from megatron.core import parallel_state +from megatron.core.tensor_parallel.mappings import _gather_along_last_dim +from megatron.core.utils import get_batch_on_this_cp_rank +from nemo.collections.llm.gpt.model.base import get_packed_seq_params +from nemo.collections.llm.gpt.model.hyena import HYENA_MODEL_OPTIONS, HyenaModel +from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer +from nemo.lightning import NeMoLogger +from nemo.lightning.data import WrappedDataLoader +from torch import Tensor + +from evo2.data.fasta_dataset import SimpleFastaDataset +from evo2.lightning import LightningPassthroughPredictionMixin +from evo2.utils.callbacks import PredictionWriter + + +CheckpointFormats = Literal["torch_dist", "zarr"] +DEFAULT_CKPT_DIR = Path( + os.environ.get( + "EVO2_CKPT_DIR", + Path(_PROJECT_ROOT) / "checkpoints" / "evo2_nemo_7b", + ) +) +DEFAULT_FASTA_PATH = Path( + os.environ.get( + "EVO2_FASTA_PATH", + Path(_PROJECT_ROOT) / "data" / "predict_example.fa", + ) +) +DEFAULT_OUTPUT_DIR = Path( + os.environ.get("EVO2_PREDICT_OUTPUT_DIR", Path(_PROJECT_ROOT) / "outputs" / "predict") +) + + +def parse_args(): + """Parse arguments for Evo2 inference.""" + ap = argparse.ArgumentParser() + + ap.add_argument( + "--fasta", + type=Path, + default=DEFAULT_FASTA_PATH, + help="Fasta path from which to generate logit predictions.", + ) + ap.add_argument( + "--ckpt-dir", + type=Path, + default=DEFAULT_CKPT_DIR, + help="NeMo2 checkpoint directory for inference.", + ) + ap.add_argument( + "--prepend-bos", + action="store_true", + help="Prepend BOS token to sequences. Defaults to False.", + ) + ap.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Order of tensor parallelism. Defaults to 1.", + ) + ap.add_argument( + "--pipeline-model-parallel-size", + type=int, + default=1, + help="Order of pipeline parallelism. Defaults to 1.", + ) + ap.add_argument( + "--context-parallel-size", + type=int, + default=1, + help="Order of context parallelism. Defaults to 1.", + ) + ap.add_argument( + "--no-sequence-parallel", + action="store_true", + help="When using TP, skip sequence parallelism. Otherwise sequence parallelism is used whenever tensor " + "parallelism is used. sequence parallelism should save a small amount of GPU memory so it's on" + " by default.", + ) + ap.add_argument( + "--batch-size", + type=int, + default=1, + help="Batch size for prediction. Defaults to 1.", + ) + ap.add_argument( + "--model-size", + type=str, + default="7b", + choices=sorted(HYENA_MODEL_OPTIONS.keys()), + help="Model size to use. Defaults to '7b'.", + ) + # output args: + ap.add_argument( + "--output-dir", + type=Path, + default=DEFAULT_OUTPUT_DIR, + help="Output dir that will contain the generated text produced by the Evo2 model. If not provided, the output will be logged.", + ) + ap.add_argument( + "--full-fp8", + action="store_true", + help="Use full FP8 precision (faster but less accurate) rather than vortex style which " + "only applies FP8 to the projection layer of the hyena mixer, when using FP8.", + ) + ap.add_argument( + "--fp8", action="store_true", help="Use FP8 precision. Defaults to BF16." + ) + # extra: + ap.add_argument( + "--ckpt-format", + type=str, + choices=["torch_dist", "zarr"], + default="torch_dist", + help="Specify checkpoint format to use. Defaults to 'torch_dist', as 'zarr' is deprecated.", + ) + ap.add_argument( + "--output-log-prob-seqs", + action="store_true", + help="Output log probability of sequences. Defaults to False.", + ) + ap.add_argument( + "--log-prob-collapse-option", + choices=["sum", "mean"], + default="mean", + help="How to collapse the log probabilities across the sequence dimension.", + ) + ap.add_argument( + "--hybrid-override-pattern", + type=str, + help="Override the hybrid override pattern in the config (specifies hyena layer ordering and type).", + ) + ap.add_argument( + "--num-layers", + type=int, + help="If set, override the number of layers specified in the requested config.", + ) + return ap.parse_args() + + +def _gather_along_cp_dim(input_, seq_dim: int = 1): + """Gather tensors and concatenate along the last dimension.""" + world_size = parallel_state.get_context_parallel_world_size() + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + + dim_size = list(input_.size()) + dim_size[0] = dim_size[0] * world_size + + output = torch.empty( + dim_size, dtype=input_.dtype, device=torch.cuda.current_device() + ) + torch.distributed.all_gather_into_tensor( + output, + input_.contiguous(), + group=parallel_state.get_tensor_model_parallel_group(), + ) + tensor_list = output.chunk(world_size, dim=0) + output = torch.cat(tensor_list, dim=seq_dim).contiguous() + + return output + + +class HyenaPredictor(LightningPassthroughPredictionMixin, HyenaModel): + """A predictor for the Hyena model. This adds in the predict step and the passthrough method.""" + + def __init__( + self, + *args, + output_log_prob_seqs: bool = False, + log_prob_collapse_option: Literal["sum", "mean"] = "mean", + **kwargs, + ): + """Initialize the predictor with our needs around computing log probabilities.""" + super().__init__(*args, **kwargs) + self.output_log_prob_seqs = output_log_prob_seqs + self.log_prob_collapse_option = log_prob_collapse_option + + def predict_step(self, batch, batch_idx: Optional[int] = None) -> Tensor: + """Alias for forward_step, also log the pad mask since sequences may not all have the same length.""" + if len(batch) == 0: + return + forward_out = self.forward_step(batch) + if not isinstance(forward_out, Tensor): + return forward_out + # Reminder: the model's predictions for input i land at output i+1. To get everything to align, we prepend the + # EOS token to the input sequences and take the outputs for all but the first token. + forward_out_tp_gathered = _gather_along_last_dim( + forward_out, group=parallel_state.get_tensor_model_parallel_group() + ) + # else: + # forward_out_tp_gathered = _collect_into_dim(forward_out, dim=-1) + forward_out_gathered = _gather_along_cp_dim(forward_out_tp_gathered) + assert self.tokenizer.vocab_size == forward_out_gathered.shape[-1] + if self.output_log_prob_seqs: + softmax_logprobs = torch.log_softmax(forward_out_gathered, dim=-1) + softmax_logprobs = softmax_logprobs[:, :-1] + input_ids = batch["tokens"][:, 1:] + assert softmax_logprobs.shape[1] == input_ids.shape[1] + + logprobs = torch.gather( + softmax_logprobs, # Gather likelihoods... + 2, # along the vocab dimension... + input_ids.unsqueeze(-1), # using the token ids to index. + ).squeeze(-1) + log_prob_seqs = torch.sum( + logprobs * batch["loss_mask"][:, 1:].float(), dim=-1 + ) + if self.log_prob_collapse_option == "mean": + log_prob_seqs = log_prob_seqs / ( + batch["loss_mask"][:, 1:].float().sum(dim=-1) + 1e-8 + ) + return { + "log_probs_seqs": log_prob_seqs.cpu(), + "seq_idx": batch["seq_idx"].cpu(), + } + else: + # If the user wants to match back to logits, then they will need to do the offsetting logic themselves. + return { + "token_logits": forward_out_gathered.cpu(), + "pad_mask": batch["loss_mask"].cpu(), + "seq_idx": batch["seq_idx"].cpu(), + } + + +def hyena_predict_forward_step(model, batch) -> torch.Tensor: + """Performs a forward step for the Hyena model. + + Args: + model: The Hyena model + batch: Dictionary containing input batch data with keys: + - tokens: Input token IDs + - position_ids: Position IDs + - labels: Labels for loss computation + - loss_mask: Mask for loss computation + + Returns: + torch.Tensor: Output from the model forward pass + """ + forward_args = { + "input_ids": batch["tokens"], + "position_ids": batch["position_ids"], + # "labels": batch["labels"], + # "loss_mask": batch["loss_mask"], + } + + forward_args["attention_mask"] = None + if "cu_seqlens" in batch: + forward_args["packed_seq_params"] = get_packed_seq_params(batch) + return model(**forward_args) + + +def hyena_predict_data_step(dataloader_iter) -> dict[str, torch.Tensor]: + """Data step for the Hyena model prediction. Modified from the original gpt data step to include the seq_idx.""" + from megatron.core import parallel_state + + # Based on: https://github.com/NVIDIA/Megatron-LM/blob/main/pretrain_gpt.py#L87 + # https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py#L828-L842 + + batch = next(dataloader_iter) + + _batch: dict + if isinstance(batch, tuple) and len(batch) == 3: + _batch = batch[0] + else: + _batch = batch + + required_device_keys = set() + required_host_keys = set() + + required_device_keys.add("attention_mask") + if "cu_seqlens" in _batch: + required_device_keys.add("cu_seqlens") + required_host_keys.add("cu_seqlens_argmin") + required_host_keys.add("max_seqlen") + + if parallel_state.is_pipeline_first_stage(): + required_device_keys.update(("tokens", "position_ids")) + if parallel_state.is_pipeline_last_stage(): + required_device_keys.update(("labels", "loss_mask", "seq_idx")) + + _batch_required_keys = {} + for key, val in _batch.items(): + if key in required_device_keys: + _batch_required_keys[key] = val.cuda(non_blocking=True) + elif key in required_host_keys: + _batch_required_keys[key] = val.cpu() + else: + _batch_required_keys[key] = None + + # slice batch along sequence dimension for context parallelism + output = get_batch_on_this_cp_rank(_batch_required_keys) + + return output + + +class PredictDataModule(LightningDataModule): + """Create a dataloader for prediction.""" + + def __init__(self, dataset: torch.utils.data.Dataset, batch_size: int = 1): + """Create a dataloader for prediction.""" + super().__init__() + self.dataset = dataset + self.batch_size = batch_size + + def setup(self, stage: Optional[str] = None) -> None: + """Set up the dataloader.""" + + def predict_dataloader(self): + """Create a dataloader for prediction.""" + # need to use this to communicate that we are in predict mode and safe to not drop last batch + return WrappedDataLoader( + mode="predict", + dataset=self.dataset, + batch_size=self.batch_size, + num_workers=8, + shuffle=False, + drop_last=False, + ) + + +def predict( + fasta_path: Path, + ckpt_dir: str, + output_dir: Path, + tensor_parallel_size: int, + pipeline_model_parallel_size: int, + context_parallel_size: int, + model_size: str = "7b", + ckpt_format: CheckpointFormats = "torch_dist", + fp8: bool = False, + full_fp8: bool = False, + work_dir: Path | None = None, + batch_size: int = 1, + output_log_prob_seqs: bool = False, + log_prob_collapse_option: Literal["sum", "mean"] = "mean", + prepend_bos: bool = False, + no_sequence_parallel: bool = False, + hybrid_override_pattern: str | None = None, + num_layers: int | None = None, +): + """Inference workflow for Evo2. + + Returns: + None + """ + if work_dir is None: + work_dir = Path(tempfile.mkdtemp()) + sequence_parallel = tensor_parallel_size > 1 and not no_sequence_parallel + output_dir.mkdir( + parents=True, exist_ok=True + ) # Make sure the output directory exists, files will be written here. + model_parallel_size = ( + tensor_parallel_size * pipeline_model_parallel_size * context_parallel_size + ) + if model_parallel_size > torch.cuda.device_count(): + raise ValueError( + f"Requested model parallel size {model_parallel_size} is greater than the " + f"number of available CUDA devices {torch.cuda.device_count()}" + ) + # Create PTL trainer. + trainer = nl.Trainer( + accelerator="gpu", + devices=model_parallel_size, + strategy=nl.MegatronStrategy( + drop_last_batch=False, + tensor_model_parallel_size=tensor_parallel_size, + pipeline_model_parallel_size=pipeline_model_parallel_size, + context_parallel_size=context_parallel_size, + pipeline_dtype=torch.bfloat16, + ckpt_load_optimizer=False, # Needs to be false for a normal model checkpoint. + ckpt_save_optimizer=False, + ckpt_async_save=False, + sequence_parallel=tensor_parallel_size > 1 and sequence_parallel, + save_ckpt_format=ckpt_format, + ckpt_load_strictness="log_all", + data_sampler=nl.MegatronDataSampler( + micro_batch_size=batch_size, + global_batch_size=batch_size, + seq_len=8192, + output_log=False, # this is needed for predict step to work + ), + ), + log_every_n_steps=1, + limit_val_batches=10, + num_sanity_val_steps=0, + callbacks=[ + PredictionWriter( + output_dir=output_dir, + write_interval="epoch", + batch_dim_key_defaults={"token_logits": 0}, + seq_dim_key_defaults={"token_logits": 1}, + ) + ], + plugins=nl.MegatronMixedPrecision( + precision="bf16-mixed", + params_dtype=torch.bfloat16, + # Only use FP8 in this plugin when using full FP8 precision and FP8. + # Otherwise use vortex_style_fp8 in the model config. + fp8="hybrid" if fp8 and full_fp8 else None, + fp8_amax_history_len=16 if fp8 and full_fp8 else 1, + fp8_amax_compute_algo="max" if fp8 and full_fp8 else "most_recent", + ), + ) + # The following two config options are really only used for testing, but may also be useful for getting output from + # specific layers of the model. + config_modifiers_init = {} + if hybrid_override_pattern is not None: + config_modifiers_init["hybrid_override_pattern"] = hybrid_override_pattern + if num_layers is not None: + config_modifiers_init["num_layers"] = num_layers + config = HYENA_MODEL_OPTIONS[model_size]( + forward_step_fn=hyena_predict_forward_step, + data_step_fn=hyena_predict_data_step, # , attention_backend=AttnBackend.fused, + distribute_saved_activations=( + False if sequence_parallel and tensor_parallel_size > 1 else True + ), + # Only use vortex style FP8 in the model config if using FP8 and not full FP8. This will only apply FP8 to + # the projection layer of the hyena mixer. + vortex_style_fp8=fp8 and not full_fp8, + **config_modifiers_init, + ) + trainer.strategy._setup_optimizers = False + + nemo_logger = NeMoLogger(log_dir=work_dir) + nemo_logger.setup(trainer, resume_if_exists=True) + resume = nl.AutoResume( + resume_if_exists=True, + resume_ignore_no_checkpoint=False, + resume_past_end=False, + restore_config=nl.RestoreConfig( + path=str(ckpt_dir), # NeMo expects a string path. + load_model_state=True, + load_optim_state=False, + load_artifacts=False, + ), + ) + tokenizer = get_nmt_tokenizer("byte-level") + model = HyenaPredictor( + config, + tokenizer=tokenizer, + output_log_prob_seqs=output_log_prob_seqs, + log_prob_collapse_option=log_prob_collapse_option, + ) + resume.setup(trainer, model) # this pulls weights from the starting checkpoint. + + dataset = SimpleFastaDataset(fasta_path, tokenizer, prepend_bos=prepend_bos) + datamodule = PredictDataModule(dataset, batch_size=batch_size) + trainer.predict(model, datamodule=datamodule) + dataset.write_idx_map( + output_dir + ) # Finally write out the index map so we can match the predictions to the original sequences. + print(f"Prediction completed. Result files are in: {output_dir}") + + +def main(): + """Entrypoint for Evo2 prediction (single inference step, no new tokens).""" + args = parse_args() + if args.ckpt_dir is None or str(args.ckpt_dir) in ("", "."): + raise SystemExit( + "ERROR: Evo2 checkpoint path is required. Put the checkpoint under " + "checkpoints/evo2_nemo_7b, set EVO2_CKPT_DIR, or pass --ckpt-dir." + ) + if not args.ckpt_dir.is_dir(): + raise SystemExit(f"ERROR: Evo2 checkpoint directory does not exist: {args.ckpt_dir}") + predict( + fasta_path=args.fasta, + ckpt_dir=args.ckpt_dir, + tensor_parallel_size=args.tensor_parallel_size, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + context_parallel_size=args.context_parallel_size, + output_dir=args.output_dir, + model_size=args.model_size, + ckpt_format=args.ckpt_format, + fp8=args.fp8, + full_fp8=args.full_fp8, + batch_size=args.batch_size, + output_log_prob_seqs=args.output_log_prob_seqs, + log_prob_collapse_option=args.log_prob_collapse_option, + prepend_bos=args.prepend_bos, + no_sequence_parallel=args.no_sequence_parallel, + hybrid_override_pattern=args.hybrid_override_pattern, + num_layers=args.num_layers, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/checkpoint_convert/convert_to_nemo.py b/scripts/tools/checkpoint_convert/convert_to_nemo.py new file mode 100644 index 0000000000000000000000000000000000000000..93c58c5cf5df8855eabc09c4df44b35ff7e29259 --- /dev/null +++ b/scripts/tools/checkpoint_convert/convert_to_nemo.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +import logging +from pathlib import Path + +from nemo.collections.llm.gpt.model.hyena import ( + HYENA_MODEL_OPTIONS, + HuggingFaceSavannaHyenaImporter, + HyenaConfig, + HyenaModel, + PyTorchHyenaImporter, +) +from nemo.lightning import io, teardown + + +def parse_args(): + """Parse command-line arguments.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--model-path", + type=str, + required=True, + help="Path to the Evo2 un-sharded (MP1) model checkpoint file, or a Hugging Face model name. Any model " + "from the Savanna Evo2 family is supported such as 'hf://arcinstitute/savanna_evo2_1b_base'.", + ) + parser.add_argument( + "--output-dir", + type=str, + required=True, + help="Output directory path for the converted model.", + ) + parser.add_argument( + "--model-size", + type=str, + choices=sorted(HYENA_MODEL_OPTIONS.keys()), + default="1b", + help="Model architecture to use, choose between 1b, 7b, 40b, or test (a sub-model of 4 layers, " + "less than 1B parameters). '*_arc_longcontext' models have GLU / FFN dimensions that support 1M " + "context length when trained with TP>>8.", + ) + parser.add_argument( + "--strip-optimizer", + action="store_true", + help="Strip the optimizer state from the model checkpoint, this works on nemo2 format checkpoints.", + ) + return parser.parse_args() + + +@io.model_importer(HyenaModel, "pytorch") +class HyenaOptimizerRemover(io.ModelConnector["HyenaModel", HyenaModel]): + """Removes the optimizer state from a nemo2 format model checkpoint.""" + + def __new__(cls, path: str, model_config=None): + """Creates a new importer instance. + + Args: + path: Path to the PyTorch model + model_config: Optional model configuration + + Returns: + PyTorchHyenaImporter instance + """ + instance = super().__new__(cls, path) + instance.model_config = model_config + return instance + + def init(self) -> HyenaModel: + """Initializes a new HyenaModel instance. + + Returns: + HyenaModel: Initialized model + """ + return HyenaModel(self.config, tokenizer=self.tokenizer) + + def get_source_model(self): + """Returns the source model.""" + model, _ = self.nemo_load(self) + return model + + def apply(self, output_path: Path, checkpoint_format: str = "torch_dist") -> Path: + """Applies the model conversion from PyTorch to NeMo format. + + Args: + output_path: Path to save the converted model + checkpoint_format: Format for saving checkpoints + + Returns: + Path: Path to the saved NeMo model + """ + source = self.get_source_model() + + target = self.init() + trainer = self.nemo_setup( + target, ckpt_async_save=False, save_ckpt_format=checkpoint_format + ) + source.to(self.config.params_dtype) + target.to(self.config.params_dtype) + self.convert_state(source, target) + self.nemo_save(output_path, trainer) + + logging.info(f"Converted Hyena model to Nemo, model saved to {output_path}") + + teardown(trainer, target) + del trainer, target + + return output_path + + def convert_state(self, source, target): + """Converts the state dictionary from source format to target format. + + Args: + source: Source model state + target: Target model + + Returns: + Result of applying state transforms + """ + mapping = {k: k for k in source.module.state_dict().keys()} + return io.apply_transforms( + source, + target, + mapping=mapping, + ) + + @property + def tokenizer(self): + """Gets the tokenizer for the model. + + Returns: + Tokenizer instance + """ + from nemo.collections.nlp.modules.common.tokenizer_utils import ( + get_nmt_tokenizer, + ) + + tokenizer = get_nmt_tokenizer( + library=self.model_config.tokenizer_library, + ) + + return tokenizer + + @property + def config(self) -> HyenaConfig: + """Gets the model configuration. + + Returns: + HyenaConfig: Model configuration + """ + return self.model_config + + +def main(): + """Convert a PyTorch Evo2 model checkpoint to a NeMo model checkpoint.""" + args = parse_args() + + evo2_config = HYENA_MODEL_OPTIONS[args.model_size]() + + if args.strip_optimizer: + importer = HyenaOptimizerRemover(args.model_path, model_config=evo2_config) + assert not args.model_path.startswith( + "hf://" + ), "Strip optimizer only works on local nemo2 format checkpoints." + elif args.model_path.startswith("hf://"): + importer = HuggingFaceSavannaHyenaImporter( + args.model_path.lstrip("hf://"), model_config=evo2_config + ) + else: + # import pdb; pdb.set_trace() + importer = PyTorchHyenaImporter(args.model_path, model_config=evo2_config) + # import pdb; pdb.set_trace() + importer.apply(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/data_process/preprocess_data_fasta.py b/scripts/tools/data_process/preprocess_data_fasta.py new file mode 100644 index 0000000000000000000000000000000000000000..fd302ea6dcc149b2774792bb2195c167e9be3cc7 --- /dev/null +++ b/scripts/tools/data_process/preprocess_data_fasta.py @@ -0,0 +1,611 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Module containing data preprocessing and splitting functions for Evo2 in BioNeMo. + +It can also be utilized as a script to dump pre-processed data to JSON. +""" + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +import multiprocessing as mp +import os +import random +import time +from contextlib import contextmanager +from pathlib import Path +from threading import Semaphore +from typing import Optional + +import numpy as np +import torch +import yaml + +from evo2.data.tokenizer import Evo2Tokenizer +from evo2.utils.config import Evo2PreprocessingConfig, Evo2TaxonomyLineage +from bionemo.noodles import back_transcribe_sequence, complement_sequence, reverse_sequence, transcribe_sequence +from bionemo.noodles.nvfaidx import NvFaidx +from megatron.core.datasets.indexed_dataset import IndexedDatasetBuilder +from nemo.utils import logging + +from evo2.data.tokenizer import Evo2Tokenizer +from evo2.utils.config import ( + Evo2PreprocessingConfig, + Evo2TaxonomyLineage, +) + + +class Evo2Preprocessor: + """Data preprocessing class for Evo2.""" + + BIN = ".bin" + IDX = ".idx" + TRAIN = "train" + VAL = "val" + TEST = "test" + + def __init__(self, params: Evo2PreprocessingConfig | None = None): + """Initialize Evo2Preprocessor. + + Args: + params (Evo2PreprocessingConfig | None): Configuration parameters for preprocessing. + """ + self.tokenizer: Evo2Tokenizer = Evo2Tokenizer(params) + + @staticmethod + @contextmanager + def preprocessing_context_manager(seed: Optional[int] = None): + """Context manager for setting and restoring the random number generator state. + + Args: + seed (int | None): Seed for the random number generator. Defaults to None. + """ + # Track current state. + current_state = random.getstate() + try: + # Set random seed. + random.seed(seed) + yield seed + finally: + # Restore random state. + random.setstate(current_state) + + @staticmethod + def _get_output_filename( + config: Evo2PreprocessingConfig, + ext: Optional[str] = None, + split: Optional[str] = None, + temp: bool = False, + ) -> Path: + """Generate the output filename for the preprocessed data. + + Args: + config (Evo2PreprocessingConfig): Configuration object containing preprocessing settings. + ext (Optional[str]): File extension for the output file. Defaults to None. + split (Optional[str]): Data split type (e.g., 'train', 'val', 'test'). Defaults to None. + temp (bool): Flag indicating whether the file is temporary. Defaults to False. + + Returns: + Path: The constructed output file path. + """ + # Get output directory. Defaults to CWD. + output_dir = config.output_dir + if output_dir is None: + output_dir = Path.cwd() + # Pickup output file prefix. + config_prefix = "{}_{}".format( + config.output_prefix, config.tokenizer_type.lower().replace(" ", "") + ) + output_filepath = Path(output_dir) / ( + config_prefix + + (f"_{split}" if split is not None else "") + + (ext if ext is not None else "") + + (".tmp" if temp else "") + ) + return output_filepath + + @staticmethod + def _subsequence_generator( + sequence: str, + subsequence_length: Optional[int] = None, + offset: Optional[int] = None, + ): + """Generate subsequences from a given sequence. + + Args: + sequence (str): The input sequence. + subsequence_length (int | None): Length of each subsequence. Defaults to the length of the sequence. + offset (int | None): Step size for generating subsequences. Defaults to subsequence_length. + + Yields: + str: Subsequences of the input sequence. + """ + subsequence_length = ( + subsequence_length if subsequence_length is not None else len(sequence) + ) + step_size = offset if offset is not None else subsequence_length + for i in range(0, len(sequence), step_size): + yield sequence[i : i + subsequence_length] + + @staticmethod + def _random_reverse_complement( + seq: str, prob: float = 0.0, seed: Optional[int] = None + ): + """Randomly reverse complements a DNA sequence based on a given probability. + + Args: + seq (str): The DNA sequence to potentially reverse complement. + prob (float): The probability of reverse complementing the sequence. Defaults to 0.0. + seed (Optional[int]): The seed for the random number generator. Defaults to None. + + Returns: + str: The original or reverse complemented DNA sequence based on the probability. + """ + with Evo2Preprocessor.preprocessing_context_manager(seed): + if random.random() < prob: + return complement_sequence(reverse_sequence(seq)) + else: + return seq + + @staticmethod + def _reverse_complement_expansion(seq: str): + """Generate a list containing the original and reverse complemented sequence. + + Args: + seq (str): The input DNA sequence. + + Returns: + list[str]: List containing the original and reverse complemented sequence. + """ + return [seq, complement_sequence(reverse_sequence(seq))] + + @staticmethod + def _train_val_test_split( + train_weight: float, + val_weight: float, + test_weight: float, + seed: Optional[int] = None, + ): + """Randomly assign a data point to train, validation, or test split based on provided weights. + + Args: + train_weight (float): The weight for the training split. + val_weight (float): The weight for the validation split. + test_weight (float): The weight for the test split. + seed (Optional[int]): The seed for the random number generator. Defaults to None. + + Returns: + str: The split assignment ('train', 'val', or 'test'). + + Raises: + ValueError: If the sum of the weights is zero or negative. + """ + with Evo2Preprocessor.preprocessing_context_manager( + seed if seed is not None else None + ): + # Generate random number. + roll = random.random() + # Rectify and normalize split ratios. + total_weight = abs(train_weight) + abs(val_weight) + abs(test_weight) + if total_weight <= 0: + raise ValueError( + "Train-validation-test split proportions cannot be zero." + ) + train_split = abs(train_weight) / total_weight + test_split = abs(test_weight) / total_weight + split = "train" + if roll > train_split: + if roll < 1 - test_split: + split = "val" + else: + split = "test" + return split + + @staticmethod + def _construct_taxonomy_token( + lineage: Evo2TaxonomyLineage, dropout: float = 0.0, seed: Optional[int] = None + ) -> Optional[str]: + """Construct a special Taxonomy token for natural language prompting of DNA generation models. + + Args: + lineage (Evo2TaxonomyLineage): The taxonomy lineage information. + dropout (float): The probability of dropping out segments of the lineage. Defaults to 0.0. + seed (Optional[int]): The seed for the random number generator. Defaults to None. + + Returns: + Optional[str]: The constructed taxonomy token or None if lineage is None. + """ + # If dropout > 0, randomly drop out segments of the lineage for training on incomplete lineages. + with Evo2Preprocessor.preprocessing_context_manager( + seed if seed is not None else None + ): + return ( + "|d__{};p__{};c__{};o__{};f__{};g__{};s__{}|".format( + lineage.domain if random.random() >= dropout else None, + lineage.phylum if random.random() >= dropout else None, + lineage.clazz if random.random() >= dropout else None, + lineage.order if random.random() >= dropout else None, + lineage.family if random.random() >= dropout else None, + lineage.genus if random.random() >= dropout else None, + lineage.species if random.random() >= dropout else None, + ) + if lineage is not None + else None + ) + + def preprocess_data( + self, + filepath: str, + seqid: str, + seq: str, + seq_idx: int, + config: Evo2PreprocessingConfig, + ): + """Preprocess fasta datapaths. + + Args: + filepath (str): Path to the .fasta file. + seqid (str): Sequence ID. + seq (str): DNA sequence. + seq_idx (int): Sequence index. + config (Evo2PreprocessingConfig): Configuration object containing preprocessing settings. + + Returns: + tuple[list[dict], float]: Preprocessed data and the time taken for preprocessing. + """ + # Timing. + start = time.time() + # Retrieve taxonomy lineage string if SeqID has associated taxonomy data. + # Note: Better implemented as a suffix tree substring dictionary, but convenient + # for identifying a large amount of sequences with identical lineages. + # Slow for extremely large dictionaries of (SeqID Substr, Taxonomy) pairs. + lineage = None + for id, tax in config.taxonomy_data.items(): + # Taxonomy ID is a substring of Seq ID. + if id in seqid: + lineage = tax + break + + # Preprocess data. + preproc_data = [] + with self.preprocessing_context_manager( + config.seed + hash(filepath) + seq_idx if config.seed is not None else None + ): + # Randomly reverse complement the sequence. + seq = self._random_reverse_complement( + seq, prob=config.random_reverse_complement + ) + seqs_to_parse = ( + self._reverse_complement_expansion(seq) + if config.embed_reverse_complement + else [seq] + ) + # 瀵圭敓鐗╁簭鍒楋紙DNA/RNA锛夎繘琛屾牸寮忓寲銆佷慨楗般€佺瓫閫夈€佹彃鍏ュ垎绫绘爣绛俱€佸垏鍒嗘垚瀛愬簭鍒椼€乼okenize锛屽苟鏈€缁堢粍缁囨垚缁撴瀯鍖栫殑璁粌鏁版嵁 + for seq in seqs_to_parse: + # Sequence Modifiers + if config.force_uppercase: + seq = seq.upper() + if config.transcribe == "transcribe": + seq = transcribe_sequence(seq) + elif config.transcribe == "back_transcribe": + seq = back_transcribe_sequence(seq) + if config.drop_empty_sequences and len(seq) == 0: + continue + if config.nnn_filter and "NNN" in seq.upper(): + continue + + # Construct taxonomy token (鍒嗙被鏍囩) with random dropout on the lineage categories per sequence. + taxonomy_token = self._construct_taxonomy_token( + lineage, dropout=config.random_lineage_dropout + ) + + # Inject taxonomy lineage tokens every prompt_spacer_length tokens in the sequence. + # If the taxonomy lineage token is not provided, then just take the original sequence. + target_length = ( + config.prompt_spacer_length - len(taxonomy_token) + if taxonomy_token is not None + else None + ) + taxonomy_injected_sequence = [ + ( + taxonomy_token + str(subseq) + if taxonomy_token is not None + else str(subseq) + ) + for subseq in self._subsequence_generator( + seq, target_length, target_length + ) + ] + # import pdb;pdb.set_trace() + # Wrap and tokenize. + preproc_data_record = { + "text": "".join(taxonomy_injected_sequence), + } + preproc_data_record["tokens"] = self.tokenizer.tokenize( + preproc_data_record["text"], + use_ftfy=config.ftfy, + enforce_sample_length=config.enforce_sample_length, + append_eod=config.append_eod, + drop_empty_sequences=config.drop_empty_sequences, + ) + preproc_data.append(preproc_data_record) + end = time.time() + return preproc_data, end - start + + def preprocess_data_task(self, file_sequence_config): + """Wrapper function to unpack args for preprocess_data. + + Args: + file_sequence_config (tuple): Tuple containing arguments for preprocess_data. + + Returns: + tuple[list[dict], float]: Preprocessed data and the time taken for preprocessing. + """ + return self.preprocess_data(*file_sequence_config) + + @staticmethod + def _yield_sequences_from_files( + config: Evo2PreprocessingConfig, semaphore: Semaphore + ): + # 鏁版嵁鐢熸垚鍣? + """Iterator over sequences within multiple input documents. Arguments for multiprocessing tasks. + + Utilized to limit the amount of sequences streamed into memory. + + Args: + config (Evo2PreprocessingConfig): Configuration object containing preprocessing settings. + semaphore (Semaphore): Semaphore to limit the number of sequences in memory. + + Yields: + tuple: Arguments for preprocess_data. + """ + + def yielder(fname, semaphore): + # Read FASTA. + index = NvFaidx(fname) + for i, (seqid, sequence) in enumerate(index.items()): + semaphore.acquire() + # Yield filename and sequence within fasta. + yield str(fname), seqid, sequence, i, config + + for fname in config.datapaths: + semaphore.acquire() + yield from yielder(fname, semaphore) + + def preprocess_generator(self, preproc_config: Evo2PreprocessingConfig): + """Main function to preprocess data for Evo2. + + Args: + preproc_config (Evo2PreprocessingConfig): Configuration object containing preprocessing settings. + + Yields: + tuple[dict, float]: Preprocessed sequence data and the time taken for preprocessing. + """ + # Track which splits have been assigned + split_assignments = { + "train": preproc_config.train_split > 0, + "val": preproc_config.valid_split > 0, + "test": preproc_config.test_split > 0, + } + splits_needed = {k for k, v in split_assignments.items() if v} + + # Instantiate multiprocessing pool. Use semaphore to limit the amount of sequences to read into memory. + semaphore = Semaphore( + preproc_config.preproc_concurrency + preproc_config.workers + ) + if preproc_config.workers > 1: + pool = mp.Pool(preproc_config.workers) + # Ordered imap for downstream seeded splitting. + preproc_tasks = pool.imap( + self.preprocess_data_task, + self._yield_sequences_from_files(preproc_config, semaphore), + chunksize=preproc_config.chunksize, + ) + else: + preproc_tasks = ( + self.preprocess_data_task(x) + for x in self._yield_sequences_from_files(preproc_config, semaphore) + ) + + # Preprocess data and split results into train, test, and split. + with self.preprocessing_context_manager( + preproc_config.seed if preproc_config.seed is not None else None + ): + for result, elapsed_time in preproc_tasks: + # Release semaphore for the task associated with the result. + semaphore.release() + # If we still need to ensure splits are assigned + if splits_needed: + # Force assign to a needed split + split = splits_needed.pop() + else: + # Regular random assignment + split = self._train_val_test_split( + preproc_config.train_split, + preproc_config.valid_split, + preproc_config.test_split, + ) + for sequence in result: + sequence["split"] = split + yield sequence, elapsed_time + + def preprocess_offline(self, preproc_config: Evo2PreprocessingConfig): + """Offline data preprocessing script for Evo2. + + Args: + preproc_config (Evo2PreprocessingConfig): Configuration object containing preprocessing settings. + """ + # Validate if binaries have already been produced for the given config and overwrite is set to False. + if any( + self._get_output_filename(preproc_config, ext, split).is_file() + for ext, split in zip( + [self.BIN, self.IDX], [self.TRAIN, self.VAL, self.TEST] + ) + ): + if not preproc_config.overwrite: + # Skip this dataset! + logging.info( + f"Skipped overwriting (overwrite: False) existing preprocessed data: {preproc_config.output_prefix}" + ) + return + else: + logging.info( + f"Overwriting (overwrite: True) existing preprocessed data: {preproc_config.output_prefix}" + ) + + # Instantiate indexed data builders. + dataset_dtype = getattr(np, preproc_config.indexed_dataset_dtype) + temp_train_bin = self._get_output_filename( + preproc_config, self.BIN, self.TRAIN, temp=True + ) + temp_val_bin = self._get_output_filename( + preproc_config, self.BIN, self.VAL, temp=True + ) + temp_test_bin = self._get_output_filename( + preproc_config, self.BIN, self.TEST, temp=True + ) + train_builder: IndexedDatasetBuilder = IndexedDatasetBuilder( + bin_path=str(temp_train_bin), dtype=dataset_dtype + ) + val_builder: IndexedDatasetBuilder = IndexedDatasetBuilder( + bin_path=str(temp_val_bin), dtype=dataset_dtype + ) + test_builder: IndexedDatasetBuilder = IndexedDatasetBuilder( + bin_path=str(temp_test_bin), dtype=dataset_dtype + ) + logging.info( + f"Created temporary binary datasets: {temp_train_bin} {temp_val_bin} {temp_test_bin}" + ) + + # Preprocess data and split results into train, validation, or test. + avg_preproc_time = 0.0 + avg_index_time = 0.0 + count = 0 + for sequence, elapsed_time in self.preprocess_generator(preproc_config): + index_start_time = time.time() + if sequence["split"] == "train": + train_builder.add_item(torch.Tensor(sequence["tokens"])) + train_builder.end_document() + elif sequence["split"] == "val": + val_builder.add_item(torch.Tensor(sequence["tokens"])) + val_builder.end_document() + elif sequence["split"] == "test": + test_builder.add_item(torch.Tensor(sequence["tokens"])) + test_builder.end_document() + index_end_time = time.time() + # Update average preprocessing and indexing time. + avg_preproc_time = (avg_preproc_time * count + elapsed_time) / (count + 1) + avg_index_time = ( + avg_index_time * count + index_end_time - index_start_time + ) / (count + 1) + count += 1 + + # Report timing. + logging.info(f"Average preprocessing time per sequence: {avg_preproc_time}") + logging.info(f"Average indexing time per sequence: {avg_index_time}") + logging.info(f"Number of sequences processed: {count}") + + # Write preprocessed index data to disk. Rename temporary binaries to denote preprocessing completion. + train_builder.finalize( + idx_path=str( + self._get_output_filename(preproc_config, self.IDX, self.TRAIN) + ) + ) + val_builder.finalize( + idx_path=str(self._get_output_filename(preproc_config, self.IDX, self.VAL)) + ) + test_builder.finalize( + idx_path=str(self._get_output_filename(preproc_config, self.IDX, self.TEST)) + ) + os.rename( + temp_train_bin, + self._get_output_filename(preproc_config, self.BIN, self.TRAIN), + ) + os.rename( + temp_val_bin, self._get_output_filename(preproc_config, self.BIN, self.VAL) + ) + os.rename( + temp_test_bin, + self._get_output_filename(preproc_config, self.BIN, self.TEST), + ) + + +def parse_args(): + """Parse arguments for preprocessing.""" + parser = argparse.ArgumentParser( + description="Preprocess FASTA files for training Evo2." + ) + parser.add_argument( + "-c", + "--config", + type=str, + required=True, + help="Path to data preprocessing config JSON.", + ) + return parser.parse_args() + + +def main(): + """Main function to execute the preprocessing script. + + This function parses command-line arguments, reads the configuration file, + and initiates the preprocessing of data as specified in the configuration. + """ + # Parse arguments. + args = parse_args() + # Read config YAML. + with open(args.config, "r") as yaml_fs: + evo2_preproc_config_batch = yaml.safe_load(yaml_fs) + # import pdb;pdb.set_trace() + for config in evo2_preproc_config_batch: + start = time.time() + # Convert into Evo2PreprocessingConfig. + evo2_preproc_config = Evo2PreprocessingConfig(**config) + if evo2_preproc_config.output_dir is not None: + evo2_preproc_config.output_dir.mkdir(parents=True, exist_ok=True) + # Instantiate Evo2Preprocessor. + evo2_preprocessor = Evo2Preprocessor(evo2_preproc_config) + # Preprocess data specified in config. + evo2_preprocessor.preprocess_offline(evo2_preproc_config) + end = time.time() + logging.info( + f"Finished preprocessing {evo2_preproc_config.output_prefix} ({evo2_preproc_config.datapaths}) in {end - start:.3f} seconds with {evo2_preproc_config.workers} workers." + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/data_process/preprocess_data_fasta.sh b/scripts/tools/data_process/preprocess_data_fasta.sh new file mode 100644 index 0000000000000000000000000000000000000000..9173feb1c40be87874fc2122a6c58a996821e719 --- /dev/null +++ b/scripts/tools/data_process/preprocess_data_fasta.sh @@ -0,0 +1,29 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +DATA_ROOT="${DATA_ROOT:-${PROJECT_ROOT}/data/data_mini}" +GENOME_DIR="${GENOME_DIR:-${DATA_ROOT}/genome_data}" + +mkdir -p "${GENOME_DIR}" +cd "${GENOME_DIR}" + +echo "EVO2_PROJECT_ROOT: ${PROJECT_ROOT}" +echo "EVO2_GENOME_DIR: ${GENOME_DIR}" + +for chr in chr20 chr21 chr22; do + if [ ! -f "${chr}.fa.gz" ]; then + wget -c "https://hgdownload.soe.ucsc.edu/goldenpath/hg38/chromosomes/${chr}.fa.gz" + fi + if [ ! -f "${chr}.fa" ]; then + zcat "${chr}.fa.gz" > "${chr}.fa" + fi +done + +cat chr20.fa chr21.fa chr22.fa > chr20_21_22.fa + +python "${PROJECT_ROOT}/scripts/tools/data_process/preprocess_data_fasta.py" \ + --config "${PROJECT_ROOT}/config/genome_preprocess_config.yaml" + +echo "FASTA preprocessing completed: ${GENOME_DIR}/preprocessed_data" \ No newline at end of file diff --git a/scripts/tools/data_process/preprocess_data_json.py b/scripts/tools/data_process/preprocess_data_json.py new file mode 100644 index 0000000000000000000000000000000000000000..609ad1ca1a65627c116196dab9620a3dee1b2d95 --- /dev/null +++ b/scripts/tools/data_process/preprocess_data_json.py @@ -0,0 +1,428 @@ +"""Processing data for pretraining.""" + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +import multiprocessing +import os +import sys + +import lm_dataformat as lmd +import numpy as np + +sys.path.append( + os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)) +) +import time +from abc import ABC, abstractmethod +from threading import Semaphore +from typing import List, Union + +import ftfy +import tqdm + +# from evo2.tokenizer import build_tokenizer +from evo2.data import indexed_dataset + + +def build_tokenizer(args): + """Initialize tokenizer.""" + if args.rank == 0: + print("> building {} tokenizer ...".format(args.tokenizer_type), flush=True) + + # Select and instantiate the tokenizer. + + if args.tokenizer_type.lower() == "CharLevelTokenizer".lower(): + tokenizer = CharLevelTokenizer(vocab_size=512) + else: + raise NotImplementedError( + "{} tokenizer is not " "implemented.".format(args.tokenizer_type) + ) + + # Add vocab size. + args.padded_vocab_size = _vocab_size_with_padding(tokenizer.vocab_size, args) + + return tokenizer + + +def _vocab_size_with_padding(orig_vocab_size, args): + """Pad vocab size so it is divisible by model parallel size and + still having GPU friendly size.""" + + after = orig_vocab_size + multiple = args.make_vocab_size_divisible_by * args.model_parallel_size + while (after % multiple) != 0: + after += 1 + if args.rank == 0: + print( + " > padded vocab (size: {}) with {} dummy tokens " + "(new size: {})".format(orig_vocab_size, after - orig_vocab_size, after), + flush=True, + ) + return after + + +class AbstractTokenizer(ABC): + """Abstract class for tokenizer.""" + + def __init__(self, name): + self.name = name + super().__init__() + + @property + @abstractmethod + def vocab_size(self): + pass + + @property + @abstractmethod + def vocab(self): + """Dictionary from vocab text token to id token.""" + + @property + @abstractmethod + def inv_vocab(self): + """Dictionary from vocab id token to text token.""" + + @abstractmethod + def tokenize(self, text): + pass + + def detokenize(self, token_ids): + raise NotImplementedError( + "detokenizer is not implemented for {} " "tokenizer".format(self.name) + ) + + @property + def cls(self): + raise NotImplementedError( + "CLS is not provided for {} " "tokenizer".format(self.name) + ) + + @property + def sep(self): + raise NotImplementedError( + "SEP is not provided for {} " "tokenizer".format(self.name) + ) + + @property + def pad(self): + raise NotImplementedError( + "PAD is not provided for {} " "tokenizer".format(self.name) + ) + + @property + def eod(self): + raise NotImplementedError( + "EOD is not provided for {} " "tokenizer".format(self.name) + ) + + @property + def mask(self): + raise NotImplementedError( + "MASK is not provided for {} " "tokenizer".format(self.name) + ) + + +class CharLevelTokenizer(AbstractTokenizer): + """Character Level Tokenizer""" + + def __init__(self, vocab_size): + name = "CharLevelTokenizer" + super().__init__(name) + self._vocab_size = vocab_size + self.eod_id = 0 + self.pad_id = 1 + + self._used_tokens = set() # 璁板綍鐢ㄨ繃鐨?token id + + def clamp(self, n): + return max(32, min(n, self.vocab_size)) + + @property + def vocab_size(self): + return self._vocab_size + + @property + def vocab(self): + raise NotImplementedError + + @property + def inv_vocab(self): + raise NotImplementedError + + def decode_token(self, token: int): + return str(chr(self.clamp(token))) + + # def tokenize(self, text: str): + # return list(np.fromstring(text, dtype=np.uint8)) + + def tokenize(self, text: str): + tokens = list(np.fromstring(text, dtype=np.uint8)) + # 璁板綍鎵€鏈?clamp 鍚庣殑 token id + clamped_tokens = [self.clamp(t) for t in tokens] + self._used_tokens.update(clamped_tokens) + return tokens + + def tokenize_batch(self, text_batch: Union[List[str], str]): + if isinstance(text_batch, list): + return [self.tokenize(s) for s in text_batch] + else: + return self.tokenize(text_batch) + + def detokenize(self, token_ids): + return "".join(list(map(self.decode_token, token_ids))) + + @property + def eod(self): + return self.eod_id + + @property + def pad(self): + return self.pad_id + + +class Encoder(object): + def __init__(self, args): + self.args = args + + def initializer(self): + # Use Encoder class as a container for global data + Encoder.tokenizer = build_tokenizer(self.args) + + def encode(self, text): + if self.args.ftfy: + text = ftfy.fix_text(text) + ids = {} + for key in self.args.jsonl_keys: + doc_ids = [] + text_ids = Encoder.tokenizer.tokenize(text) + + if ( + self.args.enforce_sample_length + and (len(text_ids) + int(self.args.append_eod)) + > self.args.enforce_sample_length + ): + raise ValueError( + "Detected input text with a length greater than the maximum " + f"possible sample length of {self.args.enforce_sample_length}.)" + ) + if len(text_ids) > 0: + doc_ids.append(text_ids) + if self.args.append_eod: + doc_ids[-1].append(Encoder.tokenizer.eod) + if self.args.enforce_sample_length: + # Pad up to max sequence length. + doc_ids[-1] += [Encoder.tokenizer.pad] * ( + self.args.enforce_sample_length - len(doc_ids[-1]) + ) + ids[key] = doc_ids + return ids, len(text) + + +def get_args(): + parser = argparse.ArgumentParser() + group = parser.add_argument_group(title="input data") + group.add_argument( + "--input", + type=str, + required=True, + help="Path to input jsonl files or lmd archive(s) - if using multiple archives, put them in a comma separated " + "list", + ) + group.add_argument( + "--jsonl-keys", + nargs="+", + default=["text"], + help="space separate listed of keys to extract from jsonl. Defa", + ) + group.add_argument( + "--num-docs", + default=None, + help="Optional: Number of documents in the input data (if known) for an accurate progress bar.", + type=int, + ) + group = parser.add_argument_group(title="tokenizer") + group.add_argument( + "--tokenizer-type", + type=str, + required=True, + choices=[ + "HFGPT2Tokenizer", + "HFTokenizer", + "GPT2BPETokenizer", + "CharLevelTokenizer", + "TiktokenTokenizer", + ], + help="What type of tokenizer to use.", + ) + group.add_argument( + "--vocab-file", type=str, default=None, help="Path to the vocab file" + ) + group.add_argument( + "--merge-file", + type=str, + default=None, + help="Path to the BPE merge file (if necessary).", + ) + group.add_argument( + "--append-eod", + action="store_true", + help="Append an token to the end of a document.", + ) + group.add_argument( + "--enforce-sample-length", + type=int, + default=None, + help="Forces all samples to have the specified length. If shorter, pads up to the length. If longer, throws an error.", + ) + group.add_argument("--ftfy", action="store_true", help="Use ftfy to clean text") + group = parser.add_argument_group(title="output data") + group.add_argument( + "--output-prefix", + type=str, + required=True, + help="Path to binary output file without suffix", + ) + group.add_argument( + "--dataset-impl", + type=str, + default="mmap", + choices=["lazy", "cached", "mmap"], + help="Dataset implementation to use. Default: mmap", + ) + + group = parser.add_argument_group(title="runtime") + group.add_argument( + "--workers", type=int, default=1, help="Number of worker processes to launch" + ) + group.add_argument( + "--log-interval", + type=int, + default=100, + help="Interval between progress updates", + ) + args = parser.parse_args() + args.keep_empty = False + + # some default/dummy values for the tokenizer + args.rank = 0 + args.make_vocab_size_divisible_by = 128 + args.model_parallel_size = 1 + + return args + + +def yield_from_files(fnames: list, semaphore): + """ + Iterator over input documents using lm_dataformat. Should be able to handle jsons / texts / + other compressed formats. Also filters out empty documents. + + :param fnames: list of filenames + """ + + def yielder(fname, semaphore): + for f in filter(lambda x: x, lmd.Reader(fname).stream_data()): + semaphore.acquire() + # import pdb;pdb.set_trace() + yield f + + for fname in fnames: + semaphore.acquire() + + yield from yielder(fname, semaphore) + + +def main(): + args = get_args() + encoder = Encoder(args) + tokenizer = build_tokenizer(args) + print(f"Vocab size: {tokenizer.vocab_size}") + print(f"Output prefix: {args.output_prefix}") + + # build a semaphore object to stop `yield_from_files` from getting ahead of encoder.encode and + # hence building up memory + semaphore = Semaphore(10000 + args.workers) + + # use multiprocessing to iterate over input documents + # import pdb;pdb.set_trace() + fin = yield_from_files(args.input.split(","), semaphore) + print(fin) + if args.workers > 1: + pool = multiprocessing.Pool(args.workers, initializer=encoder.initializer) + encoded_docs = pool.imap(encoder.encode, fin, chunksize=25) + else: + encoder.initializer() + encoded_docs = (encoder.encode(doc) for doc in fin) + # print(Encoder.tokenizer._used_tokens) + # make a dataset builder for each key in args.jsonl_keys + # each key will output to a different file beginning with args.output_prefix + output_bin_files = {} + output_idx_files = {} + builders = {} + tokenizer_name = tokenizer.name.replace(" ", "") + for key in args.jsonl_keys: + output_bin_files[key] = "{}_{}_{}_{}.bin".format( + args.output_prefix, key, tokenizer_name, "document" + ) + output_idx_files[key] = "{}_{}_{}_{}.idx".format( + args.output_prefix, key, tokenizer_name, "document" + ) + builders[key] = indexed_dataset.make_builder( + output_bin_files[key], + impl=args.dataset_impl, + vocab_size=tokenizer.vocab_size, + ) + + # actually do tokenization + proc_start = time.time() + total_bytes_processed = 0 + pbar = tqdm.tqdm() + for i, (doc, bytes_processed) in enumerate(encoded_docs, start=1): + total_bytes_processed += bytes_processed + + # release semaphore so `yield_from_files` can add another file to the buffer + semaphore.release() + + # add each tokenized document / sentence + for key, sentences in doc.items(): + for sentence in sentences: + builders[key].add_item(np.array(sentence, dtype=builders[key].dtype)) + # tell the builder that a document has finished + builders[key].end_document() + + # log progress + if i % args.log_interval == 0: + current = time.time() + elapsed = current - proc_start + mbs = total_bytes_processed / elapsed / 1024 / 1024 + pbar.set_description( + f"Processed {i}{'' if args.num_docs is None else '/' + str(args.num_docs)} documents ({i / elapsed} docs/s, {mbs} MB/s)." + ) + if i != 0: + pbar.update(args.log_interval) + + # save output file + for key in args.jsonl_keys: + builders[key].finalize(output_idx_files[key]) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/data_process/preprocess_data_json.sh b/scripts/tools/data_process/preprocess_data_json.sh new file mode 100644 index 0000000000000000000000000000000000000000..9fb0c4a9f0fbee18708b3d6130912677e89a8326 --- /dev/null +++ b/scripts/tools/data_process/preprocess_data_json.sh @@ -0,0 +1,30 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +INPUT_DIR="${INPUT_DIR:-${PROJECT_ROOT}/data/json/pretraining_or_both_phases/promoters}" +OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/data/promoters/pretraining_data_promoters}" +TOKENIZER="${TOKENIZER:-CharLevelTokenizer}" + +mkdir -p "${OUTPUT_DIR}" + +echo "EVO2_PROJECT_ROOT: ${PROJECT_ROOT}" +echo "EVO2_JSON_INPUT_DIR: ${INPUT_DIR}" +echo "EVO2_JSON_OUTPUT_DIR: ${OUTPUT_DIR}" + +for SPLIT in train test valid; do + INPUT_FILE="${INPUT_DIR}/data_promoters_${SPLIT}_chunk1.jsonl.gz" + OUTPUT_PREFIX="${OUTPUT_DIR}/data_promoters_${SPLIT}" + + echo "Processing ${INPUT_FILE} -> ${OUTPUT_PREFIX}" + python "${PROJECT_ROOT}/scripts/tools/data_process/preprocess_data_json.py" \ + --input "${INPUT_FILE}" \ + --output-prefix "${OUTPUT_PREFIX}" \ + --tokenizer-type "${TOKENIZER}" \ + --dataset-impl mmap \ + --append-eod \ + --workers 8 \ + --log-interval 100 + echo "Finished processing split: ${SPLIT}" +done \ No newline at end of file diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..be4f92a7cc7c33eaa07458c3a4f6cdcfaa787f28 --- /dev/null +++ b/scripts/train.py @@ -0,0 +1,927 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +from pathlib import Path +from typing import List, Optional + +# TODO add back support for slurm resilience. +# import nvidia_resiliency_ext.ptl_resiliency as res_module +import torch +from lightning.pytorch.callbacks import LearningRateMonitor, RichModelSummary +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.optimizer import OptimizerConfig +from nemo import lightning as nl +from nemo.collections import llm +from nemo.collections.llm.gpt.data import MockDataModule, PreTrainingDataModule +from nemo.collections.llm.gpt.data.megatron.hyena.config import parse_dataset_config +from nemo.collections.llm.gpt.data.megatron.hyena.evo2_dataset import ( + Evo2Dataset, + Evo2DatasetPadEodLossMask, +) +from nemo.collections.llm.gpt.model.hyena import HYENA_MODEL_OPTIONS +from nemo.collections.llm.recipes.tp_overlap_configs.userbuffers import ( + userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192, + userbuffers_fp8_h100_h8192_tp4_mbs1_seqlen8192, +) +from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer +from nemo.lightning.pytorch import callbacks as nl_callbacks +from nemo.lightning.pytorch.callbacks.flops_callback import FLOPsMeasurementCallback +from nemo.lightning.pytorch.callbacks.megatron_comm_overlap import ( + MegatronCommOverlapCallback, +) +from nemo.lightning.pytorch.optim import CosineAnnealingScheduler +from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule +from nemo.lightning.pytorch.strategies.utils import RestoreConfig +from nemo.utils.exp_manager import TimingCallback + +from evo2.utils.datamodule_utils import infer_global_batch_size +from evo2.utils.logger_utils import ( + WandbConfig, + setup_nemo_lightning_logger, +) + + +torch._dynamo.config.suppress_errors = True + +DEFAULT_DATASET_DIR = os.environ.get( + "EVO2_DATASET_DIR", + str(Path(_PROJECT_ROOT) / "data" / "data_mini" / "genome_data"), +) +DEFAULT_CKPT_DIR = None + + +def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace: + """Parse arguments for Evo2 model training.""" + parser = argparse.ArgumentParser( + description="Train a Hyena model using NeMo 2.0.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + data_group = parser.add_mutually_exclusive_group(required=True) + + data_group.add_argument( + "-d", + "--dataset-config", + type=str, + help="Path to the blended / weighted training dataset configuration YAML.", + ) + data_group.add_argument( + "--mock-data", + action="store_true", + help="Train with Mock data (for testing/debugging), either set this or provide a dataset config.", + ) + + parser.add_argument( + "--dataset-dir", + type=str, + default=DEFAULT_DATASET_DIR, + help="Absolute path to the dataset directory. Defaults to using the absolute or relative paths (dataset_prefix) specified in the dataset config YAML.", + ) + + parser.add_argument( + "--num-nodes", + type=int, + default=1, + help="Number of nodes to use for training, defaults to 1.", + ) + parser.add_argument( + "--devices", + type=int, + default=1, + help="Number of devices to use for training, defaults to 1.", + ) + parser.add_argument( + "--seq-length", type=int, default=8192, help="Training sequence length" + ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Order of tensor parallelism. Defaults to 1.", + ) + parser.add_argument( + "--pipeline-model-parallel-size", + type=int, + default=1, + help="Order of pipeline parallelism. Defaults to 1.", + ) + parser.add_argument( + "--context-parallel-size", + type=int, + default=1, + help="Order of context parallelism. Defaults to 1.", + ) + parser.add_argument( + "--create-tensorboard-logger", + action="store_true", + default=False, + help="Create a tensorboard logger.", + ) + parser.add_argument( + "--wandb-entity", type=str, default=None, help="The team posting this run" + ) + parser.add_argument( + "--wandb-project", type=str, default=None, help="Wandb project name " + ) + parser.add_argument( + "--wandb-tags", + nargs="+", + type=str, + default=None, + help="Tags associated with this run", + ) + parser.add_argument( + "--wandb-group", + type=str, + default=None, + help="A unique string shared by all runs in a given group", + ) + parser.add_argument( + "--wandb-job-type", + type=str, + default=None, + help="A unique string representing a type of run, which is useful when you're grouping runs together into larger experiments using group.", + ) + parser.add_argument( + "--wandb-run-name", + type=str, + default=None, + help="A unique string representing the name of the wandb run. If not provided, the name will be generated from the model and training specifications.", + ) + + parser.add_argument( + "--wandb-id", + type=str, + default=None, + help="Sets the version, mainly used to resume a previous run", + ) + parser.add_argument( + "--wandb-anonymous", + action="store_true", + help="Enable or explicitly disable anonymous logging", + ) + parser.add_argument( + "--wandb-log-model", + action="store_true", + help="Save checkpoints in wandb dir to upload on W&B servers", + ) + parser.add_argument( + "--wandb-offline", action="store_true", help="Use wandb in offline mode" + ) + parser.add_argument( + "--sequence-parallel", + action="store_true", + help="Set to enable sequence parallelism.", + ) + parser.add_argument("--fp8", action="store_true", help="Set to enable FP8") + parser.add_argument( + "--micro-batch-size", + type=int, + default=1, + help="Micro-batch size for data-parallel training.", + ) + parser.add_argument( + "--global-batch-size", + type=int, + default=None, + help="Global batch size for training. If set to None, infer it from the TP, CP, and PP parameters.", + ) + parser.add_argument( + "--grad-acc-batches", + type=int, + default=1, + help="Number of batches to accumulate gradients over.", + ) + parser.add_argument( + "--max-steps", + type=int, + help="Number of training optimizer update steps. This controls the total number of steps as well as the " + "shape of the learning rate curve.", + default=500000, + ) + parser.add_argument( + "--constant-steps", + type=int, + help="Number of steps to keep the learning rate constant before annealing. This controls the " + "shape of the learning rate curve.", + default=80000, + ) + parser.add_argument( + "--early-stop-on-step", + type=int, + help="Stop training on this step, if set. This may be useful for testing or debugging purposes.", + ) + parser.add_argument( + "--val-check-interval", + type=int, + help="Number of steps between validation measurements and model checkpoints.", + ) + parser.add_argument( + "--grad-reduce-in-fp32", + action="store_true", + default=False, + help="Gradient reduce in FP32.", + ) + parser.add_argument( + "--fp8-wgrad", + action="store_true", + default=False, + help="Faster option that is maybe less accurate (TBD) when using fp8.", + ) + parser.add_argument( + "--use-megatron-comm-overlap-llama3-8k", action="store_true", default=False + ) + parser.add_argument( + "--tp-comm-overlap-backend", + type=str, + choices=["nccl", "mpi", "gloo"], + default="nccl", + help="TP communication backend to use. Defaults to 'nccl'.", + ) + parser.add_argument("--align-param-gather", action="store_true", default=False) + # parser.add_argument("--straggler-detection", action="store_true", default=False) + parser.add_argument( + "--model-size", + type=str, + choices=sorted(HYENA_MODEL_OPTIONS.keys()), + default="7b", + help="Model architecture to use, choose between 7b, 40b, or test (a sub-model of 4 layers, less than 1B " + "parameters). '_arc_1m' models have GLU / FFN dimensions that support 1M context length when trained " + "with TP<=8.", + ) + parser.add_argument( + "--add-bias-output", + action="store_true", + default=False, + help="Add bias to the output layer to enable learning a simple prior.", + ) + parser.add_argument( + "--result-dir", + type=Path, + required=False, + default=Path("./results"), + help="Path to the result directory.", + ) + parser.add_argument( + "--experiment-name", + type=str, + required=False, + default="evo2", + help="Name of the experiment.", + ) + + parser.add_argument( + "--limit-val-batches", + type=int, + default=20, + help="Number of validation steps", + ) + parser.add_argument( + "--log-every-n-steps", + type=int, + default=1, + required=False, + help="Number of steps between logging.", + ) + parser.add_argument( + "--ckpt-dir", + type=str, + default=DEFAULT_CKPT_DIR, + help="Directory to restore an initial checkpoint from. Use this for supervised fine-tuning.", + ) + parser.add_argument( + "--wd", type=float, default=0.01, help="Weight decay for optimizer." + ) + parser.add_argument( + "--restore-optimizer-from-ckpt", + action="store_true", + help="Restore optimizer state from initial checkpoint. Defaults to False.", + ) + parser.add_argument( + "--no-average-in-collective", + action="store_true", + default=False, + help="Avaerage optimizer state in collective rather than dividing by dp size and summing.", + ) + parser.add_argument( + "--seed", type=int, default=1234, help="Set random seed for training." + ) + parser.add_argument( + "--workers", + type=int, + default=8, + help="Number of workers to use for data loading.", + ) + parser.add_argument( + "--gc-interval", + type=int, + default=0, + help="Set to a value > 0 if you want to synchronize garbage collection, will do gc every gc-interval steps.", + ) + parser.add_argument( + "--enable-preemption", + action="store_true", + default=False, + help="Enable preemption hooks. If enabled this will save a checkpoint whenever slurm exits.", + ) + parser.add_argument( + "--ckpt-async-save", + action="store_true", + default=False, + ) + parser.add_argument( + "--ckpt-format", + type=str, + choices=["torch_dist", "zarr"], + default="torch_dist", + help="Specify checkpoint format to use. Defaults to 'torch_dist', as 'zarr' is deprecated. Only use if " + "resuming training from a zarr checkpoint.", + ) + parser.add_argument( + "--eod-pad-in-loss-mask", + action="store_true", + default=False, + help="Do not predict EOD/Pad tokens (typical default, but not default in original evo2).", + ) + parser.add_argument( + "--cross-entropy-loss-fusion", + action="store_true", + default=False, + help="Use the faster, but maybe less accurate fused form of cross entropy, " + "which also has bf16 grads internally.", + ) + parser.add_argument( + "--no-fp32-residual-connection", + action="store_true", + default=False, + help="If set, turn off fp32 residual connections which may be faster but may impact accuracy.", + ) + parser.add_argument( + "--debug-ddp-parity-freq", + type=int, + default=0, + help="Set to value > 0 to debug DDP weight parity between ranks.", + ) + parser.add_argument( + "--hybrid-override-pattern", + type=str, + help="Override the hybrid override pattern in the config (specifies hyena layer ordering and type).", + ) + parser.add_argument( + "--num-layers", + type=int, + help="If set, override the number of layers specified in the requested config.", + ) + parser.add_argument( + "--create-tflops-callback", + action="store_true", + default=False, + help="Enable tflops calculation callback for Hyena / Evo2. Defaults to False.", + ) + parser.add_argument( + "--log-parameters-and-shapes", + action="store_true", + default=False, + help="Log training parameters shapes and dtypes for debugging.", + ) + parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate.") + parser.add_argument( + "--min-lr", + type=float, + default=3e-5, + help="Min learning rate in cosine annealing.", + ) + parser.add_argument( + "--warmup-steps", + type=int, + default=2500, + help="Number of warmup steps in cosine annealing", + ) + # NSYS profiling/tooling arguments + parser.add_argument( + "--nsys-profiling", + action="store_true", + default=False, + help="Enable targeted `nsys` profiling on the training loop for a defined step range. To actually get profiling" + " output you must run the whole program with `nsys`. For example: " + " `nsys profile -s none -o output_report_name -t cuda,nvtx --force-overwrite true " + "--capture-range=cudaProfilerApi --capture-range-end=stop [regular python command here]`", + ) + # start, end, rank + parser.add_argument( + "--nsys-start-step", + type=int, + required=False, + default=0, + help="Start nsys profiling after this step.", + ) + parser.add_argument( + "--nsys-end-step", + type=int, + required=False, + help="End nsys profiling after this step.", + ) + parser.add_argument( + "--no-renormalize-loss", + action="store_true", + default=False, + help="Do not renormalize the loss weights.", + ) + # rank as list of integers + parser.add_argument( + "--nsys-ranks", + type=int, + nargs="+", + required=False, + default=[0], + help="Enable nsys profiling for these ranks.", + ) + parser.add_argument( + "--activation-checkpoint-recompute-num-layers", + type=int, + help="If set, override the default value set in the config.", + ) + parser.add_argument( + "--disable-checkpointing", + action="store_false", + default=True, # False + dest="create_checkpoint_callback", + help="Disable creating a ModelCheckpoint callback.", + ) + parser.add_argument( + "--clip-grad", + type=float, + default=1.0, + help="Grad clip value. Note that when using DDP this may need to be inflated.", + ) + parser.add_argument( + "--seq-len-interpolation-factor", + type=float, + help="Adjusts the linear scaling of ROPE (Rotary Position Embedding) for context extension. " + "Set this factor relative to your base context length e.g., for an original context length of 8192 and " + "an extended context length of 524288, use 524288/8192 = 64.", + ) + parser.add_argument( + "--overlap-param-gather", + action="store_true", + default=False, + help="Overlap the parameter gather with the optimizer step. This is currently disabled due to a NeMo bug " + "when using DDP. Making this an option defaulting to False is a temporary solution until the bug is fixed.", + ) + parser.add_argument( + "--overlap-grad-reduce", + action="store_true", + default=False, + help="Overlap the gradient reduce with the optimizer step.", + ) + parser.add_argument( + "--hidden-dropout", + type=float, + default=0.0, + help="Dropout probability for the hyena layers", + ) + parser.add_argument( + "--attention-dropout", + type=float, + default=0.0, + help="Dropout probability for the attention layers.", + ) + parser.add_argument( + "--save-top-k", + type=int, + default=5, + help="Number of best checkpoints to keep. Set to -1 to save all checkpoints.", + ) + parser.add_argument( + "--metric-to-monitor-for-checkpoints", + type=str, + default="val_loss", + help="Metric to monitor for checkpoints.", + ) + parser.add_argument( + "--save-last-checkpoint", + action="store_true", + default=True, + help="Save the last checkpoint.", + ) + parser.add_argument( + "--no-save-last-checkpoint", + action="store_false", + dest="save_last_checkpoint", + default=True, + help="Disable saving the last checkpoint.", + ) + recompute_group = parser.add_mutually_exclusive_group(required=False) + recompute_group.add_argument( + "--no-activation-checkpointing", action="store_true", default=False + ) + recompute_group.add_argument( + "--selective-activation-checkpointing", action="store_true", default=False + ) + return parser.parse_args(args=args) + + +def train(args: argparse.Namespace) -> nl.Trainer: + """Main function to run Evo2 training.""" + # Instantiate tokenizer. + tokenizer = get_nmt_tokenizer( + "byte-level", + ) + + # Infer global batch size. + global_batch_size = args.global_batch_size + if global_batch_size is None: + global_batch_size = infer_global_batch_size( # batch : 8 + micro_batch_size=args.micro_batch_size, + num_nodes=args.num_nodes, + devices=args.devices, + accumulate_grad_batches=args.grad_acc_batches, + tensor_model_parallel_size=args.tensor_parallel_size, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + context_model_parallel_size=args.context_parallel_size, + ) + if args.mock_data: + data_module = MockDataModule( + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + global_batch_size=global_batch_size, + num_workers=args.workers, + tokenizer=tokenizer, + ) + else: + blended_dataset_config = parse_dataset_config( + dataset_config_path=args.dataset_config, dataset_path=args.dataset_dir + ) + Evo2DatasetPadEodLossMask if args.eod_pad_in_loss_mask else Evo2Dataset + # Instantiate pre-training module. + data_module = PreTrainingDataModule( + paths=blended_dataset_config, + # dataset_cls=dataset_cls, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + global_batch_size=global_batch_size, + seed=args.seed, # 1234 + num_workers=args.workers, + tokenizer=tokenizer, + eod_mask_loss=args.eod_pad_in_loss_mask, + ) + + if args.no_activation_checkpointing: + activation_checkpointing_args = { + "recompute_granularity": None, + "recompute_method": None, + "recompute_num_layers": None, + } + elif args.selective_activation_checkpointing: + activation_checkpointing_args = { + "recompute_granularity": "selective", + "recompute_method": None, + "recompute_num_layers": None, + } + else: + if args.activation_checkpoint_recompute_num_layers is not None: + activation_checkpointing_args = { + "recompute_num_layers": args.activation_checkpoint_recompute_num_layers, + } + else: + activation_checkpointing_args = {} + + # Retrieve model config. + config_modifiers_init = { + "tp_comm_overlap": args.use_megatron_comm_overlap_llama3_8k, + "seq_length": args.seq_length, + "hidden_dropout": args.hidden_dropout, + "attention_dropout": args.attention_dropout, + "to_upper": "weighted" if args.no_renormalize_loss else "normalized_weighted", + "distribute_saved_activations": False if args.sequence_parallel else True, + "cross_entropy_loss_fusion": args.cross_entropy_loss_fusion, + "fp32_residual_connection": not args.no_fp32_residual_connection, + "add_bias_output": args.add_bias_output, + **activation_checkpointing_args, + } + if args.hybrid_override_pattern: + config_modifiers_init["hybrid_override_pattern"] = args.hybrid_override_pattern + if args.num_layers: + config_modifiers_init["num_layers"] = args.num_layers + + if args.model_size not in HYENA_MODEL_OPTIONS: + raise ValueError(f"Invalid model size: {args.model_size}") + evo2_config = HYENA_MODEL_OPTIONS[args.model_size](**config_modifiers_init) + + # Instantiate model. + model = llm.HyenaModel(evo2_config, tokenizer=data_module.tokenizer) + + # Setup callbacks. + callbacks = [ + RichModelSummary(max_depth=4), + LearningRateMonitor(), + TimingCallback(), + ] + + if args.enable_preemption: + callbacks.append(nl_callbacks.PreemptionCallback()) + if args.debug_ddp_parity_freq > 0: + callbacks.append( + nl_callbacks.DdpParityChecker(interval=args.debug_ddp_parity_freq) + ) + if args.log_parameters_and_shapes: + callbacks.append(nl_callbacks.ParameterDebugger()) + if args.create_tflops_callback: + # Add callback that logs the tera-FLOPS per second per GPU during training. + flop_meas_callback = FLOPsMeasurementCallback( + evo2_config, + data_module, + "hyena", + ) + callbacks.append(flop_meas_callback) + + # TODO(@cye): Add this back when it works with 24.12. + # if args.straggler_detection: + # callbacks.append( + # res_module.StragglerDetectionCallback( + # report_time_interval=300, + # calc_relative_gpu_perf=True, + # calc_individual_gpu_perf=True, + # num_gpu_perf_scores_to_print=5, + # gpu_relative_perf_threshold=0.7, + # gpu_individual_perf_threshold=0.7, + # stop_if_detected=True, + # enable_ptl_logging=True, + # ) + # ) + if args.use_megatron_comm_overlap_llama3_8k: + # Pick the floating point appropriate config. + if args.fp8: + tp_comm_overlap_cfg = userbuffers_fp8_h100_h8192_tp4_mbs1_seqlen8192 + else: + tp_comm_overlap_cfg = userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192 + callbacks.append( + MegatronCommOverlapCallback( + tp_comm_overlap=evo2_config.tp_comm_overlap, + tp_comm_overlap_cfg=tp_comm_overlap_cfg, + tp_comm_bootstrap_backend=args.tp_comm_overlap_backend, + wgrad_deferral_limit=22, # default from NeMo + overlap_param_gather_with_optimizer_step=False, # Currently disabled due to an issue with checkpointing. + align_param_gather=args.align_param_gather, + ) + ) + + if args.gc_interval > 0: + callbacks.append( + nl_callbacks.GarbageCollectionCallback( + gc_interval_train=args.gc_interval, gc_interval_val=args.gc_interval + ) + ) + if args.nsys_profiling: + if args.nsys_end_step is None: + nsys_end_step = args.max_steps + else: + nsys_end_step = args.nsys_end_step + callbacks.append( + nl_callbacks.NsysCallback( + start_step=args.nsys_start_step, + end_step=nsys_end_step, + ranks=args.nsys_ranks, + gen_shape=True, + ) + ) + + wandb_run_name = ( + f"evo2-size-{args.model_size}-TP{args.tensor_parallel_size}-" + f"PP{args.pipeline_model_parallel_size}-CP{args.context_parallel_size}" + f"-GBS{global_batch_size}-MBS{args.micro_batch_size}-SkipLossRenorm{args.no_renormalize_loss}" + f"-NOAC{args.no_activation_checkpointing}-SELAC{args.selective_activation_checkpointing}" + f"-ACRNL{evo2_config.recompute_num_layers}" + f"-PAT{evo2_config.hybrid_override_pattern}" + f"-F32R{evo2_config.fp32_residual_connection}" + f"-FCE{evo2_config.cross_entropy_loss_fusion}" + f"-AIC{not args.no_average_in_collective}" + f"-PEOD{args.eod_pad_in_loss_mask}" + f"-BO{args.add_bias_output}" + f"-GCLP{args.clip_grad}" + f"-HDO{args.hidden_dropout}" + f"-ADO{args.attention_dropout}" + f"-LR{args.lr}-MINLR{args.min_lr}-WUSTEPS{args.warmup_steps}-CONSTSTEPS{args.constant_steps}-WD{args.wd}" + f"-GRFP32{args.grad_reduce_in_fp32}-FP8WG{args.fp8_wgrad and args.fp8}" + f"-OGR{args.overlap_grad_reduce}-OPG{args.overlap_param_gather}" + f"-NODES{args.num_nodes}-FP8{args.fp8}" + ) + + wandb_config: Optional[WandbConfig] = ( + None + if args.wandb_project is None + else WandbConfig( + offline=args.wandb_offline, + project=args.wandb_project, + name=( + args.wandb_run_name + if args.wandb_run_name is not None + else wandb_run_name + ), + entity=args.wandb_entity, + tags=args.wandb_tags, + group=args.wandb_group, + job_type=args.wandb_job_type, + id=args.wandb_id, + anonymous=args.wandb_anonymous, + log_model=args.wandb_log_model, + ) + ) + + nemo_logger = setup_nemo_lightning_logger( + root_dir=args.result_dir, + name=args.experiment_name, + initialize_tensorboard_logger=args.create_tensorboard_logger, + wandb_config=wandb_config, + ) + + if args.create_checkpoint_callback: + checkpoint_path = str(Path(nemo_logger.save_dir) / "checkpoints") + checkpoint_callback = nl_callbacks.ModelCheckpoint( + dirpath=checkpoint_path, + save_last=args.save_last_checkpoint, + monitor=args.metric_to_monitor_for_checkpoints, + save_top_k=args.save_top_k, + every_n_train_steps=args.val_check_interval, + always_save_context=True, + filename="{epoch}-{step}-{consumed_samples}", + save_weights_only=False, + save_optim_on_train_end=True, + save_context_on_train_end=True, + ) + callbacks.append(checkpoint_callback) + + auto_resume = nl.AutoResume( + resume_if_exists=True, + resume_ignore_no_checkpoint=True, + resume_past_end=False, + resume_from_directory=checkpoint_path, + restore_config=( + RestoreConfig( + path=args.ckpt_dir, + load_model_state=True, + load_optim_state=args.restore_optimizer_from_ckpt, + ) + if args.ckpt_dir + else None + ), + ) + else: + auto_resume = None + + ddp: DistributedDataParallelConfig = DistributedDataParallelConfig( + check_for_nan_in_grad=True, + overlap_grad_reduce=args.overlap_grad_reduce, + overlap_param_gather=args.overlap_param_gather, # Verify that this works using + grad_reduce_in_fp32=args.grad_reduce_in_fp32, + align_param_gather=args.align_param_gather, + average_in_collective=not args.no_average_in_collective, + ) + + # Initialize Megatron Strategy and Trainer. + strategy = nl.MegatronStrategy( + ddp=ddp, + tensor_model_parallel_size=args.tensor_parallel_size, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + context_parallel_size=args.context_parallel_size, + pipeline_dtype=torch.bfloat16, + sequence_parallel=args.sequence_parallel, + ckpt_load_optimizer=True, + ckpt_save_optimizer=True, + ckpt_async_save=args.ckpt_async_save, + save_ckpt_format=args.ckpt_format, + ckpt_load_strictness="log_all", # or rebasing to https://github.com/NVIDIA/NeMo/pull/11988/files#diff-7667eae242a8ef776bff78cd08e79bc81df4896a450f0a781f6ed317a3dfb7ffR139 + ) + from lightning.fabric.plugins.environments import LightningEnvironment + + + trainer = nl.Trainer( + devices=args.devices, + num_nodes=args.num_nodes, + max_steps=( + args.max_steps + if args.early_stop_on_step is None + else args.early_stop_on_step + ), + accelerator="gpu", + strategy=strategy, + callbacks=callbacks, + log_every_n_steps=args.log_every_n_steps, + limit_val_batches=args.limit_val_batches, + num_sanity_val_steps=0, + use_distributed_sampler=False, + plugins=[ + nl.MegatronMixedPrecision( + precision="bf16-mixed", + params_dtype=torch.bfloat16, + grad_reduce_in_fp32=args.grad_reduce_in_fp32, + fp8="hybrid" if args.fp8 else None, + fp8_amax_history_len=16 if args.fp8 else 1, + fp8_amax_compute_algo="max" if args.fp8 else "most_recent", + fp8_wgrad=args.fp8 + and ( + args.fp8_wgrad or args.use_megatron_comm_overlap_llama3_8k + ), # faster and less accurate when set to True, and MUST be True if using TP communication overlap + ), + LightningEnvironment(), + ], + # plugins=nl.MegatronMixedPrecision( + # precision="bf16-mixed", + # params_dtype=torch.bfloat16, + # grad_reduce_in_fp32=args.grad_reduce_in_fp32, + # fp8="hybrid" if args.fp8 else None, + # fp8_amax_history_len=16 if args.fp8 else 1, + # fp8_amax_compute_algo="max" if args.fp8 else "most_recent", + # fp8_wgrad=args.fp8 + # and ( + # args.fp8_wgrad or args.use_megatron_comm_overlap_llama3_8k + # ), # faster and less accurate when set to True, and MUST be True if using TP communication overlap + # ), + val_check_interval=args.val_check_interval, + enable_checkpointing=args.create_checkpoint_callback, + ) + + # Logger setup + nemo_logger.setup( + trainer, + resume_if_exists=True, + ) + + if auto_resume is not None: + auto_resume.setup(trainer, model) + + # Optimizer and scheduler setup + opt_config = OptimizerConfig( + optimizer="adam", + lr=args.lr, + adam_beta1=0.9, + adam_beta2=0.95, + weight_decay=args.wd, + clip_grad=args.clip_grad, + use_distributed_optimizer=True, + bf16=True, + ) + + sched = CosineAnnealingScheduler( + max_steps=trainer.max_steps, + warmup_steps=args.warmup_steps, + min_lr=args.min_lr, + constant_steps=args.constant_steps, + ) + + opt = MegatronOptimizerModule( + opt_config, + sched, + no_weight_decay_cond=evo2_config.hyena_no_weight_decay_cond_fn, + ) + opt.connect(model) + + # Start training + trainer.fit(model, data_module) + return trainer + + +def main(): + """Parsing args and running evo2 training.""" + args = parse_args() + if args.ckpt_dir and not os.path.isdir(args.ckpt_dir): + raise SystemExit(f"ERROR: Evo2 checkpoint directory does not exist: {args.ckpt_dir}") + train(args=args) + + +if __name__ == "__main__": + main() diff --git a/scripts/train_evo2_1b.sh b/scripts/train_evo2_1b.sh new file mode 100644 index 0000000000000000000000000000000000000000..e2e176d9d1539b583d6b2f0565a3b550fdfae890 --- /dev/null +++ b/scripts/train_evo2_1b.sh @@ -0,0 +1,68 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +cd "${PROJECT_ROOT}" + +DATASET_DIR="${DATASET_DIR:-${EVO2_DATASET_DIR:-${PROJECT_ROOT}/data/data_mini/genome_data}}" +CKPT_DIR="${CKPT_DIR:-}" +AVAILABLE_DEVICES="$(python -c "import torch; c=torch.cuda.device_count(); print(c if c > 0 else 1)")" +DEVICES="${DEVICES:-${AVAILABLE_DEVICES}}" +SEQ_LENGTH="${SEQ_LENGTH:-8192}" +MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-2}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-}" +MAX_STEPS="${MAX_STEPS:-1000}" +WARMUP_STEPS="${WARMUP_STEPS:-5}" +VAL_CHECK_INTERVAL="${VAL_CHECK_INTERVAL:-50}" +LIMIT_VAL_BATCHES="${LIMIT_VAL_BATCHES:-2}" +WORKERS="${WORKERS:-8}" +CHECKPOINTING="${CHECKPOINTING:-1}" + +EXTRA_ARGS=() +if [[ -n "${GLOBAL_BATCH_SIZE}" ]]; then + EXTRA_ARGS+=(--global-batch-size "${GLOBAL_BATCH_SIZE}") +fi +if [[ "${CHECKPOINTING}" != "1" ]]; then + EXTRA_ARGS+=(--disable-checkpointing --no-save-last-checkpoint) +fi + +if [[ -n "${CKPT_DIR}" ]]; then + if [[ ! -d "${CKPT_DIR}" ]]; then + echo "Error: Evo2 1B checkpoint directory not found: ${CKPT_DIR}" >&2 + echo "Set CKPT_DIR only when you want to restore from a 1B NeMo checkpoint." >&2 + exit 1 + fi + EXTRA_ARGS+=(--ckpt-dir "${CKPT_DIR}") +fi + +echo "EVO2_PROJECT_ROOT: ${PROJECT_ROOT}" +echo "EVO2_DATASET_DIR: ${DATASET_DIR}" +echo "EVO2_CKPT_DIR: ${CKPT_DIR:-}" +echo "EVO2_DEVICES: ${DEVICES}" +echo "EVO2_SEQ_LENGTH: ${SEQ_LENGTH}" +echo "EVO2_MAX_STEPS: ${MAX_STEPS}" +echo "EVO2_CHECKPOINTING: ${CHECKPOINTING}" +echo "EVO2_WORKERS: ${WORKERS}" + +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" +export MKL_NUM_THREADS="${MKL_NUM_THREADS:-1}" + +python scripts/train.py \ + -d config/genome_data_config.yaml \ + --dataset-dir "${DATASET_DIR}" \ + --model-size 1b \ + --devices "${DEVICES}" \ + --num-nodes 1 \ + --seq-length "${SEQ_LENGTH}" \ + --micro-batch-size "${MICRO_BATCH_SIZE}" \ + --lr 0.0001 \ + --warmup-steps "${WARMUP_STEPS}" \ + --max-steps "${MAX_STEPS}" \ + --clip-grad 1 \ + --wd 0.01 \ + --workers "${WORKERS}" \ + --activation-checkpoint-recompute-num-layers 1 \ + --val-check-interval "${VAL_CHECK_INTERVAL}" \ + --limit-val-batches "${LIMIT_VAL_BATCHES}" \ + "${EXTRA_ARGS[@]}" diff --git a/scripts/train_evo2_7b.sh b/scripts/train_evo2_7b.sh new file mode 100644 index 0000000000000000000000000000000000000000..4312a443f1cee4c093fc7315383b43b281254979 --- /dev/null +++ b/scripts/train_evo2_7b.sh @@ -0,0 +1,76 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +cd "${PROJECT_ROOT}" + +DATASET_DIR="${DATASET_DIR:-${EVO2_DATASET_DIR:-${PROJECT_ROOT}/data/data_mini/genome_data}}" +CKPT_DIR="${CKPT_DIR:-}" +AVAILABLE_DEVICES="$(python -c "import torch; c=torch.cuda.device_count(); print(c if c > 0 else 1)")" +DEVICES="${DEVICES:-${AVAILABLE_DEVICES}}" +TENSOR_PARALLEL_SIZE="${TENSOR_PARALLEL_SIZE:-1}" +MODEL_SIZE="${MODEL_SIZE:-7b_arc_longcontext}" +SEQ_LENGTH="${SEQ_LENGTH:-1024}" +MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-}" +MAX_STEPS="${MAX_STEPS:-1000}" +WARMUP_STEPS="${WARMUP_STEPS:-5}" +VAL_CHECK_INTERVAL="${VAL_CHECK_INTERVAL:-50}" +LIMIT_VAL_BATCHES="${LIMIT_VAL_BATCHES:-2}" +SEQUENCE_PARALLEL="${SEQUENCE_PARALLEL:-1}" +CHECKPOINTING="${CHECKPOINTING:-1}" +WORKERS="${WORKERS:-1}" + +EXTRA_ARGS=() +if [[ -n "${CKPT_DIR}" ]]; then + if [[ ! -d "${CKPT_DIR}" ]]; then + echo "ERROR: Evo2 checkpoint directory does not exist: ${CKPT_DIR}" >&2 + exit 1 + fi + EXTRA_ARGS+=(--ckpt-dir "${CKPT_DIR}") +fi +if [[ -n "${GLOBAL_BATCH_SIZE}" ]]; then + EXTRA_ARGS+=(--global-batch-size "${GLOBAL_BATCH_SIZE}") +fi +if [[ "${SEQUENCE_PARALLEL}" == "1" && "${TENSOR_PARALLEL_SIZE}" -gt 1 ]]; then + EXTRA_ARGS+=(--sequence-parallel) +fi +if [[ "${CHECKPOINTING}" != "1" ]]; then + EXTRA_ARGS+=(--disable-checkpointing --no-save-last-checkpoint) +fi + +echo "EVO2_PROJECT_ROOT: ${PROJECT_ROOT}" +echo "EVO2_DATASET_DIR: ${DATASET_DIR}" +echo "EVO2_CKPT_DIR: ${CKPT_DIR:-}" +echo "EVO2_MODEL_SIZE: ${MODEL_SIZE}" +echo "EVO2_DEVICES: ${DEVICES}" +echo "EVO2_TENSOR_PARALLEL_SIZE: ${TENSOR_PARALLEL_SIZE}" +echo "EVO2_SEQ_LENGTH: ${SEQ_LENGTH}" +echo "EVO2_MAX_STEPS: ${MAX_STEPS}" +echo "EVO2_SEQUENCE_PARALLEL: ${SEQUENCE_PARALLEL}" +echo "EVO2_CHECKPOINTING: ${CHECKPOINTING}" +echo "EVO2_WORKERS: ${WORKERS}" + +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" +export MKL_NUM_THREADS="${MKL_NUM_THREADS:-1}" + +python scripts/train_slurm.py \ + -d config/genome_data_config.yaml \ + --dataset-dir "${DATASET_DIR}" \ + --model-size "${MODEL_SIZE}" \ + --devices "${DEVICES}" \ + --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \ + --num-nodes 1 \ + --seq-length "${SEQ_LENGTH}" \ + --micro-batch-size "${MICRO_BATCH_SIZE}" \ + --lr 0.0001 \ + --warmup-steps "${WARMUP_STEPS}" \ + --max-steps "${MAX_STEPS}" \ + --clip-grad 1 \ + --wd 0.01 \ + --workers "${WORKERS}" \ + --activation-checkpoint-recompute-num-layers 1 \ + --val-check-interval "${VAL_CHECK_INTERVAL}" \ + --limit-val-batches "${LIMIT_VAL_BATCHES}" \ + "${EXTRA_ARGS[@]}" diff --git a/scripts/train_slurm.py b/scripts/train_slurm.py new file mode 100644 index 0000000000000000000000000000000000000000..eea7c32215f60d1dd8f1da11decef87efdf5f557 --- /dev/null +++ b/scripts/train_slurm.py @@ -0,0 +1,924 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Arc Institute. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Michael Poli. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024 Stanford University. All rights reserved +# SPDX-License-Identifier: LicenseRef-Apache2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import os +import sys + +_PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__)) +while _PROJECT_ROOT and not os.path.isdir(os.path.join(_PROJECT_ROOT, "model")): + _PARENT = os.path.dirname(_PROJECT_ROOT) + if _PARENT == _PROJECT_ROOT: + break + _PROJECT_ROOT = _PARENT +_MODEL_ROOT = os.path.join(_PROJECT_ROOT, "model") +_ONESCIENCE_ROOT = os.environ.get("ONESCIENCE_ROOT") +for _path in (_MODEL_ROOT, _PROJECT_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +if _ONESCIENCE_ROOT: + _ONESCIENCE_SRC = os.path.join(_ONESCIENCE_ROOT, "src") + for _path in (_ONESCIENCE_SRC, _ONESCIENCE_ROOT): + if os.path.exists(_path) and _path not in sys.path: + sys.path.insert(0, _path) +from pathlib import Path +from typing import List, Optional + +# TODO add back support for slurm resilience. +# import nvidia_resiliency_ext.ptl_resiliency as res_module +import torch +from lightning.pytorch.callbacks import LearningRateMonitor, RichModelSummary +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.optimizer import OptimizerConfig +from nemo import lightning as nl +from nemo.collections import llm +from nemo.collections.llm.gpt.data import MockDataModule, PreTrainingDataModule +from nemo.collections.llm.gpt.data.megatron.hyena.config import parse_dataset_config +from nemo.collections.llm.gpt.data.megatron.hyena.evo2_dataset import ( + Evo2Dataset, + Evo2DatasetPadEodLossMask, +) +from nemo.collections.llm.gpt.model.hyena import HYENA_MODEL_OPTIONS +from nemo.collections.llm.recipes.tp_overlap_configs.userbuffers import ( + userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192, + userbuffers_fp8_h100_h8192_tp4_mbs1_seqlen8192, +) +from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer +from nemo.lightning.pytorch import callbacks as nl_callbacks +from nemo.lightning.pytorch.callbacks.flops_callback import FLOPsMeasurementCallback +from nemo.lightning.pytorch.callbacks.megatron_comm_overlap import ( + MegatronCommOverlapCallback, +) +from nemo.lightning.pytorch.optim import CosineAnnealingScheduler +from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule +from nemo.lightning.pytorch.strategies.utils import RestoreConfig +from nemo.utils.exp_manager import TimingCallback + +from evo2.utils.datamodule_utils import infer_global_batch_size +from evo2.utils.logger_utils import ( + WandbConfig, + setup_nemo_lightning_logger, +) + + +torch._dynamo.config.suppress_errors = True + +DEFAULT_DATASET_DIR = os.environ.get( + "EVO2_DATASET_DIR", + str(Path(_PROJECT_ROOT) / "data" / "data_mini" / "genome_data"), +) +DEFAULT_CKPT_DIR = None + + +def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace: + """Parse arguments for Evo2 model training.""" + parser = argparse.ArgumentParser( + description="Train a Hyena model using NeMo 2.0.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + data_group = parser.add_mutually_exclusive_group(required=True) + + data_group.add_argument( + "-d", + "--dataset-config", + type=str, + help="Path to the blended / weighted training dataset configuration YAML.", + ) + data_group.add_argument( + "--mock-data", + action="store_true", + help="Train with Mock data (for testing/debugging), either set this or provide a dataset config.", + ) + + parser.add_argument( + "--dataset-dir", + type=str, + default=DEFAULT_DATASET_DIR, + help="Absolute path to the dataset directory. Defaults to using the absolute or relative paths (dataset_prefix) specified in the dataset config YAML.", + ) + + parser.add_argument( + "--num-nodes", + type=int, + default=1, + help="Number of nodes to use for training, defaults to 1.", + ) + parser.add_argument( + "--devices", + type=int, + default=1, + help="Number of devices to use for training, defaults to 1.", + ) + parser.add_argument( + "--seq-length", type=int, default=8192, help="Training sequence length" + ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Order of tensor parallelism. Defaults to 1.", + ) + parser.add_argument( + "--pipeline-model-parallel-size", + type=int, + default=1, + help="Order of pipeline parallelism. Defaults to 1.", + ) + parser.add_argument( + "--context-parallel-size", + type=int, + default=1, + help="Order of context parallelism. Defaults to 1.", + ) + parser.add_argument( + "--create-tensorboard-logger", + action="store_true", + default=False, + help="Create a tensorboard logger.", + ) + parser.add_argument( + "--wandb-entity", type=str, default=None, help="The team posting this run" + ) + parser.add_argument( + "--wandb-project", type=str, default=None, help="Wandb project name " + ) + parser.add_argument( + "--wandb-tags", + nargs="+", + type=str, + default=None, + help="Tags associated with this run", + ) + parser.add_argument( + "--wandb-group", + type=str, + default=None, + help="A unique string shared by all runs in a given group", + ) + parser.add_argument( + "--wandb-job-type", + type=str, + default=None, + help="A unique string representing a type of run, which is useful when you're grouping runs together into larger experiments using group.", + ) + parser.add_argument( + "--wandb-run-name", + type=str, + default=None, + help="A unique string representing the name of the wandb run. If not provided, the name will be generated from the model and training specifications.", + ) + + parser.add_argument( + "--wandb-id", + type=str, + default=None, + help="Sets the version, mainly used to resume a previous run", + ) + parser.add_argument( + "--wandb-anonymous", + action="store_true", + help="Enable or explicitly disable anonymous logging", + ) + parser.add_argument( + "--wandb-log-model", + action="store_true", + help="Save checkpoints in wandb dir to upload on W&B servers", + ) + parser.add_argument( + "--wandb-offline", action="store_true", help="Use wandb in offline mode" + ) + parser.add_argument( + "--sequence-parallel", + action="store_true", + help="Set to enable sequence parallelism.", + ) + parser.add_argument("--fp8", action="store_true", help="Set to enable FP8") + parser.add_argument( + "--micro-batch-size", + type=int, + default=1, + help="Micro-batch size for data-parallel training.", + ) + parser.add_argument( + "--global-batch-size", + type=int, + default=None, + help="Global batch size for training. If set to None, infer it from the TP, CP, and PP parameters.", + ) + parser.add_argument( + "--grad-acc-batches", + type=int, + default=1, + help="Number of batches to accumulate gradients over.", + ) + parser.add_argument( + "--max-steps", + type=int, + help="Number of training optimizer update steps. This controls the total number of steps as well as the " + "shape of the learning rate curve.", + default=500000, + ) + parser.add_argument( + "--constant-steps", + type=int, + help="Number of steps to keep the learning rate constant before annealing. This controls the " + "shape of the learning rate curve.", + default=80000, + ) + parser.add_argument( + "--early-stop-on-step", + type=int, + help="Stop training on this step, if set. This may be useful for testing or debugging purposes.", + ) + parser.add_argument( + "--val-check-interval", + type=int, + help="Number of steps between validation measurements and model checkpoints.", + ) + parser.add_argument( + "--grad-reduce-in-fp32", + action="store_true", + default=False, + help="Gradient reduce in FP32.", + ) + parser.add_argument( + "--fp8-wgrad", + action="store_true", + default=False, + help="Faster option that is maybe less accurate (TBD) when using fp8.", + ) + parser.add_argument( + "--use-megatron-comm-overlap-llama3-8k", action="store_true", default=False + ) + parser.add_argument( + "--tp-comm-overlap-backend", + type=str, + choices=["nccl", "mpi", "gloo"], + default="nccl", + help="TP communication backend to use. Defaults to 'nccl'.", + ) + parser.add_argument("--align-param-gather", action="store_true", default=False) + # parser.add_argument("--straggler-detection", action="store_true", default=False) + parser.add_argument( + "--model-size", + type=str, + choices=sorted(HYENA_MODEL_OPTIONS.keys()), + default="7b", + help="Model architecture to use, choose between 7b, 40b, or test (a sub-model of 4 layers, less than 1B " + "parameters). '_arc_1m' models have GLU / FFN dimensions that support 1M context length when trained " + "with TP<=8.", + ) + parser.add_argument( + "--add-bias-output", + action="store_true", + default=False, + help="Add bias to the output layer to enable learning a simple prior.", + ) + parser.add_argument( + "--result-dir", + type=Path, + required=False, + default=Path("./results"), + help="Path to the result directory.", + ) + parser.add_argument( + "--experiment-name", + type=str, + required=False, + default="evo2", + help="Name of the experiment.", + ) + + parser.add_argument( + "--limit-val-batches", + type=int, + default=20, + help="Number of validation steps", + ) + parser.add_argument( + "--log-every-n-steps", + type=int, + default=1, + required=False, + help="Number of steps between logging.", + ) + parser.add_argument( + "--ckpt-dir", + type=str, + default=DEFAULT_CKPT_DIR, + help="Directory to restore an initial checkpoint from. Use this for supervised fine-tuning.", + ) + parser.add_argument( + "--wd", type=float, default=0.01, help="Weight decay for optimizer." + ) + parser.add_argument( + "--restore-optimizer-from-ckpt", + action="store_true", + help="Restore optimizer state from initial checkpoint. Defaults to False.", + ) + parser.add_argument( + "--no-average-in-collective", + action="store_true", + default=False, + help="Avaerage optimizer state in collective rather than dividing by dp size and summing.", + ) + parser.add_argument( + "--seed", type=int, default=1234, help="Set random seed for training." + ) + parser.add_argument( + "--workers", + type=int, + default=8, + help="Number of workers to use for data loading.", + ) + parser.add_argument( + "--gc-interval", + type=int, + default=0, + help="Set to a value > 0 if you want to synchronize garbage collection, will do gc every gc-interval steps.", + ) + parser.add_argument( + "--enable-preemption", + action="store_true", + default=False, + help="Enable preemption hooks. If enabled this will save a checkpoint whenever slurm exits.", + ) + parser.add_argument( + "--ckpt-async-save", + action="store_true", + default=False, + ) + parser.add_argument( + "--ckpt-format", + type=str, + choices=["torch_dist", "zarr"], + default="torch_dist", + help="Specify checkpoint format to use. Defaults to 'torch_dist', as 'zarr' is deprecated. Only use if " + "resuming training from a zarr checkpoint.", + ) + parser.add_argument( + "--eod-pad-in-loss-mask", + action="store_true", + default=False, + help="Do not predict EOD/Pad tokens (typical default, but not default in original evo2).", + ) + parser.add_argument( + "--cross-entropy-loss-fusion", + action="store_true", + default=False, + help="Use the faster, but maybe less accurate fused form of cross entropy, " + "which also has bf16 grads internally.", + ) + parser.add_argument( + "--no-fp32-residual-connection", + action="store_true", + default=False, + help="If set, turn off fp32 residual connections which may be faster but may impact accuracy.", + ) + parser.add_argument( + "--debug-ddp-parity-freq", + type=int, + default=0, + help="Set to value > 0 to debug DDP weight parity between ranks.", + ) + parser.add_argument( + "--hybrid-override-pattern", + type=str, + help="Override the hybrid override pattern in the config (specifies hyena layer ordering and type).", + ) + parser.add_argument( + "--num-layers", + type=int, + help="If set, override the number of layers specified in the requested config.", + ) + parser.add_argument( + "--create-tflops-callback", + action="store_true", + default=False, + help="Enable tflops calculation callback for Hyena / Evo2. Defaults to False.", + ) + parser.add_argument( + "--log-parameters-and-shapes", + action="store_true", + default=False, + help="Log training parameters shapes and dtypes for debugging.", + ) + parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate.") + parser.add_argument( + "--min-lr", + type=float, + default=3e-5, + help="Min learning rate in cosine annealing.", + ) + parser.add_argument( + "--warmup-steps", + type=int, + default=2500, + help="Number of warmup steps in cosine annealing", + ) + # NSYS profiling/tooling arguments + parser.add_argument( + "--nsys-profiling", + action="store_true", + default=False, + help="Enable targeted `nsys` profiling on the training loop for a defined step range. To actually get profiling" + " output you must run the whole program with `nsys`. For example: " + " `nsys profile -s none -o output_report_name -t cuda,nvtx --force-overwrite true " + "--capture-range=cudaProfilerApi --capture-range-end=stop [regular python command here]`", + ) + # start, end, rank + parser.add_argument( + "--nsys-start-step", + type=int, + required=False, + default=0, + help="Start nsys profiling after this step.", + ) + parser.add_argument( + "--nsys-end-step", + type=int, + required=False, + help="End nsys profiling after this step.", + ) + parser.add_argument( + "--no-renormalize-loss", + action="store_true", + default=False, + help="Do not renormalize the loss weights.", + ) + # rank as list of integers + parser.add_argument( + "--nsys-ranks", + type=int, + nargs="+", + required=False, + default=[0], + help="Enable nsys profiling for these ranks.", + ) + parser.add_argument( + "--activation-checkpoint-recompute-num-layers", + type=int, + help="If set, override the default value set in the config.", + ) + parser.add_argument( + "--disable-checkpointing", + action="store_false", + default=True, # False + dest="create_checkpoint_callback", + help="Disable creating a ModelCheckpoint callback.", + ) + parser.add_argument( + "--clip-grad", + type=float, + default=1.0, + help="Grad clip value. Note that when using DDP this may need to be inflated.", + ) + parser.add_argument( + "--seq-len-interpolation-factor", + type=float, + help="Adjusts the linear scaling of ROPE (Rotary Position Embedding) for context extension. " + "Set this factor relative to your base context length e.g., for an original context length of 8192 and " + "an extended context length of 524288, use 524288/8192 = 64.", + ) + parser.add_argument( + "--overlap-param-gather", + action="store_true", + default=False, + help="Overlap the parameter gather with the optimizer step. This is currently disabled due to a NeMo bug " + "when using DDP. Making this an option defaulting to False is a temporary solution until the bug is fixed.", + ) + parser.add_argument( + "--overlap-grad-reduce", + action="store_true", + default=False, + help="Overlap the gradient reduce with the optimizer step.", + ) + parser.add_argument( + "--hidden-dropout", + type=float, + default=0.0, + help="Dropout probability for the hyena layers", + ) + parser.add_argument( + "--attention-dropout", + type=float, + default=0.0, + help="Dropout probability for the attention layers.", + ) + parser.add_argument( + "--save-top-k", + type=int, + default=5, + help="Number of best checkpoints to keep. Set to -1 to save all checkpoints.", + ) + parser.add_argument( + "--metric-to-monitor-for-checkpoints", + type=str, + default="val_loss", + help="Metric to monitor for checkpoints.", + ) + parser.add_argument( + "--save-last-checkpoint", + action="store_true", + default=True, + help="Save the last checkpoint.", + ) + parser.add_argument( + "--no-save-last-checkpoint", + action="store_false", + dest="save_last_checkpoint", + default=True, + help="Disable saving the last checkpoint.", + ) + recompute_group = parser.add_mutually_exclusive_group(required=False) + recompute_group.add_argument( + "--no-activation-checkpointing", action="store_true", default=False + ) + recompute_group.add_argument( + "--selective-activation-checkpointing", action="store_true", default=False + ) + return parser.parse_args(args=args) + + +def train(args: argparse.Namespace) -> nl.Trainer: + """Main function to run Evo2 training.""" + # Instantiate tokenizer. + tokenizer = get_nmt_tokenizer( + "byte-level", + ) + + # Infer global batch size. + global_batch_size = args.global_batch_size + if global_batch_size is None: + global_batch_size = infer_global_batch_size( # batch : 8 + micro_batch_size=args.micro_batch_size, + num_nodes=args.num_nodes, + devices=args.devices, + accumulate_grad_batches=args.grad_acc_batches, + tensor_model_parallel_size=args.tensor_parallel_size, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + context_model_parallel_size=args.context_parallel_size, + ) + if args.mock_data: + data_module = MockDataModule( + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + global_batch_size=global_batch_size, + num_workers=args.workers, + tokenizer=tokenizer, + ) + else: + blended_dataset_config = parse_dataset_config( + dataset_config_path=args.dataset_config, dataset_path=args.dataset_dir + ) + Evo2DatasetPadEodLossMask if args.eod_pad_in_loss_mask else Evo2Dataset + # Instantiate pre-training module. + data_module = PreTrainingDataModule( + paths=blended_dataset_config, + # dataset_cls=dataset_cls, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + global_batch_size=global_batch_size, + seed=args.seed, # 1234 + num_workers=args.workers, + tokenizer=tokenizer, + eod_mask_loss=args.eod_pad_in_loss_mask, + ) + + if args.no_activation_checkpointing: + activation_checkpointing_args = { + "recompute_granularity": None, + "recompute_method": None, + "recompute_num_layers": None, + } + elif args.selective_activation_checkpointing: + activation_checkpointing_args = { + "recompute_granularity": "selective", + "recompute_method": None, + "recompute_num_layers": None, + } + else: + if args.activation_checkpoint_recompute_num_layers is not None: + activation_checkpointing_args = { + "recompute_num_layers": args.activation_checkpoint_recompute_num_layers, + } + else: + activation_checkpointing_args = {} + + # Retrieve model config. + config_modifiers_init = { + "tp_comm_overlap": args.use_megatron_comm_overlap_llama3_8k, + "seq_length": args.seq_length, + "hidden_dropout": args.hidden_dropout, + "attention_dropout": args.attention_dropout, + "to_upper": "weighted" if args.no_renormalize_loss else "normalized_weighted", + "distribute_saved_activations": False if args.sequence_parallel else True, + "cross_entropy_loss_fusion": args.cross_entropy_loss_fusion, + "fp32_residual_connection": not args.no_fp32_residual_connection, + "add_bias_output": args.add_bias_output, + **activation_checkpointing_args, + } + if args.hybrid_override_pattern: + config_modifiers_init["hybrid_override_pattern"] = args.hybrid_override_pattern + if args.num_layers: + config_modifiers_init["num_layers"] = args.num_layers + + if args.model_size not in HYENA_MODEL_OPTIONS: + raise ValueError(f"Invalid model size: {args.model_size}") + evo2_config = HYENA_MODEL_OPTIONS[args.model_size](**config_modifiers_init) + + # Instantiate model. + model = llm.HyenaModel(evo2_config, tokenizer=data_module.tokenizer) + + # Setup callbacks. + callbacks = [ + RichModelSummary(max_depth=4), + LearningRateMonitor(), + TimingCallback(), + ] + + if args.enable_preemption: + callbacks.append(nl_callbacks.PreemptionCallback()) + if args.debug_ddp_parity_freq > 0: + callbacks.append( + nl_callbacks.DdpParityChecker(interval=args.debug_ddp_parity_freq) + ) + if args.log_parameters_and_shapes: + callbacks.append(nl_callbacks.ParameterDebugger()) + if args.create_tflops_callback: + # Add callback that logs the tera-FLOPS per second per GPU during training. + flop_meas_callback = FLOPsMeasurementCallback( + evo2_config, + data_module, + "hyena", + ) + callbacks.append(flop_meas_callback) + + # TODO(@cye): Add this back when it works with 24.12. + # if args.straggler_detection: + # callbacks.append( + # res_module.StragglerDetectionCallback( + # report_time_interval=300, + # calc_relative_gpu_perf=True, + # calc_individual_gpu_perf=True, + # num_gpu_perf_scores_to_print=5, + # gpu_relative_perf_threshold=0.7, + # gpu_individual_perf_threshold=0.7, + # stop_if_detected=True, + # enable_ptl_logging=True, + # ) + # ) + if args.use_megatron_comm_overlap_llama3_8k: + # Pick the floating point appropriate config. + if args.fp8: + tp_comm_overlap_cfg = userbuffers_fp8_h100_h8192_tp4_mbs1_seqlen8192 + else: + tp_comm_overlap_cfg = userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192 + callbacks.append( + MegatronCommOverlapCallback( + tp_comm_overlap=evo2_config.tp_comm_overlap, + tp_comm_overlap_cfg=tp_comm_overlap_cfg, + tp_comm_bootstrap_backend=args.tp_comm_overlap_backend, + wgrad_deferral_limit=22, # default from NeMo + overlap_param_gather_with_optimizer_step=False, # Currently disabled due to an issue with checkpointing. + align_param_gather=args.align_param_gather, + ) + ) + + if args.gc_interval > 0: + callbacks.append( + nl_callbacks.GarbageCollectionCallback( + gc_interval_train=args.gc_interval, gc_interval_val=args.gc_interval + ) + ) + if args.nsys_profiling: + if args.nsys_end_step is None: + nsys_end_step = args.max_steps + else: + nsys_end_step = args.nsys_end_step + callbacks.append( + nl_callbacks.NsysCallback( + start_step=args.nsys_start_step, + end_step=nsys_end_step, + ranks=args.nsys_ranks, + gen_shape=True, + ) + ) + + wandb_run_name = ( + f"evo2-size-{args.model_size}-TP{args.tensor_parallel_size}-" + f"PP{args.pipeline_model_parallel_size}-CP{args.context_parallel_size}" + f"-GBS{global_batch_size}-MBS{args.micro_batch_size}-SkipLossRenorm{args.no_renormalize_loss}" + f"-NOAC{args.no_activation_checkpointing}-SELAC{args.selective_activation_checkpointing}" + f"-ACRNL{evo2_config.recompute_num_layers}" + f"-PAT{evo2_config.hybrid_override_pattern}" + f"-F32R{evo2_config.fp32_residual_connection}" + f"-FCE{evo2_config.cross_entropy_loss_fusion}" + f"-AIC{not args.no_average_in_collective}" + f"-PEOD{args.eod_pad_in_loss_mask}" + f"-BO{args.add_bias_output}" + f"-GCLP{args.clip_grad}" + f"-HDO{args.hidden_dropout}" + f"-ADO{args.attention_dropout}" + f"-LR{args.lr}-MINLR{args.min_lr}-WUSTEPS{args.warmup_steps}-CONSTSTEPS{args.constant_steps}-WD{args.wd}" + f"-GRFP32{args.grad_reduce_in_fp32}-FP8WG{args.fp8_wgrad and args.fp8}" + f"-OGR{args.overlap_grad_reduce}-OPG{args.overlap_param_gather}" + f"-NODES{args.num_nodes}-FP8{args.fp8}" + ) + + wandb_config: Optional[WandbConfig] = ( + None + if args.wandb_project is None + else WandbConfig( + offline=args.wandb_offline, + project=args.wandb_project, + name=( + args.wandb_run_name + if args.wandb_run_name is not None + else wandb_run_name + ), + entity=args.wandb_entity, + tags=args.wandb_tags, + group=args.wandb_group, + job_type=args.wandb_job_type, + id=args.wandb_id, + anonymous=args.wandb_anonymous, + log_model=args.wandb_log_model, + ) + ) + + nemo_logger = setup_nemo_lightning_logger( + root_dir=args.result_dir, + name=args.experiment_name, + initialize_tensorboard_logger=args.create_tensorboard_logger, + wandb_config=wandb_config, + ) + + if args.create_checkpoint_callback: + checkpoint_path = str(Path(nemo_logger.save_dir) / "checkpoints") + checkpoint_callback = nl_callbacks.ModelCheckpoint( + dirpath=checkpoint_path, + save_last=args.save_last_checkpoint, + monitor=args.metric_to_monitor_for_checkpoints, + save_top_k=args.save_top_k, + every_n_train_steps=args.val_check_interval, + always_save_context=True, + filename="{epoch}-{step}-{consumed_samples}", + save_weights_only=False, + save_optim_on_train_end=True, + save_context_on_train_end=True, + ) + callbacks.append(checkpoint_callback) + + auto_resume = nl.AutoResume( + resume_if_exists=True, + resume_ignore_no_checkpoint=True, + resume_past_end=False, + resume_from_directory=checkpoint_path, + restore_config=( + RestoreConfig( + path=args.ckpt_dir, + load_model_state=True, + load_optim_state=args.restore_optimizer_from_ckpt, + ) + if args.ckpt_dir + else None + ), + ) + else: + auto_resume = None + + ddp: DistributedDataParallelConfig = DistributedDataParallelConfig( + check_for_nan_in_grad=True, + overlap_grad_reduce=args.overlap_grad_reduce, + overlap_param_gather=args.overlap_param_gather, # Verify that this works using + grad_reduce_in_fp32=args.grad_reduce_in_fp32, + align_param_gather=args.align_param_gather, + average_in_collective=not args.no_average_in_collective, + ) + + # Initialize Megatron Strategy and Trainer. + strategy = nl.MegatronStrategy( + ddp=ddp, + tensor_model_parallel_size=args.tensor_parallel_size, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + context_parallel_size=args.context_parallel_size, + pipeline_dtype=torch.bfloat16, + sequence_parallel=args.sequence_parallel, + ckpt_load_optimizer=True, + ckpt_save_optimizer=True, + ckpt_async_save=args.ckpt_async_save, + save_ckpt_format=args.ckpt_format, + ckpt_load_strictness="log_all", # or rebasing to https://github.com/NVIDIA/NeMo/pull/11988/files#diff-7667eae242a8ef776bff78cd08e79bc81df4896a450f0a781f6ed317a3dfb7ffR139 + ) + + trainer = nl.Trainer( + devices=args.devices, + num_nodes=args.num_nodes, + max_steps=( + args.max_steps + if args.early_stop_on_step is None + else args.early_stop_on_step + ), + accelerator="gpu", + strategy=strategy, + callbacks=callbacks, + log_every_n_steps=args.log_every_n_steps, + limit_val_batches=args.limit_val_batches, + num_sanity_val_steps=0, + use_distributed_sampler=False, + # plugins=[nl.MegatronMixedPrecision( + # precision="bf16-mixed", + # params_dtype=torch.bfloat16, + # grad_reduce_in_fp32=args.grad_reduce_in_fp32, + # fp8="hybrid" if args.fp8 else None, + # fp8_amax_history_len=16 if args.fp8 else 1, + # fp8_amax_compute_algo="max" if args.fp8 else "most_recent", + # fp8_wgrad=args.fp8 + # and ( + # args.fp8_wgrad or args.use_megatron_comm_overlap_llama3_8k + # ), # faster and less accurate when set to True, and MUST be True if using TP communication overlap + # ), + # LightningEnvironment(), + # ], + plugins=nl.MegatronMixedPrecision( + precision="bf16-mixed", + params_dtype=torch.bfloat16, + grad_reduce_in_fp32=args.grad_reduce_in_fp32, + fp8="hybrid" if args.fp8 else None, + fp8_amax_history_len=16 if args.fp8 else 1, + fp8_amax_compute_algo="max" if args.fp8 else "most_recent", + fp8_wgrad=args.fp8 + and ( + args.fp8_wgrad or args.use_megatron_comm_overlap_llama3_8k + ), # faster and less accurate when set to True, and MUST be True if using TP communication overlap + ), + val_check_interval=args.val_check_interval, + enable_checkpointing=args.create_checkpoint_callback, + ) + + # Logger setup + nemo_logger.setup( + trainer, + resume_if_exists=True, + ) + + if auto_resume is not None: + auto_resume.setup(trainer, model) + + # Optimizer and scheduler setup + opt_config = OptimizerConfig( + optimizer="adam", + lr=args.lr, + adam_beta1=0.9, + adam_beta2=0.95, + weight_decay=args.wd, + clip_grad=args.clip_grad, + use_distributed_optimizer=True, + bf16=True, + ) + + sched = CosineAnnealingScheduler( + max_steps=trainer.max_steps, + warmup_steps=args.warmup_steps, + min_lr=args.min_lr, + constant_steps=args.constant_steps, + ) + + opt = MegatronOptimizerModule( + opt_config, + sched, + no_weight_decay_cond=evo2_config.hyena_no_weight_decay_cond_fn, + ) + opt.connect(model) + + # Start training + trainer.fit(model, data_module) + return trainer + + +def main(): + """Parsing args and running evo2 training.""" + args = parse_args() + if args.ckpt_dir and not os.path.isdir(args.ckpt_dir): + raise SystemExit(f"ERROR: Evo2 checkpoint directory does not exist: {args.ckpt_dir}") + train(args=args) + + +if __name__ == "__main__": + main()