Delete lit-llama

lit-llama/.github/CODEOWNERS

@@ -1 +0,0 @@
-* @awaelchli @carmocca @lantiga

lit-llama/.github/azure-gpu-tests.yml

@@ -1,57 +0,0 @@
-# Python package
-# Create and test a Python package on multiple Python versions.
-# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
-# https://docs.microsoft.com/azure/devops/pipelines/languages/python
-
-trigger:
-  tags:
-    include:
-      - '*'
-  branches:
-    include:
-      - "main"
-      - "refs/tags/*"
-
-pr:
-  branches:
-    include:
-      - "main"
-
-jobs:
-  - job: testing
-    # how long to run the job before automatically cancelling
-    timeoutInMinutes: "20"
-    # how much time to give 'run always even if cancelled tasks' before stopping them
-    cancelTimeoutInMinutes: "2"
-    pool: "lit-rtx-3090"
-    variables:
-      DEVICES: $( python -c 'print("$(Agent.Name)".split("_")[-1])' )
-    container:
-      image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.0-cuda11.7.1"
-      options: "--gpus=all --shm-size=8gb"
-    workspace:
-      clean: all
-    steps:
-
-    - bash: |
-        echo "##vso[task.setvariable variable=CUDA_VISIBLE_DEVICES]$(DEVICES)"
-        cuda_ver=$(python -c "import torch ; print(''.join(map(str, torch.version.cuda.split('.')[:2])))")
-        echo "##vso[task.setvariable variable=CUDA_VERSION_MM]$cuda_ver"
-      displayName: 'set env. vars'
-
-    - bash: |
-        echo $CUDA_VISIBLE_DEVICES
-        echo $CUDA_VERSION_MM
-        lspci | egrep 'VGA|3D'
-        whereis nvidia
-        nvidia-smi
-        which python && which pip
-        python --version && pip --version && pip list
-        python -c "import torch ; mgpu = torch.cuda.device_count() ; assert mgpu == 2, f'GPU: {mgpu}'"
-      displayName: 'Image info & NVIDIA'
-
-    - script: pip install pytest -r requirements.txt
-      displayName: 'Install dependencies'
-
-    - bash: pytest -v --durations=10 --disable-pytest-warnings --strict-markers --color=yes
-      displayName: 'Testing'

lit-llama/.github/workflows/cpu-tests.yml

@@ -1,52 +0,0 @@
-name: CPU tests
-
-on:
-  push:
-    branches: [main]
-  pull_request:
-    branches: [main]
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
-  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
-
-defaults:
-  run:
-    shell: bash
-
-jobs:
-  cpu-tests:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - {os: "macOS-11", python-version: "3.10"}
-          - {os: "ubuntu-20.04", python-version: "3.10"}
-          - {os: "windows-2022", python-version: "3.10"}
-    timeout-minutes: 15
-
-    steps:
-    - uses: actions/checkout@v3
-
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
-      with:
-        python-version: ${{ matrix.python-version }}
-        cache: 'pip'
-        cache-dependency-path: |
-          requirements.txt
-          setup.py
-
-    - name: Run tests without the package installed
-      run: |
-        pip install pytest -r requirements.txt
-        pip list
-
-        pytest --disable-pytest-warnings --strict-markers --color=yes
-
-    - name: Run tests
-      run: |
-        pip install . --no-deps
-
-        pytest -v --durations=10 --disable-pytest-warnings --strict-markers --color=yes

lit-llama/.gitignore

@@ -1,16 +0,0 @@
-__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

lit-llama/LICENSE

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

lit-llama/README.md

@@ -1,193 +0,0 @@
-<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>) pretraining, finetuning, and inference code that is fully open source under the **Apache 2.0 license.**
-
-This implementation builds on [nanoGPT](<https://github.com/karpathy/nanoGPT>).
-
-The open-source code in this repository works with the original LLaMA weights that 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 the original research-only weights by Meta and the Open LLaMA weights 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/gptq.py --output_path checkpoints/lit-llama/7B/llama-gptq.4bit.pth --dtype bfloat16 --quantize gptq.int4
-```
-
-GPTQ-style int4 quantization brings GPU usage down to about ~5GB. As only the weights of the Linear layers are quantized, it is useful to also use `--dtype bfloat16` even with the quantization enabled.
-
-With the generated quantized checkpoint generation quantization then works as usual with `--quantize gptq.int4` and the newly generated checkpoint file:
-
-```bash
-python generate.py --quantize gptq.int4 --checkpoint_path checkpoints/lit-llama/7B/llama-gptq.4bit.pth
-```
-
-[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 (RTX 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.

lit-llama/evaluate/adapter.py

@@ -1,164 +0,0 @@
-# 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
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-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
-
-instruction_tuning = True
-
-
-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
-
-
-@torch.inference_mode()
-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: Path = Path("out/adapter/alpaca/lit-llama-adapter-finetuned.pth"),
-    checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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.
-        dtype: The tensor dtype for choosing the floating-point precision 
-        quantize: Whether to quantize the model and using which method:
-            ``"llm.int8"``: LLM.int8() mode,
-            ``"gptq.int4"``: GPTQ 4-bit mode.
-    """
-    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
-
-    print("Loading model ...", file=sys.stderr)
-    t0 = time.time()
-    with lazy_load(checkpoint_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()
-
-    # 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)
-
-        if instruction_tuning:
-            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
-        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)
-        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)

lit-llama/evaluate/adapter_v2.py

@@ -1,161 +0,0 @@
-# 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
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama import Tokenizer
-from lit_llama.adapter import LLaMA
-from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
-from lit_llama.adapter_v2 import add_adapter_v2_parameters_to_linear_layers
-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
-
-
-@torch.inference_mode()
-def main(
-    datasets: str = "wikitext,ptb,c4",
-    *,
-    accelerator: str = "auto",
-    adapter_path: Path = Path("out/adapter_v2/alpaca/lit-llama-adapter-finetuned.pth"),
-    checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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
-        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_v2.py`.
-        checkpoint_path: The checkpoint path to load.
-        tokenizer_path: The tokenizer path to load.
-        dtype: The tensor dtype for choosing the floating-point precision 
-        quantize: Whether to quantize the model and using which method:
-            ``"llm.int8"``: LLM.int8() mode,
-            ``"gptq.int4"``: GPTQ 4-bit mode.
-    """
-    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
-
-    print("Loading model ...", file=sys.stderr)
-    t0 = time.time()
-    with lazy_load(checkpoint_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)
-            add_adapter_v2_parameters_to_linear_layers(model)
-
-        # 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
-
-        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)
-        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)

lit-llama/evaluate/full.py

@@ -1,147 +0,0 @@
-# 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
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-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: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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.
-        dtype: The tensor dtype for choosing the floating-point precision 
-        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")
-    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)
-        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)

lit-llama/evaluate/lora.py

@@ -1,172 +0,0 @@
-# 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
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-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
-
-instruction_tuning = True
-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: Path = Path("out/lora/alpaca/lit-llama-lora-finetuned.pth"),
-    checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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.
-        dtype: The tensor dtype for choosing the floating-point precision 
-        quantize: Whether to quantize the model and using which method:
-            ``"llm.int8"``: LLM.int8() mode,
-            ``"gptq.int4"``: GPTQ 4-bit mode.
-    """
-    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()
-
-    with lazy_load(checkpoint_path) as pretrained_checkpoint, lazy_load(lora_path) as lora_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 lora weights
-            model.load_state_dict(lora_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)
-
-        if instruction_tuning:
-            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)
-        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)

lit-llama/finetune/adapter.py

@@ -1,262 +0,0 @@
-"""
-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 sys
-import time
-from pathlib import Path
-import shutil
-
-import lightning as L
-import numpy as np
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-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
-
-
-instruction_tuning = True
-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_iters = batch_size // micro_batch_size
-assert gradient_accumulation_iters > 0
-epoch_size = 50000  # train dataset size
-num_epochs = 5
-max_iters = num_epochs * (epoch_size // micro_batch_size) // devices
-weight_decay = 0.02
-max_seq_length = 256  # see scripts/prepare_alpaca.py
-warmup_iters = 2 * (epoch_size // micro_batch_size) // devices  # 2 epochs
-
-ds_config = {
-    "train_micro_batch_size_per_gpu": micro_batch_size,
-    "gradient_accumulation_steps": gradient_accumulation_iters,
-    "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_iters:
-            # linear warmup
-            lr = learning_rate * step_count / warmup_iters
-            for param_group in optimizer.param_groups:
-                param_group['lr'] = lr
-
-        t0 = time.time()
-
-        input_ids, targets = get_batch(fabric, train_data)
-        with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
-            logits = model(input_ids)
-            loss = loss_fn(logits, targets)
-            fabric.backward(loss / gradient_accumulation_iters)
-
-        if (iter_num + 1) % gradient_accumulation_iters == 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 = instruction
-    if instruction_tuning:
-        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)

lit-llama/finetune/adapter_v2.py

@@ -1,266 +0,0 @@
-"""
-Instruction-tuning with LLaMA-Adapter v2 on the Alpaca dataset following the paper
-
-LLaMA-Adapter V2: Parameter-Efficient Visual Instruction Model
-https://arxiv.org/abs/2304.15010
-
-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 sys
-import time
-from pathlib import Path
-import shutil
-
-import lightning as L
-import numpy as np
-import torch
-import torch.nn as nn
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate import generate
-from lit_llama.adapter import LLaMA, LLaMAConfig
-from lit_llama.adapter_v2 import (
-    mark_only_adapter_v2_as_trainable,
-    add_adapter_v2_parameters_to_linear_layers,
-    adapter_v2_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_iters = batch_size // micro_batch_size
-assert gradient_accumulation_iters > 0
-epoch_size = 50000  # train dataset size
-num_epochs = 5
-max_iters = num_epochs * (epoch_size // micro_batch_size) // devices
-weight_decay = 0.02
-max_seq_length = 256  # see scripts/prepare_alpaca.py
-warmup_iters = 2 * (epoch_size // micro_batch_size) // devices  # 2 epoch
-
-ds_config = {
-    "train_micro_batch_size_per_gpu": micro_batch_size,
-    "gradient_accumulation_steps": gradient_accumulation_iters,
-    "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_v2/alpaca",
-):
-
-    fabric = L.Fabric(
-        accelerator="cuda",
-        devices=1,
-        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 contained in state dict
-        model.load_state_dict(checkpoint, strict=False)
-
-    add_adapter_v2_parameters_to_linear_layers(model)
-    mark_only_adapter_v2_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_iters:
-            # linear warmup
-            lr = learning_rate * step_count / warmup_iters
-            for param_group in optimizer.param_groups:
-                param_group['lr'] = lr
-
-        t0 = time.time()
-
-        input_ids, targets = get_batch(fabric, train_data)
-        with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
-            logits = model(input_ids)
-            loss = loss_fn(logits, targets)
-            fabric.backward(loss / gradient_accumulation_iters)
-
-        if (iter_num + 1) % gradient_accumulation_iters == 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_v2_state_from_state_dict(state_dict)
-            torch.save(state_dict, file_path)
-            shutil.rmtree(tmp_path)
-    else:
-        state_dict = adapter_v2_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)

lit-llama/finetune/full.py

@@ -1,224 +0,0 @@
-"""
-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 sys
-from pathlib import Path
-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
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-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
-
-
-instruction_tuning = True
-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_iters = batch_size // micro_batch_size
-assert gradient_accumulation_iters > 0
-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_iters = 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, limit_all_gathers=True)
-
-    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, foreach=False)
-    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_iters != 0
-
-        if step_count <= warmup_iters:
-            # linear warmup
-            lr = learning_rate * step_count / warmup_iters
-            for param_group in optimizer.param_groups:
-                param_group['lr'] = lr
-
-        t0 = time.time()
-        
-        input_ids, targets = get_batch(fabric, train_data)
-        with fabric.no_backward_sync(model, enabled=is_accumulating):
-            logits = model(input_ids)
-            loss = loss_fn(logits, targets)
-            fabric.backward(loss / gradient_accumulation_iters)
-
-        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 = instruction
-    if instruction_tuning:
-        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)

lit-llama/finetune/lora.py

@@ -1,218 +0,0 @@
-"""
-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 sys
-from pathlib import Path
-import os
-import time
-
-import lightning as L
-import numpy as np
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-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
-
-
-instruction_tuning = True
-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_iters = batch_size // micro_batch_size
-assert gradient_accumulation_iters > 0
-max_iters = 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_iters = 100
-
-
-def main(
-    data_dir: str = "data/alpaca", 
-    pretrained_path: str = "checkpoints/lit-llama/7B/lit-llama.pth",
-    tokenizer_path: str = "checkpoints/lit-llama/tokenizer.model",
-    out_dir: str = "out/lora/alpaca",
-):
-
-    fabric = L.Fabric(accelerator="cuda", devices=1, 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.from_name("7B")
-    config.block_size = max_seq_length
-
-    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)
-    train(fabric, model, optimizer, train_data, val_data, tokenizer_path, out_dir)
-
-    # Save the final LoRA checkpoint at the end of training
-    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,
-    tokenizer_path: str,
-    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_iters:
-            # linear warmup
-            lr = learning_rate * step_count / warmup_iters
-            for param_group in optimizer.param_groups:
-                param_group['lr'] = lr
-
-        t0 = time.time()
-
-        input_ids, targets = get_batch(fabric, train_data)
-        with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
-            logits = model(input_ids)
-            loss = loss_fn(logits, targets)
-            fabric.backward(loss / gradient_accumulation_iters)
-
-        if (iter_num + 1) % gradient_accumulation_iters == 0:
-            optimizer.step()
-            optimizer.zero_grad()
-            step_count += 1
-                
-            if step_count % eval_interval == 0:
-                val_loss = validate(fabric, model, val_data, tokenizer_path)
-                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_path):
-    tokenizer = Tokenizer(tokenizer_path)
-    sample = {"instruction": instruction, "input": ""}
-    prompt = instruction
-    if instruction_tuning:
-        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, tokenizer_path: str) -> 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, tokenizer_path)
-    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)

lit-llama/generate.py

