Upload 58 files

evaluate/adapter.py

@@ -0,0 +1,164 @@
+# This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
+# Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
+import math
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import lightning as L
+import torch
+import tqdm
+
+# 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)

evaluate/adapter_v2.py

@@ -0,0 +1,161 @@
+# 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)

evaluate/full.py

@@ -0,0 +1,147 @@
+# 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)

evaluate/lora.py

@@ -0,0 +1,172 @@
+# 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)

finetune/adapter.py

@@ -0,0 +1,262 @@
+"""
+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)

finetune/adapter_v2.py

@@ -0,0 +1,266 @@
+"""
+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)

finetune/full.py

@@ -0,0 +1,224 @@
+"""
+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)

finetune/lora.py

@@ -0,0 +1,218 @@
+"""
+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)

generate.py

@@ -0,0 +1,170 @@
+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)

generate/adapter.py

@@ -0,0 +1,106 @@
+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)

generate/adapter_v2.py

@@ -0,0 +1,108 @@
+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)

generate/full.py

@@ -0,0 +1,103 @@
+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)

generate/lora.py

@@ -0,0 +1,118 @@
+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)

howto/convert_lora_weights.md

@@ -0,0 +1,19 @@
+# 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
+```
+

howto/customize_paths.md

@@ -0,0 +1,33 @@
+## Customize paths
+
+The project is setup to use specific paths to read the original weights and save checkpoints etc.
+
+For all scripts, you can run
+
+```shell
+python script.py -h
+```
+
+to get a list of available options. For instance, here's how you would modify the checkpoint dir:
+
+```shell
+python scripts/convert_checkpoint.py --checkpoint_dir "data/checkpoints/foo"
+```
+
+Note that this change will need to be passed along to subsequent steps, for example:
+
+```shell
+python 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.

howto/download_weights.md

@@ -0,0 +1,130 @@
+## 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).

howto/finetune_adapter.md

@@ -0,0 +1,109 @@
+# 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).

howto/finetune_adapter_v2.md

@@ -0,0 +1,114 @@
+# 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).

howto/finetune_full.md

@@ -0,0 +1,106 @@
+# 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).

howto/finetune_lora.md

@@ -0,0 +1,90 @@
+# 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).

howto/inference.md

