init

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/open-llama/7B/open_llama_7b_preview_300bt_easylm filter=lfs diff=lfs merge=lfs -text
+data/alpaca/alpaca_data_cleaned_archive_origin.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,16 @@
+__pycache__
+.idea
+.DS_Store
+*.egg-info
+build
+
+# data
+data
+checkpoints
+out
+!data/shakespeare/prepare.py
+wandb
+
+# downloaded by our tests
+original_model.py
+original_adapter.py

LICENSE

@@ -0,0 +1,201 @@
+                                 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 [2023] Lightning AI
+
+   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.

README.md

@@ -1,3 +1,187 @@
----
-license: openrail
----
+<div align="center">
+<img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/Lit_LLaMA_Badge3x.png" alt="Lit-LLaMA" width="128"/>
+
+# ⚡ Lit-LLaMA ️
+
+<!--
+<p align="center">
+  <a href="https://www.lightning.ai/">Lightning.ai</a> •
+  <a href="https://lightning.ai/docs/pytorch/stable/">PyTorch Lightning</a> •
+  <a href="https://lightning.ai/docs/fabric/stable/">Fabric</a>
+</p>
+-->
+
+![cpu-tests](https://github.com/lightning-AI/lit-llama/actions/workflows/cpu-tests.yml/badge.svg) [![Build Status](https://dev.azure.com/Lightning-AI/lit%20Models/_apis/build/status%2FLightning-AI.lit-LLaMA?branchName=main)](https://dev.azure.com/Lightning-AI/lit%20Models/_build/latest?definitionId=49&branchName=main) [![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/Lightning-AI/lit-llama/blob/master/LICENSE) [![Discord](https://img.shields.io/discord/1077906959069626439?style=plastic)](https://discord.gg/VptPCZkGNa)
+
+<img src="https://pl-public-data.s3.amazonaws.com/assets_lightning/Llama_pineapple.gif" alt="Lit-LLaMA and pineapple pizza" width="500px"/>
+
+</div>
+
+# ⚡ Lit-LLaMA ️
+Independent implementation of [LLaMA](<https://github.com/facebookresearch/llama>) that is fully open source under the **Apache 2.0 license.**
+
+This implementation builds on [nanoGPT](<https://github.com/karpathy/nanoGPT>).
+
+The original LLaMA weights are distributed by Meta under a [research-only license](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md#model-details).
+
+New Apache 2.0 licensed weights are being released as part of the [Open LLaMA project](https://github.com/openlm-research/open_llama). Both can be [loaded in Lit-LLaMA](howto/download_weights.md).
+
+## Why?
+
+We believe that AI should be fully open source and part of the collective knowledge.
+
+The original [LLaMA code](https://github.com/facebookresearch/llama) is [GPL licensed](https://github.com/facebookresearch/llama/blob/main/LICENSE) which means any project using it must also be released under GPL.
+
+This "taints" any other code and prevents integration with the rest of the ecosystem.
+
+**Lit-LLaMA solves that for good.**
+
+&nbsp;
+
+## Design principles
+**Lit-LLaMA** is:
+
+- **Simple:** Single-file implementation without boilerplate.
+- **Correct:** Numerically equivalent to the original model.
+- **Optimized:** Runs on consumer hardware or at scale.
+- **Open-source:** No strings attached.
+
+## Get involved!
+[Join our Discord](https://discord.gg/VptPCZkGNa) to build high-performance, truly open-source models for the common benefit of the community.
+
+&nbsp;
+
+## Setup
+
+Clone the repo
+
+```bash
+git clone https://github.com/Lightning-AI/lit-llama
+cd lit-llama
+```
+
+install dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+You are all set! 🎉
+
+&nbsp;
+
+## Use the model
+
+To generate text predictions, you need to download the model weights. **If you don't have them, check out our [guide](howto/download_weights.md).**
+
+Run inference:
+
+```bash
+python generate.py --prompt "Hello, my name is"
+```
+
+This will run the 7B model and require ~26 GB of GPU memory (A100 GPU).
+
+[Full guide for generating samples from the model](howto/inference.md).
+
+### Run Lit-LLaMA on consumer devices
+
+On GPUs with `bfloat16` support, the `generate.py` script will automatically convert the weights and consume about ~14 GB.
+For GPUs with less memory, or ones that don't support `bfloat16`, enable quantization (`--quantize llm.int8`):
+
+```bash
+python generate.py --quantize llm.int8 --prompt "Hello, my name is"
+```
+
+See `python generate.py --help` for more options.
+
+You can also use GPTQ-style int4 quantization, but this needs conversions of the weights first:
+
+```bash
+python quantize.py --checkpoint_path lit-llama.pth --tokenizer_path tokenizer.model --output_path llama-7b-gptq.4bit.pth --dtype bfloat16  --quantize gptq.int4
+```
+
+With the generated quantized checkpoint generation works as usual with `--quantize gptq.int4`, bringing GPU usage to about ~5GB. As only the weights of the Linear layers are quantized, it is useful to use `--dtype bfloat16` even with the quantization enabled.
+
+[Full guide for generating samples from the model](howto/inference.md).
+
+## Finetune the model
+
+We provide a simple training scripts in `finetune_lora.py` and `finetune_adapter.py` that instruction-tunes a pretrained model on the [Alpaca](https://github.com/tatsu-lab/stanford_alpaca) dataset using the techniques of [LoRA](https://arxiv.org/abs/2106.09685) and [Adapter](https://arxiv.org/abs/2303.16199).
+
+1. Download the data and generate a instruction tuning dataset:
+
+   ```bash
+   python scripts/prepare_alpaca.py
+   ```
+
+2. Run the finetuning script
+
+   ```bash
+   python finetune_lora.py
+   ```
+   or 
+   ```bash
+   python finetune_adapter.py
+   ```
+
+It is expected that you have downloaded the pretrained weights as described above.
+The finetuning requires at least one GPU with ~24 GB memory (GTX 3090). Follow the instructions in the script to efficiently fit your GPU memory.
+Note: For some GPU models you might need to set `torch.backends.cuda.enable_flash_sdp(False)` (see comments at the top of the script).
+
+More details about each finetuning method and how you can apply it to your own data can be found in our technical how-to guides.
+
+### Finetuning How-To Guides
+
+These technical tutorials illustrate how to run the finetuning code.
+
+- [Finetune with LoRA](howto/finetune_lora.md)
+- [Finetune with Adapters](howto/finetune_adapter.md)
+
+### Understanding Finetuning -- Conceptual Tutorials
+
+Looking for conceptual tutorials and explanations? We have some additional articles below:
+
+- [Understanding Parameter-Efficient Finetuning of Large Language Models: From Prefix Tuning to LLaMA-Adapters](https://lightning.ai/pages/community/article/understanding-llama-adapters/)
+
+## Pre-training
+
+We provide a simple training script based on Fabric if you want to venture into pre-training on RedPajama, a reproduction of the original LLaMA dataset.
+Conversion scripts for our optimized streaming `PackedDataset` are included.
+
+Follow this guide to start pre-training on the RedPajama dataset:
+
+- [Pretrain on RedPajama](howto/train_redpajama.md)
+
+## Get involved!
+
+We are on a quest towards fully open source AI.
+
+<img align="right" src="https://pl-public-data.s3.amazonaws.com/assets_lightning/Lit_LLaMA_Illustration3x.png" alt="Lit-LLaMA" width="128"/>
+
+Join us and start contributing, especially on the following areas:
+
+- [ ] [Pre-training](https://github.com/Lightning-AI/lit-llama/labels/pre-training)
+- [ ] [Fine-tuning (full and LoRA)](https://github.com/Lightning-AI/lit-llama/labels/fine-tuning)
+- [ ] [Quantization](https://github.com/Lightning-AI/lit-llama/labels/quantization)
+- [ ] [Sparsification](https://github.com/Lightning-AI/lit-llama/labels/sparsification)
+
+Look at `train.py` for a starting point towards pre-training / fine-tuning using [Lightning Fabric](https://lightning.ai/docs/fabric/stable/).
+
+We welcome all individual contributors, regardless of their level of experience or hardware. Your contributions are valuable, and we are excited to see what you can accomplish in this collaborative and supportive environment. 
+
+Unsure about contributing? Check out our [Contributing to Lit-LLaMA: A Hitchhiker’s Guide to the Quest for Fully Open-Source AI](https://lightning.ai/pages/community/tutorial/contributing-to-lit-llama-a-hitchhikers-guide-to-the-quest-for-fully-open-source-ai/) guide.
+
+Don't forget to [join our Discord](https://discord.gg/VptPCZkGNa)!
+
+## Acknowledgements
+
+- [@karpathy](https://github.com/karpathy) for [nanoGPT](https://github.com/karpathy/nanoGPT)
+- [@FacebookResearch](https://github.com/facebookresearch) for the original [LLaMA implementation](https://github.com/facebookresearch/llama)
+- [@TimDettmers](https://github.com/TimDettmers) for [bitsandbytes](https://github.com/TimDettmers/bitsandbytes)
+- [@Microsoft](https://github.com/microsoft) for [LoRA](https://github.com/microsoft/LoRA)
+- [@IST-DASLab](https://github.com/IST-DASLab) for [GPTQ](https://github.com/IST-DASLab/gptq)
+
+## License
+
+Lit-LLaMA is released under the [Apache 2.0](https://github.com/Lightning-AI/lightning-llama/blob/main/LICENSE) license.

checkpoints/lit-llama-bak/7B/lit-llama.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09b6a8994e8bc8ed517e600e355a19cbe41eaf8338532ff4f88d43df6b95e3cd
+size 26953750909

checkpoints/lit-llama-bak/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc820fc43f4173d6362c16658c409ed423929a807e55a984af96cce1277d39a4
+size 772031

checkpoints/lit-llama/7B/lit-llama.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09b6a8994e8bc8ed517e600e355a19cbe41eaf8338532ff4f88d43df6b95e3cd
+size 26953750909

checkpoints/lit-llama/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc820fc43f4173d6362c16658c409ed423929a807e55a984af96cce1277d39a4
+size 772031

checkpoints/open-llama/7B/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow 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
+*.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
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+open_llama_7b_preview_300bt_easylm filter=lfs diff=lfs merge=lfs -text

checkpoints/open-llama/7B/LICENSE.txt

@@ -0,0 +1,201 @@
+                                 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.

checkpoints/open-llama/7B/README.md

@@ -0,0 +1,126 @@
+---
+license: apache-2.0
+datasets:
+    - togethercomputer/RedPajama-Data-1T
+---
+
+
+# OpenLLaMA: An Open Reproduction of LLaMA
+
+In this repo, we release a permissively licensed open source reproduction of Meta AI's [LLaMA](https://ai.facebook.com/blog/large-language-model-llama-meta-ai/) large language model. In this release, we're releasing a public preview of the 7B OpenLLaMA model that has been trained with 200 billion tokens. We provide PyTorch and Jax weights of pre-trained OpenLLaMA models, as well as evaluation results and comparison against the original LLaMA models. Stay tuned for our updates.
+
+**JAX and PyTorch Weights on Huggingface Hub**
+- [200B Checkpoint](https://huggingface.co/openlm-research/open_llama_7b_preview_200bt)
+- [300B Checkpoint](https://huggingface.co/openlm-research/open_llama_7b_preview_300bt)
+
+
+## Update 5/3/2023
+We have released a new checkpoint of OpenLLaMA 7B trained on 300B tokens. In communicating
+with our users, we have realized that many existing implementations of LLaMA does not
+prepend the BOS token (id=1) at generation time. Our 200B checkpoint is sensitive
+to this and may produce degraded results without BOS token at the beginning. Hence,
+we recommend always prepending the BOS token when using our 200B checkpoint.
+
+In an effort to make our model boradly compatible with existing implementations, we have now
+released a new 300B checkpoint, which is less sensitive to BOS token and can be used
+either way.
+
+
+## Dataset and Training
+
+We train our models on the [RedPajama](https://www.together.xyz/blog/redpajama) dataset released by [Together](https://www.together.xyz/), which is a reproduction of the LLaMA training dataset containing over 1.2 trillion tokens. We follow the exactly same preprocessing steps and training hyperparameters as the original LLaMA paper, including model architecture, context length, training steps, learning rate schedule, and optimizer.  The only difference between our setting and the original one is the dataset used: OpenLLaMA employs the RedPajama dataset rather than the one utilized by the original LLaMA.
+
+We train the models on cloud TPU-v4s using [EasyLM](https://github.com/young-geng/EasyLM), a JAX based training pipeline we developed for training and fine-tuning language model. We employ a combination of normal data parallelism and [fully sharded data parallelism (also know as ZeRO stage 3)](https://engineering.fb.com/2021/07/15/open-source/fsdp/) to balance the training throughput and memory usage. Overall we reach a throughput of over 1900 tokens / second / TPU-v4 chip in our training run.
+
+
+## Evaluation
+
+We evaluated OpenLLaMA on a wide range of tasks using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).  The LLaMA results are generated by running the original LLaMA model on the same evaluation metrics. We note that our results for the LLaMA model differ slightly from the original LLaMA paper, which we believe is a result of different evaluation protocols. Similar differences have been reported in [this issue of lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness/issues/443). Additionally, we present the results of GPT-J, a 6B parameter model trained on the [Pile](https://pile.eleuther.ai/) dataset by [EleutherAI](https://www.eleuther.ai/).
+
+The original LLaMA model was trained for 1 trillion tokens and GPT-J was trained for 500 billion tokens, whereas OpenLLaMA was trained on 200 billion tokens.  We present the results in the table below. OpenLLaMA exhibits comparable performance to the original LLaMA and GPT-J across a majority of tasks, and outperforms them in some tasks. We expect that the performance of OpenLLaMA, after completing its training on 1 trillion tokens, will be enhanced even further.
+
+
+| **Task/Metric**        | **GPT-J 6B** | **LLaMA 7B** | **Open LLaMA 7B Preview 200B Tokens** |
+| ---------------------- | ------------ | ------------ | ------------------------------------- |
+| anli_r1/acc            | 0.32         | 0.35         | 0.34                                  |
+| anli_r2/acc            | 0.34         | 0.34         | 0.35                                  |
+| anli_r3/acc            | 0.35         | 0.37         | 0.34                                  |
+| arc_challenge/acc      | 0.34         | 0.39         | 0.31                                  |
+| arc_challenge/acc_norm | 0.37         | 0.41         | 0.34                                  |
+| arc_easy/acc           | 0.67         | 0.68         | 0.66                                  |
+| arc_easy/acc_norm      | 0.62         | 0.52         | 0.59                                  |
+| boolq/acc              | 0.66         | 0.75         | 0.67                                  |
+| cb/acc                 | 0.36         | 0.36         | 0.38                                  |
+| cb/f1                  | 0.26         | 0.24         | 0.29                                  |
+| hellaswag/acc          | 0.50         | 0.56         | 0.47                                  |
+| hellaswag/acc_norm     | 0.66         | 0.73         | 0.63                                  |
+| openbookqa/acc         | 0.29         | 0.29         | 0.26                                  |
+| openbookqa/acc_norm    | 0.38         | 0.41         | 0.37                                  |
+| piqa/acc               | 0.75         | 0.78         | 0.74                                  |
+| piqa/acc_norm          | 0.76         | 0.78         | 0.74                                  |
+| record/em              | 0.88         | 0.91         | 0.87                                  |
+| record/f1              | 0.89         | 0.91         | 0.88                                  |
+| rte/acc                | 0.54         | 0.56         | 0.53                                  |
+| truthfulqa_mc/mc1      | 0.20         | 0.21         | 0.21                                  |
+| truthfulqa_mc/mc2      | 0.36         | 0.34         | 0.34                                  |
+| wic/acc                | 0.50         | 0.50         | 0.50                                  |
+| winogrande/acc         | 0.64         | 0.68         | 0.62                                  |
+| wsc/acc                | 0.37         | 0.35         | 0.57                                  |
+| Average                | 0.50         | 0.52         | 0.50                                  |
+
+
+
+
+## Preview Weights Release and Usage
+
+To encourage the feedback from the community, we release a preview checkpoint of our weights. The checkpoint can be downloaded from [HuggingFace Hub](https://huggingface.co/openlm-research/open_llama_7b_preview_200bt). We release the weights in two formats: an EasyLM format to be use with our [EasyLM framework](https://github.com/young-geng/EasyLM), and a PyTorch format to be used with the [Huggingface Transformers](https://huggingface.co/docs/transformers/index) library.
+
+For using the weights in our EasyLM framework, please refer to the [LLaMA documentation of EasyLM](https://github.com/young-geng/EasyLM/blob/main/docs/llama.md). Note that unlike the original LLaMA model, our OpenLLaMA tokenizer and weights are trained completely from scratch so it is no longer needed to obtain the original LLaMA tokenizer and weights. For using the weights in the transformers library, please follow the [transformers LLaMA documentation](https://huggingface.co/docs/transformers/main/model_doc/llama). Note that we use BOS (beginning of sentence) token (id=1) during training, so it is important to prepend this token for best performance during few-shot evaluation.
+
+Both our training framework EasyLM and the preview checkpoint weights are licensed permissively under the Apache 2.0 license.
+
+
+## Future Plans
+
+The current release is only a preview of what the complete OpenLLaMA release will offer. We are currently focused on completing the training process on the entire RedPajama dataset. This can gives us a good apple-to-apple comparison between the original LLaMA and our OpenLLaMA. Other than the 7B model, we are also training a smaller 3B model in hope of facilitating language model usage in low resource use cases.  Please stay tuned for our upcoming releases.
+
+
+
+## Contact
+
+We would love to get feedback from the community. If you have any questions, please open an issue or contact us.
+
+OpenLLaMA is developed by:
+[Xinyang Geng](https://young-geng.xyz/)* and [Hao Liu](https://www.haoliu.site/)* from Berkeley AI Research.
+*Equal Contribution
+
+
+## Reference
+
+If you found OpenLLaMA useful in your research or applications, please cite using the following BibTeX:
+```
+@software{openlm2023openllama,
+  author = {Geng, Xinyang and Liu, Hao},
+  title = {OpenLLaMA: An Open Reproduction of LLaMA},
+  month = May,
+  year = 2023,
+  url = {https://github.com/openlm-research/open_llama}
+}
+```
+```
+@software{together2023redpajama,
+  author = {Together Computer},
+  title = {RedPajama-Data: An Open Source Recipe to Reproduce LLaMA training dataset},
+  month = April,
+  year = 2023,
+  url = {https://github.com/togethercomputer/RedPajama-Data}
+}
+```
+```
+@article{touvron2023llama,
+  title={Llama: Open and efficient foundation language models},
+  author={Touvron, Hugo and Lavril, Thibaut and Izacard, Gautier and Martinet, Xavier and Lachaux, Marie-Anne and Lacroix, Timoth{\'e}e and Rozi{\`e}re, Baptiste and Goyal, Naman and Hambro, Eric and Azhar, Faisal and others},
+  journal={arXiv preprint arXiv:2302.13971},
+  year={2023}
+}
+```

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_easylm

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63ab9d652aaf4e0e47f1d9a0321ef565b62c02921ce0b18a781ba0daac2ebb98
+size 13476851687

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/config.json

@@ -0,0 +1,22 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_position_embeddings": 2048,
+  "model_type": "llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.28.0.dev0",
+  "use_cache": true,
+  "vocab_size": 32000
+}

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "4.28.0.dev0"
+}

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/pytorch_model-00001-of-00002.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:adce75e45dbad20967c0e96a83b318720a767e4b8f77aabcd01cd2b38e8f0b2e
+size 9976634558

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/pytorch_model-00002-of-00002.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a80d748a6ab528f0db2249013a5d3fea17e039ad9fa1bf3e170e9070ec30f938
+size 3500315539

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/pytorch_model.bin.index.json

@@ -0,0 +1,330 @@
+{
+  "metadata": {
+    "total_size": 13476839424
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00002-of-00002.bin",
+    "model.embed_tokens.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.22.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.23.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.24.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.28.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.29.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.3.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.30.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.30.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00002.bin",
+    "model.layers.31.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.4.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.norm.weight": "pytorch_model-00002-of-00002.bin"
+  }
+}

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/special_tokens_map.json

@@ -0,0 +1 @@
+{}

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc820fc43f4173d6362c16658c409ed423929a807e55a984af96cce1277d39a4
+size 772031

checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights/tokenizer_config.json

@@ -0,0 +1 @@
+{"bos_token": "", "eos_token": "", "model_max_length": 1000000000000000019884624838656, "tokenizer_class": "LlamaTokenizer", "unk_token": ""}

checkpoints/open-llama/7B/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc820fc43f4173d6362c16658c409ed423929a807e55a984af96cce1277d39a4
+size 772031

data/alpaca/alpaca_data_cleaned_archive.json

@@ -0,0 +1,37 @@
+[
+  {
+    "instruction": "Describe DDOS attack",
+    "input": "",
+    "output": "A Distributed Denial of Service (DDoS) attack is a malicious attempt to disrupt the normal functioning of a network, service, or website by overwhelming it with a flood of internet traffic. In a DDoS attack, multiple compromised computers or devices, often referred to as a botnet, are used to generate a massive amount of traffic that exceeds the target's capacity to handle it."
+  },
+  {
+    "instruction": "What are DDOS technique?",
+    "input": "",
+    "output": "Traffic flooding: The attacker overwhelms the target with a massive volume of network traffic, consuming all available resources. \nApplication layer attacks: The attacker targets specific applications or services, attempting to exhaust their resources by exploiting vulnerabilities or by initiating resource-intensive operations. \nProtocol attacks: The attacker exploits weaknesses in network protocols, such as the TCP/IP stack, to consume server resources or disrupt communication.\nReflective/amplified attacks: The attacker uses a technique where they send requests to a large number of vulnerable systems that respond by sending a much larger response back to the target, amplifying the attack's impact."
+  },
+  {
+    "instruction": "How to avoid DDOS attack",
+    "input": "",
+    "output": "DDoS mitigation services: Consider using a reputable DDoS mitigation service that specializes in detecting and mitigating attacks. These services can help filter out malicious traffic and ensure your network remains accessible during an attack. \nNetwork monitoring: Implement robust network monitoring systems to detect any unusual traffic patterns or sudden spikes in traffic. Early detection can allow you to respond quickly and mitigate the impact of an attack."
+  },
+  {
+    "instruction": "what to do when getting ddos-ed",
+    "input": "",
+    "output": "Contact your DDoS mitigation service: If you have a DDoS mitigation service in place, immediately contact them to activate their services and inform them about the ongoing attack. They can help filter out malicious traffic and provide guidance on mitigation strategies.\nMonitor network traffic: Use network monitoring tools or services to analyze and monitor the incoming traffic. Identify the characteristics of the attack, such as the source IP addresses, traffic patterns, and attack vectors. This information can be useful for further analysis and potential legal action."
+  },
+  {
+    "instruction": "What is DDOS attack",
+    "input": "",
+    "output": "is a malicious with a flood of internet traffic. In a DDoS attack, multiple compromised computers or devices, often referred to as a botnet, are used to generate a massive amount of traffic that exceeds the target's capacity to handle it."
+  },
+  {
+    "instruction": "Which are DDOS technique?",
+    "input": "Traffic flooding, Scanning, Viewing",
+    "output": "Traffic flooding"
+  },
+  {
+    "instruction": "Avoid DDOS attack",
+    "input": "",
+    "output": "DDoS mitigation services: Consider using a reputable DDoS mitigation service that specializes in detecting and mitigating attacks. These services can help filter out malicious traffic and ensure your network remains accessible during an attack. \nNetwork monitoring: Implement robust network monitoring systems to detect any unusual traffic patterns or sudden spikes in traffic. Early detection can allow you to respond quickly and mitigate the impact of an attack."
+  }
+]

data/alpaca/alpaca_data_cleaned_archive_origin.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00c26b8da597c1aaa5a0bac023bdb8f26bbaa37a9ead7837df4aa7e51ad57459
+size 23573609

data/alpaca/test.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51be5410edd01a0351eae4535769ff441ae4278d17a6a643bed3cfcad5888c1d
+size 4607

data/alpaca/train.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4026f9d8b9f342b2d6938202e017bd5e8f81716e2c79b78a7b0de92861f15050
+size 10902

evaluate.py

@@ -0,0 +1,145 @@
+# This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
+# Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
+import math
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+import tqdm
+
+from lit_llama import LLaMA, Tokenizer
+from lit_llama.utils import EmptyInitOnDevice, llama_model_lookup
+
+from datasets import load_dataset
+
+
+def load_eval_data(dataset_name: str) -> str:
+    # this mimics gptq datautils
+    if dataset_name == "wikitext":
+        # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
+        testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+        testdata = "\n\n".join(testdata["text"])
+    elif dataset_name == "ptb":
+        testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
+        testdata = "\n\n".join(testdata["sentence"])
+    elif dataset_name == "c4":
+        testdata = load_dataset(
+            "allenai/c4",
+            "allenai--c4",
+            data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
+            split="validation",
+        )
+        testdata = " ".join(testdata[:1100]["text"])
+
+    else:
+        raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
+    return testdata
+
+
+def main(
+    datasets: str = "wikitext,ptb,c4",
+    *,
+    # compilation fails as it does not support torch.complex64 for RoPE
+    # compile: bool = False,
+    accelerator: str = "auto",
+    checkpoint_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    dtype: str = "float32",
+    quantize: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained LLaMA model and tokenizer.
+
+    Args:
+        datasets: The datasets to use as a comma separated string
+        # compile: Whether to compile the model.
+        accelerator: The hardware to run on. Possible choices are:
+            ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
+        checkpoint_path: The checkpoint path to load.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+    """
+    if not checkpoint_path:
+        checkpoint_path = Path(f"./checkpoints/lit-llama/7B/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    assert checkpoint_path.is_file()
+    assert tokenizer_path.is_file()
+
+    fabric = L.Fabric(accelerator=accelerator, devices=1)
+
+    dt = getattr(torch, dtype, None)
+    if not isinstance(dt, torch.dtype):
+        raise ValueError(f"{dtype} is not a valid dtype.")
+    dtype = dt
+
+    with EmptyInitOnDevice(
+        device=fabric.device, dtype=dtype, quantization_mode=quantize
+    ):
+        print("Loading model ...", file=sys.stderr)
+        t0 = time.time()
+        checkpoint = torch.load(checkpoint_path)
+        name = llama_model_lookup(checkpoint)
+        model = LLaMA.from_name(name)
+        model.load_state_dict(checkpoint)
+        print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+
+    # if compile:
+    #     model = torch.compile(model)
+
+    total_toks = 0
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+
+    for dsname in datasets.split(","):
+        test_string = load_eval_data(dsname)
+        encoded_text = tokenizer.encode(
+            test_string, bos=True, eos=False, device=fabric.device
+        )
+        encoded_text = encoded_text[
+            None, : 256 * model.config.block_size
+        ]  # add batch dimension, trim like gptq implementation
+        t0 = time.perf_counter()
+
+        nlls = 0
+        toks = 0
+        with torch.inference_mode():
+            block_size = 2048  # this is for compat with gptq, and indeed we get much worse beyond this (https://github.com/facebookresearch/llama/blob/57b0eb62de0636e75af471e49e2f1862d908d9d8/llama/model.py#L30)
+            for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
+                inp = encoded_text[:, i : i + block_size]
+                logits = model(inp)[0]
+                nll = torch.nn.functional.cross_entropy(
+                    logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
+                )
+                toks += inp.size(1) - 1
+                nlls += nll.item()
+
+        print(encoded_text.shape, logits.shape)
+        encoded_text = encoded_text[:, : logits.shape[0]]
+        ppl = math.exp(nlls / toks)
+        print(f"Perplexity on {dsname}: {ppl:.2f}")
+        total_toks += toks
+
+    t = time.perf_counter() - t0
+    print(
+        f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
+        file=sys.stderr,
+    )
+    print(
+        f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    CLI(main)

evaluate_adapter.py

@@ -0,0 +1,164 @@
+# This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
+# Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
+import math
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+import tqdm
+
+from lit_llama import Tokenizer
+from lit_llama.adapter import LLaMA
+from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+from scripts.prepare_alpaca import generate_prompt
+
+from datasets import load_dataset
+
+
+def load_eval_data(dataset_name: str) -> str:
+    # this mimics gptq datautils
+    if dataset_name == "wikitext":
+        # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
+        testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+        testdata = "\n\n".join(testdata["text"])
+    elif dataset_name == "ptb":
+        testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
+        testdata = "\n\n".join(testdata["sentence"])
+    elif dataset_name == "c4":
+        testdata = load_dataset(
+            "allenai/c4",
+            "allenai--c4",
+            data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
+            split="validation",
+        )
+        testdata = " ".join(testdata[:1100]["text"])
+
+    else:
+        raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
+    return testdata
+
+
+def main(
+    datasets: str = "wikitext,ptb,c4",
+    *,
+    # compilation fails as it does not support torch.complex64 for RoPE
+    # compile: bool = False,
+    accelerator: str = "auto",
+    adapter_path: Optional[Path] = None,
+    checkpoint_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    dtype: str = "float32",
+    quantize: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained LLaMA model and tokenizer.
+
+    Args:
+        datasets: The datasets to use as a comma separated string
+        # compile: Whether to compile the model.
+        accelerator: The hardware to run on. Possible choices are:
+            ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
+        adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
+            `finetune_adapter.py`.
+        checkpoint_path: The checkpoint path to load.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+    """
+    if not adapter_path:
+        adapter_path = Path("out/adapter/alpaca/lit-llama-adapter-finetuned.pth")
+    if not checkpoint_path:
+        checkpoint_path = Path(f"./checkpoints/lit-llama/7B/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    
+    assert adapter_path.is_file()
+    assert checkpoint_path.is_file()
+    assert tokenizer_path.is_file()
+
+    fabric = L.Fabric(accelerator=accelerator, devices=1)
+
+    dt = getattr(torch, dtype, None)
+    if not isinstance(dt, torch.dtype):
+        raise ValueError(f"{dtype} is not a valid dtype.")
+    dtype = dt
+
+    with EmptyInitOnDevice(
+        device=fabric.device, dtype=dtype, quantization_mode=quantize
+    ):
+        print("Loading model ...", file=sys.stderr)
+        t0 = time.time()
+        pretrained_checkpoint = lazy_load(checkpoint_path)
+        adapter_checkpoint = lazy_load(adapter_path)
+        name = llama_model_lookup(pretrained_checkpoint)
+        model = LLaMA.from_name(name)
+
+        # 1. Load the pretrained weights
+        model.load_state_dict(pretrained_checkpoint, strict=False)
+        # 2. Load the fine-tuned adapter weights
+        model.load_state_dict(adapter_checkpoint, strict=False)
+
+        print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+
+    # if compile:
+    #     model = torch.compile(model)
+
+    total_toks = 0
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+
+    for dsname in datasets.split(","):
+        test_string = load_eval_data(dsname)
+
+        sample = {"instruction": test_string, "input": input}
+        test_string = generate_prompt(sample)
+
+        encoded_text = tokenizer.encode(
+            test_string, bos=True, eos=False, device=fabric.device
+        )
+        encoded_text = encoded_text[
+            None, : 256 * model.config.block_size
+        ]  # add batch dimension, trim like gptq implementation
+        t0 = time.perf_counter()
+
+        nlls = 0
+        toks = 0
+        with torch.inference_mode():
+            block_size = 2048  # this is for compat with gptq, and indeed we get much worse beyond this (https://github.com/facebookresearch/llama/blob/57b0eb62de0636e75af471e49e2f1862d908d9d8/llama/model.py#L30)
+            for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
+                inp = encoded_text[:, i : i + block_size]
+                logits = model(inp)[0]
+                nll = torch.nn.functional.cross_entropy(
+                    logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
+                )
+                toks += inp.size(1) - 1
+                nlls += nll.item()
+
+        print(encoded_text.shape, logits.shape)
+        encoded_text = encoded_text[:, : logits.shape[0]]
+        ppl = math.exp(nlls / toks)
+        print(f"Perplexity on {dsname}: {ppl:.2f}")
+        total_toks += toks
+
+    t = time.perf_counter() - t0
+    print(
+        f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
+        file=sys.stderr,
+    )
+    print(
+        f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    CLI(main)

evaluate_full.py

@@ -0,0 +1,145 @@
+# This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
+# Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
+import math
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+import tqdm
+
+from lit_llama import LLaMA, Tokenizer
+from lit_llama.utils import EmptyInitOnDevice
+
+from datasets import load_dataset
+
+
+def load_eval_data(dataset_name: str) -> str:
+    # this mimics gptq datautils
+    if dataset_name == "wikitext":
+        # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
+        testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+        testdata = "\n\n".join(testdata["text"])
+    elif dataset_name == "ptb":
+        testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
+        testdata = "\n\n".join(testdata["sentence"])
+    elif dataset_name == "c4":
+        testdata = load_dataset(
+            "allenai/c4",
+            "allenai--c4",
+            data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
+            split="validation",
+        )
+        testdata = " ".join(testdata[:1100]["text"])
+
+    else:
+        raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
+    return testdata
+
+
+def main(
+    datasets: str = "wikitext,ptb,c4",
+    *,
+    # compilation fails as it does not support torch.complex64 for RoPE
+    # compile: bool = False,
+    accelerator: str = "auto",
+    checkpoint_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    model_size: str = "7B",
+    dtype: str = "float32",
+    quantize: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained LLaMA model and tokenizer.
+
+    Args:
+        datasets: The datasets to use as a comma separated string
+        # compile: Whether to compile the model.
+        accelerator: The hardware to run on. Possible choices are:
+            ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
+        checkpoint_path: The checkpoint path to load.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+    """
+    if not checkpoint_path:
+        checkpoint_path = Path(f"./checkpoints/lit-llama/{model_size}/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    assert checkpoint_path.is_file()
+    assert tokenizer_path.is_file()
+
+    fabric = L.Fabric(accelerator=accelerator, devices=1)
+
+    dt = getattr(torch, dtype, None)
+    if not isinstance(dt, torch.dtype):
+        raise ValueError(f"{dtype} is not a valid dtype.")
+    dtype = dt
+
+    with EmptyInitOnDevice(
+        device=fabric.device, dtype=dtype, quantization_mode=quantize
+    ):
+        print("Loading model ...", file=sys.stderr)
+        t0 = time.time()
+        model = LLaMA.from_name(model_size)
+        checkpoint = torch.load(checkpoint_path)
+        model.load_state_dict(checkpoint)
+        print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+
+    # if compile:
+    #     model = torch.compile(model)
+
+    total_toks = 0
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+
+    for dsname in datasets.split(","):
+        test_string = load_eval_data(dsname)
+        encoded_text = tokenizer.encode(
+            test_string, bos=True, eos=False, device=fabric.device
+        )
+        encoded_text = encoded_text[
+            None, : 256 * model.config.block_size
+        ]  # add batch dimension, trim like gptq implementation
+        t0 = time.perf_counter()
+
+        nlls = 0
+        toks = 0
+        with torch.inference_mode():
+            block_size = 2048  # this is for compat with gptq, and indeed we get much worse beyond this (https://github.com/facebookresearch/llama/blob/57b0eb62de0636e75af471e49e2f1862d908d9d8/llama/model.py#L30)
+            for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
+                inp = encoded_text[:, i : i + block_size]
+                logits = model(inp)[0]
+                nll = torch.nn.functional.cross_entropy(
+                    logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
+                )
+                toks += inp.size(1) - 1
+                nlls += nll.item()
+
+        print(encoded_text.shape, logits.shape)
+        encoded_text = encoded_text[:, : logits.shape[0]]
+        ppl = math.exp(nlls / toks)
+        print(f"Perplexity on {dsname}: {ppl:.2f}")
+        total_toks += toks
+
+    t = time.perf_counter() - t0
+    print(
+        f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
+        file=sys.stderr,
+    )
+    print(
+        f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    CLI(main)

evaluate_lora.py

@@ -0,0 +1,173 @@
+# This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
+# Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
+import math
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+import tqdm
+
+from lit_llama import LLaMA, Tokenizer
+from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+from lit_llama.lora import lora
+from scripts.prepare_alpaca import generate_prompt
+
+from datasets import load_dataset
+
+lora_r = 8
+lora_alpha = 16
+lora_dropout = 0.05
+
+
+def load_eval_data(dataset_name: str) -> str:
+    # this mimics gptq datautils
+    if dataset_name == "wikitext":
+        # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
+        testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+        testdata = "\n\n".join(testdata["text"])
+    elif dataset_name == "ptb":
+        testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
+        testdata = "\n\n".join(testdata["sentence"])
+    elif dataset_name == "c4":
+        testdata = load_dataset(
+            "allenai/c4",
+            "allenai--c4",
+            data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
+            split="validation",
+        )
+        testdata = " ".join(testdata[:1100]["text"])
+
+    else:
+        raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
+    return testdata
+
+
+def main(
+    datasets: str = "wikitext,ptb,c4",
+    *,
+    # compilation fails as it does not support torch.complex64 for RoPE
+    # compile: bool = False,
+    accelerator: str = "auto",
+    lora_path: Optional[Path] = None,
+    checkpoint_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    dtype: str = "float32",
+    quantize: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained LLaMA model and tokenizer
+       finetuned with LoRA.
+
+    Args:
+        datasets: The datasets to use as a comma separated string
+        # compile: Whether to compile the model.
+        accelerator: The hardware to run on. Possible choices are:
+            ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
+        lora_path: Path to the checkpoint with trained LoRA weights, which are the output of
+            `finetune_lora.py`.
+        checkpoint_path: The checkpoint path to load.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+    """
+    if not lora_path:
+        lora_path = Path("out/lora/alpaca/lit-llama-lora-finetuned.pth")
+    if not checkpoint_path:
+        checkpoint_path = Path(f"./checkpoints/lit-llama/7B/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    assert lora_path.is_file()
+    assert checkpoint_path.is_file()
+    assert tokenizer_path.is_file()
+
+    if quantize is not None:
+        raise NotImplementedError("Quantization in LoRA is not supported yet")
+
+    fabric = L.Fabric(accelerator=accelerator, devices=1)
+
+    dt = getattr(torch, dtype, None)
+    if not isinstance(dt, torch.dtype):
+        raise ValueError(f"{dtype} is not a valid dtype.")
+    dtype = dt
+
+    print("Loading model ...", file=sys.stderr)
+    t0 = time.time()
+
+    pretrained_checkpoint = lazy_load(checkpoint_path)
+    adapter_checkpoint = lazy_load(lora_path)
+    name = llama_model_lookup(pretrained_checkpoint)
+
+    with EmptyInitOnDevice(
+        device=fabric.device, dtype=dtype, quantization_mode=quantize
+    ), lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
+        model = LLaMA.from_name(name)
+
+    # 1. Load the pretrained weights
+    model.load_state_dict(pretrained_checkpoint, strict=False)
+    # 2. Load the fine-tuned adapter weights
+    model.load_state_dict(adapter_checkpoint, strict=False)
+
+    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+
+    # if compile:
+    #     model = torch.compile(model)
+
+    total_toks = 0
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+
+    for dsname in datasets.split(","):
+        test_string = load_eval_data(dsname)
+
+        sample = {"instruction": test_string, "input": input}
+        test_string = generate_prompt(sample)
+        
+        encoded_text = tokenizer.encode(
+            test_string, bos=True, eos=False, device=fabric.device
+        )
+        encoded_text = encoded_text[
+            None, : 256 * model.config.block_size
+        ]  # add batch dimension, trim like gptq implementation
+        t0 = time.perf_counter()
+
+        nlls = 0
+        toks = 0
+        with torch.inference_mode():
+            block_size = 2048  # this is for compat with gptq, and indeed we get much worse beyond this (https://github.com/facebookresearch/llama/blob/57b0eb62de0636e75af471e49e2f1862d908d9d8/llama/model.py#L30)
+            for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
+                inp = encoded_text[:, i : i + block_size]
+                logits = model(inp)[0]
+                nll = torch.nn.functional.cross_entropy(
+                    logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
+                )
+                toks += inp.size(1) - 1
+                nlls += nll.item()
+
+        print(encoded_text.shape, logits.shape)
+        encoded_text = encoded_text[:, : logits.shape[0]]
+        ppl = math.exp(nlls / toks)
+        print(f"Perplexity on {dsname}: {ppl:.2f}")
+        total_toks += toks
+
+    t = time.perf_counter() - t0
+    print(
+        f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
+        file=sys.stderr,
+    )
+    print(
+        f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    CLI(main)

finetune_adapter.py

@@ -0,0 +1,253 @@
+"""
+Instruction-tuning with LLaMA-Adapter on the Alpaca dataset following the paper
+
+LLaMA-Adapter: Efficient Fine-tuning of Language Models with Zero-init Attention
+https://arxiv.org/abs/2303.16199
+
+This script runs on a single GPU by default. You can adjust the `micro_batch_size` to fit your GPU memory.
+You can finetune within 1 hour as done in the original paper using DeepSpeed Zero-2 on 8 A100 GPUs by setting the
+devices variable to `devices = 8` and `micro_batch_size = 8` (or higher).
+
+Note: If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
+`torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).
+"""
+import os
+import time
+from pathlib import Path
+import shutil
+
+import lightning as L
+import numpy as np
+import torch
+
+from generate import generate
+from lit_llama.adapter import LLaMA, LLaMAConfig, mark_only_adapter_as_trainable, adapter_state_from_state_dict
+from lit_llama.tokenizer import Tokenizer
+from scripts.prepare_alpaca import generate_prompt
+from lightning.fabric.strategies import DeepSpeedStrategy
+
+
+eval_interval = 600
+save_interval = 1000
+eval_iters = 100
+log_interval = 1
+devices = 1
+
+# Hyperparameters
+learning_rate = 9e-3
+batch_size = 64 / devices
+micro_batch_size = 4
+gradient_accumulation_steps = batch_size // micro_batch_size
+epoch_size = 50000  # train dataset size
+num_epochs = 5
+max_iters = num_epochs * epoch_size // devices
+weight_decay = 0.02
+max_seq_length = 256  # see scripts/prepare_alpaca.py
+warmup_steps = epoch_size * 2 // micro_batch_size // devices  # 2 epochs
+
+ds_config = {
+    "train_micro_batch_size_per_gpu": micro_batch_size,
+    "gradient_accumulation_steps": gradient_accumulation_steps,
+    "zero_optimization": {"stage": 2},
+}
+
+
+def main(
+    data_dir: str = "data/alpaca", 
+    pretrained_path: str = "checkpoints/lit-llama/7B/lit-llama.pth",
+    out_dir: str = "out/adapter/alpaca",
+):
+
+    fabric = L.Fabric(
+        accelerator="cuda", 
+        devices=devices, 
+        strategy=(DeepSpeedStrategy(config=ds_config) if devices > 1 else "auto"), 
+        precision="bf16-true",
+    )
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if fabric.global_rank == 0:
+        os.makedirs(out_dir, exist_ok=True)
+
+    train_data, val_data = load_datasets(data_dir=data_dir)
+
+    config = LLaMAConfig(block_size=max_seq_length)
+
+    if not os.path.isfile(pretrained_path):
+        raise FileNotFoundError(
+            f"Can't find the pretrained weights at {pretrained_path}."
+            " Please follow the instructions in the README to download them."
+        )
+    checkpoint = torch.load(pretrained_path)
+
+    with fabric.init_module():
+        model = LLaMA(config)
+        # strict=False because missing keys due to adapter weights not containted in state dict
+        model.load_state_dict(checkpoint, strict=False)
+
+    mark_only_adapter_as_trainable(model)
+
+    num_params = sum([p.numel() for p in model.parameters() if p.requires_grad])
+    print(f"Number of trainable parameters: {num_params}")
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+    model, optimizer = fabric.setup(model, optimizer)
+    train(fabric, model, optimizer, train_data, val_data, out_dir)
+
+    # Save the final checkpoint at the end of training
+    save_model_checkpoint(fabric, model, os.path.join(out_dir, "lit-llama-adapter-finetuned.pth"))
+
+
+def train(
+    fabric: L.Fabric,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    train_data: np.ndarray,
+    val_data: np.ndarray,
+    out_dir: str,
+) -> None:
+    """The training loop.
+
+    Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
+    """
+    step_count = 0
+
+    for iter_num in range(max_iters):
+
+        if step_count <= warmup_steps:
+            # linear warmup
+            lr = learning_rate * step_count / warmup_steps
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = lr
+
+        t0 = time.time()
+
+        input_ids, targets = get_batch(fabric, train_data)
+        logits = model(input_ids)
+        loss = loss_fn(logits, targets)
+        with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_steps != 0)):
+            fabric.backward(loss / gradient_accumulation_steps)
+
+        if (iter_num + 1) % gradient_accumulation_steps == 0:
+            optimizer.step()
+            optimizer.zero_grad()
+            step_count += 1
+                
+            if step_count % eval_interval == 0:
+                val_loss = validate(fabric, model, val_data)
+                fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
+                fabric.barrier()
+
+            if step_count % save_interval == 0:
+                print(f"Saving adapter weights to {out_dir}")
+                # TODO: Provide a function/script to merge the adapter weights with pretrained weights
+                save_model_checkpoint(fabric, model, os.path.join(out_dir, f"iter-{iter_num:06d}.pth"))
+
+        dt = time.time() - t0
+        if iter_num % log_interval == 0:
+            fabric.print(f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms")
+
+
+def generate_response(model, instruction, input=""):
+    tokenizer = Tokenizer("checkpoints/lit-llama/tokenizer.model")
+    sample = {"instruction": instruction, "input": input}
+    prompt = generate_prompt(sample)
+    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
+
+    output = generate(
+        model,
+        idx=encoded,
+        max_seq_length=max_seq_length,
+        max_new_tokens=100,
+        temperature=0.8,
+    )
+    output = tokenizer.decode(output)
+    return output # output.split("### Response:")[1].strip()
+
+
+@torch.no_grad()
+def validate(fabric: L.Fabric, model: torch.nn.Module, val_data: np.ndarray) -> torch.Tensor:
+    fabric.print("Validating ...")
+    model.eval()
+    losses = torch.zeros(eval_iters)
+    for k in range(eval_iters):
+        input_ids, targets = get_batch(fabric, val_data)
+        logits = model(input_ids)
+        loss = loss_fn(logits, targets)
+        losses[k] = loss.item()
+    val_loss = losses.mean()
+
+    # produce an example:
+    instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
+    output = generate_response(model, instruction)
+    fabric.print(instruction)
+    fabric.print(output)
+
+    model.train()
+    return val_loss.item()
+
+def loss_fn(logits, targets):
+    # shift the targets such that output n predicts token n+1
+    logits = logits[..., :-1, :].contiguous()
+    targets = targets[..., 1:].contiguous()
+    loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+    return loss
+    
+
+def get_batch(fabric: L.Fabric, data: list):
+    ix = torch.randint(len(data), (micro_batch_size,))
+
+    input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
+    labels = [data[i]["labels"].type(torch.int64) for i in ix]
+
+    max_len = max(len(s) for s in input_ids)
+
+    def pad_right(x, pad_id):
+        # pad right based on the longest sequence
+        n = max_len - len(x)
+        return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
+
+    x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
+    y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
+    x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
+    return x, y
+
+
+def load_datasets(data_dir):
+    train_data = torch.load(os.path.join(data_dir, "train.pt"))
+    val_data = torch.load(os.path.join(data_dir, "test.pt"))
+    return train_data, val_data
+
+
+def save_model_checkpoint(fabric, model, file_path):
+    file_path = Path(file_path)
+
+    if isinstance(fabric.strategy, DeepSpeedStrategy):
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+
+        tmp_path = file_path.with_suffix(".tmp")
+        fabric.save(tmp_path, {"model": model})
+        fabric.barrier()
+        if fabric.global_rank == 0:
+            # Create a consolidated checkpoint with the same name next to the deepspeed checkpoint
+            # and only keep the adapter weights
+            state_dict = get_fp32_state_dict_from_zero_checkpoint(tmp_path)
+            state_dict = adapter_state_from_state_dict(state_dict)
+            torch.save(state_dict, file_path)
+            shutil.rmtree(tmp_path)
+    else:
+        state_dict = adapter_state_from_state_dict(model.state_dict())
+        if fabric.global_rank == 0:
+            torch.save(state_dict, file_path)
+        fabric.barrier()
+
+
+if __name__ == "__main__":
+    # Uncomment this line if you see an error: "Expected is_sm80 to be true, but got false"
+    # torch.backends.cuda.enable_flash_sdp(False)
+    torch.set_float32_matmul_precision("high")
+
+    from jsonargparse.cli import CLI
+
+    CLI(main)

finetune_full.py

@@ -0,0 +1,214 @@
+"""
+Instruction-tuning on the Alpaca dataset using a regular finetuning procedure (updating all layers).
+
+Note: If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
+`torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).
+"""
+import os
+import time
+from functools import partial
+
+import lightning as L
+from lightning.fabric.strategies import FSDPStrategy
+import numpy as np
+import torch
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+from generate import generate
+from lit_llama.model import Block, LLaMA, LLaMAConfig
+from lit_llama.tokenizer import Tokenizer
+from lit_llama.utils import save_model_checkpoint
+from scripts.prepare_alpaca import generate_prompt
+
+
+eval_interval = 1000
+save_interval = 1000
+eval_iters = 100
+log_interval = 100
+devices = 4
+
+# Hyperparameters
+learning_rate = 3e-5
+batch_size = 128 / devices
+micro_batch_size = 4
+gradient_accumulation_steps = batch_size // micro_batch_size
+epoch_size = 50000  # train dataset size
+num_epochs = 5
+max_iters = num_epochs * epoch_size // micro_batch_size // devices
+weight_decay = 0.0
+block_size = 512
+warmup_steps = 100
+
+
+def main(
+    data_dir: str = "data/alpaca",
+    pretrained_path: str = "checkpoints/lit-llama/7B/lit-llama.pth",
+    out_dir: str = "out/full/alpaca",
+):
+
+    auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
+    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block)
+
+    fabric = L.Fabric(accelerator="cuda", devices=devices, precision="bf16-mixed", strategy=strategy)
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if fabric.global_rank == 0:
+        os.makedirs(out_dir, exist_ok=True)
+
+    train_data, val_data = load_datasets(data_dir=data_dir)
+
+    config = LLaMAConfig.from_name("7B")
+    config.block_size = block_size
+
+    checkpoint = torch.load(pretrained_path)
+
+    with fabric.device:
+        torch.set_default_tensor_type(torch.HalfTensor)
+        model = LLaMA(config).bfloat16()
+        torch.set_default_tensor_type(torch.FloatTensor)
+        model.load_state_dict(checkpoint, strict=False) 
+
+    model = fabric.setup_module(model)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+    optimizer = fabric.setup_optimizers(optimizer)
+
+    train(fabric, model, optimizer, train_data, val_data, out_dir)
+
+    # Save the final checkpoint at the end of training
+    save_model_checkpoint(fabric, model, os.path.join(out_dir, "lit-llama-full-finetuned.pth"))
+
+
+def train(
+    fabric: L.Fabric,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    train_data: np.ndarray,
+    val_data: np.ndarray,
+    out_dir: str,
+) -> None:
+    """The training loop.
+
+    Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
+    """
+    step_count = 0
+    model.train()
+
+    for iter_num in range(max_iters):
+
+        is_accumulating = (iter_num + 1) % gradient_accumulation_steps == 0
+
+        if step_count <= warmup_steps:
+            # linear warmup
+            lr = learning_rate * step_count / warmup_steps
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = lr
+
+        t0 = time.time()
+
+        with fabric.no_backward_sync(model, enabled=is_accumulating):
+            input_ids, targets = get_batch(fabric, train_data)
+            logits = model(input_ids)
+            loss = loss_fn(logits, targets)
+            fabric.backward(loss)
+
+        if not is_accumulating:
+            optimizer.step()
+            optimizer.zero_grad()
+            step_count += 1
+
+            if step_count % eval_interval == 0:
+                val_loss = validate(fabric, model, val_data)
+                fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
+                fabric.barrier()
+
+            if step_count % save_interval == 0:
+                print(f"Saving weights to {out_dir}")
+                save_model_checkpoint(fabric, model, os.path.join(out_dir, f"iter-{iter_num:06d}-ckpt.pth"))
+
+        dt = time.time() - t0
+        if iter_num % log_interval == 0:
+            fabric.print(f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms")
+
+
+def generate_response(model, instruction):
+    tokenizer = Tokenizer("checkpoints/lit-llama/tokenizer.model")
+    sample = {"instruction": instruction, "input": ""}
+    prompt = generate_prompt(sample)
+    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
+
+    output = generate(
+        model,
+        idx=encoded,
+        max_seq_length=block_size,
+        max_new_tokens=100,
+    )
+    output = tokenizer.decode(output)
+    return output # output.split("### Response:")[1].strip()
+
+
+@torch.no_grad()
+def validate(fabric: L.Fabric, model: torch.nn.Module, val_data: np.ndarray) -> torch.Tensor:
+    fabric.print("Validating ...")
+    model.eval()
+    losses = torch.zeros(eval_iters)
+    for k in range(eval_iters):
+        input_ids, targets = get_batch(fabric, val_data)
+        logits = model(input_ids)
+        loss = loss_fn(logits, targets)
+        losses[k] = loss.item()
+    out = losses.mean()
+
+    # produce an example:
+    instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
+    
+    output = generate_response(model, instruction)
+    fabric.print(instruction)
+    fabric.print(output)
+
+    model.train()
+    return out.item()
+
+
+def loss_fn(logits, targets):
+    # shift the targets such that output n predicts token n+1
+    logits = logits[..., :-1, :].contiguous()
+    targets = targets[..., 1:].contiguous()
+    loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+    return loss
+
+
+def get_batch(fabric: L.Fabric, data: list):
+    ix = torch.randint(len(data), (micro_batch_size,))
+
+    input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
+    labels = [data[i]["labels"].type(torch.int64) for i in ix]
+
+    max_len = max(len(s) for s in input_ids)
+
+    def pad_right(x, pad_id):
+        # pad right based on the longest sequence
+        n = max_len - len(x)
+        return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
+
+    x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
+    y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
+    x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
+    return x, y
+
+
+def load_datasets(data_dir):
+    train_data = torch.load(os.path.join(data_dir, "train.pt"))
+    val_data = torch.load(os.path.join(data_dir, "test.pt"))
+    return train_data, val_data
+
+
+if __name__ == "__main__":
+    # Uncomment this line if you see an error: "Expected is_sm80 to be true, but got false"
+    # torch.backends.cuda.enable_flash_sdp(False)
+    torch.set_float32_matmul_precision("high")
+
+    from jsonargparse.cli import CLI
+
+    CLI(main)

finetune_lora.py

@@ -0,0 +1,211 @@
+"""
+Instruction-tuning with LoRA on the Alpaca dataset.
+
+Note: If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
+`torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).
+"""
+import os
+import time
+
+import lightning as L
+import numpy as np
+import torch
+
+from generate import generate
+from lit_llama.lora import mark_only_lora_as_trainable, lora, lora_state_dict
+from lit_llama.model import LLaMA, LLaMAConfig
+from lit_llama.tokenizer import Tokenizer
+from scripts.prepare_alpaca import generate_prompt
+
+
+eval_interval = 100
+save_interval = 100
+eval_iters = 100
+log_interval = 1
+
+# Hyperparameters
+learning_rate = 3e-4
+batch_size = 128
+micro_batch_size = 4
+gradient_accumulation_steps = batch_size // micro_batch_size
+max_iters = 2 #50000 * 3 // micro_batch_size
+weight_decay = 0.0
+max_seq_length = 256  # see scripts/prepare_alpaca.py
+lora_r = 8
+lora_alpha = 16
+lora_dropout = 0.05
+warmup_steps = 100
+
+
+def main(
+    data_dir: str = "data/alpaca", 
+    pretrained_path: str = "checkpoints/lit-llama/7B/lit-llama.pth",
+    out_dir: str = "out/lora/alpaca",
+):
+
+    #fabric = L.Fabric(accelerator="cuda", precision="bf16-true")
+    fabric = L.Fabric(accelerator="cpu", devices=2, precision="bf16-true")
+    fabric.launch()
+    fabric.seed_everything(1337 + fabric.global_rank)
+
+    if fabric.global_rank == 0:
+        os.makedirs(out_dir, exist_ok=True)
+    print("loading dataset ", data_dir)
+    train_data, val_data = load_datasets(data_dir=data_dir)
+    print("train data: ", len(train_data))
+    print("val data: ", len(val_data))
+    config = LLaMAConfig.from_name("7B")
+    config.block_size = max_seq_length
+    print("loading pretrained model ", pretrained_path)
+    checkpoint = torch.load(pretrained_path)
+
+    with fabric.init_module(), lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
+        model = LLaMA(config)
+        # strict=False because missing keys due to LoRA weights not contained in checkpoint state
+        model.load_state_dict(checkpoint, strict=False)
+
+    mark_only_lora_as_trainable(model)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+    model, optimizer = fabric.setup(model, optimizer)
+    print("start training")
+    train(fabric, model, optimizer, train_data, val_data, out_dir)
+
+    # Save the final LoRA checkpoint at the end of training
+    print(f"Saving LoRA weights to {out_dir}")
+    checkpoint = lora_state_dict(model)
+    fabric.save(os.path.join(out_dir, "lit-llama-lora-finetuned.pth"), checkpoint)
+
+
+def train(
+    fabric: L.Fabric,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    train_data: np.ndarray,
+    val_data: np.ndarray,
+    out_dir: str,
+) -> None:
+    """The training loop.
+
+    Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
+    """
+    step_count = 0
+    print("max iters:", max_iters )
+
+    for iter_num in range(max_iters):
+        print("iter_num", iter_num)
+        if step_count <= warmup_steps:
+            # linear warmup
+            lr = learning_rate * step_count / warmup_steps
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = lr
+
+        t0 = time.time()
+
+        input_ids, targets = get_batch(fabric, train_data)
+        logits = model(input_ids)
+        print("calculate loss")
+        loss = loss_fn(logits, targets)
+        print("backward")
+        fabric.backward(loss)
+
+        if (iter_num + 1) % gradient_accumulation_steps == 0:
+            print("step optimizer")
+            optimizer.step()
+            optimizer.zero_grad()
+            step_count += 1
+            if step_count % eval_interval == 0:
+                val_loss = validate(fabric, model, val_data)
+                fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
+                fabric.barrier()
+
+            if step_count % save_interval == 0:
+                print(f"Saving LoRA weights to {out_dir}")
+                # We are only saving the LoRA weights
+                # TODO: Provide a function/script to merge the LoRA weights with pretrained weights
+                checkpoint = lora_state_dict(model)
+                fabric.save(os.path.join(out_dir, f"iter-{iter_num:06d}-ckpt.pth"), checkpoint)
+
+        dt = time.time() - t0
+        if iter_num % log_interval == 0:
+            fabric.print(f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms")
+
+
+def generate_response(model, instruction):
+    tokenizer = Tokenizer("checkpoints/lit-llama/tokenizer.model")
+    sample = {"instruction": instruction, "input": ""}
+    prompt = generate_prompt(sample)
+    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
+
+    output = generate(
+        model,
+        idx=encoded,
+        max_seq_length=max_seq_length,
+        max_new_tokens=100,
+    )
+    output = tokenizer.decode(output)
+    return output # output.split("### Response:")[1].strip()
+
+
+@torch.no_grad()
+def validate(fabric: L.Fabric, model: torch.nn.Module, val_data: np.ndarray) -> torch.Tensor:
+    fabric.print("Validating ...")
+    model.eval()
+    losses = torch.zeros(eval_iters)
+    for k in range(eval_iters):
+        input_ids, targets = get_batch(fabric, val_data)
+        logits = model(input_ids)
+        loss = loss_fn(logits, targets)
+        losses[k] = loss.item()
+    out = losses.mean()
+
+    # produce an example:
+    instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
+    output = generate_response(model, instruction)
+    fabric.print(instruction)
+    fabric.print(output)
+
+    model.train()
+    return out.item()
+
+def loss_fn(logits, targets):
+    # shift the targets such that output n predicts token n+1
+    logits = logits[..., :-1, :].contiguous()
+    targets = targets[..., 1:].contiguous()
+    loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+    return loss
+    
+
+def get_batch(fabric: L.Fabric, data: list):
+    ix = torch.randint(len(data), (micro_batch_size,))
+
+    input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
+    labels = [data[i]["labels"].type(torch.int64) for i in ix]
+
+    max_len = max(len(s) for s in input_ids)
+
+    def pad_right(x, pad_id):
+        # pad right based on the longest sequence
+        n = max_len - len(x)
+        return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
+
+    x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
+    y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
+    x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
+    return x, y
+
+
+def load_datasets(data_dir):
+    train_data = torch.load(os.path.join(data_dir, "train.pt"))
+    val_data = torch.load(os.path.join(data_dir, "test.pt"))
+    return train_data, val_data
+
+
+if __name__ == "__main__":
+    # Uncomment this line if you see an error: "Expected is_sm80 to be true, but got false"
+    # torch.backends.cuda.enable_flash_sdp(False)
+    torch.set_float32_matmul_precision("high")
+    
+    from jsonargparse.cli import CLI
+
+    CLI(main)

generate.py

@@ -0,0 +1,162 @@
+import sys
+import time
+import warnings
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+
+from lit_llama import LLaMA, Tokenizer
+from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+
+@torch.no_grad()
+def generate(
+    model: torch.nn.Module,
+    idx: torch.Tensor,
+    max_new_tokens: int,
+    max_seq_length: int,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    eos_id: Optional[int] = None,
+) -> torch.Tensor:
+    """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
+
+    The implementation of this function is modified from A. Karpathy's nanoGPT.
+
+    Args:
+        model: The model to use.
+        idx: Tensor of shape (T) with indices of the prompt sequence.
+        max_new_tokens: The number of new tokens to generate.
+        max_seq_length: The maximum sequence length allowed.
+        temperature: Scales the predicted logits by 1 / temperature
+        top_k: If specified, only sample among the tokens with the k highest probabilities
+        eos_id: If specified, stop generating any more token once the <eos> token is triggered
+    """
+    # create an empty tensor of the expected final shape and fill in the current tokens
+    T = idx.size(0)
+    T_new = T + max_new_tokens
+    empty = torch.empty(T_new, dtype=idx.dtype, device=idx.device)
+    empty[:T] = idx
+    idx = empty
+
+    # generate max_new_tokens tokens
+    for t in range(T, T_new):
+        # ignore the not-filled-yet tokens
+        idx_cond = idx[:t]
+        # if the sequence context is growing too long we must crop it at max_seq_length
+        idx_cond = idx_cond if T <= max_seq_length else idx_cond[-max_seq_length:]
+
+        # forward
+        logits = model(idx_cond.view(1, -1))
+        logits = logits[0, -1] / temperature
+
+        # optionally crop the logits to only the top k options
+        if top_k is not None:
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[[-1]]] = -float("Inf")
+
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        idx_next = torch.multinomial(probs, num_samples=1)
+
+        # concatenate the new generation
+        idx[t] = idx_next
+
+        # if <eos> token is triggered, return the output (stop generation)
+        if idx_next == eos_id:
+            return idx[:t + 1]  # include the EOS token
+
+    return idx
+
+
+def main(
+    prompt: str = "Hello, my name is",
+    *,
+    num_samples: int = 1,
+    max_new_tokens: int = 50,
+    top_k: int = 200,
+    temperature: float = 0.8,
+    checkpoint_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    quantize: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained LLaMA model and tokenizer.
+
+    Args:
+        prompt: The prompt string to use for generating the samples.
+        num_samples: The number of text samples to generate.
+        max_new_tokens: The number of generation steps to take.
+        top_k: The number of top most probable tokens to consider in the sampling process.
+        temperature: A value controlling the randomness of the sampling process. Higher values result in more random
+            samples.
+        checkpoint_path: The checkpoint path to load.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+    """
+    if not checkpoint_path:
+        checkpoint_path = Path(f"./checkpoints/lit-llama/7B/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    assert checkpoint_path.is_file(), checkpoint_path
+    assert tokenizer_path.is_file(), tokenizer_path
+
+    fabric = L.Fabric(devices="auto", accelerator="cuda")
+    #fabric = L.Fabric(accelerator="cpu")
+    fabric.launch()
+    dtype = torch.bfloat16 if fabric.device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
+
+    print("Loading model ...", file=sys.stderr)
+    t0 = time.time()
+    with lazy_load(checkpoint_path) as checkpoint:
+        name = llama_model_lookup(checkpoint)
+
+        with EmptyInitOnDevice(
+                device=fabric.device, dtype=dtype, quantization_mode=quantize
+        ):
+            model = LLaMA.from_name(name)
+
+        model.load_state_dict(checkpoint)
+    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+    encoded_prompt = tokenizer.encode(prompt, bos=True, eos=False, device=fabric.device)
+
+    L.seed_everything(1234)
+    for i in range(num_samples):
+        t0 = time.perf_counter()
+        y = generate(
+            model,
+            encoded_prompt,
+            max_new_tokens,
+            model.config.block_size,  # type: ignore[union-attr,arg-type]
+            temperature=temperature,
+            top_k=top_k,
+        )
+        t = time.perf_counter() - t0
+        print(tokenizer.decode(y))
+        print(f"Time for inference {i + 1}: {t:.02f} sec total, {max_new_tokens / t:.02f} tokens/sec", file=sys.stderr)
+    if fabric.device.type == "cuda":
+        print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+    torch.backends.cuda.max_split_size_mb = 16
+    torch.quantization.quantize_dynamic
+    torch.set_float32_matmul_precision("high")
+    warnings.filterwarnings(
+        # Triggered internally at ../aten/src/ATen/EmptyTensor.cpp:31
+        "ignore", 
+        message="ComplexHalf support is experimental and many operators don't support it yet"
+    )
+    warnings.filterwarnings(
+        # Triggered in bitsandbytes/autograd/_functions.py:298
+        "ignore",
+        message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization",
+    )
+    CLI(main)

generate_adapter.py

@@ -0,0 +1,117 @@
+import sys
+import time
+import warnings
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+
+from generate import generate
+from lit_llama import Tokenizer
+from lit_llama.adapter import LLaMA
+from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+from scripts.prepare_alpaca import generate_prompt
+
+
+def main(
+    prompt: str = "What food do lamas eat?",
+    input: str = "",
+    adapter_path: Optional[Path] = None,
+    pretrained_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    quantize: Optional[str] = None,
+    max_new_tokens: int = 100,
+    top_k: int = 200,
+    temperature: float = 0.8,
+) -> None:
+    """Generates a response based on a given instruction and an optional input.
+    This script will only work with checkpoints from the instruction-tuned LLaMA-Adapter model.
+    See `finetune_adapter.py`.
+
+    Args:
+        prompt: The prompt/instruction (Alpaca style).
+        adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
+            `finetune_adapter.py`.
+        input: Optional input (Alpaca style).
+        pretrained_path: The path to the checkpoint with pretrained LLaMA weights.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+        max_new_tokens: The number of generation steps to take.
+        top_k: The number of top most probable tokens to consider in the sampling process.
+        temperature: A value controlling the randomness of the sampling process. Higher values result in more random
+            samples.
+    """
+    if not adapter_path:
+        adapter_path = Path("out/adapter/alpaca/lit-llama-adapter-finetuned.pth")
+    if not pretrained_path:
+        pretrained_path = Path(f"./checkpoints/lit-llama/7B/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    
+    assert adapter_path.is_file()
+    assert pretrained_path.is_file()
+    assert tokenizer_path.is_file()
+
+    fabric = L.Fabric(devices=1)
+    dtype = torch.bfloat16 if fabric.device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
+
+    print("Loading model ...", file=sys.stderr)
+    t0 = time.time()
+    with (lazy_load(pretrained_path) as pretrained_checkpoint,
+          lazy_load(adapter_path) as adapter_checkpoint):
+        name = llama_model_lookup(pretrained_checkpoint)
+
+        with EmptyInitOnDevice(
+                device=fabric.device, dtype=dtype, quantization_mode=quantize
+        ):
+            model = LLaMA.from_name(name)
+
+        # 1. Load the pretrained weights
+        model.load_state_dict(pretrained_checkpoint, strict=False)
+        # 2. Load the fine-tuned adapter weights
+        model.load_state_dict(adapter_checkpoint, strict=False)
+
+    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+    sample = {"instruction": prompt, "input": input}
+    prompt = generate_prompt(sample)
+    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
+
+    t0 = time.perf_counter()
+    output = generate(
+        model,
+        idx=encoded,
+        max_seq_length=max_new_tokens,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_k=top_k,
+        eos_id=tokenizer.eos_id
+    )
+    t = time.perf_counter() - t0
+
+    output = tokenizer.decode(output)
+    output = output.split("### Response:")[1].strip()
+    print(output)
+
+    print(f"\n\nTime for inference: {t:.02f} sec total, {max_new_tokens / t:.02f} tokens/sec", file=sys.stderr)
+    if fabric.device.type == "cuda":
+        print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    warnings.filterwarnings(
+        # Triggered internally at ../aten/src/ATen/EmptyTensor.cpp:31
+        "ignore", 
+        message="ComplexHalf support is experimental and many operators don't support it yet"
+    )
+    CLI(main)

generate_full.py

@@ -0,0 +1,160 @@
+import sys
+import time
+import warnings
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+
+from lit_llama import LLaMA, Tokenizer
+from lit_llama.utils import EmptyInitOnDevice
+
+
+@torch.no_grad()
+def generate(
+    model: torch.nn.Module,
+    idx: torch.Tensor,
+    max_new_tokens: int,
+    max_seq_length: int,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    eos_id: Optional[int] = None,
+) -> torch.Tensor:
+    """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
+
+    The implementation of this function is modified from A. Karpathy's nanoGPT.
+
+    Args:
+        model: The model to use.
+        idx: Tensor of shape (T) with indices of the prompt sequence.
+        max_new_tokens: The number of new tokens to generate.
+        max_seq_length: The maximum sequence length allowed.
+        temperature: Scales the predicted logits by 1 / temperature
+        top_k: If specified, only sample among the tokens with the k highest probabilities
+        eos_id: If specified, stop generating any more token once the <eos> token is triggered
+    """
+    # create an empty tensor of the expected final shape and fill in the current tokens
+    T = idx.size(0)
+    T_new = T + max_new_tokens
+    empty = torch.empty(T_new, dtype=idx.dtype, device=idx.device)
+    empty[:T] = idx
+    idx = empty
+
+    # generate max_new_tokens tokens
+    for t in range(T, T_new):
+        # ignore the not-filled-yet tokens
+        idx_cond = idx[:t]
+        # if the sequence context is growing too long we must crop it at max_seq_length
+        idx_cond = idx_cond if T <= max_seq_length else idx_cond[-max_seq_length:]
+
+        # forward
+        logits = model(idx_cond.view(1, -1))
+        logits = logits[0, -1] / temperature
+
+        # optionally crop the logits to only the top k options
+        if top_k is not None:
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[[-1]]] = -float("Inf")
+
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        idx_next = torch.multinomial(probs, num_samples=1)
+
+        # concatenate the new generation
+        idx[t] = idx_next
+
+        # if <eos> token is triggered, return the output (stop generation)
+        if idx_next == eos_id:
+            return idx[:t + 1]  # include the EOS token
+
+    return idx
+
+
+def main(
+    prompt: str = "Hello, my name is",
+    *,
+    num_samples: int = 1,
+    max_new_tokens: int = 50,
+    top_k: int = 200,
+    temperature: float = 0.8,
+    checkpoint_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    model_size: str = "7B",
+    quantize: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained LLaMA model and tokenizer.
+
+    Args:
+        prompt: The prompt string to use for generating the samples.
+        num_samples: The number of text samples to generate.
+        max_new_tokens: The number of generation steps to take.
+        top_k: The number of top most probable tokens to consider in the sampling process.
+        temperature: A value controlling the randomness of the sampling process. Higher values result in more random
+            samples.
+        checkpoint_path: The checkpoint path to load.
+        tokenizer_path: The tokenizer path to load.
+        model_size: The model size to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+    """
+    if not checkpoint_path:
+        checkpoint_path = Path(f"./checkpoints/lit-llama/{model_size}/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    assert checkpoint_path.is_file(), checkpoint_path
+    assert tokenizer_path.is_file(), tokenizer_path
+
+    fabric = L.Fabric(devices=1)
+    dtype = torch.bfloat16 if fabric.device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
+
+    print("Loading model ...", file=sys.stderr)
+    t0 = time.time()
+    with EmptyInitOnDevice(
+        device=fabric.device, dtype=dtype, quantization_mode=quantize
+    ):
+        model = LLaMA.from_name(model_size)
+
+    checkpoint = torch.load(checkpoint_path)
+    model.load_state_dict(checkpoint)
+    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+    encoded_prompt = tokenizer.encode(prompt, bos=True, eos=False, device=fabric.device)
+
+    L.seed_everything(1234)
+    for i in range(num_samples):
+        t0 = time.perf_counter()
+        y = generate(
+            model,
+            encoded_prompt,
+            max_new_tokens,
+            model.config.block_size,  # type: ignore[union-attr,arg-type]
+            temperature=temperature,
+            top_k=top_k,
+        )
+        t = time.perf_counter() - t0
+        print(tokenizer.decode(y))
+        print(f"Time for inference {i + 1}: {t:.02f} sec total, {max_new_tokens / t:.02f} tokens/sec", file=sys.stderr)
+    if fabric.device.type == "cuda":
+        print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    warnings.filterwarnings(
+        # Triggered internally at ../aten/src/ATen/EmptyTensor.cpp:31
+        "ignore", 
+        message="ComplexHalf support is experimental and many operators don't support it yet"
+    )
+    warnings.filterwarnings(
+        # Triggered in bitsandbytes/autograd/_functions.py:298
+        "ignore", 
+        message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization",
+    )
+    CLI(main)

generate_lora.py

@@ -0,0 +1,131 @@
+import sys
+import time
+import warnings
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+
+from generate import generate
+from lit_llama import Tokenizer, LLaMA
+from lit_llama.lora import lora
+from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
+from scripts.prepare_alpaca import generate_prompt
+
+lora_r = 8
+lora_alpha = 16
+lora_dropout = 0.05
+
+
+def main(
+    prompt: str = "What food do lamas eat?",
+    input: str = "",
+    lora_path: Optional[Path] = None,
+    pretrained_path: Optional[Path] = None,
+    tokenizer_path: Optional[Path] = None,
+    quantize: Optional[str] = None,
+    dtype: str = "float32",
+    max_new_tokens: int = 100,
+    top_k: int = 200,
+    temperature: float = 0.8,
+) -> None:
+    """Generates a response based on a given instruction and an optional input.
+    This script will only work with checkpoints from the instruction-tuned LoRA model.
+    See `finetune_lora.py`.
+
+    Args:
+        prompt: The prompt/instruction (Alpaca style).
+        lora_path: Path to the checkpoint with trained LoRA weights, which are the output of
+            `finetune_lora.py`.
+        input: Optional input (Alpaca style).
+        pretrained_path: The path to the checkpoint with pretrained LLaMA weights.
+        tokenizer_path: The tokenizer path to load.
+        quantize: Whether to quantize the model and using which method:
+            ``"llm.int8"``: LLM.int8() mode,
+            ``"gptq.int4"``: GPTQ 4-bit mode.
+        dtype: The dtype to use during generation.
+        max_new_tokens: The number of generation steps to take.
+        top_k: The number of top most probable tokens to consider in the sampling process.
+        temperature: A value controlling the randomness of the sampling process. Higher values result in more random
+            samples.
+    """
+    if not lora_path:
+        lora_path = Path("out/lora/alpaca/lit-llama-lora-finetuned.pth")
+    if not pretrained_path:
+        pretrained_path = Path(f"./checkpoints/lit-llama/7B/lit-llama.pth")
+    if not tokenizer_path:
+        tokenizer_path = Path("./checkpoints/lit-llama/tokenizer.model")
+    
+    assert lora_path.is_file()
+    assert pretrained_path.is_file()
+    assert tokenizer_path.is_file()
+
+    if quantize is not None:
+        raise NotImplementedError("Quantization in LoRA is not supported yet")
+
+    fabric = L.Fabric(devices=1)
+
+    dt = getattr(torch, dtype, None)
+    if not isinstance(dt, torch.dtype):
+        raise ValueError(f"{dtype} is not a valid dtype.")
+    dtype = dt
+
+    print("Loading model ...", file=sys.stderr)
+    t0 = time.time()
+
+    with (lazy_load(pretrained_path) as pretrained_checkpoint,
+          lazy_load(lora_path) as adapter_checkpoint):
+        name = llama_model_lookup(pretrained_checkpoint)
+
+        with EmptyInitOnDevice(
+                device=fabric.device, dtype=dtype, quantization_mode=quantize
+        ), lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
+            model = LLaMA.from_name(name)
+
+            # 1. Load the pretrained weights
+            model.load_state_dict(pretrained_checkpoint, strict=False)
+            # 2. Load the fine-tuned adapter weights
+            model.load_state_dict(adapter_checkpoint, strict=False)
+
+    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(tokenizer_path)
+    sample = {"instruction": prompt, "input": input}
+    prompt = generate_prompt(sample)
+    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
+
+    t0 = time.perf_counter()
+    output = generate(
+        model,
+        idx=encoded,
+        max_seq_length=max_new_tokens,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_k=top_k,
+        eos_id=tokenizer.eos_id
+    )
+    t = time.perf_counter() - t0
+
+    output = tokenizer.decode(output)
+    output = output.split("### Response:")[1].strip()
+    print(output)
+
+    print(f"\n\nTime for inference: {t:.02f} sec total, {max_new_tokens / t:.02f} tokens/sec", file=sys.stderr)
+    if fabric.device.type == "cuda":
+        print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    warnings.filterwarnings(
+        # Triggered internally at ../aten/src/ATen/EmptyTensor.cpp:31
+        "ignore", 
+        message="ComplexHalf support is experimental and many operators don't support it yet"
+    )
+    CLI(main)

howto/customize_paths.md

@@ -0,0 +1,33 @@
+## Customize paths
+
+The project is setup to use specific paths to read the original weights and save checkpoints etc.
+
+For all scripts, you can run
+
+```shell
+python script.py -h
+```
+
+to get a list of available options. For instance, here's how you would modify the checkpoint dir:
+
+```shell
+python scripts/convert_checkpoint.py --checkpoint_dir "data/checkpoints/foo"
+```
+
+Note that this change will need to be passed along to subsequent steps, for example:
+
+```shell
+python scripts/generate.py \
+  --checkpoint_path "data/checkpoints/foo/7B/lit-llama.pth" \
+  --tokenizer_path "data/checkpoints/foo/tokenizer.model"
+```
+
+and
+
+```shell
+python scripts/quantize.py \
+  --checkpoint_path "data/checkpoints/foo/7B/lit-llama.pth" \
+  --tokenizer_path "data/checkpoints/foo/tokenizer.model"
+```
+
+To avoid this, you can use symbolic links to create shortcuts and avoid passing different paths.

howto/download_weights.md

@@ -0,0 +1,131 @@
+## Downloading pretrained weights
+
+Except for when you are training from scratch, you will need the pretrained weights from Meta.
+
+### Original Meta weights
+
+Download the model weights following the instructions on the official [LLaMA repository](https://github.com/facebookresearch/llama).
+
+Once downloaded, you should have a folder like this:
+
+```text
+checkpoints/llama
+├── 7B
+│   ├── ...
+│   └── consolidated.00.pth
+├── 13B
+│   ...
+└── tokenizer.model
+```
+
+Convert the weights to the Lit-LLaMA format:
+
+```bash
+python scripts/convert_checkpoint.py --model_size 7B
+```
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+### OpenLLaMA
+
+OpenLM Research has released **Apache 2.0 licensed** weights obtained by training LLaMA on the 1.2 trillion token open-source [RedPajama](https://github.com/togethercomputer/RedPajama-Data) dataset.
+
+Weights were released in preview on intermediate number of tokens (200B, 300B at the time of writing). In order to get them do:
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com): git lfs install
+git clone https://huggingface.co/openlm-research/open_llama_7b_preview_300bt checkpoints/open-llama/7B
+```
+
+Or if you don't have `git-lfs` installed:
+
+```bash
+python scripts/download.py --repo_id openlm-research/open_llama_7b_preview_300bt --local_dir checkpoints/open-llama/7B
+```
+
+Once downloaded, you should have a folder like this:
+
+```text
+checkpoints/open-llama/
+└── 7B
+    └── open_llama_7b_preview_300bt_transformers_weights
+        ├── ...
+        ├── pytorch_model-00001-of-00002.bin
+        ├── pytorch_model-00002-of-00002.bin
+        ├── pytorch_model.bin.index.json
+        └── tokenizer.model
+```
+
+Convert the weights to the Lit-LLaMA format:
+
+```bash
+python scripts/convert_hf_checkpoint.py --checkpoint_dir checkpoints/open-llama/7B/open_llama_7b_preview_300bt_transformers_weights --model_size 7B
+```
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+Once converted, you should have a folder like this:
+
+```text
+checkpoints/lit-llama/
+├── 7B
+│   └── lit-llama.pth
+└── tokenizer.model
+```
+
+You are all set. Now you can continue with inference or finetuning.
+
+Try running [`generate.py` to test the imported weights](inference.md).
+
+
+### Alternative sources
+
+You might find LLaMA weights hosted online in the HuggingFace hub. Beware that this infringes the original weight's license.
+You could try downloading them by running the following command with a specific repo id:
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com): git lfs install
+git clone REPO_ID checkpoints/hf-llama/7B
+```
+
+Or if you don't have `git-lfs` installed:
+
+```bash
+python scripts/download.py --repo_id REPO_ID --local_dir checkpoints/hf-llama/7B
+```
+
+Once downloaded, you should have a folder like this:
+
+```text
+checkpoints/hf-llama/
+└── 7B
+    ├── ...
+    ├── pytorch_model-00001-of-00002.bin
+    ├── pytorch_model-00002-of-00002.bin
+    ├── pytorch_model.bin.index.json
+    └── tokenizer.model
+```
+
+Convert the weights to the Lit-LLaMA format:
+
+```bash
+python scripts/convert_hf_checkpoint.py --model_size 7B
+```
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+Once converted, you should have a folder like this:
+
+```text
+checkpoints/lit-llama/
+├── 7B
+│   └── lit-llama.pth
+└── tokenizer.model
+```
+
+You are all set. Now you can continue with inference or finetuning.
+
+Try running [`generate.py` to test the imported weights](inference.md).

howto/finetune_adapter.md

@@ -0,0 +1,102 @@
+# Finetuning with Adapter
+
+[LLaMA-Adapter](https://arxiv.org/abs/2303.16199) is a form of prefix-tuning that prepends a learnable adaption-prompt to the inputs of the attention blocks in LLaMA. In total, there are only 1.2M parameters to update during finetuning, which significantly reduces the memory footprint and speeds up training.
+
+We are able to demonstrate instruction-finetuning Lit-LLaMA 7B on the [Alpaca](https://github.com/tatsu-lab/stanford_alpaca) dataset on a **single GTX 3090 (24GB) GPU**. If using 8 GPUs, finetuning can be completed in under 1 hour.
+
+If you are new to LLaMA-Adapter and are interest to learn more about how it works before proceeding with the finetuning guide below, you might find our article [Understanding Parameter-Efficient Finetuning of Large Language Models: From Prefix Tuning to LLaMA-Adapters](https://lightning.ai/pages/community/article/understanding-llama-adapters/) helpful.
+
+## Preparation
+
+The steps here only need to be done once:
+
+1. Follow the instructions in the [README](README.md) to install the dependencies.
+2. Download and convert the weights and save them in the `./checkpoints` folder as described [here](download_weights.md).
+3. If you want to utilize more than one GPU, you should `pip install deepspeed`.
+4. Download the data and generate the Alpaca instruction tuning dataset:
+
+   ```bash
+   python scripts/prepare_alpaca.py
+   ```
+
+   or [prepare your own dataset](#tune-on-your-own-dataset).
+
+## Running the finetuning
+
+```bash
+python finetune_adapter.py
+```
+
+The finetuning requires at least one GPU with ~24 GB memory (GTX 3090).
+You can speed up training by setting the `devices` variable in the script to utilize more GPUs if available.
+Depending on the available GPU memory, you can also tune the `micro_batch_size` parameter to utilize the GPU efficiently.
+
+For example, the following settings will let you finetune the model in under 1 hour using DeepSpeed Zero-2:
+```python
+devices = 8
+micro_batch_size = 8
+```
+
+This script will save checkpoints periodically to the folder `out/`.
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+## Test the model
+
+You can test the finetuned model with your own instructions by running:
+
+```bash
+python generate_adapter.py \
+    --prompt "Recommend a movie to watch on the weekend." \
+    --quantize llm.int8
+```
+Output:
+```
+A good movie to watch on the weekend would be The Lion King, since it's a classic family film that everyone can enjoy...
+```
+If your GPU supports `bfloat16`, the script will automatically use it. Together with `--quantize llm.int8`, this brings the memory consumption down to ~8 GB.
+
+## Tune on your dataset
+
+With only a few modifications, you can prepare and train on your own instruction dataset.
+
+1. Create a json file in which each row holds one instruction-response pair. 
+   A row has an entry for 'instruction', 'input', and 'output', where 'input' is optional an can be 
+   the empty string if the instruction doesn't require a context. Below is an example json file:
+
+    ```
+    [
+        {
+            "instruction": "Arrange the given numbers in ascending order.",
+            "input": "2, 4, 0, 8, 3",
+            "output": "0, 2, 3, 4, 8"
+        },
+        ...
+    ]
+    ```
+
+2. Make a copy of `scripts/prepare_alpaca.py` and name it what you want:
+
+    ```bash
+    cp scripts/prepare_alpaca.py scripts/prepare_mydata.py
+    ```
+
+3. Modify `scripts/prepare_mydata.py` to read the json data file.
+4. Run the script to generate the preprocessed, tokenized train-val split:
+
+    ```bash
+    python scripts/prepare_mydata.py --destination_path data/mydata/
+    ```
+
+5. Run `finetune_adapter.py` by passing in the location of your data (and optionally other parameters):
+   
+    ```bash
+    python finetune_adapter.py --data_dir data/mydata/ --out_dir out/myexperiment
+    ```
+
+
+## Troubleshooting
+
+If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
+`torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).

howto/finetune_full.md

@@ -0,0 +1,104 @@
+# Full Finetuning
+
+Full finetuning updates all layers in the pretrained LLaMA model. This *regular* finetuning procedure is typically considered as the baseline for parameter-efficient alternatives such as Low-Rank Adaptation (LoRA) or LLaMA-Adapter.
+
+The current  [finetune_full.py](../scripts/finetune_full.py) we provide uses 4 A100 GPUs with a fully-sharded data parallel strategy to finetune Lit-LLaMA 7B on [Alpaca](https://github.com/tatsu-lab/stanford_alpaca) dataset. The A100 GPUs have 40 GB each, but it may require less memory to finetune this model.
+
+
+
+## Preparation
+
+The steps here only need to be done once:
+
+1. Follow the instructions in the [README](README.md) to install the dependencies.
+
+2. Download and convert the weights and save them in the `./checkpoints` folder as described [here](download_weights.md).
+
+4. Download the data and generate the Alpaca instruction tuning dataset:
+
+   ```bash
+   python scripts/prepare_alpaca.py
+   ```
+
+   or [prepare your own dataset](#tune-on-your-own-dataset).
+
+## Running the finetuning
+
+```bash
+python finetune_full.py
+```
+
+
+You can speed up training by setting the `devices` variable in the script to utilize more GPUs if available or increase the `batch_size`.
+Depending on the available GPU memory, you can also tune the `micro_batch_size` parameter to utilize the GPU efficiently.
+
+For example, the following settings will let you finetune the model in 32 hours using a fully-sharded data parallel strategy:
+```python
+devices = 4
+batch_size = 128 // devices
+micro_batch_size = 4
+```
+
+This script will save checkpoints periodically to the folder `out/`.
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+## Test the model
+
+You can test the finetuned model with your own instructions by running:
+
+```bash
+python generate_full.py \
+    --prompt "Recommend a movie to watch on the weekend." \
+    --quantize llm.int8
+```
+Output:
+```
+A good movie to watch on the weekend would be The Lion King, since it's a classic family film that everyone can enjoy...
+```
+If your GPU supports `bfloat16`, the script will automatically use it. Together with `--quantize llm.int8`, this brings the memory consumption down to ~8 GB.
+
+## Tune on your dataset
+
+With only a few modifications, you can prepare and train on your own instruction dataset.
+
+1. Create a json file in which each row holds one instruction-response pair. 
+   A row has an entry for 'instruction', 'input', and 'output', where 'input' is optional an can be 
+   the empty string if the instruction doesn't require a context. Below is an example json file:
+
+    ```
+    [
+        {
+            "instruction": "Arrange the given numbers in ascending order.",
+            "input": "2, 4, 0, 8, 3",
+            "output": "0, 2, 3, 4, 8"
+        },
+        ...
+    ]
+    ```
+
+2. Make a copy of `scripts/prepare_alpaca.py` and name it what you want:
+
+    ```bash
+    cp scripts/prepare_alpaca.py scripts/prepare_mydata.py
+    ```
+
+3. Modify `scripts/prepare_mydata.py` to read the json data file.
+4. Run the script to generate the preprocessed, tokenized train-val split:
+
+    ```bash
+    python scripts/prepare_mydata.py --destination_path data/mydata/
+    ```
+
+5. Run `finetune_full.py` by passing in the location of your data (and optionally other parameters):
+   
+    ```bash
+    python finetune_full.py --data_dir data/mydata/ --out_dir out/myexperiment
+    ```
+
+
+## Troubleshooting
+
+If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
+`torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).

howto/finetune_lora.md

@@ -0,0 +1,88 @@
+# Finetuning with LoRA
+
+[Low-rank adaption (LoRA)](https://arxiv.org/abs/2106.09685) is a technique to approximate the update to the linear layers in a LLM with a low-rank matrix factorization. This significantly reduces the number of trainable parameters and speeds up training with little impact on the final performance of the model.
+We demonstrate this method by instruction-finetuning LLaMA 7B on the [Alpaca](https://github.com/tatsu-lab/stanford_alpaca) dataset on a **single GTX 3090 (24GB) GPU**.
+
+## Preparation
+
+The steps here only need to be done once:
+
+1. Follow the instructions in the [README](README.md) to install the dependencies.
+2. Download and convert the weights and save them in the `./checkpoints` folder as described [here](download_weights.md).
+3. Download the data and generate the instruction tuning dataset:
+
+   ```bash
+   python scripts/prepare_alpaca.py
+   ```
+
+## Running the finetuning
+
+```bash
+python finetune_lora.py
+```
+
+The finetuning requires at least one GPU with ~24 GB memory (GTX 3090).
+
+This script will save checkpoints periodically to the folder `out/`.
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+
+## Test the model
+
+You can test the finetuned model with your own instructions by running:
+
+```bash
+python generate_lora.py --prompt "Recommend a movie to watch on the weekend."
+```
+Output:
+```
+I would recommend the movie The Martian (2015). It is a sci-fi movie starring Matt Damon that follows the story of...
+```
+
+If your GPU supports `bfloat16`, you can additionally pass `--dtype bfloat16` to bring the memory consumption down to ~14 GB.
+
+## Tune on your dataset
+
+With only a few modifications, you can prepare and train on your own instruction dataset.
+
+1. Create a json file in which each row holds one instruction-response pair. 
+   A row has an entry for 'instruction', 'input', and 'output', where 'input' is optional an can be 
+   the empty string if the instruction doesn't require a context. Below is an example json file:
+
+    ```
+    [
+        {
+            "instruction": "Arrange the given numbers in ascending order.",
+            "input": "2, 4, 0, 8, 3",
+            "output": "0, 2, 3, 4, 8"
+        },
+        ...
+    ]
+    ```
+
+2. Make a copy of `scripts/prepare_alpaca.py` and name it what you want:
+
+    ```bash
+    cp scripts/prepare_alpaca.py scripts/prepare_mydata.py
+    ```
+
+3. Modify `scripts/prepare_mydata.py` to read the json data file.
+4. Run the script to generate the preprocessed, tokenized train-val split:
+
+    ```bash
+    python scripts/prepare_mydata.py --destination_path data/mydata/
+    ```
+
+5. Run `finetune_lora.py` by passing in the location of your data (and optionally other parameters):
+    
+    ```bash
+    python finetune_lora.py --data_dir data/mydata/ --out_dir out/myexperiment
+    ```
+
+
+## Troubleshooting
+
+If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
+`torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).

howto/inference.md

@@ -0,0 +1,37 @@
+# Inference
+
+We demonstrate how to run inference (next token prediction) with the LLaMA base model in the [`generate.py`](generate.py) script:
+
+```bash
+python generate.py --prompt "Hello, my name is"
+```
+Output:
+```
+Hello my name is TJ. I have a passion for the outdoors, love hiking and exploring. I also enjoy traveling and learning new things. I especially enjoy long walks, good conversation and a friendly smile.
+```
+
+The script assumes you have downloaded and converted the weights and saved them in the `./checkpoints` folder as described [here](download_weights.md).
+
+> **Note**
+> All scripts support argument [customization](customize_paths.md)
+
+With the default settings, this will run the 7B model and require ~26 GB of GPU memory (A100 GPU).
+
+## Run Lit-LLaMA on consumer devices
+
+On GPUs with `bfloat16` support, the `generate.py` script will automatically convert the weights and consume about ~14 GB.
+For GPUs with less memory, or ones that don't support `bfloat16`, enable quantization (`--quantize llm.int8`):
+
+```bash
+python generate.py --quantize llm.int8 --prompt "Hello, my name is"
+```
+This will consume about ~10 GB of GPU memory or ~8 GB if also using `bfloat16`.
+See `python generate.py --help` for more options.
+
+You can also use GPTQ-style int4 quantization, but this needs conversions of the weights first:
+
+```bash
+python quantize.py --checkpoint_path lit-llama.pth --tokenizer_path tokenizer.model --output_path llama-7b-gptq.4bit.pt --dtype bfloat16  --quantize gptq.int4
+```
+
+With the generated quantized checkpoint generation works as usual with `--quantize gptq.int4`, bringing GPU usage to about ~5GB. As only the weights of the Linear layers are quantized, it is useful to use `--dtype bfloat16` even with the quantization enabled.

howto/tpus.md

@@ -0,0 +1,51 @@
+# TPU support
+
+Lit-LLaMA used `lightning.Fabric` under the hood, which itself supports TPUs (via [PyTorch XLA](https://github.com/pytorch/xla)).
+
+The following commands will allow you to set up a `Google Cloud` instance with a [TPU v4](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm) VM:
+
+```shell
+gcloud compute tpus tpu-vm create lit-llama --version=tpu-vm-v4-pt-2.0 --accelerator-type=v4-8 --zone=us-central2-b
+gcloud compute tpus tpu-vm ssh lit-llama --zone=us-central2-b
+```
+
+Now that you are in the machine, let's clone the repository and install the dependencies
+
+```shell
+git clone https://github.com/Lightning-AI/lit-llama
+cd lit-llama
+pip install -r requirements.txt
+```
+
+By default, computations will run using the new (and experimental) PjRT runtime. Still, it's recommended that you set the following environment variables
+
+```shell
+export PJRT_DEVICE=TPU
+export ALLOW_MULTIPLE_LIBTPU_LOAD=1
+```
+
+> **Note**
+> You can find an extensive guide on how to get set-up and all the available options [here](https://cloud.google.com/tpu/docs/v4-users-guide).
+
+Since you created a new machine, you'll probably need to download the weights. You could scp them into the machine with `gcloud compute tpus tpu-vm scp` or you can follow the steps described in our [downloading guide](download_weights.md).
+
+## Inference
+
+Generation works out-of-the-box with TPUs:
+
+```shell
+python3 generate.py --prompt "Hello, my name is" --num_samples 2
+```
+
+This command will take a long time as XLA needs to compile the graph (~13 min) before running the model.
+In fact, you'll notice that the second sample takes considerable less time (~12 sec).
+
+## Finetuning
+
+Coming soon.
+
+> **Warning**
+> When you are done, remember to delete your instance 
+> ```shell
+> gcloud compute tpus tpu-vm delete lit-llama --zone=us-central2-b
+> ```

howto/train_redpajama.md

@@ -0,0 +1,133 @@
+# Pre-train LLaMA on RedPajama
+
+This howto will walk you through setting up the RedPajama dataset and launching the pre-training script.
+
+## What's RedPajama
+
+[RedPajama](https://github.com/togethercomputer/RedPajama-Data) is an open-source reproduction of the original LLaMA training dataset.
+
+It contains a total of 1.2 trillion tokens, divided into
+
+```text
+Commoncrawl   878B
+C4            175B
+GitHub         59B
+Books          26B
+ArXiv          28B
+Wikipedia      24B
+StackExchange  20B
+```
+
+The [RedPajama repo](https://github.com/togethercomputer/RedPajama-Data) contains the source code for collecting and preparing
+the dataset, and it is Apache 2.0 licensed.
+
+The data itself is licensed according to the original licenses with which its invidivdual parts were released.
+The GitHub datasets are limited to MIT, BSD, or Apache 2.0 repositories.
+
+Along with the full [RedPajama-1T dataset](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T),
+the [RedPajama-1T-Sample](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample) 1B sample dataset
+is also available for development.
+
+You can download the data using git lfs:
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com): git lfs install
+git clone https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T data/RedPajama-Data-1T
+```
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com): git lfs install
+git clone https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample data/RedPajama-Data-1T-Sample
+```
+
+## Prepare RedPajama for training
+
+The dataset consists of 2084 `jsonl` files (the sample dataset contains 11). In order to start pre-training lit-llama
+on it, you need to read, tokenize, and write the data in binary chunks. This will leverage the `PackedDataset`
+streaming dataset that comes with lit-llama.
+
+Do to so, run
+
+```bash
+python scripts/prepare_redpajama.py --source_path data/RedPajama-Data-1T --tokenizer_path checkpoints/lit-llama/tokenizer.model --destination_path data/lit-redpajama
+```
+
+or
+
+```bash
+python scripts/prepare_redpajama.py --source_path data/RedPajama-Data-1T-Sample --tokenizer_path checkpoints/lit-llama/tokenizer.model --destination_path data/lit-redpajama-sample --sample True
+```
+
+for the sample dataset.
+
+In the above we are assuming that you will be using the same tokenizer as used in LLaMA, but any trained [SentencePiece](https://github.com/google/sentencepiece) tokenizer with a 32000 vocabulary size will do here.
+
+The script will take a while to run, so time for :tea:
+
+## Pre-training
+
+Running the pre-training script requires at least 4 GPUs with 40GB+ each (A100).
+
+```bash
+python train_redpajama.py --devices 4 --train_data_dir data/lit-redpajama
+```
+
+For running on the sample dataset:
+
+```bash
+python train_redpajama.py --devices 4 --train_data_dir data/lit-redpajama-sample
+```
+
+The script will save checkpoints periodically to the folder `out/`.
+
+The `train_redpajama.py` script will pre-train the LLaMA 7B model with FSDP in
+`bfloat16` precision and gradient accumulation.
+
+You can easily change the size of the model by passing a different string to
+
+```python
+config = LLaMAConfig.from_name("7B")
+```
+
+in the `main` function.
+
+Keep in mind that the original LLaMA training for the 7B model required 83k A100 80GB
+hours, so you'll need access to a cluster.
+
+Once you're in a cluster, you can follow [these instructions](https://lightning.ai/docs/fabric/stable/guide/multi_node/other.html)
+to launch the script across machines:
+
+- [SLURM cluster](https://lightning.ai/docs/fabric/stable/guide/multi_node/slurm.html)
+- [Barebones cluster](https://lightning.ai/docs/fabric/stable/guide/multi_node/barebones.html)
+- [MPI](https://lightning.ai/docs/fabric/stable/guide/multi_node/other.html)
+
+The script contains several configurations and hyperparameters you can tweak:
+
+```python
+out_dir = "out/training"
+save_interval = 1000
+eval_interval = 1000
+eval_iters = 100
+log_interval = 1
+
+# Hyperparameters
+learning_rate = 6e-4
+batch_size = 125
+micro_batch_size = 5
+max_iters = 600000  # num_epochs * epoch_size // devices
+weight_decay = 1e-1
+beta1 = 0.9
+beta2 = 0.95
+grad_clip = 1.0
+decay_lr = True
+warmup_iters = 2000
+lr_decay_iters = max_iters
+min_lr = 6e-5
+```
+
+In particular, `micro_batch_size` should be adjusted so the process will use the available
+GPU memory.
+
+Last, logging is kept minimal in the script. In order to use a particular logger
+please refer to <https://lightning.ai/docs/fabric/stable/api/loggers.html> or
+call a logging client library like `wandb` directly.

lit_llama/__init__.py

@@ -0,0 +1,2 @@
+from lit_llama.model import LLaMAConfig, LLaMA, RMSNorm, build_rope_cache, apply_rope
+from lit_llama.tokenizer import Tokenizer