@@ -1,170 +0,0 @@
-import sys
-import time
-import warnings
-from pathlib import Path
-from typing import Optional
-
-import lightning as L
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama import LLaMA, Tokenizer
-from lit_llama.utils import lazy_load, llama_model_lookup, quantization
-
-
-@torch.no_grad()
-def generate(
-    model: LLaMA,
-    idx: torch.Tensor,
-    max_new_tokens: int,
-    *,
-    max_seq_length: Optional[int] = None,
-    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
-    if max_seq_length is None:
-        max_seq_length = min(T_new, model.config.block_size)
-
-    device, dtype = idx.device, idx.dtype
-    # create an empty tensor of the expected final shape and fill in the current tokens
-    empty = torch.empty(T_new, dtype=dtype, device=device)
-    empty[:T] = idx
-    idx = empty
-    input_pos = torch.arange(0, T, device=device)
-
-    if idx.device.type == "xla":
-        import torch_xla.core.xla_model as xm
-
-        xm.mark_step()
-
-    # generate max_new_tokens tokens
-    for _ in range(max_new_tokens):
-        x = idx.index_select(0, input_pos).view(1, -1)
-
-        # forward
-        logits = model(x, max_seq_length, input_pos)
-        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 = torch.where(logits < v[[-1]], -float("Inf"), logits)
-
-        probs = torch.nn.functional.softmax(logits, dim=-1)
-        idx_next = torch.multinomial(probs, num_samples=1).to(dtype=dtype)
-
-        # advance
-        input_pos = input_pos[-1:] + 1
-
-        if idx.device.type == "xla":
-            xm.mark_step()
-
-        # concatenate the new generation
-        idx = idx.index_copy(0, input_pos, idx_next)
-
-        # if <eos> token is triggered, return the output (stop generation)
-        if idx_next == eos_id:
-            return idx[:input_pos]  # 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: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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.
-    """
-    assert checkpoint_path.is_file(), checkpoint_path
-    assert tokenizer_path.is_file(), tokenizer_path
-
-    precision = "bf16-true" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "32-true"
-    fabric = L.Fabric(devices=1, precision=precision)
-
-    print("Loading model ...", file=sys.stderr)
-    t0 = time.time()
-    with lazy_load(checkpoint_path) as checkpoint:
-        name = llama_model_lookup(checkpoint)
-
-        with fabric.init_module(empty_init=True), 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(model)
-
-    tokenizer = Tokenizer(tokenizer_path)
-    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=fabric.device)
-    prompt_length = encoded.size(0)
-
-    L.seed_everything(1234)
-    for i in range(num_samples):
-        t0 = time.perf_counter()
-        y = generate(model, encoded, max_new_tokens, temperature=temperature, top_k=top_k)
-        t = time.perf_counter() - t0
-
-        model.reset_cache()
-        print(tokenizer.decode(y))
-        tokens_generated = y.size(0) - prompt_length
-        print(f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_generated / 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)

lit-llama/generate/adapter.py

@@ -1,106 +0,0 @@
-import sys
-import time
-import warnings
-from pathlib import Path
-from typing import Optional
-
-import lightning as L
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate import generate
-from lit_llama import Tokenizer
-from lit_llama.adapter import LLaMA
-from lit_llama.utils import lazy_load, llama_model_lookup, quantization
-from scripts.prepare_alpaca import generate_prompt
-
-
-def main(
-    prompt: str = "What food do lamas eat?",
-    input: str = "",
-    adapter_path: Path = Path("out/adapter/alpaca/lit-llama-adapter-finetuned.pth"),
-    pretrained_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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.
-    """
-    assert adapter_path.is_file()
-    assert pretrained_path.is_file()
-    assert tokenizer_path.is_file()
-
-    precision = "bf16-true" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "32-true"
-    fabric = L.Fabric(devices=1, precision=precision)
-
-    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 fabric.init_module(empty_init=True), 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(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)
-    prompt_length = encoded.size(0)
-
-    t0 = time.perf_counter()
-    y = generate(model, encoded, max_new_tokens, temperature=temperature, top_k=top_k, eos_id=tokenizer.eos_id)
-    t = time.perf_counter() - t0
-
-    model.reset_cache()
-    output = tokenizer.decode(y)
-    output = output.split("### Response:")[1].strip()
-    print(output)
-
-    tokens_generated = y.size(0) - prompt_length
-    print(f"\n\nTime for inference: {t:.02f} sec total, {tokens_generated / 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)

lit-llama/generate/adapter_v2.py

@@ -1,108 +0,0 @@
-import sys
-import time
-import warnings
-from pathlib import Path
-from typing import Optional
-
-import lightning as L
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate import generate
-from lit_llama import Tokenizer
-from lit_llama.adapter import LLaMA
-from lit_llama.utils import lazy_load, llama_model_lookup, quantization
-from lit_llama.adapter_v2 import add_adapter_v2_parameters_to_linear_layers
-from scripts.prepare_alpaca import generate_prompt
-
-
-def main(
-    prompt: str = "What food do lamas eat?",
-    input: str = "",
-    adapter_path: Path = Path("out/adapter_v2/alpaca/lit-llama-adapter-finetuned.pth"),
-    pretrained_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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_v2.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_v2.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.
-    """
-    assert adapter_path.is_file()
-    assert pretrained_path.is_file()
-    assert tokenizer_path.is_file()
-
-    precision = "bf16-true" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "32-true"
-    fabric = L.Fabric(devices=1, precision=precision)
-
-    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 fabric.init_module(empty_init=True), quantization(mode=quantize):
-            model = LLaMA.from_name(name)
-            add_adapter_v2_parameters_to_linear_layers(model)
-
-        # 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(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)
-    prompt_length = encoded.size(0)
-
-    t0 = time.perf_counter()
-    y = generate(model, encoded, max_new_tokens, temperature=temperature, top_k=top_k, eos_id=tokenizer.eos_id)
-    t = time.perf_counter() - t0
-
-    model.reset_cache()
-    output = tokenizer.decode(y)
-    output = output.split("### Response:")[1].strip()
-    print(output)
-
-    tokens_generated = y.size(0) - prompt_length
-    print(f"\n\nTime for inference: {t:.02f} sec total, {tokens_generated / 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)

lit-llama/generate/full.py

@@ -1,103 +0,0 @@
-import sys
-import time
-import warnings
-from pathlib import Path
-from typing import Optional
-
-import lightning as L
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).absolute().parent.parent
-sys.path.append(str(wd))
-
-from lit_llama import LLaMA, Tokenizer
-from lit_llama.utils import quantization
-from scripts.prepare_alpaca import generate_prompt
-from generate import generate
-
-
-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: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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")
-    assert checkpoint_path.is_file(), checkpoint_path
-    assert tokenizer_path.is_file(), tokenizer_path
-
-    precision = "bf16-true" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "32-true"
-    fabric = L.Fabric(devices=1, precision=precision)
-
-    print("Loading model ...", file=sys.stderr)
-    t0 = time.time()
-    
-    with fabric.init_module(empty_init=True), 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(model)
-
-    tokenizer = Tokenizer(tokenizer_path)
-    sample = {"instruction": prompt, "input": input}
-    prompt = generate_prompt(sample)
-    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=fabric.device)
-    prompt_length = encoded.size(0)
-
-    L.seed_everything(1234)
-    for i in range(num_samples):
-        t0 = time.perf_counter()
-        y = generate(model, encoded, max_new_tokens, temperature=temperature, top_k=top_k)
-        t = time.perf_counter() - t0
-
-        model.reset_cache()
-        print(tokenizer.decode(y))
-        tokens_generated = y.size(0) - prompt_length
-        print(f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_generated / 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)

lit-llama/generate/lora.py

@@ -1,118 +0,0 @@
-import sys
-import time
-import warnings
-from pathlib import Path
-from typing import Optional
-
-import lightning as L
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate import generate
-from lit_llama import Tokenizer, LLaMA
-from lit_llama.lora import lora
-from lit_llama.utils import 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: Path = Path("out/lora/alpaca/lit-llama-lora-finetuned.pth"),
-    pretrained_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    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 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.
-        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.
-    """
-    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")
-
-    precision = "bf16-true" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "32-true"
-    fabric = L.Fabric(devices=1, precision=precision)
-
-    print("Loading model ...", file=sys.stderr)
-    t0 = time.time()
-
-    with lazy_load(pretrained_path) as pretrained_checkpoint, lazy_load(lora_path) as lora_checkpoint:
-        name = llama_model_lookup(pretrained_checkpoint)
-
-        with fabric.init_module(empty_init=True), 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 lora weights
-            model.load_state_dict(lora_checkpoint, strict=False)
-
-    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
-
-    model.eval()
-    model = fabric.setup(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_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)

lit-llama/howto/convert_lora_weights.md

@@ -1,19 +0,0 @@
-# Merging LoRA weights into base model weights
-
-Purpose: By merging our selected LoRA weights into the base model weights, we can benefit from all base model optimisation such as quantisation (available in this repo), pruning, caching, etc.
-
-
-## How to run?
-
-After you have finish finetuning using LoRA, select your weight and run the converter script:
-
-```bash
-python scripts/convert_lora_weights.py --lora_path out/lora/your-folder/your-weight-name.pth
-```
-
-The converted base weight file will be saved into the same folder with the name `{your-weight-name}-lora-merged-weights.pth`. Now you can run `generate.py` with the merged weights and apply quantisation:
-
-```bash
-python generate.py --checkpoint_path out/lora/your-folder/your-weight-name-lora-merged-weights.pth --quantize llm.int8
-```
-

lit-llama/howto/customize_paths.md

@@ -1,33 +0,0 @@
-## 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 generate.py \
-  --checkpoint_path "data/checkpoints/foo/7B/lit-llama.pth" \
-  --tokenizer_path "data/checkpoints/foo/tokenizer.model"
-```
-
-and
-
-```shell
-python quantize/gptq.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.

lit-llama/howto/download_weights.md

@@ -1,130 +0,0 @@
-## 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 (1T 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 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 --local_dir checkpoints/open-llama/7B
-```
-
-Once downloaded, you should have a folder like this:
-
-```text
-checkpoints/open-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 --checkpoint_dir checkpoints/open-llama/7B --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).

lit-llama/howto/finetune_adapter.md

@@ -1,109 +0,0 @@
-# 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 RTX 3090 (24GB) GPU**. If using 8 GPUs, finetuning can be completed in under 1 hour.
-
-If you are new to LLaMA-Adapter and are interested 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.
-
-## LLaMA-Adapter v2
-
-The LLaMA-Adapter authors developed a newer adapter method called LLaMA-Adapter v2, which is related to this LLaMA-Adapter method but includes more trainable parameters. LLaMA-Adapter v2 is also available via Lit-LLaMA; you can read more about it in [the related how-to doc here](./finetune_adapter_v2.md).
-
-## 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-dataset).
-
-See also: [Finetuning on an unstructured dataset](unstructured_dataset.md)
-
-## Running the finetuning
-
-```bash
-python finetune/adapter.py
-```
-
-The finetuning requires at least one GPU with ~24 GB memory (RTX 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).

lit-llama/howto/finetune_adapter_v2.md

@@ -1,114 +0,0 @@
-# Finetuning with Adapter v2
-
-[LLaMA-Adapter v2](https://arxiv.org/abs/2304.15010) 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 ~4 M 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 RTX 3090 (24GB) GPU**. If using 8 GPUs, finetuning can be completed in under 1 hour.
-
-If you are new to LLaMA-Adapter and are interested 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.
-
-## LLaMA-Adapter v1 versus LLaMA-Adapter v2
-
-LLaMA-Adapter v2 extends the original LLaMA-Adapter idea by adding trainable bias and scale parameters to each linear layer in the transformer. Furthermore, LLaMA-Adapter v2 makes the normalization layers trainable. Where the 7B LLaMA model has 1.2M trainable parameters with LLaMA v1, LLaMA-Adapter v2 adds 2.8 M trainable parameters for the bias and scale parameters and ~300k trainable parameters for the normalization layers. So, adapter v2 has ~4.3 M trainable parameters in total.
-
-If you are interested in using the more lightweight LLaMA-Adapter v1 approach, see [the related LLaMA Adapter how-to doc here](./finetune_adapter.md).
-
-While LLaMA-Adapter v2 increases the number of trainable parameters from 1.2 M (from LLaMA-Apdapter v1) to 4.3 M, the inference cost is not significantly impacted. This is because the additional bias and scale parameters are cheap to compute in the forward pass, and the RMSNorm parameters are already included in the base model. In LLaMA-Adapter v1, the RMSNorm parameters are not trainable.
-
-
-## 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-dataset).
-
-See also: [Finetuning on an unstructured dataset](unstructured_dataset.md)
-
-## Running the finetuning
-
-```bash
-python finetune/adapter_v2.py
-```
-
-The finetuning requires at least one GPU with ~24 GB memory (RTX 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_v2.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_v2.py` by passing in the location of your data (and optionally other parameters):
-
-    ```bash
-    python finetune/adapter_v2.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).

lit-llama/howto/finetune_full.md

@@ -1,106 +0,0 @@
-# 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](../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).
-
-See also: [Finetuning on an unstructured dataset](unstructured_dataset.md)
-
-## 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).

lit-llama/howto/finetune_lora.md

@@ -1,90 +0,0 @@
-# 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 RTX 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
-   ```
-
-See also: [Finetuning on an unstructured dataset](unstructured_dataset.md)
-
-## Running the finetuning
-
-```bash
-python finetune/lora.py
-```
-
-The finetuning requires at least one GPU with ~24 GB memory (RTX 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).

lit-llama/howto/inference.md

@@ -1,43 +0,0 @@
-# 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/gptq.py --output_path checkpoints/lit-llama/7B/llama-gptq.4bit.pth --dtype bfloat16 --quantize gptq.int4
-```
-
-GPTQ-style int4 quantization brings GPU usage down to about ~5GB. As only the weights of the Linear layers are quantized, it is useful to also use `--dtype bfloat16` even with the quantization enabled.
-
-With the generated quantized checkpoint generation quantization then works as usual with `--quantize gptq.int4` and the newly generated checkpoint file:
-
-```bash
-python generate.py --quantize gptq.int4 --checkpoint_path checkpoints/lit-llama/7B/llama-gptq.4bit.pth
-```

lit-llama/howto/tpus.md

@@ -1,51 +0,0 @@
-# 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 3
-```
-
-This command will take take ~20s for the first generation time as XLA needs to compile the graph.
-You'll notice that afterwards, generation times drop to ~5s.
-
-## 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
-> ```

lit-llama/howto/train_redpajama.md

@@ -1,133 +0,0 @@
-# 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 pretrain/redpajama.py --devices 4 --train_data_dir data/lit-redpajama
-```
-
-For running on the sample dataset:
-
-```bash
-python pretrain/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 // micro_batch_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/howto/unstructured_dataset.md

@@ -1,18 +0,0 @@
-# Finetuning on an unstructured dataset
-
-While most scripts were made to finetune on instruction datasets, it is possible to finetune on any dataset. This is useful for experimentation while not being as expensive as training a full model.
-
-This guide is only to prepare the finetuning, as either LoRA or Adapter-v1 methods support this dataset type!
-
-## Preparation
-
-1. Gather your text into an input file named `input.txt`
-2. Divide the data into training and validation sets using the following script:
-
-    ```bash
-    python scripts/prepare_any_text.py
-    ```
-
-3. Modify relevant scripts for your finetuning method under `finetune/` and `evaluate/`, setting the `instruction_tuning` variable to `False`
-
-And then you're set! Proceed to run the [LoRA guide](./finetune_lora.md) or [Adapter v1 guide](./finetune_adapter.md).

lit-llama/lit_llama/__init__.py

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

lit-llama/lit_llama/adapter.py

@@ -1,313 +0,0 @@
-"""Implementation of the paper:
-
-LLaMA-Adapter: Efficient Fine-tuning of Language Models with Zero-init Attention
-https://arxiv.org/abs/2303.16199
-
-                                                                             |              Prefix cross-attention
-                                                                             |
-  ┌─────────────────┐                                                        |               ┌──────────────────┐
-  ┆        x        ┆                                                        |               ┆      prefix      ┆
-  └─────────────────┘                                                        |               └──────────────────┘
-           |                                                                 |                        |
-           ▼                                                                 |                        ▼
-  ┌──────────────────┐                                                       |              ┌─────────────────────┐
-  ┆  self-attention  ┆ --------------------------------------------------------------┐      ┆  linear projection  ┆
-  └──────────────────┘                                                       |       ┆      └─────────────────────┘
-           |                                                                 |       ┆                |         \
-           ▼                                                                 |       ▼                ▼          ▼
-         ╭───╮     ┌────────────────┐ ╭───╮ ┌──────────────────────────┐     |  ┌─────────┐    ┌──────────────┐  ┌────────────────┐
-         ┆ + ┆ ◀── ┆  gating factor ┆-┆ x ┆-┆  prefix cross-attention  ┆     |  ┆  query  ┆    ┆  prefix key  ┆  ┆  prefix value  ┆
-         ╰───╯     └────────────────┘ ╰───╯ └──────────────────────────┘     |  └─────────┘    └──────────────┘  └────────────────┘
-           |                                                                 |          \             |           /
-           ▼                                                                 |           ▼            ▼          ▼
-                                                                             |         ┌────────────────────────────────┐
-                                                                             |         ┆  scaled dot-product attention  ┆
-                                                                             |         └────────────────────────────────┘
-
-
-In order to inject learnable information from the prefix to pretrained weights we need to sum outputs from
-self-attention and prefix cross-attention (times gating factor). For prefix cross-attention we need `query` (from
-self-attention as a result of linear projection), `prefix key` and `prefix value` (from cross-attention as a result of
-linear projection).
-The output of prefix cross-attention is multiplied by gating factor, which is a learnable parameter that is needed to
-avoid potential disruption of pretrained weights caused by incorporating randomly initialized tensors. This factor is
-initialized with zeros to avoid noise from the adaption prompts at the early training stage.
-More about it: https://lightning.ai/pages/community/article/understanding-llama-adapters/
-
-Notes about implementation: as per paper adapter's prefix is concatenated with the input, while here outputs of
-self-attention and prefix cross-attention are summed. Both variants are mathematically equivalent:
-https://github.com/ZrrSkywalker/LLaMA-Adapter/issues/47
-"""
-# mypy: ignore-errors
-from dataclasses import dataclass
-from typing import Optional, Tuple, List, Union
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-
-import lit_llama.model as llama
-from lit_llama.model import build_rope_cache, apply_rope, RMSNorm, MLP, KVCache, RoPECache
-
-
-@dataclass
-class LLaMAConfig(llama.LLaMAConfig):
-    adapter_prompt_length: int = 10
-    adapter_start_layer: int = 2
-
-
-class CausalSelfAttention(nn.Module):
-    """A modification of `lit_llama.model.CausalSelfAttention` that adds the attention
-    over the adaption prompt."""
-
-    def __init__(self, config: LLaMAConfig, block_idx: int) -> None:
-        super().__init__()
-        assert config.n_embd % config.n_head == 0
-
-        # key, query, value projections for all heads, but in a batch
-        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
-        # output projection
-        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
-
-        if block_idx >= config.adapter_start_layer:
-            # adapter embedding layer
-            self.adapter_wte = nn.Embedding(config.adapter_prompt_length, config.n_embd)
-            # a learnable gating factor (to avoid potential disruption of pretrained weights) initialized with zeros (to
-            # avoid noise from adaption prompts at the early training stage)
-            self.gating_factor = torch.nn.Parameter(torch.zeros(1, config.n_head, 1, 1))
-
-        self.n_head = config.n_head
-        self.n_embd = config.n_embd
-        self.block_size = config.block_size
-        self.block_idx = block_idx
-        self.adapter_prompt_length = config.adapter_prompt_length
-        self.adapter_start_layer = config.adapter_start_layer
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope: RoPECache,
-        mask: torch.Tensor,
-        max_seq_length: int,
-        input_pos: Optional[torch.Tensor] = None,
-        kv_cache: Optional[KVCache] = None,
-        adapter_kv_cache: Optional[KVCache] = None,
-    ) -> Tuple[torch.Tensor, Optional[KVCache], Optional[KVCache]]:
-        # notation:
-        # - B  | batch
-        # - T  | time-step (sequence length)
-        # - C  | embeddings size (n_embd) = head size * num heads
-        # - hs | head size
-        # - nh | number of heads
-
-        B, T, C = x.size()
-
-        # instead of calculating `query`, `key` and `value` by separately multiplying input `x` with corresponding
-        # weight matrices do it (for all heads) in a single multiplication with a matrix of 3x size (concatenated
-        # weights for q, k, v) and then split the result along `embedding size` dimension
-        q, k, v = self.c_attn(x).split(self.n_embd, dim=2) # (B, T, 3 * C) --> 3 * (B, T, C)
-
-        # in order to move head_size (hs) dimension right after batch (B) dimension, we need to first split
-        # embedding size (C) dimension into num_heads (nh) and head_size (hs)
-        head_size = C // self.n_head
-        k = k.view(B, T, self.n_head, head_size)
-        q = q.view(B, T, self.n_head, head_size)
-        v = v.view(B, T, self.n_head, head_size)
-
-        # "Unlike standard positional embeddings rotary embeddings must be applied at every layer"
-        q = apply_rope(q, rope) # (B, T, nh, hs)
-        k = apply_rope(k, rope) # (B, T, nh, hs)
-
-        # now `key`, 'query` and `value` tensors are correctly represented: for each element in a batch (B)
-        # there is a number of heads (nh) and for each head there is a sequence of elements (T), each of them is
-        # represented by a vector of size `hs`
-        k = k.transpose(1, 2)  # (B, nh, T, hs)
-        q = q.transpose(1, 2)  # (B, nh, T, hs)
-        v = v.transpose(1, 2)  # (B, nh, T, hs)
-
-        if kv_cache is not None:
-            cache_k, cache_v = kv_cache # 2 * (B, nh, max_seq_length, hs)
-            # check if reached token limit
-            if input_pos[-1] >= max_seq_length:
-                # if we reached token limit and thus there is no space to put newly calculated `key` and `value`
-                # right next to cached ones, we need to rotate cache tensor along `max_seq_length` dimension by one
-                # element to the left: this will free up space for new `key` and `value`
-                input_pos = torch.tensor(max_seq_length - 1, device=input_pos.device)
-                # shift 1 position to the left
-                cache_k = torch.roll(cache_k, -1, dims=2)
-                cache_v = torch.roll(cache_v, -1, dims=2)
-            k = cache_k.index_copy(2, input_pos, k) # (B, nh, max_seq_length, hs)
-            v = cache_v.index_copy(2, input_pos, v) # (B, nh, max_seq_length, hs)
-            kv_cache = k, v
-
-        # efficient attention using Flash Attention CUDA kernels
-        # ↓ (B, nh, T, hs) @ (B, nh, T, hs).mT --> (B, nh, T, T) @ (B, nh, T, hs) --> (B, nh, T, hs)
-        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0) # (B, nh, T, hs)
-
-        # "Adapters are applied to the topmost layers to better tune the language
-        # representations with higher-level semantics".
-        if self.block_idx >= self.adapter_start_layer:
-            if adapter_kv_cache is not None:
-                ak, av = adapter_kv_cache # 2 * (B, nh, aT, hs)
-            else:
-                prefix = self.adapter_wte.weight.reshape(1, self.adapter_prompt_length, self.n_embd)
-                aT = prefix.size(1)
-                _, ak, av = self.c_attn(prefix).split(self.n_embd, dim=2) # (1, aT, 3 * C) --> 3 * (1, aT, C)
-                ak = ak.view(1, aT, self.n_head, head_size).repeat(B, 1, 1, 1).transpose(1, 2) # (B, nh, aT, hs)
-                av = av.view(1, aT, self.n_head, head_size).repeat(B, 1, 1, 1).transpose(1, 2) # (B, nh, aT, hs)
-                adapter_kv_cache = (ak, av)
-
-            # Apply cross-attention with `query`, `adapter_key`, `adapter_value` and sum the output with the output
-            # obtained from self-attention step. This is mathematically equivalent to concatenation of prefix and input as per paper.
-            amask = torch.ones(q.shape[-2], ak.shape[-2], dtype=torch.bool, device=x.device) # (T, aT)
-            # ↓ (B, nh, T, hs) @ (B, nh, aT, hs).mT --> (B, nh, T, aT) @ (B, nh, aT, hs) --> (B, nh, T, hs)
-            ay = F.scaled_dot_product_attention(q, ak, av, attn_mask=amask, dropout_p=0.0, is_causal=False) # (B, nh, T, hs)
-            y = y + self.gating_factor * ay
-
-        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
-
-        # output projection
-        y = self.c_proj(y) # (B, T, C)
-
-        return y, kv_cache, adapter_kv_cache
-
-    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
-        """For backward compatibility with old checkpoints that have a single gating value for all heads."""
-        name = prefix + "gating_factor"
-        if name in state_dict:
-            tensor = state_dict[name]
-            # in case we are loading with `utils.lazy_load()`
-            tensor = tensor._load_tensor() if hasattr(tensor, "_load_tensor") else tensor
-
-            if len(tensor.shape) < 4:
-                # For old checkpoints with unified gating value
-                state_dict[name] = tensor.reshape(1, 1, 1, 1).repeat(1, self.n_head, 1, 1)
-            else:
-                state_dict[name] = tensor
-
-        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
-
-
-class Block(nn.Module):
-    """The implementation is identical to `lit_llama.model.Block` with the exception that
-    we replace the attention layer where adaption is implemented."""
-
-    def __init__(self, config: LLaMAConfig, block_idx: int) -> None:
-        super().__init__()
-        self.rms_1 = RMSNorm(config.n_embd)
-        self.attn = CausalSelfAttention(config, block_idx)
-        self.rms_2 = RMSNorm(config.n_embd)
-        self.mlp = MLP(config)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope: RoPECache,
-        mask: torch.Tensor,
-        max_seq_length: int,
-        input_pos: Optional[torch.Tensor] = None,
-        kv_cache: Optional[KVCache] = None,
-        adapter_kv_cache: Optional[KVCache] = None,
-    ) -> Tuple[torch.Tensor, Optional[KVCache], Optional[KVCache]]:
-        h, new_kv_cache, new_adapter_kv_cache = self.attn(
-            self.rms_1(x), rope, mask, max_seq_length, input_pos, kv_cache, adapter_kv_cache
-        )
-        x = x + h
-        x = x + self.mlp(self.rms_2(x))
-        return x, new_kv_cache, new_adapter_kv_cache
-
-
-class LLaMA(llama.LLaMA):
-    """The implementation is identical to `lit_llama.model.LLaMA` with the exception that
-    the `Block` saves the layer index and passes it down to the attention layer."""
-
-    def __init__(self, config: LLaMAConfig) -> None:
-        nn.Module.__init__(self)
-        assert config.vocab_size is not None
-        assert config.block_size is not None
-        self.config = config
-
-        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-        self.transformer = nn.ModuleDict(
-            dict(
-                wte=nn.Embedding(config.vocab_size, config.n_embd),
-                h=nn.ModuleList(Block(config, i) for i in range(config.n_layer)),
-                ln_f=RMSNorm(config.n_embd),
-            )
-        )
-
-        self.rope_cache: Optional[RoPECache] = None
-        self.mask_cache: Optional[torch.Tensor] = None
-        self.kv_caches: List[KVCache] = []
-        self.adapter_kv_caches: List[KVCache] = []
-
-    @classmethod
-    def from_name(cls, name: str):
-        return cls(LLaMAConfig.from_name(name))
-
-    def reset_cache(self) -> None:
-        super().reset_cache()
-        self.adapter_kv_caches.clear()
-
-    def forward(
-        self, idx: torch.Tensor, max_seq_length: Optional[int] = None, input_pos: Optional[torch.Tensor] = None
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[KVCache]]]:
-        B, T = idx.size()
-
-        block_size = self.config.block_size
-        if max_seq_length is None:
-            max_seq_length = block_size
-        assert T <= max_seq_length, f"Cannot forward sequence of length {T}, max seq length is only {max_seq_length}"
-        assert max_seq_length <= block_size, f"Cannot attend to {max_seq_length}, block size is only {block_size}"
-        assert T <= block_size, f"Cannot forward sequence of length {T}, block size is only {block_size}"
-
-        if self.rope_cache is None:
-            self.rope_cache = self.build_rope_cache(idx) # (block_size, head_size / 2, 2)
-        if self.mask_cache is None:
-            self.mask_cache = self.build_mask_cache(idx) # (1, 1, block_size, block_size)
-
-        if input_pos is not None:
-            rope = self.rope_cache.index_select(0, input_pos)
-            mask = self.mask_cache.index_select(2, input_pos)
-            mask = mask[:, :, :, :max_seq_length]
-        else:
-            rope = self.rope_cache[:T]
-            mask = self.mask_cache[:, :, :T, :T]
-
-        # forward the model itself
-        x = self.transformer.wte(idx)  # token embeddings of shape (B, T, n_embd)
-
-        if input_pos is None:  # proxy for use_cache=False
-            for block in self.transformer.h:
-                x, *_ = block(x, rope, mask, max_seq_length)
-        else:
-            if not self.kv_caches:
-                head_size = self.config.n_embd // self.config.n_head
-                cache_shape = (B, self.config.n_head, max_seq_length, head_size)
-                self.kv_caches = [
-                    (torch.zeros(cache_shape, device=x.device, dtype=x.dtype), torch.zeros(cache_shape, device=x.device, dtype=x.dtype))
-                    for _ in range(self.config.n_layer)
-                ]
-            if not self.adapter_kv_caches:
-                self.adapter_kv_caches = [None for _ in range(self.config.n_layer)]
-            for i, block in enumerate(self.transformer.h):
-                x, self.kv_caches[i], self.adapter_kv_caches[i] = block(
-                    x, rope, mask, max_seq_length, input_pos, self.kv_caches[i], self.adapter_kv_caches[i]
-                )
-
-        x = self.transformer.ln_f(x) # (B, T, n_embd)
-
-        logits = self.lm_head(x)  # (B, T, vocab_size)
-
-        return logits
-
-
-def mark_only_adapter_as_trainable(model: LLaMA) -> None:
-    """Sets `requires_grad=False` for all non-adapter weights."""
-    for name, param in model.named_parameters():
-        param.requires_grad = "adapter_wte" in name or "gating_factor" in name
-
-
-def adapter_state_from_state_dict(state_dict: dict) -> dict:
-    """Returns the model state dict with only the adapter weights for saving."""
-    return {name: param for name, param in state_dict.items() if "adapter_wte" in name or "gating_factor" in name}