@@ -0,0 +1,43 @@
+# 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
+```

howto/tpus.md

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

howto/train_redpajama.md

@@ -0,0 +1,133 @@
+# Pre-train LLaMA on RedPajama
+
+This howto will walk you through setting up the RedPajama dataset and launching the pre-training script.
+
+## What's RedPajama
+
+[RedPajama](https://github.com/togethercomputer/RedPajama-Data) is an open-source reproduction of the original LLaMA training dataset.
+
+It contains a total of 1.2 trillion tokens, divided into
+
+```text
+Commoncrawl   878B
+C4            175B
+GitHub         59B
+Books          26B
+ArXiv          28B
+Wikipedia      24B
+StackExchange  20B
+```
+
+The [RedPajama repo](https://github.com/togethercomputer/RedPajama-Data) contains the source code for collecting and preparing
+the dataset, and it is Apache 2.0 licensed.
+
+The data itself is licensed according to the original licenses with which its invidivdual parts were released.
+The GitHub datasets are limited to MIT, BSD, or Apache 2.0 repositories.
+
+Along with the full [RedPajama-1T dataset](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T),
+the [RedPajama-1T-Sample](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample) 1B sample dataset
+is also available for development.
+
+You can download the data using git lfs:
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com): git lfs install
+git clone https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T data/RedPajama-Data-1T
+```
+
+```bash
+# Make sure you have git-lfs installed (https://git-lfs.com): git lfs install
+git clone https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample data/RedPajama-Data-1T-Sample
+```
+
+## Prepare RedPajama for training
+
+The dataset consists of 2084 `jsonl` files (the sample dataset contains 11). In order to start pre-training lit-llama
+on it, you need to read, tokenize, and write the data in binary chunks. This will leverage the `PackedDataset`
+streaming dataset that comes with lit-llama.
+
+Do to so, run
+
+```bash
+python scripts/prepare_redpajama.py --source_path data/RedPajama-Data-1T --tokenizer_path checkpoints/lit-llama/tokenizer.model --destination_path data/lit-redpajama
+```
+
+or
+
+```bash
+python scripts/prepare_redpajama.py --source_path data/RedPajama-Data-1T-Sample --tokenizer_path checkpoints/lit-llama/tokenizer.model --destination_path data/lit-redpajama-sample --sample True
+```
+
+for the sample dataset.
+
+In the above we are assuming that you will be using the same tokenizer as used in LLaMA, but any trained [SentencePiece](https://github.com/google/sentencepiece) tokenizer with a 32000 vocabulary size will do here.
+
+The script will take a while to run, so time for :tea:
+
+## Pre-training
+
+Running the pre-training script requires at least 4 GPUs with 40GB+ each (A100).
+
+```bash
+python 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.

howto/unstructured_dataset.md

@@ -0,0 +1,18 @@
+# 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/__init__.py

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

lit_llama/adapter.py

@@ -0,0 +1,313 @@
+"""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/adapter_v2.py

@@ -0,0 +1,45 @@
+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/lora.py

@@ -0,0 +1,476 @@
+# 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/model.py

@@ -0,0 +1,321 @@
+"""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/packed_dataset.py

@@ -0,0 +1,260 @@
+# 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/quantization.py

@@ -0,0 +1,614 @@
+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/tokenizer.py

@@ -0,0 +1,49 @@
+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/utils.py

@@ -0,0 +1,496 @@
+"""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()

pretrain/redpajama.py

@@ -0,0 +1,321 @@
+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)

pretrain/shakespeare.py

@@ -0,0 +1,166 @@
+"""
+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()

quantize/gptq.py

@@ -0,0 +1,238 @@
+# 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)

scripts/convert_checkpoint.py

@@ -0,0 +1,141 @@
+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)

scripts/convert_hf_checkpoint.py

@@ -0,0 +1,167 @@
+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)
+

scripts/convert_lora_weights.py

@@ -0,0 +1,95 @@
+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)

scripts/download.py

@@ -0,0 +1,34 @@
+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)

scripts/prepare_alpaca.py

@@ -0,0 +1,131 @@
+"""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)

scripts/prepare_any_text.py

@@ -0,0 +1,97 @@
+"""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)

scripts/prepare_dolly.py

@@ -0,0 +1,133 @@
+"""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)

scripts/prepare_redpajama.py

@@ -0,0 +1,181 @@
+import json
+import glob
+import os
+from pathlib import Path
+import sys
+
+# support running without installing as a package
+wd = Path(__file__).parent.parent.resolve()
+sys.path.append(str(wd))
+
+import numpy as np
+from tqdm import tqdm
+
+from lit_llama import Tokenizer
+import lit_llama.packed_dataset as packed_dataset
+
+
+filenames_sample = [
+    "arxiv_sample.jsonl",
+    "book_sample.jsonl",
+    "c4_sample.jsonl",
+    "cc_2019-30_sample.jsonl",
+    "cc_2020-05_sample.jsonl",
+    "cc_2021-04_sample.jsonl",
+    "cc_2022-05_sample.jsonl",
+    "cc_2023-06_sample.jsonl",
+    "github_sample.jsonl",
+    "stackexchange_sample.jsonl",
+    "wikipedia_sample.jsonl",
+]
+
+filename_sets = {
+    "arxiv": "arxiv/arxiv*",
+    "book": "book/book*",
+    "c4": "c4/c4-train*",
+    "common_crawl": "common_crawl/*",
+    "github": "github/filtered*",
+    "stackexchange": "stackexchange/stackexchange*",
+    "wikipedia": "wikipedia/wiki*",
+}
+
+
+def prepare_sample(
+    source_path: Path,
+    tokenizer_path: Path,
+    destination_path: Path,
+    chunk_size: int,
+    match = ""
+) -> None:
+    """Prepare the "Red Pajama" dataset. We assume tokenizer has been trained (i.e. we reuse LLaMA's tokenizer model)."""
+    destination_path.mkdir(parents=True, exist_ok=True)
+
+    tokenizer = Tokenizer(tokenizer_path)
+
+    for name in filenames_sample:
+        if match and match not in name:
+            continue
+
+        filepath = source_path / name
+
+        if not filepath.is_file():
+            raise RuntimeError(
+                f"Input file not found at {filepath}. \n"
+                "Make sure you download the data, e.g. wget -i https://data.together.xyz/redpajama-data-1T/v1.0.0/urls.txt or through \n"
+                "https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T \n"
+                "https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample \n"
+            )
+
+        prefix, _ = os.path.splitext(name)
+
+        builder = packed_dataset.PackedDatasetBuilder(
+            outdir=destination_path,
+            prefix=prefix,
+            chunk_size=chunk_size,
+            sep_token=tokenizer.bos_id,
+            dtype="auto",
+            vocab_size=tokenizer.vocab_size,
+        )
+
+        print(f"Processing {name}")
+
+        with open(filepath, encoding="utf-8") as f:
+            for row in tqdm(f):
+                text = json.loads(row)["text"]
+                text_ids = tokenizer.encode(text)
+                builder.add_array(np.array(text_ids, dtype=builder.dtype))
+
+        builder.write_reminder()
+
+
+def prepare_full(
+    source_path: Path,
+    tokenizer_path: Path,
+    destination_path: Path,
+    chunk_size: int,
+    match: str = ""
+) -> None:
+    """Prepare the "Red Pajama" dataset. We assume tokenizer has been trained (i.e. we reuse LLaMA's tokenizer model)."""
+    import zstandard as zstd
+
+    destination_path.mkdir(parents=True, exist_ok=True)
+
+    tokenizer = Tokenizer(tokenizer_path)
+
+    for set_name, pattern in filename_sets.items():
+        if match and match not in set_name:
+            continue
+
+        is_cc = set_name == "common_crawl"
+
+        filenames = glob.glob(os.path.join(source_path, pattern), recursive=True)
+
+        if not filenames:
+            raise RuntimeError(
+                f"No files matching {pattern} found at {source_path}. \n"
+                "Make sure you download the data, e.g. wget -i https://data.together.xyz/redpajama-data-1T/v1.0.0/urls.txt or through \n"
+                "https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T \n"
+                "https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample \n"
+            )
+
+        builder = packed_dataset.PackedDatasetBuilder(
+            outdir=destination_path,
+            prefix=set_name,
+            chunk_size=chunk_size,
+            sep_token=tokenizer.bos_id,
+            dtype="auto",
+            vocab_size=tokenizer.vocab_size,
+        )
+
+        for name in filenames:
+            filepath = source_path / name
+
+            print(f"Processing {name}")
+
+            if is_cc:
+                with zstd.open(open(filepath, "rb"), "rt", encoding="utf-8") as f:
+                    for row in tqdm(f):
+                        text = json.loads(row)["text"]
+                        text_ids = tokenizer.encode(text)
+                        builder.add_array(np.array(text_ids, dtype=builder.dtype))
+            else:
+                with open(filepath, encoding="utf-8") as f:
+                    for row in tqdm(f):
+                        text = json.loads(row)["text"]
+                        text_ids = tokenizer.encode(text)
+                        builder.add_array(np.array(text_ids, dtype=builder.dtype))
+
+        builder.write_reminder()
+
+
+def prepare(
+    source_path: Path = Path("data/RedPajama-Data-1T-Sample"),
+    tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
+    destination_path: Path = Path("data/red_pajama_sample"),
+    chunk_size: int = 2049 * 1024,  # 2048 block size + 1 for causal (from LLama), 1024 blocks
+    sample: bool = False,
+    match: str = "",
+) -> None:
+    """Prepare the "Red Pajama" dataset. We assume tokenizer has been trained (i.e. we reuse LLaMA's tokenizer model)."""
+    if sample:
+        prepare_sample(
+            source_path=source_path,
+            tokenizer_path=tokenizer_path,
+            destination_path=destination_path,
+            chunk_size=chunk_size,
+            match=match,
+        )
+    else:
+        prepare_full(
+            source_path=source_path,
+            tokenizer_path=tokenizer_path,
+            destination_path=destination_path,
+            chunk_size=chunk_size,
+            match=match,
+        )
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    CLI(prepare)

scripts/prepare_shakespeare.py

@@ -0,0 +1,69 @@
+# MIT License
+
+# Copyright (c) 2022 Andrej Karpathy
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+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 numpy as np
+import requests
+
+
+def prepare(destination_path: Path = Path("data/shakespeare")) -> None:
+    """Prepare the "Tiny Shakespeare" dataset."""
+    destination_path.mkdir(parents=True, exist_ok=True)
+
+    # download the tiny shakespeare dataset
+    input_file_path = destination_path / "input.txt"
+    if not input_file_path.exists():
+        data_url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
+        with open(input_file_path, "w") as f:
+            f.write(requests.get(data_url).text)
+
+    with open(input_file_path) as f:
+        data = f.read()
+    n = len(data)
+    train_data = data[: int(n * 0.9)]
+    val_data = data[int(n * 0.9) :]
+
+    from lit_llama import Tokenizer
+
+    Tokenizer.train(input=input_file_path, destination=destination_path, vocab_size=100)
+    tokenizer = Tokenizer(destination_path / "tokenizer.model")
+    train_ids = tokenizer.encode(train_data)
+    val_ids = tokenizer.encode(val_data)
+    print(f"train has {len(train_ids):,} tokens")
+    print(f"val has {len(val_ids):,} tokens")
+
+    # export to bin files
+    train_ids = np.array(train_ids, dtype=np.uint16)
+    val_ids = np.array(val_ids, dtype=np.uint16)
+    train_ids.tofile(destination_path / "train.bin")
+    val_ids.tofile(destination_path / "val.bin")
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    CLI(prepare)

setup.py

@@ -0,0 +1,26 @@
+import os
+
+from setuptools import setup, find_packages
+
+
+_PATH_ROOT = os.path.dirname(__file__)
+
+with open(os.path.join(_PATH_ROOT, "README.md"), encoding="utf-8") as fo:
+    readme = fo.read()
+
+setup(
+    name='lit-llama',
+    version='0.1.0',
+    description='Implementation of the LLaMA language model',
+    author='Lightning AI',
+    url='https://github.com/lightning-AI/lit-llama',
+    install_requires=[
+        "torch>=2.0.0",
+        "lightning @ git+https://github.com/Lightning-AI/lightning@master",
+        "sentencepiece",
+        "bitsandbytes",
+    ],
+    packages=find_packages(),
+    long_description=readme,
+    long_description_content_type="text/markdown",
+)

tests/conftest.py

@@ -0,0 +1,42 @@
+import sys
+from pathlib import Path
+
+import pytest
+
+wd = Path(__file__).parent.parent.absolute()
+
+
+@pytest.fixture()
+def orig_llama():
+    sys.path.append(str(wd))
+
+    from scripts.download import download_original
+
+    download_original(wd)
+
+    import original_model
+
+    return original_model
+
+
+@pytest.fixture()
+def orig_llama_adapter():
+    sys.path.append(str(wd))
+
+    from scripts.download import download_original
+
+    download_original(wd)
+
+    import original_adapter
+
+    return original_adapter
+
+
+@pytest.fixture()
+def lit_llama():
+    # this adds support for running tests without the package installed
+    sys.path.append(str(wd))
+
+    import lit_llama
+
+    return lit_llama

tests/test_adapter.py

@@ -0,0 +1,55 @@
+from dataclasses import asdict
+import pytest
+import sys
+import torch
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="EmptyInitOnDevice on CPU not working for Windows.")
+@pytest.mark.parametrize("model_size", ["7B", "13B", "30B", "65B"])
+def test_config_identical(model_size, lit_llama):
+    import lit_llama.adapter as llama_adapter
+    import lit_llama.model as llama
+    from lit_llama.utils import EmptyInitOnDevice
+
+    llama_config = asdict(llama.LLaMAConfig.from_name(model_size))
+    adapter_config = asdict(llama_adapter.LLaMAConfig.from_name(model_size))
+
+    del adapter_config["adapter_prompt_length"]
+    del adapter_config["adapter_start_layer"]
+    assert adapter_config == llama_config
+
+    with EmptyInitOnDevice():
+        llama_model = llama.LLaMA.from_name(model_size)
+        adapter_model = llama_adapter.LLaMA.from_name(model_size)
+        assert llama_model.lm_head.weight.shape == adapter_model.lm_head.weight.shape
+
+
+def test_adapter_load_gating_factor(lit_llama):
+    """Tests backward-compatible loading of checkpoints after the `gating_factor` was extended per-head
+    in PR #297.
+    """
+    import lit_llama.adapter as llama_adapter
+    from lit_llama.utils import lazy_load
+
+    config = llama_adapter.LLaMAConfig(n_head=4, block_size=100, n_embd=16)
+    attn = llama_adapter.CausalSelfAttention(config=config, block_idx=3)
+
+    # Old checkpoint format
+    state_dict={
+        "gating_factor": torch.tensor(0.42),  # in old checkpoints, this was a scalar
+        "c_attn.weight": torch.zeros(3 * 16, 16),
+        "c_proj.weight": torch.zeros(16, 16),
+        "adapter_wte.weight": torch.zeros(10, 16),
+    }
+    attn.load_state_dict(state_dict=state_dict)
+    assert torch.equal(attn.gating_factor, torch.full((1, 4, 1, 1), 0.42))
+
+    # New checkpoint format
+    state_dict={
+        "gating_factor": torch.tensor([0.42, 0.42, 0.42, 0.42]).reshape(1, 4, 1, 1),
+        "c_attn.weight": torch.zeros(3 * 16, 16),
+        "c_proj.weight": torch.zeros(16, 16),
+        "adapter_wte.weight": torch.zeros(10, 16),
+    }
+    attn.load_state_dict(state_dict=state_dict)
+    assert torch.equal(attn.gating_factor, torch.full((1, 4, 1, 1), 0.42))

tests/test_adapter_v2.py

@@ -0,0 +1,26 @@
+import pytest
+import sys
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="EmptyInitOnDevice on CPU not working for Windows.")
+@pytest.mark.parametrize("model_size", ["7B", "13B", "30B", "65B"])
+def test_config_identical(model_size, lit_llama):
+    import torch.nn as nn
+    import lit_llama.adapter as llama_adapter
+    from lit_llama.adapter_v2 import adapter_v2_linear_with_bias_and_scale
+    import lit_llama.model as llama
+    from lit_llama.utils import EmptyInitOnDevice
+
+    with EmptyInitOnDevice():
+        llama_model = llama.LLaMA.from_name(model_size)
+        adapter_model = llama_adapter.LLaMA.from_name(model_size)
+
+        for module in adapter_model.modules():
+            if isinstance(module, nn.Linear):
+                adapter_v2_linear_with_bias_and_scale(module)
+
+        print(adapter_model.transformer.h[2].attn.c_attn.adapter_bias)
+        assert not hasattr(llama_model.transformer.h[2].attn.c_attn, 'adapter_bias')
+        assert not hasattr(llama_model.transformer.h[2].attn.c_attn, 'adapter_scale')
+        assert hasattr(adapter_model.transformer.h[2].attn.c_attn, 'adapter_bias')
+        assert hasattr(adapter_model.transformer.h[2].attn.c_attn, 'adapter_scale')

tests/test_generate.py

@@ -0,0 +1,117 @@
+import functools
+import subprocess
+import sys
+from contextlib import contextmanager, redirect_stdout
+from io import StringIO
+from pathlib import Path
+from unittest import mock
+from unittest.mock import Mock, call, ANY
+
+import torch
+
+wd = Path(__file__).parent.parent.absolute()
+
+
+@functools.lru_cache(maxsize=1)
+def load_generate_script():
+    sys.path.append(str(wd))
+
+    import generate as generate
+
+    return generate
+
+
+def test_generate():
+    generate = load_generate_script()
+
+    from lit_llama.model import LLaMA, LLaMAConfig
+
+    T, C = 5, 3
+    logits = torch.randn(T, C)
+    input_idx = torch.randint(10, size=(T,))
+
+    config = LLaMAConfig(block_size=128, vocab_size=16, n_layer=1, n_head=4, n_embd=8)
+    model = LLaMA(config)
+    max_new_tokens = 20
+
+    multinomial_results = []
+    original_multinomial = torch.multinomial
+
+    def multinomial(*args, **kwargs):
+        out = original_multinomial(*args, **kwargs)
+        multinomial_results.append(out)
+        return out
+
+    with mock.patch("torch.multinomial", multinomial):
+        out = generate.generate(model, input_idx, max_new_tokens, max_seq_length=10, top_k=4)
+
+    assert out.size(0) == T + max_new_tokens
+    multinomial_results = torch.hstack(multinomial_results)
+    expected = torch.cat((input_idx, multinomial_results))
+    assert out.shape == expected.shape
+    torch.testing.assert_close(out, expected)
+
+
+@mock.patch("torch.cuda.is_bf16_supported", return_value=False)
+def test_main(tmp_path, monkeypatch):
+    generate = load_generate_script()
+
+    checkpoint_path = tmp_path / "ckpt"
+    checkpoint_path.touch()
+    tokenizer_path = tmp_path / "tokenizer"
+    tokenizer_path.touch()
+
+    class FabricMock(Mock):
+        @property
+        def device(self):
+            return torch.device("cpu")
+
+        @contextmanager
+        def init_module(self, empty_init):
+            yield
+
+    monkeypatch.setattr(generate.L, "Fabric", FabricMock)
+    model_mock = Mock()
+    monkeypatch.setattr(generate.LLaMA, "from_name", model_mock)
+    lookup_mock = Mock(return_value="1T")
+    monkeypatch.setattr(generate, "llama_model_lookup", lookup_mock)
+    load_mock = Mock()
+    load_mock.return_value = load_mock
+    load_mock.__enter__ = Mock()
+    load_mock.__exit__ = Mock()
+    monkeypatch.setattr(generate.torch, "load", load_mock)
+    monkeypatch.setattr(generate, "lazy_load", load_mock)
+    tokenizer_mock = Mock()
+    tokenizer_mock.return_value.encode.return_value = torch.tensor([[1, 2, 3]])
+    tokenizer_mock.return_value.decode.return_value = "foo bar baz"
+    monkeypatch.setattr(generate, "Tokenizer", tokenizer_mock)
+    generate_mock = Mock()
+    generate_mock.return_value = torch.tensor([[3, 2, 1]])
+    monkeypatch.setattr(generate, "generate", generate_mock)
+
+    num_samples = 2
+    out = StringIO()
+    with redirect_stdout(out):
+        generate.main(
+            checkpoint_path=checkpoint_path,
+            tokenizer_path=tokenizer_path,
+            temperature=2.0,
+            top_k=2,
+            num_samples=num_samples,
+        )
+
+    model_mock.assert_called_once_with("1T")
+    load_mock.assert_called_once_with(checkpoint_path)
+    tokenizer_mock.assert_called_once_with(tokenizer_path)
+    assert len(tokenizer_mock.return_value.decode.mock_calls) == num_samples
+    assert torch.allclose(tokenizer_mock.return_value.decode.call_args[0][0], generate_mock.return_value)
+    assert generate_mock.mock_calls == [call(ANY, ANY, 50, temperature=2.0, top_k=2)] * num_samples
+    # only the generated result is printed to stdout
+    assert out.getvalue() == "foo bar baz\n" * num_samples
+
+
+def test_cli():
+    cli_path = wd / "generate.py"
+    output = subprocess.check_output([sys.executable, cli_path, "-h"])
+    output = str(output.decode())
+    assert "Generates text samples" in output