+
+
+# ⚡ Lit-LLaMA ️
+
+
+
+ [](https://dev.azure.com/Lightning-AI/lit%20Models/_build/latest?definitionId=49&branchName=main) [](https://github.com/Lightning-AI/lit-llama/blob/master/LICENSE) [](https://discord.gg/VptPCZkGNa)
+
+
+
+
+
+Join us and start contributing, especially on the following areas:
+
+- [ ] [Pre-training](https://github.com/Lightning-AI/lit-llama/labels/pre-training)
+- [ ] [Fine-tuning (full and LoRA)](https://github.com/Lightning-AI/lit-llama/labels/fine-tuning)
+- [ ] [Quantization](https://github.com/Lightning-AI/lit-llama/labels/quantization)
+- [ ] [Sparsification](https://github.com/Lightning-AI/lit-llama/labels/sparsification)
+
+Look at `train.py` for a starting point towards pre-training / fine-tuning using [Lightning Fabric](https://lightning.ai/docs/fabric/stable/).
+
+We welcome all individual contributors, regardless of their level of experience or hardware. Your contributions are valuable, and we are excited to see what you can accomplish in this collaborative and supportive environment.
+
+Unsure about contributing? Check out our [Contributing to Lit-LLaMA: A Hitchhiker’s Guide to the Quest for Fully Open-Source AI](https://lightning.ai/pages/community/tutorial/contributing-to-lit-llama-a-hitchhikers-guide-to-the-quest-for-fully-open-source-ai/) guide.
+
+Don't forget to [join our Discord](https://discord.gg/VptPCZkGNa)!
+
+## Acknowledgements
+
+- [@karpathy](https://github.com/karpathy) for [nanoGPT](https://github.com/karpathy/nanoGPT)
+- [@FacebookResearch](https://github.com/facebookresearch) for the original [LLaMA implementation](https://github.com/facebookresearch/llama)
+- [@TimDettmers](https://github.com/TimDettmers) for [bitsandbytes](https://github.com/TimDettmers/bitsandbytes)
+- [@Microsoft](https://github.com/microsoft) for [LoRA](https://github.com/microsoft/LoRA)
+- [@IST-DASLab](https://github.com/IST-DASLab) for [GPTQ](https://github.com/IST-DASLab/gptq)
+
+## License
+
+Lit-LLaMA is released under the [Apache 2.0](https://github.com/Lightning-AI/lightning-llama/blob/main/LICENSE) license.
diff --git a/lit-llama/evaluate/adapter.py b/lit-llama/evaluate/adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..9aa378279cbb216ac120f9c7da29b14a761947df
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/evaluate/adapter_v2.py b/lit-llama/evaluate/adapter_v2.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1337a7790d91859bb21fd85717f69161f3a0e97
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/evaluate/full.py b/lit-llama/evaluate/full.py
new file mode 100644
index 0000000000000000000000000000000000000000..48d5fb89a9ccfb1ba6be31f2132095a461012726
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/evaluate/lora.py b/lit-llama/evaluate/lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c2192b3f8e4e4ed646c32645b47337e88b518c1
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/finetune/adapter.py b/lit-llama/finetune/adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc815830437cce0a53f3d9d7a80145977d2866a1
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/finetune/adapter_v2.py b/lit-llama/finetune/adapter_v2.py
new file mode 100644
index 0000000000000000000000000000000000000000..c686cd157c77ac264722579dc9de97b588db564c
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/finetune/full.py b/lit-llama/finetune/full.py
new file mode 100644
index 0000000000000000000000000000000000000000..9248e8de696e4894703cbb5bbe8795dd5d4e71b2
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/finetune/lora.py b/lit-llama/finetune/lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..18737015c5d2290406a4558248d8b8b311cb5bf4
--- /dev/null
+++ b/lit-llama/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)
diff --git a/lit-llama/generate.py b/lit-llama/generate.py
new file mode 100644
index 0000000000000000000000000000000000000000..91a7a6e4233485b704449bc7ef7421719850bdeb
--- /dev/null
+++ b/lit-llama/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