lit-llama/lit_llama/adapter_v2.py

@@ -1,45 +0,0 @@
-import torch
-from torch import Tensor
-import torch.nn as nn
-from torch.nn import functional as F
-
-from lit_llama.adapter import LLaMA
-
-
-def get_adapter_substrings():
-    substrings = ["adapter_wte", "gating_factor"]  # regular adapter v1 parameters
-    substrings.extend(["adapter_scale", "adapter_bias"])  # adapter v2: new bias and scale used in Linear
-    substrings.extend(["rms_1", "rms_2", "ln_f"])  # adapter v2: RMSNorm parameters are now trainable
-    return substrings
-
-
-def mark_only_adapter_v2_as_trainable(model: LLaMA) -> None:
-    """Sets `requires_grad=False` for all non-adapter weights."""
-    for name, param in model.named_parameters():
-        param.requires_grad = any(s in name for s in get_adapter_substrings())
-
-
-def adapter_v2_state_from_state_dict(state_dict: dict) -> dict:
-    """Returns the model state dict with only the adapter weights for saving."""
-    return {name: param for name, param in state_dict.items()
-            if any(s in name for s in get_adapter_substrings())}
-
-
-def adapter_v2_new_forward(self, input: Tensor) -> Tensor:
-    return self.adapter_scale * (
-        F.linear(input, self.weight, self.bias) + self.adapter_bias
-    )
-
-
-def adapter_v2_linear_with_bias_and_scale(layer):
-    layer.adapter_bias = torch.nn.Parameter(torch.zeros(layer.weight.shape[0]), requires_grad=True)
-    layer.adapter_scale = torch.nn.Parameter(torch.ones(layer.weight.shape[0]), requires_grad=True)
-    bound_method = adapter_v2_new_forward.__get__(layer, layer.__class__)
-    setattr(layer, 'forward', bound_method)
-    return layer
-
-
-def add_adapter_v2_parameters_to_linear_layers(model):
-    for module in model.modules():
-        if isinstance(module, nn.Linear):
-            adapter_v2_linear_with_bias_and_scale(module)

lit-llama/lit_llama/lora.py

@@ -1,476 +0,0 @@
-# Derived from https://github.com/microsoft/LoRA
-#  ------------------------------------------------------------------------------------------
-#  Copyright (c) Microsoft Corporation. All rights reserved.
-#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
-#  ------------------------------------------------------------------------------------------
-
-r"""
-    Low Ranking Adaptation for LLMs scheme.
-
-             ┌───────────────────┐
-             ┆         h         ┆
-             └───────────────────┘
-                       ▲
-                       |
-                       +
-                    /     \
-    ┌─────────────────┐    ╭───────────────╮     Matrix initialization:
-    ┆                 ┆     \      B      /      B = 0
-    ┆   pretrained    ┆      \    r*d    /       A = N(0, sigma^2)
-    ┆    weights      ┆       ╰─────────╯
-    ┆                 ┆       |    r    |        r - rank
-    ┆   W e R^(d*d)   ┆       | ◀─────▶ |
-    ┆                 ┆       ╭─────────╮
-    └─────────────────┘      /     A     \
-              ▲             /     d*r     \
-               \           ╰───────────────╯
-                \                ▲
-                 \              /
-                  \            /
-             ┌───────────────────┐
-             ┆         x         ┆
-             └───────────────────┘
-
-With LoRA (Low Ranking Adaptation: https://arxiv.org/abs/2106.09685) instead of learning weights of size d*d,
-we can freeze the pretrained weights and instead learn two matrices of size d*r and r*d (they will store weight updates
-for the pretrained weights): the number of parameters in this case will be reduced drastically (depending on the rank of
-course) yet after multiplication of matrices d*r and r*d we will get a matrix d*d which we can sum with frozen
-pretrained weights and thus fine-tune the model.
-
-The goal of this approach is to move weight updates into a separate matrix which is decomposed with
-two matrices of a lower rank.
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-import math
-from typing import Dict, List
-
-import lit_llama.model as llama
-
-from contextlib import contextmanager
-from dataclasses import dataclass
-
-
-class LoRALayer():
-    def __init__(
-        self, 
-        r: int, 
-        lora_alpha: int, 
-        lora_dropout: float,
-        merge_weights: bool,
-    ):
-        """Store LoRA specific attributes in a class.
-
-        Args:
-            r: rank of the weight update matrices. To make sense of using LoRA the rank should be smaller than the rank of
-                the weights of the model.  The rank can be as low as 1: https://arxiv.org/pdf/2106.09685.pdf (section 7.2)
-            lora_alpha: alpha is needed for scaling updates as alpha/r
-                "This scaling helps to reduce the need to retune hyperparameters when we vary r"
-                https://arxiv.org/pdf/2106.09685.pdf (section 4.1)
-            lora_dropout: dropout that is applied on the input in the LoRA branch (before multiplying by matrix A)
-            merge_weights: whether we want to merge pretrained weights and LoRA weight updates. This is useful if one wants to use
-                fine-tuned model as a standalone one (without storing LoRA weights separately) plus it helps to reduce
-                overhead during inference.
-        """
-        self.r = r
-        self.lora_alpha = lora_alpha
-        # Optional dropout
-        if lora_dropout > 0.:
-            self.lora_dropout = nn.Dropout(p=lora_dropout)
-        else:
-            self.lora_dropout = lambda x: x
-        # Mark the weight as unmerged
-        self.merged = False
-        self.merge_weights = merge_weights
-
-
-class MergedLinear(nn.Linear, LoRALayer):
-    # LoRA implemented in a dense layer
-    def __init__(
-        self, 
-        # ↓ this part is for pretrained weights
-        in_features: int, 
-        out_features: int, 
-        # ↓ the remaining part is for LoRA
-        r: int = 0, 
-        lora_alpha: int = 1, 
-        lora_dropout: float = 0.,
-        enable_lora: List[bool] = [False],
-        fan_in_fan_out: bool = False,
-        merge_weights: bool = True,
-        **kwargs
-    ):
-        """LoRA wrapper around linear class that is used for calculation of q, k and v matrices.
-
-        This class has three weight matrices:
-            1. Pretrained weights are stored as `self.weight` (because of the nn.Linear inheritance)
-            2. LoRA A matrix as `self.lora_A`
-            3. LoRA B matrix as `self.lora_B`
-        Only LoRA's A and B matrices are updated, pretrained weights stay frozen.
-
-        Args:
-            in_features: number of input features of the pretrained weights
-            out_features: number of output features of the pretrained weights
-            r: rank of the weight update matrices. To make sense of using LoRA the rank should be smaller than the rank of
-                the weights of the model.  The rank can be as low as 1: https://arxiv.org/pdf/2106.09685.pdf (section 7.2)
-            lora_alpha: alpha is needed for scaling updates as alpha/r
-                "This scaling helps to reduce the need to retune hyperparameters when we vary r"
-                https://arxiv.org/pdf/2106.09685.pdf (section 4.1)
-            lora_dropout: dropout that is applied on the input in the LoRA branch (before multiplying by matrix A)
-            enable_lora: MergeLinear class is for attention mechanism where qkv are calculated with a single weight matrix. If we
-                don't want to apply LoRA for all three (query, key and value) we can set it as False. For example if we want
-                to apply LoRA only to `query` and `value` but keep `key` without weight updates we should pass `[True,
-                False, True]`
-            fan_in_fan_out: set this to True if the layer to replace stores weight like (fan_in, fan_out).  For example, gpt-2 uses
-                `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`
-                https://github.com/huggingface/peft/blob/main/src/peft/tuners/lora.py#LL53C9-L53C112
-            merge_weights: whether we want to merge pretrained weights and LoRA weight updates. This is useful if one wants to use
-                fine-tuned model as a standalone one (without storing LoRA weight separately) plus it helps to reduce
-                overhead during inference.
-        """
-        nn.Linear.__init__(self, in_features, out_features, **kwargs)
-        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
-                           merge_weights=merge_weights)
-        assert out_features % len(enable_lora) == 0, \
-            'The length of enable_lora must divide out_features'
-        self.enable_lora = enable_lora
-        self.fan_in_fan_out = fan_in_fan_out
-
-        # Actual trainable parameters
-        # To better understand initialization let's imagine that we have such parameters:
-        # ⚬ in_features: 128 (embeddings_size)
-        # ⚬ out_features: 384 (3 * embedding_size)
-        # ⚬ r: 2
-        # ⚬ enable_lora: [True, False, True]
-        if r > 0 and any(enable_lora):
-            self.lora_A = nn.Parameter(
-                self.weight.new_zeros((r * sum(enable_lora), in_features)))  # (4, 128)
-            self.lora_B = nn.Parameter(
-                self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r))  # (256, 2)
-            ) # weights for Conv1D with groups=sum(enable_lora)
-            # Notes about shapes above
-            # - self.lora_A has shape (4, 128): 4 because rank is 2 and LoRA is applied only to two matrices;
-            # 128 is the input size of the x (embedding size). (4, 128) and not (128, 4) because later on in
-            # F.linear function weights are automatically transposed. In addition conv1d requires channels to
-            # be before seq length
-            # - self.lora_B has shape (256, 2): 256 because LoRA is applied only to two matrices, so the output is
-            # 128*2; 2 tells to have two channels per group for group convolution
-
-            # Scaling:
-            # This balances the pretrained model`s knowledge and the new task-specific adaptation
-            # https://lightning.ai/pages/community/tutorial/lora-llm/
-            # So, set alpha to 1.0 to fully add LoRA. If the LoRA seems to have too much effect (i.e., overfitted), set
-            # alpha to lower value. If the LoRA seems to have too little effect, set alpha to higher than 1.0. You can
-            # tune these values to your needs. This value can be even slightly greater than 1.0!
-            # https://github.com/cloneofsimo/lora
-            self.scaling = self.lora_alpha / self.r
-
-            # Freezing the pre-trained weight matrix
-            self.weight.requires_grad = False # (384, 128)
-
-            # Compute the indices
-            # Indices are needed to properly pad weight updates with zeros. If we want to fine-tune queries and values,
-            # but not keys, then the weights update should be:
-            #
-            # [[ΔW,ΔW,ΔW, ..., 0,0,0, ..., ΔW,ΔW,ΔW,],
-            #  [....................................],
-            #  [ΔW,ΔW,ΔW, ..., 0,0,0, ..., ΔW,ΔW,ΔW,]]
-            #      ↑              ↑            ↑
-            # ________________________________________
-            # | query         | key       | value    |
-            # ----------------------------------------
-            self.lora_ind = self.weight.new_zeros(
-                (out_features, ), dtype=torch.bool
-            ).view(len(enable_lora), -1)  # (3, 128)
-            self.lora_ind[enable_lora, :] = True  # (3, 128)
-            self.lora_ind = self.lora_ind.view(-1)  # (384,)
-        self.reset_parameters()
-        if fan_in_fan_out:
-            self.weight.data = self.weight.data.T
-
-    def reset_parameters(self):
-        """Reset all the weights, even including pretrained ones."""
-        nn.Linear.reset_parameters(self)
-        if hasattr(self, 'lora_A'):
-            # initialize A the same way as the default for nn.Linear and B to zero
-            # Wondering why 'a' is equal to math.sqrt(5)?: https://github.com/pytorch/pytorch/issues/15314
-            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            nn.init.zeros_(self.lora_B)
-
-    def zero_pad(self, x: torch.Tensor) -> torch.Tensor:
-        """Properly pad weight updates with zeros.
-
-        If, based on `self.enable_lora`, we want to fine-tune queries and values, but not keys,
-        then the weights update should be:
-
-        [[ΔW,ΔW,ΔW, ..., 0,0,0, ..., ΔW,ΔW,ΔW,],
-         [....................................],
-         [ΔW,ΔW,ΔW, ..., 0,0,0, ..., ΔW,ΔW,ΔW,]]
-            ↑              ↑            ↑
-        ________________________________________
-        | query         | key       | value    |
-        ----------------------------------------
-
-        Args:
-            x: tensor with weights update that will be padded with zeros if necessary
-
-        Returns:
-            A tensor with weight updates and zeros for deselected q, k or v
-        """
-        # Let's image that:
-        # ⚬ input x has shape (64, 64, 256): (batch_size, sequence_length, embeddings_size)
-        # ⚬ embeddings_size: 128
-        # ⚬ self.out_features: 384 (3 * embeddings_size)
-        # ⚬ enable_lora: [True, False, True]
-        # Then x has embeddings_size of 256 (2 * 128 as enable_lora only for query and value, not keys) and expected
-        # embeddings_size is 384 (self.out_features), so that means that we need to pad from 256 to 384 with zeros, but
-        # only for key updates (this is where self.lora_ind comes in handy)
-        # Note: double transpose (in the beginning and in the end) is basically a guard for two-dimensional tensors
-        # for example when we want to merge/unmerge LoRA weights and pretrained weights
-        x = x.transpose(0, 1)
-        result = x.new_zeros((*x.shape[:-1], self.out_features))  # (64, 64, 384)
-        result = result.view(-1, self.out_features)  # (4096, 384)
-        result[:, self.lora_ind] = x.reshape(
-            -1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)
-        )  # (4096, 256)
-        return result.view((*x.shape[:-1], self.out_features)).transpose(0, 1)  # (64, 64, 384)
-
-    def train(self, mode: bool = True):
-        """Set the module into train or eval mode if `mode` is True of False respectively.
-
-        For train mode (train(True)) if weights are merged we need to subtract weights updates (LoRA_A @ LoRA_B) from
-        pretrained weights so we can continue training LoRA's matrices A and B and keep pretrained weights frozen.
-
-        For eval mode (train(False)) if weights are not merged we need to add weight updates to pretrained weights in
-        order to reduce computational overhead during inference.
-
-        Args:
-            mode: if True the module will be set into train mode (affects Dropout and BatchNorm), if False - eval mode.
-
-        """
-        def T(w):
-            return w.T if self.fan_in_fan_out else w
-        # despite being called from nn.Linear this method will put all layers into train mode, including nn.Dropout
-        # of course except parameters (such as self.lora_A, self.lora_B)
-        nn.Linear.train(self, mode)
-
-        # if train(True) -> unmerge unless we already have them unmerged
-        # if train(False) -> merge unless we already have them merged
-        should = self.merged if mode else not self.merged
-
-        # Let's assume that:
-        # ⚬ self.weight.data: (384, 128) or (3 * embedding_size, embedding_size)
-        # ⚬ self.lora_A.data: (4, 128)
-        # ⚬ self.lora_B.data: (256, 2)
-        if self.merge_weights and should:
-            if self.r > 0 and any(self.enable_lora):
-                delta_w = F.conv1d(
-                    self.lora_A.data.unsqueeze(0),   # (4, 128) -> (1, 4, 128)
-                    self.lora_B.data.unsqueeze(-1),  # (256, 2) -> (256, 2, 1)
-                    groups=sum(self.enable_lora)
-                ).squeeze(0) # (1, 4, 128) @ (256, 2, 1) -> (1, 256, 128) -> (256, 128)
-                # -1: W = W - delta_W (unmerge), +1: W = W + delta_W (merge)
-                sign = -1 if mode else 1
-                self.weight.data += sign * self.zero_pad(T(delta_w * self.scaling)) # (256, 128) after zero_pad (384, 128)
-            self.merged = not mode
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Do the forward pass.
-
-        If LoRA's weights are merged with pretrained ones then it's a simple matrix multiplication.
-        If not, then multiply pretrained weights with input, apply LoRA on input and do summation.
-
-        Args:
-            x: input tensor of shape (batch_size, context_length, embedding_size)
-
-        Returns:
-            Output tensor of shape (batch_size, context_length, 3 * embedding_size)
-        """
-        def T(w):
-            return w.T if self.fan_in_fan_out else w
-
-        # Let's assume that:
-        # ⚬ x: (64, 64, 128) or (batch_size, context_length, embedding_size)
-        # ⚬ self.weight: (384, 128) or (3 * embedding_size, embedding_size)
-        # ⚬ self.lora_A.data: (4, 128)
-        # ⚬ self.lora_B.data: (256, 2)
-
-        # the logic here is that the weights are merged only during inference
-        # so if they are merged we don't need to do anything with LoRA's A and B matrices
-        # but if the weights are not merged that means that the forward method is called during
-        # training and we need to forward pass input through pretrained weights, LoRA A and B matrices
-        # and do the summation (as per scheme at the top of the file)
-        if self.merged:
-            return F.linear(x, T(self.weight), bias=self.bias)
-        else:
-            # `F.linear` automatically transposes the second argument (T(self.weight) in our case)
-            result = F.linear(x, T(self.weight), bias=self.bias)  # (64, 64, 128) @ (384, 128) -> (64, 64, 384)
-            if self.r > 0:
-                after_A = F.linear(self.lora_dropout(x), self.lora_A)  # (64, 64, 128) @ (4, 128) -> (64, 64, 4)
-                # For F.conv1d:
-                # ⚬ input: input tensor of shape (mini-batch, in_channels, iW)
-                # ⚬ weight: filters of shape (out_channels, in_channels/groups, kW)
-                # ⚬ groups: split input into groups, in_channels should be divisible by the number of groups. Default: 1
-                # presumably iW - sequence width/length, kW - kernel width
-                after_B = F.conv1d(
-                    after_A.transpose(-2, -1),  # (64, 64, 4) -> (64, 4, 64)
-                    self.lora_B.unsqueeze(-1),  # (256, 2) -> (256, 2, 1)
-                    groups=sum(self.enable_lora)
-                ).transpose(-2, -1)  # (64, 4, 64) @ (256, 2, 1) -> (64, 256, 64) -> (64, 64, 256)
-                result += self.zero_pad(after_B) * self.scaling  # (64, 64, 256) after zero_pad (64, 64, 384)
-            return result
-
-
-def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None:
-    """Freeze all modules except LoRA's and depending on 'bias' value unfreezes bias weights.
-
-    Args:
-        model: model with LoRA layers
-        bias: 
-            ``"none"``: all bias weights will be frozen,
-            ``"lora_only"``: only bias weight for LoRA layers will be unfrozen,
-            ``"all"``: all bias weights will be unfrozen.
-
-    Raises:
-        NotImplementedError: if `bias` not in ["none", "lora_only", "all"]
-    """
-    # freeze all layers except LoRA's
-    for n, p in model.named_parameters():
-        if 'lora_' not in n:
-            p.requires_grad = False
-
-    # depending on the `bias` value unfreeze bias weights
-    if bias == 'none':
-        return
-    elif bias == 'all':
-        for n, p in model.named_parameters():
-            if 'bias' in n:
-                p.requires_grad = True
-    elif bias == 'lora_only':
-        for m in model.modules():
-            if isinstance(m, LoRALayer) and \
-                hasattr(m, 'bias') and \
-                m.bias is not None:
-                    m.bias.requires_grad = True
-    else:
-        raise NotImplementedError
-
-
-def lora_state_dict(model: nn.Module, bias: str = 'none') -> Dict[str, torch.Tensor]:
-    """Return state_dict with weights of LoRA's A and B matrices and with biases depending on the `bias` value.
-
-    Args:
-        model: model with LoRA layers
-        bias: 
-            ``"none"``: state dict will not store bias weights,
-            ``"lora_only"``: state dict will store bias weights only from LoRA layers,
-            ``"all"``: state dict will store all bias weights.
-
-    Returns:
-        Weights and biases of LoRA layers
-
-    Raises:
-        NotImplementedError: if `bias` not in ["none", "lora_only", "all"]
-    """
-    my_state_dict = model.state_dict()
-    if bias == 'none':
-        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k}
-    elif bias == 'all':
-        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k or 'bias' in k}
-    elif bias == 'lora_only':
-        to_return = {}
-        for k in my_state_dict:
-            if 'lora_' in k:
-                to_return[k] = my_state_dict[k]
-                bias_name = k.split('lora_')[0]+'bias'
-                if bias_name in my_state_dict:
-                    to_return[bias_name] = my_state_dict[bias_name]
-        return to_return
-    else:
-        raise NotImplementedError
-
-
-@dataclass
-class LoRAConfig:
-    r: float = 0.0
-    alpha: float = 1.0
-    dropout: float = 0.0
-
-
-class CausalSelfAttention(llama.CausalSelfAttention):
-    lora_config = None
-
-    def __init__(self, config: llama.LLaMAConfig) -> None:
-        """Causal self-attention with calculating qkv matrices with a single matrix* and Low Ranking Adaptation for
-        parameter-efficient fine-tuning.
-
-        *Instead of creating multiple heads and concatenating the result (in addition to creating separate matrices for
-        query, key and value for each head) we can do this in a single pass with a single weight matrix.
-
-        Args:
-            config: 
-                ``"block_size"``: size of the context of the model,
-                ``"vocab_size"``: number of unique tokens,
-                ``"padded_vocab_size"``: padded size of the vocabulary to the nearest multiple of 64 (leads to a greater performance),
-                ``"n_layer"``: number of transformer blocks (self-attention + MLP),
-                ``"n_head"``: number of heads in multi-head attention mechanism,
-                ``"n_embd"``: size of the embedding: vector representation of each token.
-        """
-        # Skip the parent class __init__ altogether and replace it to avoid
-        # useless allocations
-        nn.Module.__init__(self)
-        assert config.n_embd % config.n_head == 0
-
-        # key, query, value projections for all heads, but in a batch
-        self.c_attn = MergedLinear(
-            in_features=config.n_embd,
-            out_features=3 * config.n_embd,
-            r=self.lora_config.r,
-            lora_alpha=self.lora_config.alpha,
-            lora_dropout=self.lora_config.dropout,
-            enable_lora=[True, False, True],
-            fan_in_fan_out = False,
-            merge_weights=True,
-            bias=False)
-        # output projection
-        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
-        # regularization
-        self.n_head = config.n_head
-        self.n_embd = config.n_embd
-        self.block_size = config.block_size
-        self.rope_cache = None
-
-
-@contextmanager
-def lora(r, alpha, dropout, enabled: bool = True):
-    """Apply context manager under which you can instantiate the model with LoRA.
-
-    In a nutshell the code inside this function forces to use LoRA variant of causal self-attention
-    instead of the original one (without LoRA).
-
-    Args:
-        r: rank of the weight update matrices. To make sense of using LoRA the rank should be smaller than the rank of
-            the weights of the model.  The rank can be as low as 1: https://arxiv.org/pdf/2106.09685.pdf (section 7.2)
-        alpha: alpha is needed for scaling updates as alpha/r
-            "This scaling helps to reduce the need to retune hyperparameters when we vary r"
-            https://arxiv.org/pdf/2106.09685.pdf (section 4.1)
-        dropout: dropout that is applied on the input in the LoRA branch (before multiplying by matrix A)
-        enabled: enables/disables LoRA
-    """
-    if not enabled:
-        yield
-        return
-
-    CausalSelfAttention.lora_config = LoRAConfig(r=r, alpha=alpha, dropout=dropout)
-    # when entering context manager replace link to causal self-attention class from original
-    # to a variant with LoRA
-    causal_self_attention = llama.CausalSelfAttention
-    llama.CausalSelfAttention = CausalSelfAttention
-    yield
-    # when exiting context manager - restore link to original causal self-attention class
-    llama.CausalSelfAttention = causal_self_attention
-
-    CausalSelfAttention.lora_config = None

lit-llama/lit_llama/model.py

@@ -1,321 +0,0 @@
-"""Full definition of a LLaMA Language Model, all of it in this single file.
-
-Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
-"""
-# mypy: ignore-errors
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from typing_extensions import Self
-
-from lit_llama.utils import find_multiple
-
-
-MaskCache = torch.Tensor
-RoPECache = torch.Tensor
-KVCache = Tuple[torch.Tensor, torch.Tensor]
-
-
-@dataclass
-class LLaMAConfig:
-    block_size: int = 2048
-    vocab_size: int = 32000
-    padded_vocab_size: Optional[int] = None
-    n_layer: int = 32
-    n_head: int = 32
-    n_embd: int = 4096
-
-    def __post_init__(self):
-        if self.padded_vocab_size is None:
-            self.padded_vocab_size = find_multiple(self.vocab_size, 64)
-
-    @classmethod
-    def from_name(cls, name: str) -> Self:
-        return cls(**llama_configs[name])
-
-
-llama_configs = {
-    "7B": dict(n_layer=32, n_head=32, n_embd=4096),
-    "13B": dict(n_layer=40, n_head=40, n_embd=5120),
-    "30B": dict(n_layer=60, n_head=52, n_embd=6656),
-    "65B": dict(n_layer=80, n_head=64, n_embd=8192),
-}
-
-
-class LLaMA(nn.Module):
-    def __init__(self, config: LLaMAConfig) -> None:
-        super().__init__()
-        assert config.padded_vocab_size is not None
-        self.config = config
-
-        self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
-        self.transformer = nn.ModuleDict(
-            dict(
-                wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
-                h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
-                ln_f=RMSNorm(config.n_embd),
-            )
-        )
-
-        self.rope_cache: Optional[RoPECache] = None
-        self.mask_cache: Optional[MaskCache] = None
-        self.kv_caches: List[KVCache] = []
-
-    def _init_weights(self, module: nn.Module) -> None:
-        if isinstance(module, nn.Linear):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(2 * self.config.n_layer))
-        elif isinstance(module, nn.Embedding):
-            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(2 * self.config.n_layer))
-
-    def forward(
-        self, idx: torch.Tensor, max_seq_length: Optional[int] = None, input_pos: Optional[torch.Tensor] = None
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[KVCache]]]:
-        B, T = idx.size()
-
-        block_size = self.config.block_size
-        if max_seq_length is None:
-            max_seq_length = block_size
-        assert T <= max_seq_length, f"Cannot forward sequence of length {T}, max seq length is only {max_seq_length}"
-        assert max_seq_length <= block_size, f"Cannot attend to {max_seq_length}, block size is only {block_size}"
-        assert T <= block_size, f"Cannot forward sequence of length {T}, block size is only {block_size}"
-
-        if self.rope_cache is None:
-            self.rope_cache = self.build_rope_cache(idx)
-        if self.mask_cache is None:
-            self.mask_cache = self.build_mask_cache(idx)
-
-        if input_pos is not None:
-            rope = self.rope_cache.index_select(0, input_pos)
-            mask = self.mask_cache.index_select(2, input_pos)
-            mask = mask[:, :, :, :max_seq_length]
-        else:
-            rope = self.rope_cache[:T]
-            mask = self.mask_cache[:, :, :T, :T]
-
-        # forward the model itself
-        x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
-
-        if input_pos is None:  # proxy for use_cache=False
-            for block in self.transformer.h:
-                x, _ = block(x, rope, mask, max_seq_length)
-        else:
-            if not self.kv_caches:
-                head_size = self.config.n_embd // self.config.n_head
-                cache_shape = (B, self.config.n_head, max_seq_length, head_size)
-                self.kv_caches = [
-                    (torch.zeros(cache_shape, device=x.device, dtype=x.dtype), torch.zeros(cache_shape, device=x.device, dtype=x.dtype))
-                    for _ in range(self.config.n_layer)
-                ]
-            for i, block in enumerate(self.transformer.h):
-                x, self.kv_caches[i] = block(x, rope, mask, max_seq_length, input_pos, self.kv_caches[i])
-
-        x = self.transformer.ln_f(x)
-
-        logits = self.lm_head(x)  # (b, t, vocab_size)
-
-        return logits
-
-    @classmethod
-    def from_name(cls, name: str) -> Self:
-        return cls(LLaMAConfig.from_name(name))
-
-    def build_rope_cache(self, idx: torch.Tensor) -> RoPECache:
-        return build_rope_cache(
-            seq_len=self.config.block_size,
-            n_elem=self.config.n_embd // self.config.n_head,
-            dtype=idx.dtype,
-            device=idx.device,
-        )
-
-    def build_mask_cache(self, idx: torch.Tensor) -> MaskCache:
-        ones = torch.ones((self.config.block_size, self.config.block_size), device=idx.device, dtype=torch.bool)
-        return torch.tril(ones).unsqueeze(0).unsqueeze(0)
-
-    def reset_cache(self) -> None:
-        self.kv_caches.clear()
-        if self.mask_cache.device.type == "xla":
-            # https://github.com/Lightning-AI/lit-parrot/pull/83#issuecomment-1558150179
-            self.rope_cache = None
-            self.mask_cache = None
-
-
-class Block(nn.Module):
-    def __init__(self, config: LLaMAConfig) -> None:
-        super().__init__()
-        self.rms_1 = RMSNorm(config.n_embd)
-        self.attn = CausalSelfAttention(config)
-        self.rms_2 = RMSNorm(config.n_embd)
-        self.mlp = MLP(config)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope: RoPECache,
-        mask: MaskCache,
-        max_seq_length: int,
-        input_pos: Optional[torch.Tensor] = None,
-        kv_cache: Optional[KVCache] = None,
-    ) -> Tuple[torch.Tensor, Optional[KVCache]]:
-        h, new_kv_cache = self.attn(self.rms_1(x), rope, mask, max_seq_length, input_pos, kv_cache)
-        x = x + h
-        x = x + self.mlp(self.rms_2(x))
-        return x, new_kv_cache
-
-
-class CausalSelfAttention(nn.Module):
-    def __init__(self, config: LLaMAConfig) -> None:
-        super().__init__()
-        assert config.n_embd % config.n_head == 0
-
-        # key, query, value projections for all heads, but in a batch
-        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
-        # output projection
-        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
-
-        self.n_head = config.n_head
-        self.n_embd = config.n_embd
-        self.block_size = config.block_size
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        rope: RoPECache,
-        mask: MaskCache,
-        max_seq_length: int,
-        input_pos: Optional[torch.Tensor] = None,
-        kv_cache: Optional[KVCache] = None,
-    ) -> Tuple[torch.Tensor, Optional[KVCache]]:
-        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
-
-        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
-        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
-
-        head_size = C // self.n_head
-        k = k.view(B, T, self.n_head, head_size)
-        q = q.view(B, T, self.n_head, head_size)
-        v = v.view(B, T, self.n_head, head_size)
-
-        q = apply_rope(q, rope)
-        k = apply_rope(k, rope)
-
-        k = k.transpose(1, 2)  # (B, nh, T, hs)
-        q = q.transpose(1, 2)  # (B, nh, T, hs)
-        v = v.transpose(1, 2)  # (B, nh, T, hs)
-
-        if kv_cache is not None:
-            cache_k, cache_v = kv_cache
-            # check if reached token limit
-            if input_pos[-1] >= max_seq_length:
-                input_pos = torch.tensor(max_seq_length - 1, device=input_pos.device)
-                # shift 1 position to the left
-                cache_k = torch.roll(cache_k, -1, dims=2)
-                cache_v = torch.roll(cache_v, -1, dims=2)
-            k = cache_k.index_copy(2, input_pos, k)
-            v = cache_v.index_copy(2, input_pos, v)
-            kv_cache = k, v
-
-        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
-        #  att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
-        #  att = att.masked_fill(mask[:,:,:T,:T] == 0, float('-inf'))
-        #  att = F.softmax(att, dim=-1)
-        #  y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
-
-        # efficient attention using Flash Attention CUDA kernels
-        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
-
-        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
-
-        # output projection
-        y = self.c_proj(y)
-
-        return y, kv_cache
-
-
-class MLP(nn.Module):
-    def __init__(self, config: LLaMAConfig) -> None:
-        super().__init__()
-        hidden_dim = 4 * config.n_embd
-        n_hidden = int(2 * hidden_dim / 3)
-        n_hidden = find_multiple(n_hidden, 256)
-
-        self.c_fc1 = nn.Linear(config.n_embd, n_hidden, bias=False)
-        self.c_fc2 = nn.Linear(config.n_embd, n_hidden, bias=False)
-        self.c_proj = nn.Linear(n_hidden, config.n_embd, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
-        x = self.c_proj(x)
-        return x
-
-
-class RMSNorm(nn.Module):
-    """Root Mean Square Layer Normalization.
-
-    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
-    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
-    """
-
-    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:
-        super().__init__()
-        self.scale = nn.Parameter(torch.ones(size))
-        self.eps = eps
-        self.dim = dim
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # NOTE: the original RMSNorm paper implementation is not equivalent
-        # norm_x = x.norm(2, dim=self.dim, keepdim=True)
-        # rms_x = norm_x * d_x ** (-1. / 2)
-        # x_normed = x / (rms_x + self.eps)
-        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
-        x_normed = x * torch.rsqrt(norm_x + self.eps)
-        return self.scale * x_normed
-
-
-def build_rope_cache(
-    seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
-) -> RoPECache:
-    """Enhanced Transformer with Rotary Position Embedding.
-
-    Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
-    transformers/rope/__init__.py. MIT License:
-    https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
-    """
-    # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
-    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))
-
-    # Create position indexes `[0, 1, ..., seq_len - 1]`
-    seq_idx = torch.arange(seq_len, dtype=dtype, device=device)
-
-    # Calculate the product of position index and $\theta_i$
-    idx_theta = torch.outer(seq_idx, theta).float()
-
-    cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
-
-    # this is to mimic the behaviour of complex32, else we will get different results
-    if dtype in (torch.float16, torch.bfloat16, torch.int8):
-        cache = cache.half()
-    return cache
-
-
-def apply_rope(x: torch.Tensor, rope_cache: RoPECache) -> torch.Tensor:
-    # truncate to support variable sizes
-    T = x.size(1)
-    rope_cache = rope_cache[:T]
-
-    # cast because the reference does
-    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
-    rope_cache = rope_cache.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
-    x_out2 = torch.stack(
-        [
-            xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
-            xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
-        ],
-        -1,
-    )
-
-    x_out2 = x_out2.flatten(3)
-    return x_out2.type_as(x)

lit-llama/lit_llama/packed_dataset.py

@@ -1,260 +0,0 @@
-# Very loosely inspired by indexed_dataset in Fairseq, Megatron
-# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py
-
-
-import os
-import struct
-import random
-
-import numpy as np
-import torch
-from torch.utils.data import IterableDataset, get_worker_info
-
-
-dtypes = {
-    1: np.uint8,
-    2: np.int8,
-    3: np.int16,
-    4: np.int32,
-    5: np.int64,
-    6: np.float32,
-    7: np.float64,
-    8: np.uint16,
-}
-
-
-def code(dtype):
-    for k in dtypes.keys():
-        if dtypes[k] == dtype:
-            return k
-    raise ValueError(dtype)
-
-
-HDR_MAGIC = b"LITPKDS"
-HDR_SIZE = 24  # bytes
-
-
-class PackedDataset(IterableDataset):
-    def __init__(self, filenames, n_chunks, block_size, seed=12345, shuffle=True, wrap=False, num_processes=1, process_rank=0):
-        self._filenames = filenames
-        self._n_chunks = n_chunks
-        self._block_size = block_size
-        self._seed = seed
-        self._shuffle = shuffle
-        self._wrap = wrap
-        self._num_processes = num_processes
-        self._process_rank = process_rank
-
-    def __iter__(self):
-        worker_info = get_worker_info()
-        num_workers = worker_info.num_workers if worker_info is not None else 1
-        worker_id = worker_info.id if worker_info is not None else 0
-        num_shards = num_workers * self._num_processes
-        shard_id = self._process_rank * num_workers + worker_id
-
-        max_num_files = len(self._filenames) // num_shards * num_shards
-        filenames = self._filenames[shard_id : max_num_files : num_shards]
-
-        return PackedDatasetIterator(
-            filenames=filenames,
-            n_chunks=self._n_chunks,
-            block_size=self._block_size,
-            seed=self._seed,
-            shuffle=self._shuffle,
-            wrap=self._wrap,
-        )
-
-
-class PackedDatasetBuilder(object):
-    def __init__(
-        self,
-        outdir,
-        prefix,
-        chunk_size,
-        sep_token,
-        dtype="auto",
-        vocab_size=None,
-    ):
-        if dtype == "auto":
-            if vocab_size is None:
-                raise ValueError("vocab_size cannot be None when dtype='auto'")
-            if vocab_size is not None and vocab_size < 65500:
-                self._dtype = np.uint16
-            else:
-                self._dtype = np.int32
-        else:
-            self._dtype = dtype
-        self._counter = 0
-        self._chunk_size = chunk_size
-        self._outdir = outdir
-        self._prefix = prefix
-        self._sep_token = sep_token
-        self._arr = np.zeros(self._chunk_size, dtype=self._dtype)
-        self._arr.fill(self._sep_token)
-        self._idx = 0
-        self._version = 1
-        self._filenames = []
-
-    def _write_chunk(self):
-        filename = f"{self._prefix}_{self._counter:010d}.bin"
-        filename = os.path.join(self._outdir, filename)
-
-        with open(filename, "wb") as f:
-            f.write(HDR_MAGIC)
-            f.write(struct.pack("<Q", self._version))
-            f.write(struct.pack("<B", code(self._dtype)))
-            f.write(struct.pack("<Q", self._chunk_size))
-            f.write(self._arr.tobytes(order="C"))
-
-        self._filenames.append(filename)
-        self._counter += 1
-        self._arr.fill(self._sep_token)
-        self._idx = 0
-
-    @property
-    def dtype(self):
-        return self._dtype
-
-    @property
-    def filenames(self):
-        return self._filenames.copy()
-
-    def add_array(self, arr):
-        while self._idx + arr.shape[0] > self._chunk_size:
-            part_len = self._chunk_size - self._idx
-            self._arr[self._idx : self._idx + part_len] = arr[:part_len]
-            self._write_chunk()
-            arr = arr[part_len:]
-
-        arr_len = arr.shape[0]
-        self._arr[self._idx : self._idx + arr_len] = arr
-        self._idx += arr_len
-
-    def write_reminder(self):
-        self._write_chunk()
-
-
-class PackedDatasetIterator:
-    def __init__(self, filenames, n_chunks, block_size, seed, shuffle, wrap):
-        self._seed = seed
-        self._shuffle = shuffle
-        self._rng = np.random.default_rng(seed) if shuffle else None
-        self._block_idxs = None
-
-        self._wrap = wrap
-
-        # TODO: instead of filenames, we could have a single text stream
-        #       (or text file) with the sequence of all files to be
-        #       fetched/loaded.
-        self._filenames = filenames
-        self._file_idx = 0
-
-        self._n_chunks = n_chunks
-
-        self._dtype = None
-        self._block_size = block_size
-        self._n_blocks = None
-
-        self._mmaps = []
-        self._buffers = []
-
-        self._block_idxs = []
-        self._curr_idx = 0
-
-        self._load_n_chunks()
-
-    def _read_header(self, path):
-        with open(path, "rb") as f:
-            magic = f.read(len(HDR_MAGIC))
-            assert magic == HDR_MAGIC, "File doesn't match expected format."
-            version = struct.unpack("<Q", f.read(8))
-            assert (1,) == version
-            (dtype_code,) = struct.unpack("<B", f.read(1))
-            dtype = dtypes[dtype_code]
-            (chunk_size,) = struct.unpack("<Q", f.read(8))
-        return dtype, chunk_size
-
-    def _close_mmaps(self):
-        for mmap in self._mmaps:
-            mmap._mmap.close()
-
-    def _load_n_chunks(self):
-        self._close_mmaps()
-        self._mmaps = []
-        self._buffers = []
-
-        if self._n_chunks > len(self._filenames[self._file_idx:]):
-            if not self._wrap:
-                raise StopIteration
-            else:
-                self._file_idx = 0
-
-        for i in range(self._n_chunks):
-            filename = self._filenames[self._file_idx + i]
-            if self._dtype is None:
-                self._dtype, self._chunk_size = self._read_header(
-                    filename
-                )
-                self._n_blocks = self._chunk_size // self._block_size
-            # TODO: check header matches with previous files
-            mmap = np.memmap(filename, mode="r", order="C", offset=HDR_SIZE)
-            self._mmaps.append(mmap)
-            self._buffers.append(memoryview(mmap))
-
-        self._file_idx += self._n_chunks
-        n_all_blocks = self._n_chunks * self._n_blocks
-
-        self._block_idxs = (
-            self._rng.permutation(n_all_blocks)
-            if self._shuffle
-            else range(n_all_blocks)
-        )
-
-        self._curr_idx = 0
-
-    def __del__(self):
-        self._close_mmaps()
-        del self._mmaps
-        del self._buffers
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        if self._curr_idx >= len(self._block_idxs):
-            self._load_n_chunks()
-            # TODO: trigger fetching next next n_chunks if remote
-        block_idx = self._block_idxs[self._curr_idx]
-        chunk_id = block_idx // self._n_blocks
-        buffer = self._buffers[chunk_id]
-        elem_id = (block_idx % self._n_blocks) * self._block_size
-        offset = np.dtype(self._dtype).itemsize * elem_id
-        arr = np.frombuffer(
-            buffer, dtype=self._dtype, count=self._block_size, offset=offset
-        )
-        self._curr_idx += 1
-        return torch.from_numpy(arr.astype(np.int64))
-
-
-class CombinedDataset(IterableDataset):
-    def __init__(self, datasets, seed, weights=None):
-        self._seed = seed
-        self._datasets = datasets
-        self._weights = weights
-        n_datasets = len(datasets)
-        if weights is None:
-            self._weights = [1 / n_datasets] * n_datasets
-
-    def __iter__(self):
-        return CombinedDatasetIterator(self._datasets, self._seed, self._weights)
-
-
-class CombinedDatasetIterator:
-    def __init__(self, datasets, seed, weights):
-        self._datasets = [iter(el) for el in datasets]
-        self._weights = weights
-        self._rng = random.Random(seed)
-
-    def __next__(self):
-        dataset, = self._rng.choices(self._datasets, weights=self._weights, k=1)
-        return next(dataset)

lit-llama/lit_llama/quantization.py

@@ -1,614 +0,0 @@
-import os
-from contextlib import contextmanager
-import warnings
-import math
-
-import torch
-
-# configuration for bitsandbytes before import
-os.environ["BITSANDBYTES_NOWELCOME"] = "1"
-warnings.filterwarnings(
-    "ignore",
-    message="MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization",
-)
-warnings.filterwarnings(
-    "ignore",
-    message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization",
-)
-warnings.filterwarnings(
-    "ignore",
-    message="The installed version of bitsandbytes was compiled without GPU support. 8-bit optimizers and GPU quantization are unavailable.",
-)
-
-try:
-    import bitsandbytes as bnb  # noqa: E402
-except:
-    bnb = None
-
-try:
-    import triton  # noqa: E402
-    import triton.language as tl  # noqa: E402
-except:
-    triton = None
-
-if bnb is not None:
-
-    class Linear8bitLt(bnb.nn.Linear8bitLt):
-        """Wraps `bnb.nn.Linear8bitLt` and enables instantiation directly on the device and
-        re-quantizaton when loading the state dict.
-
-
-        This should only be used for inference. For training, use `bnb.nn.Linear8bitLt` directly.
-        """
-
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs, has_fp16_weights=False, threshold=6.0)
-            # We quantize the initial weight here so we don't end up filling the device
-            # memory with float32 weights which could lead to OOM.
-            self._quantize_weight(self.weight.data)
-
-        def _load_from_state_dict(self, local_state_dict, *args, **kwargs):
-            # There is only one key that ends with `*.weight`, the other one is the bias
-            weight_key = next(
-                (name for name in local_state_dict.keys() if name.endswith("weight")),
-                None,
-            )
-            if weight_key is None:
-                return
-
-            # Load the weight from the state dict and re-quantize it
-            weight = local_state_dict.pop(weight_key)
-            self._quantize_weight(weight)
-
-            # If there is a bias, let nn.Module load it
-            if local_state_dict:
-                super()._load_from_state_dict(local_state_dict, *args, **kwargs)
-
-        def _quantize_weight(self, weight: torch.Tensor) -> None:
-            # This code is taken and adapted from `bnb.nn.Int8Params.cuda()`
-            B = weight.contiguous().half().cuda()
-            CB, CBt, SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B)
-            del CBt
-            del SCBt
-            self.weight.data = CB
-            setattr(self.weight, "CB", CB)
-            setattr(self.weight, "SCB", SCB)
-
-
-if triton is not None:
-    # This is adapted from the OpenAI Triton matmul example.
-    @triton.autotune(
-        configs=[
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 128,
-                    "BLOCK_SIZE_N": 256,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=3,
-                num_warps=8,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 256,
-                    "BLOCK_SIZE_N": 128,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=3,
-                num_warps=8,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 256,
-                    "BLOCK_SIZE_N": 64,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=4,
-                num_warps=4,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 64,
-                    "BLOCK_SIZE_N": 256,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=4,
-                num_warps=4,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 128,
-                    "BLOCK_SIZE_N": 128,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=4,
-                num_warps=4,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 128,
-                    "BLOCK_SIZE_N": 64,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=4,
-                num_warps=4,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 64,
-                    "BLOCK_SIZE_N": 128,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=4,
-                num_warps=4,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 128,
-                    "BLOCK_SIZE_N": 32,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=4,
-                num_warps=4,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 64,
-                    "BLOCK_SIZE_N": 32,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=5,
-                num_warps=2,
-            ),
-            triton.Config(
-                {
-                    "BLOCK_SIZE_M": 32,
-                    "BLOCK_SIZE_N": 64,
-                    "BLOCK_SIZE_K": 32,
-                    "GROUP_SIZE_M": 8,
-                },
-                num_stages=5,
-                num_warps=2,
-            ),
-        ],
-        key=["M", "N", "K"],
-    )
-    @triton.jit
-    def linear_kernel_4bit_weight(
-        # Pointers to matrices
-        a_ptr,
-        b_ptr,
-        c_ptr,
-        bscales_ptr,
-        bzeros_ptr,
-        # bdequant,
-        # Matrix dimensions
-        M,
-        N,
-        K,
-        # The stride variables represent how much to increase the ptr by when moving by 1
-        # element in a particular dimension. E.g. stride_am is how much to increase a_ptr
-        # by to get the element one row down (A has M rows)
-        stride_am,
-        stride_ak,
-        stride_bk,
-        stride_bn,
-        stride_cm,
-        stride_cn,
-        # Meta-parameters
-        BLOCK_SIZE_M: tl.constexpr,
-        BLOCK_SIZE_N: tl.constexpr,
-        BLOCK_SIZE_K: tl.constexpr,
-        GROUP_SIZE_M: tl.constexpr,
-    ):
-        """Kernel for computing the matmul C = A x B.T.
-        A has shape (M, K), B has shape (N, K) and C has shape (M, N)
-        """
-        # -----------------------------------------------------------
-        # Map program ids `pid` to the block of C it should compute.
-        # This is done in a grouped ordering to promote L2 data reuse
-        # See above `L2 Cache Optimizations` section for details
-        pid = tl.program_id(axis=0)
-        num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-        num_pid_in_group = GROUP_SIZE_M * num_pid_n
-        group_id = pid // num_pid_in_group
-        first_pid_m = group_id * GROUP_SIZE_M
-        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
-        pid_m = first_pid_m + (pid % group_size_m)
-        pid_n = (pid % num_pid_in_group) // group_size_m
-
-        # ----------------------------------------------------------
-        # Create pointers for the first blocks of A and B.
-        # We will advance this pointer as we move in the K direction
-        # and accumulate
-        # a_ptrs is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
-        # b_ptrs is a block of [BLOCK_SIZE_K, BLOCK_SIZE_n] pointers
-        # see above `Pointer Arithmetics` section for details
-        offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-        offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-        a_mask = offs_am[:, None] < M
-        b_mask = offs_bn[None, :] < N
-        offs_k = tl.arange(0, BLOCK_SIZE_K)
-        a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
-        b_ptrs = b_ptr + (
-            (offs_k[:, None] // 2) * stride_bk + offs_bn[None, :] * stride_bn
-        )
-
-        bscales_ptrs = bscales_ptr + offs_bn[None, :]
-        bzeros_ptrs = bzeros_ptr + offs_bn[None, :]
-
-        scale = tl.load(bscales_ptrs)
-        zero = tl.load(bzeros_ptrs)
-        # -----------------------------------------------------------
-        # Iterate to compute a block of the C matrix
-        # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
-        # of fp32 values for higher accuracy.
-        # `accumulator` will be converted back to fp16 after the loop
-        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-        for k in range(0, K, BLOCK_SIZE_K):
-            # wasteful as it is to load everything twice, my attempts at avoiding it lead to slower code
-            b12 = tl.load(b_ptrs, mask=b_mask)
-            # Note that for simplicity, we don't apply a mask in K here.
-            a = tl.load(a_ptrs, mask=a_mask).to(tl.float32)
-            b = (
-                ((b12.to(tl.uint8) >> ((offs_k[:, None] % 2) * 4)) & 0xF).to(tl.float32)
-                - zero
-            ) * scale
-            accumulator += tl.dot(a, b)
-
-            # Advance the ptrs to the next K block
-            a_ptrs += BLOCK_SIZE_K * stride_ak
-            b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk
-        c = accumulator
-
-        # -----------------------------------------------------------
-        # Write back the block of the output matrix C
-        offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-        offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-        c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-        tl.store(c_ptrs, c, mask=c_mask)
-
-    def qlinear_4bit_weight(inp, weight, scales, zeros):
-        weight = weight.t().contiguous()
-        c_shape = inp.shape[:-1] + weight.shape[-1:]
-        inp = inp.reshape(-1, inp.shape[-1]).contiguous()
-        # we pad the input to amortize triton compilation cost better
-        PAD_TO = 256
-        if inp.shape[0] % PAD_TO != 0:
-            c_crop = inp.shape[0]
-            new_inp_shape0 = inp.shape[0] + PAD_TO - inp.shape[0] % PAD_TO
-            inp2 = inp.new_empty((new_inp_shape0, inp.shape[1]))
-            inp2[: inp.shape[0]] = inp
-            inp2[inp.shape[0] :].zero_()
-            inp = inp2
-        else:
-            c_crop = None
-
-        assert inp.shape[1] == weight.shape[0] * 2, "incompatible dimensions"
-
-        assert scales.shape == (weight.shape[1], 1)
-        assert zeros.shape == (weight.shape[1], 1)
-        scales = scales.contiguous()
-        zeros = zeros.contiguous()
-        K, N = weight.shape
-        M, K = inp.shape
-        assert (
-            K % 32 == 0
-        ), "We don't check memory-out-of-bounds with K so K must be divisible by BLOCK_SIZE_K"
-        # allocates output
-        c = torch.empty((M, N), device=inp.device, dtype=inp.dtype)
-        # 1D launch kernel where each block gets its own program.
-        grid = lambda META: (
-            triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
-        )
-        linear_kernel_4bit_weight[grid](
-            inp,
-            weight,
-            c,
-            scales,
-            zeros,
-            M,
-            N,
-            K,
-            inp.stride(0),
-            inp.stride(1),
-            weight.stride(0),
-            weight.stride(1),
-            c.stride(0),
-            c.stride(1),
-        )
-        return c[:c_crop].reshape(c_shape)
-
-else:
-    qlinear_4bit_weight = None
-
-
-# for correctness but with terrible perf
-class ColBlockQuantizedLinear(torch.nn.Module):
-    def __init__(self, in_features, out_features, bias: bool, *, bits, tile_cols):
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.tile_cols = tile_cols if tile_cols != -1 else self.in_features
-        self.bits = bits
-        self.entries_per_byte = 8 // bits
-        assert self.entries_per_byte > 0 and self.entries_per_byte * self.bits == 8
-        assert in_features % self.entries_per_byte == 0
-        self.register_buffer(
-            "quant_weight",
-            torch.empty(
-                (self.out_features, self.in_features // self.entries_per_byte),
-                dtype=torch.uint8,
-            )
-            .t()
-            .contiguous()
-            .t(),
-        )
-        self.register_buffer(
-            "scales",
-            torch.empty(
-                (
-                    self.out_features,
-                    (self.in_features + self.tile_cols - 1) // self.tile_cols,
-                )
-            ),
-        )
-        self.register_buffer("zeros", torch.empty_like(self.scales))
-        assert isinstance(bias, bool)
-        if bias:
-            self.register_buffer("bias", torch.empty((self.out_features,)))
-        else:
-            self.register_buffer("bias", None)
-
-    def pack_weight(self, weight):
-        weight = weight.to(device=self.quant_weight.device, copy=True)
-        for j in range(self.scales.size(1)):
-            weight[:, j * self.tile_cols : (j + 1) * self.tile_cols] /= self.scales[
-                :, j : j + 1
-            ]
-            weight[:, j * self.tile_cols : (j + 1) * self.tile_cols] += self.zeros[
-                :, j : j + 1
-            ]
-        weight = weight.clamp_(min=0, max=2**self.bits - 1).to(dtype=torch.uint8)
-        self.quant_weight.zero_()
-        for nr in range(self.entries_per_byte):
-            self.quant_weight += weight[:, nr :: self.entries_per_byte] << (
-                nr * self.bits
-            )
-
-    def get_weight(self, dtype=torch.float):
-        weight = torch.empty(
-            (self.out_features, self.in_features),
-            device=self.quant_weight.device,
-            dtype=dtype,
-        )
-        mask = (1 << self.bits) - 1
-        for nr in range(self.entries_per_byte):
-            weight[:, nr :: self.entries_per_byte] = (
-                (self.quant_weight >> (nr * self.bits)) & mask
-            ).float()
-        self.quant_weight.to(dtype)
-        for j in range(self.scales.size(1)):
-            weight[:, j * self.tile_cols : (j + 1) * self.tile_cols] -= self.zeros[
-                :, j : j + 1
-            ]
-            weight[:, j * self.tile_cols : (j + 1) * self.tile_cols] *= self.scales[
-                :, j : j + 1
-            ]
-        return weight
-
-    def forward(self, inp):
-        if (
-            triton is not None
-            and self.bits == 4
-            and self.quant_weight.device.type == "cuda"
-            and self.zeros.shape[1] == 1
-            and self.quant_weight.shape[1] % 32 == 0
-        ):
-            return qlinear_4bit_weight(inp, self.quant_weight, self.scales, self.zeros)
-        weight = self.get_weight(dtype=inp.dtype)
-        return torch.nn.functional.linear(inp, weight, self.bias)
-
-
-class GPTQQuantizer:
-    # The algorithm and code has been taken from  https://github.com/IST-DASLab/gptq/
-    # E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
-    # portions copyright by the authors licensed under the Apache License 2.0
-    # All errors are our own.
-
-    def __init__(
-        self,
-        linear_module,
-        *,
-        bits,
-        perchannel=True,
-        sym=False,
-        blocksize=128,
-        percdamp=0.01,
-        groupsize=-1,
-        actorder=False
-    ):
-        assert isinstance(linear_module, torch.nn.Linear)
-
-        self.linear_module = linear_module
-        self.dev = self.linear_module.weight.device
-        self.rows = linear_module.weight.shape[0]
-        self.columns = linear_module.weight.shape[1]
-        self.H = torch.zeros((self.columns, self.columns), device=self.dev)
-        self.nsamples = 0
-        self.bits = bits
-        self.maxq = 2**bits - 1
-        self.perchannel = perchannel
-        self.sym = sym
-        self.blocksize = blocksize
-        self.percdamp = percdamp
-        self.groupsize = groupsize
-        self.actorder = actorder
-        self.tile_cols = self.columns if groupsize == -1 else groupsize
-        self.scales = torch.zeros(
-            (self.rows, (self.columns + self.tile_cols - 1) // self.tile_cols),
-            dtype=self.linear_module.weight.dtype,
-            device=self.dev,
-        )
-        self.zeros = torch.zeros_like(self.scales)
-        assert not (
-            self.actorder and self.groupsize != -1
-        ), "The permutation trick does not work for grouped quantization"
-
-    @staticmethod
-    def quantize_weight(x, scale, zero, maxq):
-        q = torch.clamp(torch.round(x / scale) + zero, 0, maxq)
-        x_rec = scale * (q - zero)
-        return x_rec
-
-    def find_params_weight(self, x):
-        dev = x.device
-
-        shape = x.shape
-        if self.perchannel:
-            x = x.flatten(1)
-        else:
-            x = x.flatten().unsqueeze(0)
-
-        tmp = torch.zeros(x.shape[0], device=dev)
-        xmin = torch.minimum(x.min(1)[0], tmp)
-        xmax = torch.maximum(x.max(1)[0], tmp)
-
-        if self.sym:
-            xmax = torch.maximum(torch.abs(xmin), xmax)
-            tmp = xmin < 0
-            if torch.any(tmp):
-                xmin[tmp] = -xmax[tmp]
-        tmp = (xmin == 0) & (xmax == 0)
-        xmin[tmp] = -1
-        xmax[tmp] = +1
-
-        scale = (xmax - xmin) / self.maxq
-        if self.sym:
-            zero = torch.full_like(scale, (self.maxq + 1) / 2)
-        else:
-            zero = torch.round(-xmin / scale)
-
-        if not self.perchannel:
-            tmp = shape[0]
-            scale = scale.repeat(tmp)
-            zero = zero.repeat(tmp)
-
-        shape = [-1] + [1] * (len(shape) - 1)
-        scale = scale.reshape(shape)
-        zero = zero.reshape(shape)
-        return scale, zero
-
-    def collect_input_stats(self, _1, inp, _2):
-        inp = inp[0].detach()
-        self.last_inp = inp
-        if len(inp.shape) == 2:
-            inp = inp.unsqueeze(0)
-        tmp = inp.shape[0]
-        if len(inp.shape) == 3:
-            inp = inp.reshape((-1, inp.shape[-1]))
-        inp = inp.t()
-        self.H *= self.nsamples / (self.nsamples + tmp)
-        self.nsamples += tmp
-        # inp = inp.float()
-        inp = math.sqrt(2 / self.nsamples) * inp.float()
-        # self.H += 2 / self.nsamples * inp.matmul(inp.t())
-        self.H += inp.matmul(inp.t())
-
-    def quantize(self):
-        W = self.linear_module.weight.detach().to(dtype=torch.float, copy=True)
-
-        scale, zero = self.find_params_weight(W)
-        self.scales[:] = scale
-        self.zeros[:] = zero
-
-        H = self.H
-        del self.H
-        dead = torch.diag(H) == 0
-        H[dead, dead] = 1
-        W[:, dead] = 0
-        if self.actorder:
-            perm = torch.argsort(torch.diag(H), descending=True)
-            W = W[:, perm]
-            H = H[perm][:, perm]
-
-        Losses = torch.zeros_like(W)
-        Q = torch.zeros_like(W)
-
-        damp = self.percdamp * torch.mean(torch.diag(H))
-        diag = torch.arange(self.columns, device=self.dev)
-        H[diag, diag] += damp
-        H = torch.linalg.cholesky(H)
-        H = torch.cholesky_inverse(H)
-        H = torch.linalg.cholesky(H, upper=True)
-        Hinv = H
-
-        for i1 in range(0, self.columns, self.blocksize):
-            i2 = min(i1 + self.blocksize, self.columns)
-            count = i2 - i1
-
-            W1 = W[:, i1:i2].clone()
-            Q1 = torch.zeros_like(W1)
-            Err1 = torch.zeros_like(W1)
-            Losses1 = torch.zeros_like(W1)
-            Hinv1 = Hinv[i1:i2, i1:i2]
-
-            for i in range(count):
-                w = W1[:, i]
-                d = Hinv1[i, i]
-
-                if self.groupsize != -1:
-                    if (i1 + i) % self.groupsize == 0:
-                        scale, zero = self.find_params_weight(
-                            W[:, (i1 + i) : (i1 + i + self.groupsize)]
-                        )
-                        self.scales[:, (i1 + i) // self.groupsize] = scale
-                        self.zeros[:, (i1 + i) // self.groupsize] = zero
-
-                q = self.quantize_weight(w.unsqueeze(1), scale, zero, self.maxq)
-                q = q.squeeze(1)
-                assert q.dim() == 1
-                Q1[:, i] = q
-                Losses1[:, i] = (w - q) ** 2 / d**2
-
-                err1 = (w - q) / d
-                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
-                Err1[:, i] = err1
-
-            Q[:, i1:i2] = Q1
-            Losses[:, i1:i2] = Losses1 / 2
-
-            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
-
-        if self.actorder:
-            invperm = torch.argsort(perm)
-            Q = Q[:, invperm]
-
-        weight = Q.reshape(self.linear_module.weight.shape).to(
-            self.linear_module.weight.data.dtype
-        )
-        error = torch.sum(Losses).item()
-
-        q_module = ColBlockQuantizedLinear(
-            self.linear_module.in_features,
-            self.linear_module.out_features,
-            self.linear_module.bias is not None,
-            bits=self.bits,
-            tile_cols=self.groupsize,
-        ).to(self.dev)
-        q_module.scales = self.scales
-        q_module.zeros = self.zeros
-        q_module.pack_weight(weight)
-        q_module.bias = self.linear_module.bias
-        return q_module, error

lit-llama/lit_llama/tokenizer.py

@@ -1,49 +0,0 @@
-import os
-from pathlib import Path
-from typing import Optional
-
-import torch
-from sentencepiece import SentencePieceProcessor, SentencePieceTrainer
-
-
-class Tokenizer:
-    """Tokenizer for LLaMA."""
-
-    def __init__(self, model_path: Path) -> None:
-        self.processor = SentencePieceProcessor(model_file=str(model_path))
-        self.bos_id = self.processor.bos_id()
-        self.eos_id = self.processor.eos_id()
-        self.pad_id = self.processor.pad_id()
-
-    @property
-    def vocab_size(self) -> int:
-        return self.processor.vocab_size()
-
-    def encode(
-        self,
-        string: str,
-        bos: bool = True,
-        eos: bool = False,
-        max_length: int = -1,
-        pad: bool = False,
-        device: Optional[torch.device] = None
-    ) -> torch.Tensor:
-        tokens = self.processor.encode(string)
-        if bos:
-            tokens = [self.bos_id] + tokens
-        if eos:
-            tokens = tokens + [self.eos_id]
-        if max_length > 0:
-            tokens = tokens[:max_length]
-        if pad and len(tokens) < max_length:
-            tokens += [self.pad_id] * (max_length - len(tokens))
-
-        return torch.tensor(tokens, dtype=torch.int, device=device)
-
-    def decode(self, tokens: torch.Tensor) -> str:
-        return self.processor.decode(tokens.tolist())
-
-    @staticmethod
-    def train(input: str, destination: str, vocab_size=32000) -> None:
-        model_prefix = os.path.join(destination, "tokenizer")
-        SentencePieceTrainer.Train(input=input, model_prefix=model_prefix, vocab_size=vocab_size)

lit-llama/lit_llama/utils.py

@@ -1,496 +0,0 @@
-"""Utility functions for training and inference."""
-
-import functools
-import pickle
-import warnings
-from io import BytesIO
-from pathlib import Path
-from contextlib import contextmanager
-
-import torch
-import torch.utils._device
-from lightning.fabric.strategies import DeepSpeedStrategy, FSDPStrategy
-from torch.distributed.fsdp import FullStateDictConfig
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import StateDictType
-from torch.serialization import normalize_storage_type
-
-llama_model_sizes = {
-    4096: "7B",  # 7B n_embd=4096
-    5120: "13B",  # 13B n_embd=5120
-    6656: "30B",  # 30B n_embd=6656
-    8192: "65B",  # 65B n_embd=8192
-}
-
-
-def llama_model_lookup(checkpoint: dict) -> str:
-    """Returns the LLaMA model name from the checkpoint.
-    
-    Checks the width of the lm_head.weight matrix, as these uniquely identify the model.
-    """
-    embedding_size = checkpoint['transformer.wte.weight'].shape[1]
-    return llama_model_sizes[embedding_size]
-
-
-def find_multiple(n: int, k: int) -> int:
-    if n % k == 0:
-        return n
-    return n + k - (n % k)
-
-
-def save_model_checkpoint(fabric, model, file_path):
-    """Handles boilerplate logic for retrieving and saving the state_dict.
-    
-    This will be upstreamed to Fabric soon.
-    """
-    file_path = Path(file_path)
-
-    if isinstance(fabric.strategy, DeepSpeedStrategy):
-        from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
-
-        fabric.save(file_path, {"model": model})
-        fabric.barrier()
-        if fabric.global_rank == 0:
-            # Create a consolidated checkpoint with the same name next to the deepspeed checkpoint
-            convert_zero_checkpoint_to_fp32_state_dict(file_path, file_path.with_suffix(".pth"))
-        return
-
-    if isinstance(fabric.strategy, FSDPStrategy):
-        save_policy = FullStateDictConfig(offload_to_cpu=(fabric.world_size > 1), rank0_only=True)
-        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
-            state_dict = model._forward_module.state_dict()
-    else:
-        state_dict = model.state_dict()
-
-    if fabric.global_rank == 0:
-        torch.save(state_dict, file_path)
-    fabric.barrier()
-
-
-class EmptyInitOnDevice(torch.overrides.TorchFunctionMode):
-    def __init__(self, device=None, dtype=None, quantization_mode=None):
-        """
-        Create tensors with given device and dtype and don't run initialization
-           (but instead use "empty tensors", i.e. uninitialized memory).
-
-            device: `torch.device` to work with
-            dtype: `torch.dtype` to work with
-            quantization_mode: optional string, quantization mode to work with, default `None`.
-                 Available modes: `llm.int8` bitsnbytes LLM.int8 quantization (only on GPU)
-                                  `gptq.int4`, `gptq.int8`: GPTQ pre-quantized models
-
-        Example::
-            with EmptyInitOnDevice("cuda", dtype=torch.bfloat16):
-               model = LLaMA.from_name('7B')
-            model.load_state_dict(torch.load('llama-lit/7B/lit-llama.pth'))"""
-
-        self.quantization_mode = quantization_mode
-        self.quantized_linear_cls = None
-        if self.quantization_mode == 'llm.int8':
-            if device.type != "cuda":
-                raise ValueError("Quantization is only supported on the GPU.")
-            from .quantization import Linear8bitLt
-            self.quantized_linear_cls = Linear8bitLt
-        elif self.quantization_mode == 'gptq.int4':
-            from .quantization import ColBlockQuantizedLinear
-            self.quantized_linear_cls = functools.partial(ColBlockQuantizedLinear, bits=4, tile_cols=-1)
-        elif self.quantization_mode == 'gptq.int8':
-            from .quantization import ColBlockQuantizedLinear
-            self.quantized_linear_cls = functools.partial(ColBlockQuantizedLinear, bits=8, tile_cols=-1)
-        elif self.quantization_mode is not None:
-            raise RuntimeError(f"unknown quantization mode {self.quantization_mode}")
-        self.device = device
-        self.dtype = dtype
-
-    def __enter__(self):
-        if self.quantized_linear_cls != None:
-            self.torch_linear_cls = torch.nn.Linear
-            torch.nn.Linear = self.quantized_linear_cls
-        return super().__enter__()
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.quantized_linear_cls != None:
-            torch.nn.Linear = self.torch_linear_cls
-        return super().__exit__(exc_type, exc_val, exc_tb)
-
-    def __torch_function__(self, func, types, args=(), kwargs=None):
-        kwargs = kwargs or {}
-        if getattr(func, "__module__", None) == "torch.nn.init":
-            if "tensor" in kwargs:
-                return kwargs["tensor"]
-            else:
-                return args[0]
-        if (
-            self.device is not None
-            and func in torch.utils._device._device_constructors()
-            and kwargs.get("device") is None
-        ):
-            kwargs["device"] = self.device
-        if (
-            self.dtype is not None
-            and func in torch.utils._device._device_constructors()
-            and kwargs.get("dtype") is None
-        ):
-            kwargs["dtype"] = self.dtype
-        return func(*args, **kwargs)
-
-
-@contextmanager
-def quantization(mode: str = None):
-    quantized_linear_cls = None
-    if mode == 'llm.int8':
-        from .quantization import Linear8bitLt
-        quantized_linear_cls = Linear8bitLt
-    elif mode == 'gptq.int4':
-        from .quantization import ColBlockQuantizedLinear
-        quantized_linear_cls = functools.partial(ColBlockQuantizedLinear, bits=4, tile_cols=-1)
-    elif mode == 'gptq.int8':
-        from .quantization import ColBlockQuantizedLinear
-        quantized_linear_cls = functools.partial(ColBlockQuantizedLinear, bits=8, tile_cols=-1)
-    elif mode is not None:
-        raise ValueError(f"Unknown quantization mode: {mode}")
-
-    enabled = mode is not None
-    torch_linear_cls = torch.nn.Linear
-    if enabled:
-        torch.nn.Linear = quantized_linear_cls
-    yield
-    if enabled:
-        torch.nn.Linear = torch_linear_cls
-
-
-# this is taken from torchhacks https://github.com/lernapparat/torchhacks
-
-
-class NotYetLoadedTensor:
-    def __init__(self, metatensor, archiveinfo, storageinfo, rebuild_args):
-        self.metatensor = metatensor
-        self.archiveinfo = archiveinfo
-        self.storageinfo = storageinfo
-        self.rebuild_args = rebuild_args
-
-    @classmethod
-    def rebuild_from_type_v2(cls, func, new_type, args, state, *, archiveinfo=None):
-        ret = func(*args)
-        if isinstance(ret, NotYetLoadedTensor):
-            old_lt = ret._load_tensor
-
-            def _load_tensor():
-                t = old_lt()
-                return torch._tensor._rebuild_from_type_v2(
-                    lambda: t, new_type, (), state
-                )
-
-            ret._load_tensor = _load_tensor
-            return ret
-        return torch._tensor._rebuild_from_type_v2(func, new_type, args, state)
-
-    @classmethod
-    def rebuild_parameter(
-        cls, data, requires_grad, backward_hooks, *, archiveinfo=None
-    ):
-        if isinstance(data, NotYetLoadedTensor):
-            old_lt = data._load_tensor
-
-            def _load_tensor():
-                t = old_lt()
-                return torch._utils._rebuild_parameter(t, requires_grad, backward_hooks)
-
-            data._load_tensor = _load_tensor
-            return data
-        return torch._utils._rebuild_parameter(data, requires_grad, backward_hooks)
-
-    @classmethod
-    def rebuild_tensor_v2(
-        cls,
-        storage,
-        storage_offset,
-        size,
-        stride,
-        requires_grad,
-        backward_hooks,
-        metadata=None,
-        *,
-        archiveinfo=None,
-    ):
-        rebuild_args = (
-            storage_offset,
-            size,
-            stride,
-            requires_grad,
-            backward_hooks,
-            metadata,
-        )
-        metatensor = torch._utils._rebuild_tensor_v2(
-            storage,
-            storage_offset,
-            size,
-            stride,
-            requires_grad,
-            backward_hooks,
-            metadata,
-        )
-        storageinfo = storage.archiveinfo
-        return NotYetLoadedTensor(metatensor, archiveinfo, storageinfo, rebuild_args)
-
-    def _load_tensor(self):
-        name, storage_cls, fn, device, size = self.storageinfo
-        dtype = self.metatensor.dtype
-
-        uts = (
-            self.archiveinfo.zipfile_context.zf.get_storage_from_record(
-                f"data/{fn}",
-                size * torch._utils._element_size(dtype),
-                torch.UntypedStorage,
-            )
-            ._typed_storage()
-            ._untyped_storage
-        )
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            storage = torch.storage.TypedStorage(
-                wrap_storage=uts, dtype=self.metatensor.dtype, _internal=True
-            )
-        tensor = torch._utils._rebuild_tensor_v2(storage, *self.rebuild_args)
-        return tensor
-
-    @classmethod
-    def __torch_function__(cls, func, types, args=(), kwargs=None):
-        if kwargs is None:
-            kwargs = {}
-        loaded_args = [
-            (a._load_tensor() if isinstance(a, NotYetLoadedTensor) else a) for a in args
-        ]
-        res = func(*loaded_args, **kwargs)
-        # gc.collect would be costly here, maybe do it optionally
-        return res
-
-    def __getattr__(self, name):
-        # properties
-        ## TODO: device, is_...??
-        ## TODO: mH, mT, H, T, data, imag, real
-        ## name ???
-        if name in {
-            "dtype",
-            "grad",
-            "grad_fn",
-            "layout",
-            "names",
-            "ndim",
-            "output_nr",
-            "requires_grad",
-            "retains_grad",
-            "shape",
-            "volatile",
-        }:
-            return getattr(self.metatensor, name)
-        if name in {"size"}:
-            return getattr(self.metatensor, name)
-        # materializing with contiguous is needed for quantization
-        if name in {"contiguous"}:
-            return getattr(self._load_tensor(), name)
-
-        raise AttributeError(f"{type(self)} does not have {name}")
-
-    def __repr__(self):
-        return f"NotYetLoadedTensor({repr(self.metatensor)})"
-
-
-class LazyLoadingUnpickler(pickle.Unpickler):
-    def __init__(self, file, zipfile_context):
-        super().__init__(file)
-        self.zipfile_context = zipfile_context
-
-    def find_class(self, module, name):
-        res = super().find_class(module, name)
-        if module == "torch._utils" and name == "_rebuild_tensor_v2":
-            return functools.partial(
-                NotYetLoadedTensor.rebuild_tensor_v2, archiveinfo=self
-            )
-        elif module == "torch._tensor" and name == "_rebuild_from_type_v2":
-            return functools.partial(
-                NotYetLoadedTensor.rebuild_from_type_v2, archiveinfo=self
-            )
-        elif module == "torch._utils" and name == "_rebuild_parameter":
-            return functools.partial(
-                NotYetLoadedTensor.rebuild_parameter, archiveinfo=self
-            )
-        return res
-
-    def persistent_load(self, pid):
-        name, cls, fn, device, size = pid
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            s = torch.storage.TypedStorage(dtype=cls().dtype, device="meta")
-        s.archiveinfo = pid
-        return s
-
-
-class lazy_load:
-    def __init__(self, fn):
-        self.zf = torch._C.PyTorchFileReader(str(fn))
-        with BytesIO(self.zf.get_record("data.pkl")) as pkl:
-            mup = LazyLoadingUnpickler(pkl, self)
-            self.sd = mup.load()
-
-    def __enter__(self):
-        return self.sd
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        del self.zf  # I don't think there is a way to force closing...
-        self.zf = None
-
-
-class SavingProxyForStorage:
-    def __init__(self, obj, saver, protocol_version=5):
-        self.protocol_version = protocol_version
-        self.saver = saver
-        if not (isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj)):
-            raise TypeError(f"expected storage, not {type(obj)}")
-
-        # this logic is taken from PyTorch 2.0+ torch/serialization.py
-        if isinstance(obj, torch.storage.TypedStorage):
-            # PT upstream wants to deprecate this eventually...
-            storage = obj._untyped_storage
-            storage_type_str = obj._pickle_storage_type()
-            storage_type = getattr(torch, storage_type_str)
-            storage_numel = obj._size()
-        else:
-            storage = obj
-            storage_type = normalize_storage_type(type(obj))
-            storage_numel = storage.nbytes()
-
-        storage_key = saver._write_storage_and_return_key(storage)
-        location = torch.serialization.location_tag(storage)
-
-        self.storage_info = (
-            "storage",
-            storage_type,
-            storage_key,
-            location,
-            storage_numel,
-        )
-
-    def __reduce_ex__(self, protocol_version):
-        assert False, "this should be handled with out of band"
-
-
-class SavingProxyForTensor:
-    def __init__(self, tensor, saver, protocol_version=5):
-        self.protocol_version = protocol_version
-        self.reduce_ret_fn, (storage, *other_reduce_args) = tensor.__reduce_ex__(
-            protocol_version
-        )
-        assert isinstance(
-            storage, torch.storage.TypedStorage
-        ), "Please check for updates"
-        storage_proxy = SavingProxyForStorage(
-            storage, saver, protocol_version=protocol_version
-        )
-        self.reduce_args = (storage_proxy, *other_reduce_args)
-
-    def __reduce_ex__(self, protocol_version):
-        if protocol_version != self.protocol_version:
-            raise RuntimeError(
-                f"Unexpected protocol version: expected {self.protocol_version}, got {protocol_version}"
-            )
-        return self.reduce_ret_fn, self.reduce_args
-
-
-class IncrementalPyTorchPickler(pickle.Pickler):
-    def __init__(self, saver, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.storage_dtypes = {}
-        self.saver = saver
-        self.id_map = {}
-
-    # this logic is taken from PyTorch 2.0+ torch/serialization.py
-    def persistent_id(self, obj):
-        # FIXME: the docs say that persistent_id should only return a string
-        # but torch store returns tuples. This works only in the binary protocol
-        # see
-        # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects
-        # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537
-        if isinstance(obj, SavingProxyForStorage):
-            return obj.storage_info
-
-        if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj):
-            if isinstance(obj, torch.storage.TypedStorage):
-                # TODO: Once we decide to break serialization FC, this case
-                # can be deleted
-                storage = obj._untyped_storage
-                storage_dtype = obj.dtype
-                storage_type_str = obj._pickle_storage_type()
-                storage_type = getattr(torch, storage_type_str)
-                storage_numel = obj._size()
-
-            else:
-                storage = obj
-                storage_dtype = torch.uint8
-                storage_type = normalize_storage_type(type(obj))
-                storage_numel = storage.nbytes()
-
-            # If storage is allocated, ensure that any other saved storages
-            # pointing to the same data all have the same dtype. If storage is
-            # not allocated, don't perform this check
-            if storage.data_ptr() != 0:
-                if storage.data_ptr() in self.storage_dtypes:
-                    if storage_dtype != self.storage_dtypes[storage.data_ptr()]:
-                        raise RuntimeError(
-                            "Cannot save multiple tensors or storages that "
-                            "view the same data as different types"
-                        )
-                else:
-                    self.storage_dtypes[storage.data_ptr()] = storage_dtype
-
-            storage_key = self.id_map.get(storage._cdata)
-            if storage_key is None:
-                storage_key = self.saver._write_storage_and_return_key(storage)
-                self.id_map[storage._cdata] = storage_key
-            location = torch.serialization.location_tag(storage)
-
-            return ("storage", storage_type, storage_key, location, storage_numel)
-
-        return None
-
-
-class incremental_save:
-    def __init__(self, name):
-        self.name = name
-        self.zipfile = torch._C.PyTorchFileWriter(str(name))
-        self.has_saved = False
-        self.next_key = 0
-
-    def __enter__(self):
-        return self
-
-    def store_early(self, tensor):
-        if isinstance(tensor, torch.Tensor):
-            return SavingProxyForTensor(tensor, self)
-        raise TypeError(f"can only store tensors early, not {type(tensor)}")
-
-    def save(self, obj):
-        if self.has_saved:
-            raise RuntimeError("have already saved")
-        # Write the pickle data for `obj`
-        data_buf = BytesIO()
-        pickler = IncrementalPyTorchPickler(self, data_buf, protocol=5)
-        pickler.dump(obj)
-        data_value = data_buf.getvalue()
-        self.zipfile.write_record("data.pkl", data_value, len(data_value))
-        self.has_saved = True
-
-    def _write_storage_and_return_key(self, storage):
-        if self.has_saved:
-            raise RuntimeError("have already saved")
-        key = self.next_key
-        self.next_key += 1
-        name = f"data/{key}"
-        if storage.device.type != "cpu":
-            storage = storage.cpu()
-        num_bytes = storage.nbytes()
-        self.zipfile.write_record(name, storage.data_ptr(), num_bytes)
-        return key
-
-    def __exit__(self, type, value, traceback):
-        self.zipfile.write_end_of_file()

lit-llama/pretrain/redpajama.py

@@ -1,321 +0,0 @@
-import os
-import sys
-import math
-import glob
-import time
-from functools import partial
-from pathlib import Path
-from typing import Tuple, Optional
-
-import lightning as L
-from lightning.fabric.strategies import FSDPStrategy
-
-import torch
-from torch.utils.data import DataLoader
-from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
-
-import numpy as np
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama.model import Block, LLaMA, LLaMAConfig
-from lit_llama.packed_dataset import PackedDataset, CombinedDataset
-from lit_llama.utils import save_model_checkpoint
-
-
-out_dir = "out/training"
-save_interval = 1000
-eval_interval = 1000
-eval_iters = 100
-log_interval = 1
-
-# compile = False
-
-# Hyperparameters
-learning_rate = 6e-4
-batch_size = 125
-micro_batch_size = 5
-max_iters = 600000  # num_epochs * (epoch_size // micro_batch_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
-
-
-# Data proportions from https://arxiv.org/pdf/2302.13971.pdf Table 1
-data_config = [
-    ("arxiv", 2.5),
-    ("book", 4.5),
-    ("c4", 15.0),
-    ("cc", 67.0),
-    ("github", 4.5),
-    ("stackexchange", 2.0),
-    ("wikipedia", 4.5),
-]
-
-
-def main(
-    devices: int = 4,
-    train_data_dir: Path = "data/lit-redpajama",
-    val_data_dir: Optional[Path] = None,
-) -> None:
-    auto_wrap_policy = partial(
-        transformer_auto_wrap_policy, transformer_layer_cls={Block}
-    )
-    strategy = FSDPStrategy(
-        auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True
-    )
-
-    fabric = L.Fabric(
-        accelerator="cuda", devices=devices, precision="bf16-mixed", strategy=strategy
-    )
-    fabric.launch()
-    fabric.seed_everything(1337)
-
-    if fabric.global_rank == 0:
-        os.makedirs(out_dir, exist_ok=True)
-
-    config = LLaMAConfig.from_name("7B")
-
-    train_dataloader, val_dataloader = create_dataloaders(
-        batch_size=micro_batch_size,
-        block_size=config.block_size,
-        fabric=fabric,
-        train_data_dir=train_data_dir,
-        val_data_dir=val_data_dir,
-        seed=1338,
-    )
-    if val_dataloader is None:
-        train_dataloader = fabric.setup_dataloaders(train_dataloader)
-    else:
-        train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)
-
-    with fabric.device:
-        torch.set_default_dtype(torch.bfloat16)
-        model = LLaMA(config)
-        model.apply(model._init_weights)
-        torch.set_default_dtype(torch.float32)
-
-    # if compile:
-    #     model = torch.compile(model)
-
-    optimizer = torch.optim.AdamW(
-        model.parameters(),
-        lr=learning_rate,
-        weight_decay=weight_decay,
-        betas=(beta1, beta2),
-        foreach=False,
-    )
-
-    model, optimizer = fabric.setup(model, optimizer)
-
-    process_batch_size = batch_size // devices
-    gradient_accumulation_iters = process_batch_size // micro_batch_size
-
-    train(fabric, model, optimizer, train_dataloader, val_dataloader, gradient_accumulation_iters, devices)
-
-
-def train(
-    fabric: L.Fabric,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    train_dataloader: DataLoader,
-    val_dataloader: Optional[DataLoader],
-    grad_accum_steps: int,
-    devices: int,
-) -> None:
-    """The training loop.
-
-    Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
-    """
-
-    step_count = 0
-
-    step_time = 0.0
-    tokens = 0
-    tokens_sec = 0.0
-    prev_t1 = time.time()
-
-    for iter_num, train_data in enumerate(train_dataloader):
-        t0 = time.time()
-
-        # determine and set the learning rate for this iteration
-        lr = get_lr(iter_num) if decay_lr else learning_rate
-        for param_group in optimizer.param_groups:
-            param_group["lr"] = lr
-
-
-        input_ids = train_data[:, 0 : model.config.block_size].contiguous()
-        targets = train_data[:, 1 : model.config.block_size + 1].contiguous()
-        
-        is_accumulating = (iter_num + 1) % grad_accum_steps != 0
-
-        with fabric.no_backward_sync(model, enabled=is_accumulating):
-            logits = model(input_ids)
-            loss = torch.nn.functional.cross_entropy(
-                logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
-            )
-            fabric.backward(loss / grad_accum_steps)
-
-        t1 = time.time()
-
-        if not is_accumulating:
-            fabric.clip_gradients(model, optimizer, max_norm=grad_clip)
-
-            optimizer.step()
-            optimizer.zero_grad()
-            step_count += 1
-
-            t1 = time.time()
-
-            if val_dataloader is not None and step_count % eval_interval == 0:
-                val_loss = validate(fabric, model, val_dataloader)
-                fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
-                fabric.barrier()
-                fabric.log_dict(
-                    {"iter": iter_num, "val_loss": val_loss, "step": step_count, "lr": lr}
-                )
-
-            if step_count % save_interval == 0:
-                fabric.print(f"Saving checkpoint to {out_dir}")
-                save_model_checkpoint(
-                    fabric, model, os.path.join(out_dir, f"iter-{iter_num:06d}-ckpt.pth")
-                )
-
-        dt = t1 - t0
-
-        tokens += micro_batch_size * model.config.block_size
-        step_time += t1 - prev_t1
-        prev_t1 = t1
-
-        if iter_num % log_interval == 0:
-            tokens_sec_str = f"{tokens / step_time:.0f}" if not is_accumulating else "-"
-
-            fabric.log_dict(
-                {"iter": iter_num, "train_loss": loss, "step": step_count, "lr": lr}
-            )
-            fabric.print(
-                    f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms, speed: {tokens_sec_str} toks/s/device"
-            )
-
-        if not is_accumulating:
-            tokens = 0
-            step_time = 0.0
-
-        if iter_num > max_iters:
-            break
-
-
-@torch.no_grad()
-def validate(
-    fabric: L.Fabric, model: torch.nn.Module, val_dataloader: DataLoader
-) -> torch.Tensor:
-    fabric.print("Validating ...")
-    model.eval()
-    losses = torch.zeros(eval_iters)
-    for k, val_data in enumerate(val_dataloader):
-        input_ids = val_data[:, 0 : model.config.block_size].contiguous()
-        targets = val_data[:, 1 : model.config.block_size + 1].contiguous()
-        logits = model(input_ids)
-        loss = torch.nn.functional.cross_entropy(
-            logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
-        )
-        losses[k] = loss.item()
-    out = losses.mean()
-    model.train()
-    return out
-
-
-def create_dataloader(
-    batch_size: int,
-    block_size: int,
-    data_dir: str,
-    fabric,
-    shuffle: bool = True,
-    seed: int = 12345,
-) -> DataLoader:
-    datasets = []
-    for prefix, _ in data_config:
-        filenames = glob.glob(os.path.join(data_dir, prefix + "*"))
-        dataset = PackedDataset(
-            filenames, n_chunks=4, block_size=block_size, shuffle=shuffle, seed=seed,
-            num_processes=fabric.world_size, process_rank=fabric.global_rank,
-        )
-        datasets.append(dataset)
-
-    if not datasets:
-        raise RuntimeError(
-            f"No data found at {data_dir}. Make sure you ran prepare_redpajama.py to create the dataset."
-        )
-
-    weights = [weight for _, weight in data_config]
-    sum_weights = sum(weights)
-    weights = [el / sum_weights for el in weights]
-
-    combined_dataset = CombinedDataset(datasets=datasets, seed=seed, weights=weights)
-
-    return DataLoader(combined_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)
-
-
-def create_dataloaders(
-    batch_size: int,
-    block_size: int,
-    fabric,
-    train_data_dir: str = "data/lit-redpajama",
-    val_data_dir: Optional[str] = None,
-    seed: int = 12345,
-) -> Tuple[DataLoader, DataLoader]:
-    # Increase by one because we need the next word as well
-    effective_block_size = block_size + 1
-    train_dataloader = create_dataloader(
-        batch_size=batch_size,
-        block_size=effective_block_size,
-        fabric=fabric,
-        data_dir=train_data_dir,
-        shuffle=True,
-        seed=seed,
-    )
-    val_dataloader = (
-        create_dataloader(
-            batch_size=batch_size,
-            block_size=effective_block_size,
-            fabric=fabric,
-            data_dir=val_data_dir,
-            shuffle=False,
-            seed=seed,
-        )
-        if val_data_dir
-        else None
-    )
-    return train_dataloader, val_dataloader
-
-
-# learning rate decay scheduler (cosine with warmup)
-def get_lr(it):
-    # 1) linear warmup for warmup_iters steps
-    if it < warmup_iters:
-        return learning_rate * it / warmup_iters
-    # 2) if it > lr_decay_iters, return min learning rate
-    if it > lr_decay_iters:
-        return min_lr
-    # 3) in between, use cosine decay down to min learning rate
-    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
-    assert 0 <= decay_ratio <= 1
-    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1
-    return min_lr + coeff * (learning_rate - min_lr)
-
-
-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)

lit-llama/pretrain/shakespeare.py

@@ -1,166 +0,0 @@
-"""
-This script is a placeholder for training LLaMA from scratch.
-Currently, it just trains on the Shakespeare dataset.
-"""
-from pathlib import Path
-import sys
-import os
-import time
-from functools import partial
-from typing import Tuple
-
-import lightning as L
-from lightning.fabric.strategies import FSDPStrategy
-
-import torch
-from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
-
-import numpy as np
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama.model import Block, LLaMA, LLaMAConfig
-from lit_llama.utils import save_model_checkpoint
-
-
-out_dir = "out/training"
-eval_interval = 2000
-eval_iters = 200
-log_interval = 1
-# compilation fails as it does not support torch.complex64 for RoPE
-# compile = False
-
-# Hyperparameters
-learning_rate = 6e-4
-batch_size = 2
-max_iters = 600000
-weight_decay = 1e-1
-beta1 = 0.9
-beta2 = 0.95
-grad_clip = 1.0
-
-# For shakespeare, choose smaller block size than vanilla LLaMA
-block_size = 1024
-
-
-def main() -> None:
-    auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
-    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True)
-
-    fabric = L.Fabric(accelerator="cuda", devices=4, 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()
-
-    config = LLaMAConfig.from_name("7B")
-    config.block_size = block_size
-    config.vocab_size = 100  # from prepare_shakespeare.py
-
-    with fabric.device:
-        model = LLaMA(config)
-
-    # if compile:
-    #     model = torch.compile(model)
-
-    model = fabric.setup_module(model)
-
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False)
-    optimizer = fabric.setup_optimizers(optimizer)
-
-    train(fabric, model, optimizer, train_data, val_data)
-
-
-def train(
-    fabric: L.Fabric,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    train_data: np.ndarray,
-    val_data: np.ndarray,
-) -> None:
-    """The training loop.
-
-    Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
-    """
-
-    iter_num = 0
-
-    while True:
-        # TODO: add learning rate scheduling
-
-        # evaluate the loss on train/val sets and write checkpoints
-        if iter_num > 0 and iter_num % eval_interval == 0:
-            val_loss = validate(fabric, model, val_data)
-            fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
-            fabric.print(f"Saving checkpoint to {out_dir}")
-            save_model_checkpoint(fabric, model, os.path.join(out_dir, f"iter-{iter_num:06d}-ckpt.pth"))
-
-        t0 = time.time()
-
-        input_ids, targets = get_batch(
-            fabric,
-            train_data,
-            block_size=model.config.block_size,  # type: ignore[union-attr,arg-type]
-        )
-        logits = model(input_ids)
-        loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
-
-        fabric.backward(loss)
-
-        # TODO: Gradient clipping
-        # if grad_clip != 0.0:
-        #     fabric.clip_gradients(model, optimizer, max_norm=grad_clip)
-
-        optimizer.step()
-        optimizer.zero_grad()
-
-        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")
-        iter_num += 1
-
-        if iter_num > max_iters:
-            break
-
-
-@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,
-            block_size=model.config.block_size,  # type: ignore[union-attr,arg-type]
-        )
-        logits = model(input_ids)
-        loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
-        losses[k] = loss.item()
-    out = losses.mean()
-    model.train()
-    return out
-
-
-def get_batch(fabric: L.Fabric, data: np.ndarray, block_size: int) -> Tuple[torch.Tensor, torch.Tensor]:
-    ix = torch.randint(len(data) - block_size, (batch_size,))
-    x = torch.stack([torch.from_numpy((data[i : i + block_size]).astype(np.int64)) for i in ix])
-    y = torch.stack([torch.from_numpy((data[i + 1 : i + 1 + block_size]).astype(np.int64)) for i in ix])
-    x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
-    return x, y
-
-
-def load_datasets(data_dir: str = "data/shakespeare") -> Tuple[np.ndarray, np.ndarray]:
-    train_data = np.memmap(os.path.join(data_dir, "train.bin"), dtype=np.uint16, mode="r")
-    val_data = np.memmap(os.path.join(data_dir, "val.bin"), dtype=np.uint16, mode="r")
-    return train_data, val_data
-
-
-if __name__ == "__main__":
-    torch.set_float32_matmul_precision("high")
-    main()

lit-llama/quantize/gptq.py

@@ -1,238 +0,0 @@
-# This adapts GPTQ's quantization process: https://github.com/IST-DASLab/gptq/
-# E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
-# portions copyright by the authors licensed under the Apache License 2.0
-import gc
-import sys
-import time
-from pathlib import Path
-from typing import Optional
-
-import torch
-from datasets import load_dataset
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama import LLaMA, Tokenizer
-from lit_llama.quantization import GPTQQuantizer
-from lit_llama.utils import EmptyInitOnDevice, llama_model_lookup
-
-
-def get_sample_data():
-    traindata = load_dataset(
-        "allenai/c4",
-        "allenai--c4",
-        data_files={"train": "en/c4-train.00000-of-01024.json.gz"},
-        split="train",
-    )
-    # heuristic for the data size?
-    txt = "\n".join(
-        traindata[i]["text"] for i in torch.randperm(len(traindata))[:1000].tolist()
-    )
-    return txt
-
-
-@torch.no_grad()
-def llama_blockwise_quantization(
-    model, sample_inputs, working_device, *, bits=4, groupsize=-1
-):
-    """
-    This is the classic post-training quantization of all linear layers.
-    We quantize in order, i.e. when observing the inputs, we use the outputs of the previously quantized layers rather
-    than doing them all at once.
-    """
-    print(model)
-    print(model.config)
-
-    print("Getting inputs for first block")
-    model.transformer.wte.to(working_device)
-    sample_inputs = sample_inputs.to(working_device)
-    inps = model.transformer.wte(sample_inputs)
-    model.transformer.wte.to("cpu")
-    torch.cuda.empty_cache()
-
-    rope_cache = model.build_rope_cache(sample_inputs)
-    mask_cache = model.build_mask_cache(sample_inputs)
-
-    print("Starting to quantize blocks")
-    outs = torch.zeros_like(inps)
-
-    # better than relying on enumeration? originally the code bundled
-    # the two mlp fc layers
-    # we could automate this with a lot of hooks and another iteration
-    submodules_to_process = [
-        "attn.c_attn",
-        "attn.c_proj",
-        "mlp.c_fc1",
-        "mlp.c_fc2",
-        "mlp.c_proj",
-    ]
-
-    for i, block in enumerate(model.transformer.h):
-        block.to(working_device)
-
-        for name in submodules_to_process:
-            print(i, name, end=" ")
-            t0 = time.perf_counter()
-            print("collecting stats", end=" ")
-            sys.stdout.flush()
-            module = block.get_submodule(name)
-
-            gptq = GPTQQuantizer(
-                module,
-                bits=bits,
-                groupsize=groupsize,
-                actorder=(groupsize == -1),
-            )
-            handle = module.register_forward_hook(gptq.collect_input_stats)
-            for j in range(inps.size(0)):
-                outs[j : j + 1], _ = block(
-                    inps[j : j + 1],
-                    rope=rope_cache,
-                    mask=mask_cache,
-                    max_seq_length=model.config.block_size
-                )
-
-            handle.remove()
-
-            print("quantizing", end=" ")
-            sys.stdout.flush()
-            q_module, error = gptq.quantize()
-
-            # replace the linear module with the quantized module
-            pname, dname = name.rsplit(".", 1)
-            setattr(block.get_submodule(pname), dname, q_module)
-
-            # cleanup in an attempt to not run out of memory
-            del gptq
-            gc.collect()
-            torch.cuda.empty_cache()
-            t1 = time.perf_counter()
-            print(f"time {int(t1 - t0 + 0.5)}s quantization error {error:.1f}")
-
-        for j in range(inps.size(0)):
-            outs[j : j + 1], _ = block(
-                inps[j : j + 1],
-                rope=rope_cache,
-                mask=mask_cache,
-                max_seq_length=model.config.block_size
-            )
-
-        block.cpu()
-        gc.collect()
-        torch.cuda.empty_cache()
-
-        # the outputs are the next block's inputs and we'll reuse the old inputs
-        inps, outs = outs, inps
-
-    model.transformer.ln_f.to(working_device)
-    for j in range(inps.size(0)):
-        outs[j : j + 1] = model.transformer.ln_f(inps[j : j + 1])
-    model.transformer.ln_f.to("cpu")
-    inps, outs = outs, inps
-
-    model.lm_head.to(working_device)
-    gptq = GPTQQuantizer(
-        model.lm_head,
-        bits=bits,
-        groupsize=groupsize,
-        actorder=(groupsize == -1),
-    )
-    handle = model.lm_head.register_forward_hook(gptq.collect_input_stats)
-    for j in range(inps.size(0)):
-        model.lm_head(inps[j : j + 1])
-    handle.remove()
-    q_module, error = gptq.quantize()
-    model.lm_head = q_module
-    model.lm_head.to("cpu")
-
-
-def main(
-    *,
-    checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
-    output_path: Optional[Path] = None,
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    n_samples: int = 128,
-    dtype: str = "float32",
-    quantize: Optional[str] = None,
-) -> None:
-    """Generates text samples based on a pre-trained LLaMA model and tokenizer.
-
-    Args:
-        checkpoint_path: The checkpoint path to load.
-        output_path: Path to write the quantized model's state dict to.
-        tokenizer_path: The tokenizer path to load.
-        n_samples: Number of example inputs to use for statistics (default: 128)
-        dtype: The dtype to use to load the model.
-        quantize: Mode to quantize the model to:
-            ``"gptq.int4"``: GPTQ 4-bit mode.
-            Note that ``"llm.int8"```does not need a quantization step.
-    """
-    assert checkpoint_path.is_file()
-    assert tokenizer_path.is_file()
-    if output_path is None:
-        output_path = checkpoint_path.parent / "llama-gptq.4bit.pth"
-    assert output_path.parent.is_dir() and (not output_path.exists() or output_path.is_file())
-
-    device = "cuda"
-
-    dt = getattr(torch, dtype, None)
-    if not isinstance(dt, torch.dtype):
-        raise ValueError(f"{dtype} is not a valid dtype.")
-    dtype = dt
-
-    if quantize == "gptq.int4":
-        bits = 4
-    elif quantize == "gptq.int8":
-        bits = 8
-    else:
-        raise RuntimeError(f"unknown/unsupported quantization mode {quantize}")
-
-    # we avoid loading the entire model on the GPU and do this block by block
-    with EmptyInitOnDevice(
-        device="cpu",
-        dtype=dtype,
-    ):
-        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()
-
-    tokenizer = Tokenizer(tokenizer_path)
-
-    test_string = get_sample_data()
-    encoded_text = tokenizer.encode(
-        test_string,
-        bos=True,
-        eos=False,
-    )
-    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)
-    encoded_text = encoded_text[: n_samples * block_size].reshape(n_samples, block_size)
-
-    t0 = time.perf_counter()
-    llama_blockwise_quantization(model, encoded_text, device, bits=bits)
-    t = time.perf_counter() - t0
-
-    print(
-        f"\n\nTime for quantization: {t:.02f} sec total",
-        file=sys.stderr,
-    )
-    print(
-        f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
-        file=sys.stderr,
-    )
-
-    torch.save(model.state_dict(), output_path)
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    torch.set_float32_matmul_precision("high")
-    CLI(main)

lit-llama/requirements.txt

@@ -1,9 +0,0 @@
-torch>=2.0.0
-lightning @ git+https://github.com/Lightning-AI/lightning@master
-sentencepiece
-tqdm  # convert_checkpoint.py
-numpy  # train.py dataset memmap
-jsonargparse[signatures]  # generate.py, convert_checkpoint.py CLI
-bitsandbytes  # quantization.py
-datasets  # evaluate.py
-zstandard  # prepare_redpajama.py

lit-llama/scripts/convert_checkpoint.py

@@ -1,141 +0,0 @@
-import gc
-import shutil
-from pathlib import Path
-from typing import Dict
-
-import torch
-from tqdm import tqdm
-
-"""
-Sample usage:
-
-```bash
-python -m scripts.convert_checkpoint -h
-
-python -m scripts.convert_checkpoint converted
-```
-"""
-
-
-def convert_state_dict(state_dict: Dict[str, torch.Tensor], dtype: torch.dtype = torch.float32) -> Dict[str, torch.Tensor]:
-    converted = {}
-    converted["transformer.wte.weight"] = state_dict["tok_embeddings.weight"].to(dtype)
-    converted["lm_head.weight"] = state_dict["output.weight"].to(dtype)
-    converted["transformer.ln_f.scale"] = state_dict["norm.weight"].to(dtype)
-
-    for layer_idx in sorted(set([k.split(".")[1] for k in state_dict if k.startswith("layers")])):
-        # attention
-        # the wq, wk, wv from the FB model are stacked in our model as c_attn
-        converted[f"transformer.h.{layer_idx}.attn.c_attn.weight"] = torch.cat(
-            (
-                state_dict[f"layers.{layer_idx}.attention.wq.weight"].to(dtype),
-                state_dict[f"layers.{layer_idx}.attention.wk.weight"].to(dtype),
-                state_dict[f"layers.{layer_idx}.attention.wv.weight"].to(dtype),
-            )
-        )
-        converted[f"transformer.h.{layer_idx}.attn.c_proj.weight"] = state_dict[
-            f"layers.{layer_idx}.attention.wo.weight"
-        ].to(dtype)
-        # mlp
-        converted[f"transformer.h.{layer_idx}.mlp.c_fc1.weight"] = state_dict[
-            f"layers.{layer_idx}.feed_forward.w1.weight"
-        ].to(dtype)
-        converted[f"transformer.h.{layer_idx}.mlp.c_proj.weight"] = state_dict[
-            f"layers.{layer_idx}.feed_forward.w2.weight"
-        ].to(dtype)
-        converted[f"transformer.h.{layer_idx}.mlp.c_fc2.weight"] = state_dict[
-            f"layers.{layer_idx}.feed_forward.w3.weight"
-        ].to(dtype)
-        # rms norm
-        converted[f"transformer.h.{layer_idx}.rms_1.scale"] = state_dict[f"layers.{layer_idx}.attention_norm.weight"].to(dtype)
-        converted[f"transformer.h.{layer_idx}.rms_2.scale"] = state_dict[f"layers.{layer_idx}.ffn_norm.weight"].to(dtype)
-    return converted
-
-
-shard_dims = {
-    "lm_head.weight": 0,
-    "wte.weight": 1,
-    "attn.c_attn.weight": 0,
-    "attn.c_proj.weight": 1,
-    "mlp.c_fc1.weight": 0,
-    "mlp.c_fc2.weight": 0,
-    "mlp.c_proj.weight": 1
-}
-
-
-def meta_weights_for_nano_model(
-    *,
-    output_dir: Path = Path("checkpoints/lit-llama"),
-    checkpoint_dir: Path = Path("checkpoints/llama/"),
-    model_size: str = "7B",
-    dtype: str = "float32",
-) -> None:
-    output_dir = output_dir / model_size
-    checkpoint_dir = checkpoint_dir / model_size
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # the tokenizer is the same for all model sizes, so we store it in the parent dir
-    shutil.copy(checkpoint_dir.parent / "tokenizer.model", output_dir.parent)
-
-    dt = getattr(torch, dtype, None)
-    if not isinstance(dt, torch.dtype):
-        raise ValueError(f"{dtype} is not a valid dtype.")
-    dtype = dt
-
-    checkpoint_files = sorted(checkpoint_dir.glob("*.pth"))
-    checkpoint_files.sort()
-    n_checkpoints = len(checkpoint_files)
-
-    if n_checkpoints == 0:
-        raise RuntimeError(f"No checkpoints were found at checkpoint_dir {checkpoint_dir}. `consolidated.0*.pth` files expected at that location.")
-
-    # for the bigger models, there are multiple model-parallel checkpoints
-    # and we combine them into one single file
-    combined = None
-    for file in tqdm(checkpoint_files, total=n_checkpoints):
-        checkpoint = torch.load(file, map_location="cpu")
-        converted = convert_state_dict(checkpoint, dtype=dtype)
-        if combined is None:
-            combined = converted
-            continue
-        for name, param in converted.items():
-            dim = None
-            for k, d in shard_dims.items():
-                if k in name:
-                    dim = d
-                    break
-            if dim is None:
-                # Extra check: assert that tensors are the same if not sharded
-                # assert torch.allclose(combined[name], param)
-                continue
-            combined[name] = torch.cat((combined[name], param), dim=dim)
-
-        del checkpoint
-        del converted
-        gc.collect()
-
-    for name, param in combined.items():
-        if "c_attn" not in name:
-            continue
-
-        # Turn [Q1, K1, V1, Q2, K2, V2, ...] into [Q1, Q2, ..., K1, K2, .., V1, V2, ...]
-
-        src_chunk_len = param.shape[0] // n_checkpoints
-        mat_len = src_chunk_len // 3
-        dst_chunk_len = mat_len * n_checkpoints
-        attn = torch.clone(param)
-        for i in range(n_checkpoints):
-            for j in range(3):
-                param[j * dst_chunk_len + i * mat_len: j * dst_chunk_len + (i+1) * mat_len] = \
-                    attn[i * src_chunk_len + j * mat_len: i * src_chunk_len + (j+1) * mat_len]
-
-        del attn
-        gc.collect()
-
-    torch.save(combined, output_dir / "lit-llama.pth")
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(meta_weights_for_nano_model)

lit-llama/scripts/convert_hf_checkpoint.py

@@ -1,167 +0,0 @@
-import collections
-import contextlib
-import gc
-import json
-import shutil
-import sys
-from pathlib import Path
-
-import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama.model import LLaMA, LLaMAConfig
-from lit_llama.utils import EmptyInitOnDevice, lazy_load, incremental_save
-
-
-@torch.no_grad()
-def convert_hf_checkpoint(
-    *,
-    output_dir: Path = Path("checkpoints/lit-llama/7B"),
-    checkpoint_dir: Path = Path("checkpoints/hf-llama/7B"),
-    model_size: str = "7B",
-    dtype: str = "float32",
-    verify: bool = False,
-) -> None:
-    """
-    Perform the reverse operation of: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py
-    """
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # the tokenizer is the same for all model sizes, so we store it in the parent dir
-    shutil.copy(checkpoint_dir / "tokenizer.model", output_dir.parent)
-
-    dt = getattr(torch, dtype, None)
-    if not isinstance(dt, torch.dtype):
-        raise ValueError(f"{dtype} is not a valid dtype.")
-    dtype = dt
-
-    print("Initializing lit-llama")
-    config = LLaMAConfig.from_name(model_size)
-
-    with EmptyInitOnDevice(device="meta", dtype=dtype):
-        model = LLaMA(config)
-
-    qkv_size = model.transformer.h[0].attn.c_attn.weight.shape[0] // 3
-
-    # initialize a new empty state dict to hold our new weights
-    sd_meta = model.state_dict()
-    sd = {}
-
-    # Load the json file containing weight mapping
-    pytorch_bin_map_json_path = checkpoint_dir / "pytorch_model.bin.index.json"
-    with open(pytorch_bin_map_json_path) as json_map:
-        bin_index = json.load(json_map)
-    bin_files = set(checkpoint_dir / bin for bin in bin_index["weight_map"].values())
-    if not bin_files:
-        raise ValueError(f"Expected {str(checkpoint_dir)!r} to contain .bin files")
-
-    def permute(w):
-        dim = config.n_embd
-        w = w._load_tensor().to(dtype)
-        return (
-            w.view(config.n_head, 2, dim // config.n_head // 2, dim)
-            .transpose(1, 2)
-            .reshape(dim, dim)
-        )
-
-    weight_map = {
-        "self_attn.o_proj.weight": "attn.c_proj.weight",
-        "self_attn.q_proj.weight": "attn.c_attn.weight",
-        "self_attn.k_proj.weight": "attn.c_attn.weight",
-        "self_attn.v_proj.weight": "attn.c_attn.weight",
-        "mlp.gate_proj.weight": "mlp.c_fc1.weight",
-        "mlp.up_proj.weight": "mlp.c_fc2.weight",
-        "mlp.down_proj.weight": "mlp.c_proj.weight",
-        "input_layernorm.weight": "rms_1.scale",
-        "post_attention_layernorm.weight": "rms_2.scale",
-        "model.embed_tokens.weight": "transformer.wte.weight",
-        "model.norm.weight": "transformer.ln_f.scale",
-        "lm_head.weight": "lm_head.weight",
-    }
-
-    print(f"Saving to disk at {output_dir}")
-    unprocessed_weights = collections.defaultdict(dict)
-
-    with incremental_save(output_dir / "lit-llama.pth") as saver:
-        # for checkpoints that split the QKV across several files, we need to keep all the bin files
-        # open, so we use `ExitStack` to close them all together at the end
-        with contextlib.ExitStack() as stack:
-            for bin_file in bin_files:
-                print("Processing", bin_file)
-                hf_weights = stack.enter_context(lazy_load(bin_file))
-                for name, param in hf_weights.items():
-                    skip = False
-                    if "rotary_emb.inv_freq" in name:
-                        continue
-                    if "model.layers" in name:
-                        block_id = int(name.split(".")[2])
-                        from_name = ".".join(name.split(".")[3:])
-                        to_name = weight_map[from_name]
-                        sd_key = f"transformer.h.{block_id}.{to_name}"
-
-                        if "q_proj" in name:
-                            unprocessed_weights[sd_key]["q_proj"] = param
-                            skip = True
-                        elif "k_proj" in name:
-                            unprocessed_weights[sd_key]["k_proj"] = param
-                            skip = True
-                        elif "v_proj" in name:
-                            unprocessed_weights[sd_key]["v_proj"] = param
-                            skip = True
-                        if skip and len(unprocessed_weights[sd_key]) == 3:
-                            w = torch.empty(
-                                sd_meta[sd_key].shape, dtype=sd_meta[sd_key].dtype
-                            )
-                            w[:qkv_size] = permute(unprocessed_weights[sd_key]["q_proj"])
-                            w[qkv_size:-qkv_size] = permute(
-                                unprocessed_weights[sd_key]["k_proj"]
-                            )
-                            w[-qkv_size:] = (
-                                unprocessed_weights[sd_key]["v_proj"]
-                                ._load_tensor()
-                                .to(dtype)
-                            )
-                            sd[sd_key] = w
-                            del unprocessed_weights[sd_key]
-                            skip = False
-                        else:
-                            sd[sd_key] = param._load_tensor().to(dtype)
-                    else:
-                        sd_key = weight_map[name]
-                        sd[sd_key] = param._load_tensor().to(dtype)
-                    if not skip:
-                        sd[sd_key] = saver.store_early(sd[sd_key])
-                gc.collect()
-        saver.save(sd)
-
-    assert len(unprocessed_weights) == 0, f"unexpected partial weights {list(unprocessed_weights)}"
-    if verify:
-        try:
-            from transformers import LlamaForCausalLM
-        except ImportError:
-            raise ImportError("verify=True requires transformers to be installed, please `pip install transformers`")
-        print("Verifying...")
-
-        token_sample = torch.randint(0, config.vocab_size, size=(1, config.block_size), dtype=torch.int64)
-        out = model(token_sample)
-        del model
-        gc.collect()
-
-        print("Loading original model for comparison")
-        model_hf = LlamaForCausalLM.from_pretrained(checkpoint_dir)
-        out_hf = model_hf(token_sample)["logits"]
-
-        print("Comparing outputs")
-        assert out.device.type == out_hf.device.type
-        assert out.dtype == out_hf.dtype
-        assert torch.testing.assert_close(out, out_hf)
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(convert_hf_checkpoint)
-

lit-llama/scripts/convert_lora_weights.py

@@ -1,95 +0,0 @@
-import sys
-import time
-from pathlib import Path
-from typing import Optional
-
-import lightning as L
-import torch
-import torch.nn as nn
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from lit_llama import LLaMA
-from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
-from lit_llama.lora import lora
-
-def del_lora_state_dict(model: nn.Module):
-    base_model_dict = model.state_dict()
-    key_to_delete = [k for k in base_model_dict if "lora_" in k]
-    for del_key in key_to_delete:
-        del base_model_dict[del_key]
-    return base_model_dict
-
-
-def lora_model_lookup(checkpoint: dict) -> int:
-    """Returns the LoRA rank from the adapter checkpoint.
-
-    """
-    return checkpoint["transformer.h.0.attn.c_attn.lora_B"].shape[1]
-     
-
-def main(
-    accelerator: str = "auto",
-    lora_path: Optional[Path] = None,
-    checkpoint_path: Optional[Path] = None,
-    dtype: str = "bfloat16",
-) -> None:
-    """Merges lora weights to base model.
-
-    Args:
-        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.
-        dtype: `torch.dtype` to work with
-    """
-    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")
-
-    assert lora_path.is_file()
-    assert checkpoint_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
-
-    print("Loading model ...", file=sys.stderr)
-    t0 = time.time()
-
-    with (lazy_load(checkpoint_path) as pretrained_checkpoint,
-          lazy_load(lora_path) as lora_checkpoint):
-        name = llama_model_lookup(pretrained_checkpoint)
-        rank = lora_model_lookup(lora_checkpoint)
-
-        with EmptyInitOnDevice(
-                device=fabric.device, dtype=dtype
-        ), lora(r=rank, alpha=16, dropout=0.05, 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 lora weights
-            model.load_state_dict(lora_checkpoint, strict=False)
-
-    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
-
-    model.eval()
-    base_model_dict = del_lora_state_dict(model)
-    save_path = lora_path.with_stem(f"{lora_path.stem}-lora-merged-weights")
-    print("Saving LoRA to base model weights ...")
-    torch.save(base_model_dict, save_path)
-    print(f"Model saved at {save_path}")
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(main)

lit-llama/scripts/download.py

@@ -1,34 +0,0 @@
-import os
-from typing import Optional
-from urllib.request import urlretrieve
-
-files = {
-    "original_model.py": "https://gist.githubusercontent.com/lantiga/fd36849fb1c498da949a0af635318a7b/raw/7dd20f51c2a1ff2886387f0e25c1750a485a08e1/llama_model.py",
-    "original_adapter.py": "https://gist.githubusercontent.com/awaelchli/546f33fcdb84cc9f1b661ca1ca18418d/raw/e81d8f35fb1fec53af1099349b0c455fc8c9fb01/original_adapter.py",
-}
-
-
-def download_original(wd: str) -> None:
-    for file, url in files.items():
-        filepath = os.path.join(wd, file)
-        if not os.path.isfile(filepath):
-            print(f"Downloading original implementation to {filepath!r}")
-            urlretrieve(url=url, filename=file)
-            print("Done")
-        else:
-            print("Original implementation found. Skipping download.")
-
-
-def download_from_hub(repo_id: Optional[str] = None, local_dir: str = "checkpoints/hf-llama/7B") -> None:
-    if repo_id is None:
-        raise ValueError("Please pass `--repo_id=...`. You can try googling 'huggingface hub llama' for options.")
-
-    from huggingface_hub import snapshot_download
-
-    snapshot_download(repo_id, local_dir=local_dir, local_dir_use_symlinks=False)
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(download_from_hub)

lit-llama/scripts/prepare_alpaca.py

@@ -1,131 +0,0 @@
-"""Implementation derived from https://github.com/tloen/alpaca-lora"""
-import sys
-from pathlib import Path
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-import torch
-import requests
-import json
-from torch.utils.data import random_split
-from lit_llama.tokenizer import Tokenizer
-from tqdm import tqdm
-
-
-DATA_FILE = "https://raw.githubusercontent.com/tloen/alpaca-lora/main/alpaca_data_cleaned_archive.json"
-DATA_FILE_NAME = "alpaca_data_cleaned_archive.json"
-IGNORE_INDEX = -1
-
-
-def prepare(
-    destination_path: Path = Path("data/alpaca"), 
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    test_split_size: int = 2000,
-    max_seq_length: int = 256,
-    seed: int = 42,
-    mask_inputs: bool = False,  # as in alpaca-lora
-    data_file_name: str = DATA_FILE_NAME
-) -> None:
-    """Prepare the Alpaca dataset for instruction tuning.
-    
-    The output is a training and validation dataset saved as `train.pt` and `val.pt`,
-    which stores the preprocessed and tokenized prompts and labels.
-    """
-    
-    destination_path.mkdir(parents=True, exist_ok=True)
-    file_path = destination_path / data_file_name
-    download(file_path)
-
-    # TODO: If we don't have the Meta weights, where do we get the tokenizer from?
-    tokenizer = Tokenizer(tokenizer_path)
-    
-    with open(file_path, "r") as file:
-        data = json.load(file)
-
-    # Partition the dataset into train and test
-    train_split_size = len(data) - test_split_size
-    train_set, test_set = random_split(
-        data, 
-        lengths=(train_split_size, test_split_size),
-        generator=torch.Generator().manual_seed(seed),
-    )
-    train_set, test_set = list(train_set), list(test_set)
-
-    print(f"train has {len(train_set):,} samples")
-    print(f"val has {len(test_set):,} samples")
-
-    print("Processing train split ...")
-    train_set = [prepare_sample(sample, tokenizer, max_seq_length, mask_inputs) for sample in tqdm(train_set)]
-    torch.save(train_set, file_path.parent / "train.pt")
-
-    print("Processing test split ...")
-    test_set = [prepare_sample(sample, tokenizer, max_seq_length, mask_inputs) for sample in tqdm(test_set)]
-    torch.save(test_set, file_path.parent / "test.pt")
-
-
-def download(file_path: Path):
-    """Downloads the raw json data file and saves it in the given destination."""
-    if file_path.exists():
-        return
-    with open(file_path, "w") as f:
-        f.write(requests.get(DATA_FILE).text)
-
-
-def prepare_sample(example: dict, tokenizer: Tokenizer, max_length: int, mask_inputs: bool = True):
-    """Processes a single sample.
-    
-    Each sample in the dataset consists of:
-    - instruction: A string describing the task
-    - input: A string holding a special input value for the instruction.
-        This only applies to some samples, and in others this is empty.
-    - output: The response string
-
-    This function processes this data to produce a prompt text and a label for
-    supervised training. The input text is formed as a single message including all
-    the instruction, the input (optional) and the response.
-    The label/target is the same message but can optionally have the instruction + input text
-    masked out (mask_inputs=True).
-
-    Finally, both the prompt and the label get tokenized. If desired, all tokens
-    in the label that correspond to the original input prompt get masked out (default).
-    """
-    full_prompt = generate_prompt(example)
-    full_prompt_and_response = full_prompt + example["output"]
-    encoded_full_prompt = tokenize(tokenizer, full_prompt, max_length=max_length, eos=False)
-    encoded_full_prompt_and_response = tokenize(tokenizer, full_prompt_and_response, eos=True, max_length=max_length)
-
-    # The labels are the full prompt with response, but with the prompt masked out
-    labels = encoded_full_prompt_and_response.clone()
-    if mask_inputs:
-        labels[:len(encoded_full_prompt)] = IGNORE_INDEX
-
-    return {**example, "input_ids": encoded_full_prompt_and_response, "input_ids_no_response": encoded_full_prompt, "labels": labels}
-
-
-def tokenize(tokenizer: Tokenizer, string: str, max_length: int, eos=True) -> torch.Tensor:
-    return tokenizer.encode(string, bos=True, eos=eos, max_length=max_length)
-
-
-def generate_prompt(example):
-    """Generates a standardized message to prompt the model with an instruction, optional input and a
-    'response' field."""
-
-    if example["input"]:
-        return (
-            "Below is an instruction that describes a task, paired with an input that provides further context. "
-            "Write a response that appropriately completes the request.\n\n"
-            f"### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:"
-        )
-    return (
-        "Below is an instruction that describes a task. "
-        "Write a response that appropriately completes the request.\n\n"
-        f"### Instruction:\n{example['instruction']}\n\n### Response:"
-    )
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(prepare)

lit-llama/scripts/prepare_any_text.py

@@ -1,97 +0,0 @@
-"""Implementation derived from https://github.com/tloen/alpaca-lora"""
-import sys
-from pathlib import Path
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-import torch
-import requests
-import json
-from torch.utils.data import random_split
-from lit_llama.tokenizer import Tokenizer
-from tqdm import tqdm
-
-
-IGNORE_INDEX = -1
-
-DATA_FILE_NAME = "input.txt"
-
-
-def prepare(
-    destination_path: Path = Path("data/any"),
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    test_split_ratio: float = 0.9,  # default 90% train, 10% validation
-    max_seq_length: int = 256,
-    seed: int = 42,
-    data_file_name: str = DATA_FILE_NAME,
-) -> None:
-    """Prepare any dataset for finetuning (akin to Shakespheare full tuning).
-
-    The output is a training and validation dataset saved as `train.pt` and `val.pt`,
-    which stores the preprocessed and tokenized prompts and labels.
-    """
-
-    destination_path.mkdir(parents=True, exist_ok=True)
-    file_path = destination_path / data_file_name
-    if not file_path.exists():
-        raise AssertionError(f"{data_file_name} is provided by the user")
-
-    # TODO: If we don't have the Meta weights, where do we get the tokenizer from?
-    tokenizer = Tokenizer(tokenizer_path)
-
-    data = []
-
-    with open(file_path, "r") as input_file:
-        for line in input_file.readlines():
-            data.append(line)
-
-    # Partition the dataset into train and test
-    train_split_size = int(len(data) * test_split_ratio)
-    test_split_size = len(data) - train_split_size
-    train_set, test_set = random_split(
-        data,
-        lengths=(train_split_size, test_split_size),
-        generator=torch.Generator().manual_seed(seed),
-    )
-    train_set, test_set = list(train_set), list(test_set)
-
-    print(f"train has {len(train_set):,} samples")
-    print(f"val has {len(test_set):,} samples")
-
-    print("Processing train split ...")
-    train_set = [
-        prepare_line(line, tokenizer, max_seq_length) for line in tqdm(train_set)
-    ]
-    torch.save(train_set, file_path.parent / "train.pt")
-
-    print("Processing test split ...")
-    test_set = [
-        prepare_line(line, tokenizer, max_seq_length) for line in tqdm(test_set)
-    ]
-    torch.save(test_set, file_path.parent / "test.pt")
-
-
-def prepare_line(line: str, tokenizer: Tokenizer, max_length: int):
-    """Processes a single sample.
-
-    This function processes the line to produce the tokenized version of it.
-    """
-    encoded_full_prompt = tokenize(tokenizer, line, max_length=max_length, eos=False)
-    return {
-        "input_ids": encoded_full_prompt,
-        "labels": encoded_full_prompt,
-    }
-
-
-def tokenize(
-    tokenizer: Tokenizer, string: str, max_length: int, eos=True
-) -> torch.Tensor:
-    return tokenizer.encode(string, bos=True, eos=eos, max_length=max_length)
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(prepare)

lit-llama/scripts/prepare_dolly.py

@@ -1,133 +0,0 @@
-"""Implementation derived from https://github.com/tloen/alpaca-lora"""
-import sys
-from pathlib import Path
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-import torch
-import requests
-import json
-from torch.utils.data import random_split
-from lit_llama.tokenizer import Tokenizer
-from tqdm import tqdm
-
-
-DATA_FILE = "https://huggingface.co/datasets/databricks/databricks-dolly-15k/resolve/main/databricks-dolly-15k.jsonl"
-DATA_FILE_NAME = "dolly_data_cleaned.json"
-IGNORE_INDEX = -1
-
-
-def prepare(
-    destination_path: Path = Path("data/dolly"), 
-    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
-    test_split_size: int = 2000,
-    max_seq_length: int = 1024,
-    seed: int = 42,
-    mask_inputs: bool = False,  # as in alpaca-lora
-) -> None:
-    """Prepare the Dolly dataset for instruction tuning.
-    
-    The output is a training and validation dataset saved as `train.pt` and `val.pt`,
-    which stores the preprocessed and tokenized prompts and labels.
-    """
-    
-    destination_path.mkdir(parents=True, exist_ok=True)
-    file_path = destination_path / DATA_FILE_NAME
-    download(file_path)
-
-    # TODO: If we don't have the Meta weights, where do we get the tokenizer from?
-    tokenizer = Tokenizer(tokenizer_path)
-
-    with open(file_path, "r") as file:
-        data = file.readlines()
-        data = [json.loads(line) for line in data]
-    for item in data:
-        item["input"] = item.pop("context")
-        item["output"] = item.pop("response")
-
-    # Partition the dataset into train and test
-    train_split_size = len(data) - test_split_size
-    train_set, test_set = random_split(
-        data, 
-        lengths=(train_split_size, test_split_size),
-        generator=torch.Generator().manual_seed(seed),
-    )
-    train_set, test_set = list(train_set), list(test_set)
-
-    print(f"train has {len(train_set):,} samples")
-    print(f"val has {len(test_set):,} samples")
-
-    print("Processing train split ...")
-    train_set = [prepare_sample(sample, tokenizer, max_seq_length, mask_inputs) for sample in tqdm(train_set)]
-    torch.save(train_set, file_path.parent / "train.pt")
-
-    print("Processing test split ...")
-    test_set = [prepare_sample(sample, tokenizer, max_seq_length, mask_inputs) for sample in tqdm(test_set)]
-    torch.save(test_set, file_path.parent / "test.pt")
-
-
-def download(file_path: Path):
-    """Downloads the raw json data file and saves it in the given destination."""
-    if file_path.exists():
-        return
-    with open(file_path, "w") as f:
-        f.write(requests.get(DATA_FILE).text)
-
-
-def prepare_sample(example: dict, tokenizer: Tokenizer, max_length: int, mask_inputs: bool = True):
-    """Processes a single sample.
-    
-    Each sample in the dataset consists of:
-    - instruction: A string describing the task
-    - input: A string holding a special input value for the instruction.
-        This only applies to some samples, and in others this is empty.
-    - output: The response string
-
-    This function processes this data to produce a prompt text and a label for
-    supervised training. The prompt text is formed as a single message including both
-    the instruction and the input. The label/target is the same message but with the
-    response attached.
-
-    Finally, both the prompt and the label get tokenized. If desired, all tokens
-    in the label that correspond to the original input prompt get masked out (default).
-    """
-    full_prompt = generate_prompt(example)
-    full_prompt_and_response = full_prompt + example["output"]
-    encoded_full_prompt = tokenize(tokenizer, full_prompt, max_length=max_length, eos=False)
-    encoded_full_prompt_and_response = tokenize(tokenizer, full_prompt_and_response, eos=True, max_length=max_length)
-
-    # The labels are the full prompt with response, but with the prompt masked out
-    labels = encoded_full_prompt_and_response.clone()
-    if mask_inputs:
-        labels[:len(encoded_full_prompt)] = IGNORE_INDEX
-
-    return {**example, "input_ids": encoded_full_prompt_and_response, "input_ids_no_response": encoded_full_prompt, "labels": labels}
-
-
-def tokenize(tokenizer: Tokenizer, string: str, max_length: int, eos=True) -> torch.Tensor:
-    return tokenizer.encode(string, bos=True, eos=eos, max_length=max_length)
-
-
-def generate_prompt(example):
-    """Generates a standardized message to prompt the model with an instruction, optional input and a
-    'response' field."""
-
-    if example["input"]:
-        return (
-            f"Below is an instruction that describes a task, paired with an input that provides further context. "
-            "Write a response that appropriately completes the request.\n\n"
-            f"### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:"
-        )
-    return (
-        f"Below is an instruction that describes a task. "
-        "Write a response that appropriately completes the request.\n\n"
-        f"### Instruction:\n{example['instruction']}\n\n### Response:"
-    )
-
-
-if __name__ == "__main__":
-    from jsonargparse import CLI
-
-    CLI(prepare)