2 |
3 |
4 | # ⚡ Lit-LLaMA ️
5 |
6 |  [](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)
7 |
8 |
9 | ⚠️ Warning: Not Actively Maintained
10 |
11 | This repository is no longer actively maintained. For a more up-to-date alternative, please visit the LitGPT project:
12 | https://github.com/Lightning-AI/litgpt, which serves as the successor to this repository.
13 |
14 | Feel free to explore, reuse, or fork, but be aware that no further updates or support will be provided.
15 |
16 |
17 |
18 |
19 |
172 |
173 | Join us and start contributing, especially on the following areas:
174 |
175 | - [ ] [Pre-training](https://github.com/Lightning-AI/lit-llama/labels/pre-training)
176 | - [ ] [Fine-tuning (full and LoRA)](https://github.com/Lightning-AI/lit-llama/labels/fine-tuning)
177 | - [ ] [Quantization](https://github.com/Lightning-AI/lit-llama/labels/quantization)
178 | - [ ] [Sparsification](https://github.com/Lightning-AI/lit-llama/labels/sparsification)
179 |
180 | Look at `train.py` for a starting point towards pre-training / fine-tuning using [Lightning Fabric](https://lightning.ai/docs/fabric/stable/).
181 |
182 | 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.
183 |
184 | 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.
185 |
186 | Don't forget to [join our Discord](https://discord.gg/VptPCZkGNa)!
187 |
188 | ## Acknowledgements
189 |
190 | - [@karpathy](https://github.com/karpathy) for [nanoGPT](https://github.com/karpathy/nanoGPT)
191 | - [@FacebookResearch](https://github.com/facebookresearch) for the original [LLaMA implementation](https://github.com/facebookresearch/llama)
192 | - [@TimDettmers](https://github.com/TimDettmers) for [bitsandbytes](https://github.com/TimDettmers/bitsandbytes)
193 | - [@Microsoft](https://github.com/microsoft) for [LoRA](https://github.com/microsoft/LoRA)
194 | - [@IST-DASLab](https://github.com/IST-DASLab) for [GPTQ](https://github.com/IST-DASLab/gptq)
195 |
196 | ## License
197 |
198 | Lit-LLaMA is released under the [Apache 2.0](https://github.com/Lightning-AI/lightning-llama/blob/main/LICENSE) license.
199 |
--------------------------------------------------------------------------------
/evaluate/adapter.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | # This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
4 | # Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
5 | import math
6 | import sys
7 | import time
8 | from pathlib import Path
9 | from typing import Optional
10 |
11 | import lightning as L
12 | import torch
13 | import tqdm
14 |
15 | # support running without installing as a package
16 | wd = Path(__file__).parent.parent.resolve()
17 | sys.path.append(str(wd))
18 |
19 | from lit_llama import Tokenizer
20 | from lit_llama.adapter import LLaMA
21 | from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
22 | from scripts.prepare_alpaca import generate_prompt
23 |
24 | from datasets import load_dataset
25 |
26 | instruction_tuning = True
27 |
28 |
29 | def load_eval_data(dataset_name: str) -> str:
30 | # this mimics gptq datautils
31 | if dataset_name == "wikitext":
32 | # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
33 | testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
34 | testdata = "\n\n".join(testdata["text"])
35 | elif dataset_name == "ptb":
36 | testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
37 | testdata = "\n\n".join(testdata["sentence"])
38 | elif dataset_name == "c4":
39 | testdata = load_dataset(
40 | "allenai/c4",
41 | "allenai--c4",
42 | data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
43 | split="validation",
44 | )
45 | testdata = " ".join(testdata[:1100]["text"])
46 |
47 | else:
48 | raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
49 | return testdata
50 |
51 |
52 | @torch.inference_mode()
53 | def main(
54 | datasets: str = "wikitext,ptb,c4",
55 | *,
56 | # compilation fails as it does not support torch.complex64 for RoPE
57 | # compile: bool = False,
58 | accelerator: str = "auto",
59 | adapter_path: Path = Path("out/adapter/alpaca/lit-llama-adapter-finetuned.pth"),
60 | checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
61 | tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
62 | dtype: str = "float32",
63 | quantize: Optional[str] = None,
64 | ) -> None:
65 | """Generates text samples based on a pre-trained LLaMA model and tokenizer.
66 |
67 | Args:
68 | datasets: The datasets to use as a comma separated string
69 | # compile: Whether to compile the model.
70 | accelerator: The hardware to run on. Possible choices are:
71 | ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
72 | adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
73 | `finetune_adapter.py`.
74 | checkpoint_path: The checkpoint path to load.
75 | tokenizer_path: The tokenizer path to load.
76 | dtype: The tensor dtype for choosing the floating-point precision
77 | quantize: Whether to quantize the model and using which method:
78 | ``"llm.int8"``: LLM.int8() mode,
79 | ``"gptq.int4"``: GPTQ 4-bit mode.
80 | """
81 | assert adapter_path.is_file()
82 | assert checkpoint_path.is_file()
83 | assert tokenizer_path.is_file()
84 |
85 | fabric = L.Fabric(accelerator=accelerator, devices=1)
86 |
87 | dt = getattr(torch, dtype, None)
88 | if not isinstance(dt, torch.dtype):
89 | raise ValueError(f"{dtype} is not a valid dtype.")
90 | dtype = dt
91 |
92 | print("Loading model ...", file=sys.stderr)
93 | t0 = time.time()
94 | with lazy_load(checkpoint_path) as pretrained_checkpoint, lazy_load(adapter_path) as adapter_checkpoint:
95 | name = llama_model_lookup(pretrained_checkpoint)
96 |
97 | with EmptyInitOnDevice(
98 | device=fabric.device, dtype=dtype, quantization_mode=quantize
99 | ):
100 | model = LLaMA.from_name(name)
101 |
102 | # 1. Load the pretrained weights
103 | model.load_state_dict(pretrained_checkpoint, strict=False)
104 | # 2. Load the fine-tuned adapter weights
105 | model.load_state_dict(adapter_checkpoint, strict=False)
106 |
107 | print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
108 |
109 | model.eval()
110 |
111 | # if compile:
112 | # model = torch.compile(model)
113 |
114 | total_toks = 0
115 | model = fabric.setup_module(model)
116 |
117 | tokenizer = Tokenizer(tokenizer_path)
118 |
119 | for dsname in datasets.split(","):
120 | test_string = load_eval_data(dsname)
121 |
122 | if instruction_tuning:
123 | sample = {"instruction": test_string, "input": input}
124 | test_string = generate_prompt(sample)
125 |
126 | encoded_text = tokenizer.encode(
127 | test_string, bos=True, eos=False, device=fabric.device
128 | )
129 | encoded_text = encoded_text[
130 | None, : 256 * model.config.block_size
131 | ] # add batch dimension, trim like gptq implementation
132 | t0 = time.perf_counter()
133 |
134 | nlls = 0
135 | toks = 0
136 | 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)
137 | for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
138 | inp = encoded_text[:, i : i + block_size]
139 | logits = model(inp)[0]
140 | nll = torch.nn.functional.cross_entropy(
141 | logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
142 | )
143 | toks += inp.size(1) - 1
144 | nlls += nll.item()
145 |
146 | print(encoded_text.shape, logits.shape)
147 | ppl = math.exp(nlls / toks)
148 | print(f"Perplexity on {dsname}: {ppl:.2f}")
149 | total_toks += toks
150 |
151 | t = time.perf_counter() - t0
152 | print(
153 | f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
154 | file=sys.stderr,
155 | )
156 | print(
157 | f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
158 | file=sys.stderr,
159 | )
160 |
161 |
162 | if __name__ == "__main__":
163 | from jsonargparse import CLI
164 |
165 | torch.set_float32_matmul_precision("high")
166 | CLI(main)
167 |
--------------------------------------------------------------------------------
/evaluate/adapter_v2.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | # This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
4 | # Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
5 | import math
6 | import sys
7 | import time
8 | from pathlib import Path
9 | from typing import Optional
10 |
11 | import lightning as L
12 | import torch
13 | import tqdm
14 |
15 | # support running without installing as a package
16 | wd = Path(__file__).parent.parent.resolve()
17 | sys.path.append(str(wd))
18 |
19 | from lit_llama import Tokenizer
20 | from lit_llama.adapter import LLaMA
21 | from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
22 | from lit_llama.adapter_v2 import add_adapter_v2_parameters_to_linear_layers
23 | from scripts.prepare_alpaca import generate_prompt
24 |
25 | from datasets import load_dataset
26 |
27 |
28 | def load_eval_data(dataset_name: str) -> str:
29 | # this mimics gptq datautils
30 | if dataset_name == "wikitext":
31 | # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
32 | testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
33 | testdata = "\n\n".join(testdata["text"])
34 | elif dataset_name == "ptb":
35 | testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
36 | testdata = "\n\n".join(testdata["sentence"])
37 | elif dataset_name == "c4":
38 | testdata = load_dataset(
39 | "allenai/c4",
40 | "allenai--c4",
41 | data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
42 | split="validation",
43 | )
44 | testdata = " ".join(testdata[:1100]["text"])
45 |
46 | else:
47 | raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
48 | return testdata
49 |
50 |
51 | @torch.inference_mode()
52 | def main(
53 | datasets: str = "wikitext,ptb,c4",
54 | *,
55 | accelerator: str = "auto",
56 | adapter_path: Path = Path("out/adapter_v2/alpaca/lit-llama-adapter-finetuned.pth"),
57 | checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
58 | tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
59 | dtype: str = "float32",
60 | quantize: Optional[str] = None,
61 | ) -> None:
62 | """Generates text samples based on a pre-trained LLaMA model and tokenizer.
63 |
64 | Args:
65 | datasets: The datasets to use as a comma separated string
66 | accelerator: The hardware to run on. Possible choices are:
67 | ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
68 | adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
69 | `finetune_adapter_v2.py`.
70 | checkpoint_path: The checkpoint path to load.
71 | tokenizer_path: The tokenizer path to load.
72 | dtype: The tensor dtype for choosing the floating-point precision
73 | quantize: Whether to quantize the model and using which method:
74 | ``"llm.int8"``: LLM.int8() mode,
75 | ``"gptq.int4"``: GPTQ 4-bit mode.
76 | """
77 | assert adapter_path.is_file()
78 | assert checkpoint_path.is_file()
79 | assert tokenizer_path.is_file()
80 |
81 | fabric = L.Fabric(accelerator=accelerator, devices=1)
82 |
83 | dt = getattr(torch, dtype, None)
84 | if not isinstance(dt, torch.dtype):
85 | raise ValueError(f"{dtype} is not a valid dtype.")
86 | dtype = dt
87 |
88 | print("Loading model ...", file=sys.stderr)
89 | t0 = time.time()
90 | with lazy_load(checkpoint_path) as pretrained_checkpoint, lazy_load(adapter_path) as adapter_checkpoint:
91 | name = llama_model_lookup(pretrained_checkpoint)
92 |
93 | with EmptyInitOnDevice(
94 | device=fabric.device, dtype=dtype, quantization_mode=quantize
95 | ):
96 | model = LLaMA.from_name(name)
97 | add_adapter_v2_parameters_to_linear_layers(model)
98 |
99 | # 1. Load the pretrained weights
100 | model.load_state_dict(pretrained_checkpoint, strict=False)
101 | # 2. Load the fine-tuned adapter weights
102 | model.load_state_dict(adapter_checkpoint, strict=False)
103 |
104 | print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
105 |
106 | model.eval()
107 |
108 | # if compile:
109 | # model = torch.compile(model)
110 |
111 | total_toks = 0
112 | model = fabric.setup_module(model)
113 |
114 | tokenizer = Tokenizer(tokenizer_path)
115 |
116 | for dsname in datasets.split(","):
117 | test_string = load_eval_data(dsname)
118 |
119 | sample = {"instruction": test_string, "input": input}
120 | test_string = generate_prompt(sample)
121 |
122 | encoded_text = tokenizer.encode(
123 | test_string, bos=True, eos=False, device=fabric.device
124 | )
125 | encoded_text = encoded_text[
126 | None, : 256 * model.config.block_size
127 | ] # add batch dimension, trim like gptq implementation
128 | t0 = time.perf_counter()
129 |
130 | nlls = 0
131 | toks = 0
132 |
133 | 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)
134 | for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
135 | inp = encoded_text[:, i : i + block_size]
136 | logits = model(inp)[0]
137 | nll = torch.nn.functional.cross_entropy(
138 | logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
139 | )
140 | toks += inp.size(1) - 1
141 | nlls += nll.item()
142 |
143 | print(encoded_text.shape, logits.shape)
144 | ppl = math.exp(nlls / toks)
145 | print(f"Perplexity on {dsname}: {ppl:.2f}")
146 | total_toks += toks
147 |
148 | t = time.perf_counter() - t0
149 | print(
150 | f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
151 | file=sys.stderr,
152 | )
153 | print(
154 | f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
155 | file=sys.stderr,
156 | )
157 |
158 |
159 | if __name__ == "__main__":
160 | from jsonargparse import CLI
161 |
162 | torch.set_float32_matmul_precision("high")
163 | CLI(main)
164 |
--------------------------------------------------------------------------------
/evaluate/full.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | # This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
4 | # Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
5 | import math
6 | import sys
7 | import time
8 | from pathlib import Path
9 | from typing import Optional
10 |
11 | import lightning as L
12 | import torch
13 | import tqdm
14 |
15 | # support running without installing as a package
16 | wd = Path(__file__).parent.parent.resolve()
17 | sys.path.append(str(wd))
18 |
19 | from lit_llama import LLaMA, Tokenizer
20 | from lit_llama.utils import EmptyInitOnDevice
21 |
22 | from datasets import load_dataset
23 |
24 |
25 | def load_eval_data(dataset_name: str) -> str:
26 | # this mimics gptq datautils
27 | if dataset_name == "wikitext":
28 | # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
29 | testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
30 | testdata = "\n\n".join(testdata["text"])
31 | elif dataset_name == "ptb":
32 | testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
33 | testdata = "\n\n".join(testdata["sentence"])
34 | elif dataset_name == "c4":
35 | testdata = load_dataset(
36 | "allenai/c4",
37 | "allenai--c4",
38 | data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
39 | split="validation",
40 | )
41 | testdata = " ".join(testdata[:1100]["text"])
42 |
43 | else:
44 | raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
45 | return testdata
46 |
47 |
48 | def main(
49 | datasets: str = "wikitext,ptb,c4",
50 | *,
51 | # compilation fails as it does not support torch.complex64 for RoPE
52 | # compile: bool = False,
53 | accelerator: str = "auto",
54 | checkpoint_path: Optional[Path] = None,
55 | tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
56 | model_size: str = "7B",
57 | dtype: str = "float32",
58 | quantize: Optional[str] = None,
59 | ) -> None:
60 | """Generates text samples based on a pre-trained LLaMA model and tokenizer.
61 |
62 | Args:
63 | datasets: The datasets to use as a comma separated string
64 | # compile: Whether to compile the model.
65 | accelerator: The hardware to run on. Possible choices are:
66 | ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
67 | checkpoint_path: The checkpoint path to load.
68 | tokenizer_path: The tokenizer path to load.
69 | dtype: The tensor dtype for choosing the floating-point precision
70 | quantize: Whether to quantize the model and using which method:
71 | ``"llm.int8"``: LLM.int8() mode,
72 | ``"gptq.int4"``: GPTQ 4-bit mode.
73 | """
74 | if not checkpoint_path:
75 | checkpoint_path = Path(f"checkpoints/lit-llama/{model_size}/lit-llama.pth")
76 | assert checkpoint_path.is_file()
77 | assert tokenizer_path.is_file()
78 |
79 | fabric = L.Fabric(accelerator=accelerator, devices=1)
80 |
81 | dt = getattr(torch, dtype, None)
82 | if not isinstance(dt, torch.dtype):
83 | raise ValueError(f"{dtype} is not a valid dtype.")
84 | dtype = dt
85 |
86 | with EmptyInitOnDevice(
87 | device=fabric.device, dtype=dtype, quantization_mode=quantize
88 | ):
89 | print("Loading model ...", file=sys.stderr)
90 | t0 = time.time()
91 | model = LLaMA.from_name(model_size)
92 | checkpoint = torch.load(checkpoint_path)
93 | model.load_state_dict(checkpoint)
94 | print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
95 |
96 | model.eval()
97 |
98 | # if compile:
99 | # model = torch.compile(model)
100 |
101 | total_toks = 0
102 | model = fabric.setup_module(model)
103 |
104 | tokenizer = Tokenizer(tokenizer_path)
105 |
106 | for dsname in datasets.split(","):
107 | test_string = load_eval_data(dsname)
108 | encoded_text = tokenizer.encode(
109 | test_string, bos=True, eos=False, device=fabric.device
110 | )
111 | encoded_text = encoded_text[
112 | None, : 256 * model.config.block_size
113 | ] # add batch dimension, trim like gptq implementation
114 | t0 = time.perf_counter()
115 |
116 | nlls = 0
117 | toks = 0
118 | with torch.inference_mode():
119 | 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)
120 | for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
121 | inp = encoded_text[:, i : i + block_size]
122 | logits = model(inp)[0]
123 | nll = torch.nn.functional.cross_entropy(
124 | logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
125 | )
126 | toks += inp.size(1) - 1
127 | nlls += nll.item()
128 |
129 | print(encoded_text.shape, logits.shape)
130 | ppl = math.exp(nlls / toks)
131 | print(f"Perplexity on {dsname}: {ppl:.2f}")
132 | total_toks += toks
133 |
134 | t = time.perf_counter() - t0
135 | print(
136 | f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
137 | file=sys.stderr,
138 | )
139 | print(
140 | f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
141 | file=sys.stderr,
142 | )
143 |
144 |
145 | if __name__ == "__main__":
146 | from jsonargparse import CLI
147 |
148 | torch.set_float32_matmul_precision("high")
149 | CLI(main)
150 |
--------------------------------------------------------------------------------
/evaluate/lora.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | # This mimics GPTQ's evaluation metrics: https://github.com/IST-DASLab/gptq/
4 | # Thanks to E. Frantar et al GPTQ: Accurate Post-training Compression for GPT, arXiv:2210.17323
5 | import math
6 | import sys
7 | import time
8 | from pathlib import Path
9 | from typing import Optional
10 |
11 | import lightning as L
12 | import torch
13 | import tqdm
14 |
15 | # support running without installing as a package
16 | wd = Path(__file__).parent.parent.resolve()
17 | sys.path.append(str(wd))
18 |
19 | from lit_llama import LLaMA, Tokenizer
20 | from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup
21 | from lit_llama.lora import lora
22 | from scripts.prepare_alpaca import generate_prompt
23 |
24 | from datasets import load_dataset
25 |
26 | instruction_tuning = True
27 | lora_r = 8
28 | lora_alpha = 16
29 | lora_dropout = 0.05
30 |
31 |
32 | def load_eval_data(dataset_name: str) -> str:
33 | # this mimics gptq datautils
34 | if dataset_name == "wikitext":
35 | # traindata = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
36 | testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
37 | testdata = "\n\n".join(testdata["text"])
38 | elif dataset_name == "ptb":
39 | testdata = load_dataset("ptb_text_only", "penn_treebank", split="test")
40 | testdata = "\n\n".join(testdata["sentence"])
41 | elif dataset_name == "c4":
42 | testdata = load_dataset(
43 | "allenai/c4",
44 | "allenai--c4",
45 | data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
46 | split="validation",
47 | )
48 | testdata = " ".join(testdata[:1100]["text"])
49 |
50 | else:
51 | raise ValueError("invalid dataset name (wikitext, ptb, c4 are allowed)")
52 | return testdata
53 |
54 |
55 | def main(
56 | datasets: str = "wikitext,ptb,c4",
57 | *,
58 | # compilation fails as it does not support torch.complex64 for RoPE
59 | # compile: bool = False,
60 | accelerator: str = "auto",
61 | lora_path: Path = Path("out/lora/alpaca/lit-llama-lora-finetuned.pth"),
62 | checkpoint_path: Path = Path("checkpoints/lit-llama/7B/lit-llama.pth"),
63 | tokenizer_path: Path = Path("checkpoints/lit-llama/tokenizer.model"),
64 | dtype: str = "float32",
65 | quantize: Optional[str] = None,
66 | ) -> None:
67 | """Generates text samples based on a pre-trained LLaMA model and tokenizer
68 | finetuned with LoRA.
69 |
70 | Args:
71 | datasets: The datasets to use as a comma separated string
72 | # compile: Whether to compile the model.
73 | accelerator: The hardware to run on. Possible choices are:
74 | ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
75 | lora_path: Path to the checkpoint with trained LoRA weights, which are the output of
76 | `finetune_lora.py`.
77 | checkpoint_path: The checkpoint path to load.
78 | tokenizer_path: The tokenizer path to load.
79 | dtype: The tensor dtype for choosing the floating-point precision
80 | quantize: Whether to quantize the model and using which method:
81 | ``"llm.int8"``: LLM.int8() mode,
82 | ``"gptq.int4"``: GPTQ 4-bit mode.
83 | """
84 | assert lora_path.is_file()
85 | assert checkpoint_path.is_file()
86 | assert tokenizer_path.is_file()
87 |
88 | if quantize is not None:
89 | raise NotImplementedError("Quantization in LoRA is not supported yet")
90 |
91 | fabric = L.Fabric(accelerator=accelerator, devices=1)
92 |
93 | dt = getattr(torch, dtype, None)
94 | if not isinstance(dt, torch.dtype):
95 | raise ValueError(f"{dtype} is not a valid dtype.")
96 | dtype = dt
97 |
98 | print("Loading model ...", file=sys.stderr)
99 | t0 = time.time()
100 |
101 | with lazy_load(checkpoint_path) as pretrained_checkpoint, lazy_load(lora_path) as lora_checkpoint:
102 | name = llama_model_lookup(pretrained_checkpoint)
103 |
104 | with EmptyInitOnDevice(
105 | device=fabric.device, dtype=dtype, quantization_mode=quantize
106 | ), lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
107 | model = LLaMA.from_name(name)
108 |
109 | # 1. Load the pretrained weights
110 | model.load_state_dict(pretrained_checkpoint, strict=False)
111 | # 2. Load the fine-tuned lora weights
112 | model.load_state_dict(lora_checkpoint, strict=False)
113 |
114 | print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)
115 |
116 | model.eval()
117 |
118 | # if compile:
119 | # model = torch.compile(model)
120 |
121 | total_toks = 0
122 | model = fabric.setup_module(model)
123 |
124 | tokenizer = Tokenizer(tokenizer_path)
125 |
126 | for dsname in datasets.split(","):
127 | test_string = load_eval_data(dsname)
128 |
129 | if instruction_tuning:
130 | sample = {"instruction": test_string, "input": input}
131 | test_string = generate_prompt(sample)
132 |
133 | encoded_text = tokenizer.encode(
134 | test_string, bos=True, eos=False, device=fabric.device
135 | )
136 | encoded_text = encoded_text[
137 | None, : 256 * model.config.block_size
138 | ] # add batch dimension, trim like gptq implementation
139 | t0 = time.perf_counter()
140 |
141 | nlls = 0
142 | toks = 0
143 | with torch.inference_mode():
144 | 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)
145 | for i in tqdm.tqdm(range(0, encoded_text.shape[1], block_size)):
146 | inp = encoded_text[:, i : i + block_size]
147 | logits = model(inp)[0]
148 | nll = torch.nn.functional.cross_entropy(
149 | logits[:-1], inp[0, 1:].to(dtype=torch.long), reduction="sum"
150 | )
151 | toks += inp.size(1) - 1
152 | nlls += nll.item()
153 |
154 | print(encoded_text.shape, logits.shape)
155 | ppl = math.exp(nlls / toks)
156 | print(f"Perplexity on {dsname}: {ppl:.2f}")
157 | total_toks += toks
158 |
159 | t = time.perf_counter() - t0
160 | print(
161 | f"\n\nTime for inference: {t:.02f} sec total, {total_toks / t:.02f} tokens/sec",
162 | file=sys.stderr,
163 | )
164 | print(
165 | f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB",
166 | file=sys.stderr,
167 | )
168 |
169 |
170 | if __name__ == "__main__":
171 | from jsonargparse import CLI
172 |
173 | torch.set_float32_matmul_precision("high")
174 | CLI(main)
175 |
--------------------------------------------------------------------------------
/finetune/full.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | """
4 | Instruction-tuning on the Alpaca dataset using a regular finetuning procedure (updating all layers).
5 |
6 | Note: If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
7 | `torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).
8 | """
9 | import sys
10 | from pathlib import Path
11 | import os
12 | import time
13 | from functools import partial
14 |
15 | import lightning as L
16 | from lightning.fabric.strategies import FSDPStrategy
17 | import numpy as np
18 | import torch
19 | from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
20 |
21 | # support running without installing as a package
22 | wd = Path(__file__).parent.parent.resolve()
23 | sys.path.append(str(wd))
24 |
25 | from generate import generate
26 | from lit_llama.model import Block, LLaMA, LLaMAConfig
27 | from lit_llama.tokenizer import Tokenizer
28 | from lit_llama.utils import save_model_checkpoint
29 | from scripts.prepare_alpaca import generate_prompt
30 |
31 |
32 | instruction_tuning = True
33 | eval_interval = 1000
34 | save_interval = 1000
35 | eval_iters = 100
36 | log_interval = 100
37 | devices = 4
38 |
39 | # Hyperparameters
40 | learning_rate = 3e-5
41 | batch_size = 128 / devices
42 | micro_batch_size = 4
43 | gradient_accumulation_iters = batch_size // micro_batch_size
44 | assert gradient_accumulation_iters > 0
45 | epoch_size = 50000 # train dataset size
46 | num_epochs = 5
47 | max_iters = num_epochs * (epoch_size // micro_batch_size) // devices
48 | weight_decay = 0.0
49 | block_size = 512
50 | warmup_iters = 100
51 |
52 |
53 | def main(
54 | data_dir: str = "data/alpaca",
55 | pretrained_path: str = "checkpoints/lit-llama/7B/lit-llama.pth",
56 | out_dir: str = "out/full/alpaca",
57 | ):
58 |
59 | auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
60 | strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True)
61 |
62 | fabric = L.Fabric(accelerator="cuda", devices=devices, precision="bf16-mixed", strategy=strategy)
63 | fabric.launch()
64 | fabric.seed_everything(1337 + fabric.global_rank)
65 |
66 | if fabric.global_rank == 0:
67 | os.makedirs(out_dir, exist_ok=True)
68 |
69 | train_data, val_data = load_datasets(data_dir=data_dir)
70 |
71 | config = LLaMAConfig.from_name("7B")
72 | config.block_size = block_size
73 |
74 | checkpoint = torch.load(pretrained_path)
75 |
76 | with fabric.device:
77 | torch.set_default_tensor_type(torch.HalfTensor)
78 | model = LLaMA(config).bfloat16()
79 | torch.set_default_tensor_type(torch.FloatTensor)
80 | model.load_state_dict(checkpoint, strict=False)
81 |
82 | model = fabric.setup_module(model)
83 |
84 | optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, foreach=False)
85 | optimizer = fabric.setup_optimizers(optimizer)
86 |
87 | train(fabric, model, optimizer, train_data, val_data, out_dir)
88 |
89 | # Save the final checkpoint at the end of training
90 | save_model_checkpoint(fabric, model, os.path.join(out_dir, "lit-llama-full-finetuned.pth"))
91 |
92 |
93 | def train(
94 | fabric: L.Fabric,
95 | model: torch.nn.Module,
96 | optimizer: torch.optim.Optimizer,
97 | train_data: np.ndarray,
98 | val_data: np.ndarray,
99 | out_dir: str,
100 | ) -> None:
101 | """The training loop.
102 |
103 | Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
104 | """
105 | step_count = 0
106 | model.train()
107 |
108 | for iter_num in range(max_iters):
109 |
110 | is_accumulating = (iter_num + 1) % gradient_accumulation_iters != 0
111 |
112 | if step_count <= warmup_iters:
113 | # linear warmup
114 | lr = learning_rate * step_count / warmup_iters
115 | for param_group in optimizer.param_groups:
116 | param_group['lr'] = lr
117 |
118 | t0 = time.time()
119 |
120 | input_ids, targets = get_batch(fabric, train_data)
121 | with fabric.no_backward_sync(model, enabled=is_accumulating):
122 | logits = model(input_ids)
123 | loss = loss_fn(logits, targets)
124 | fabric.backward(loss / gradient_accumulation_iters)
125 |
126 | if not is_accumulating:
127 | optimizer.step()
128 | optimizer.zero_grad()
129 | step_count += 1
130 |
131 | if step_count % eval_interval == 0:
132 | val_loss = validate(fabric, model, val_data)
133 | fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
134 | fabric.barrier()
135 |
136 | if step_count % save_interval == 0:
137 | print(f"Saving weights to {out_dir}")
138 | save_model_checkpoint(fabric, model, os.path.join(out_dir, f"iter-{iter_num:06d}-ckpt.pth"))
139 |
140 | dt = time.time() - t0
141 | if iter_num % log_interval == 0:
142 | fabric.print(f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms")
143 |
144 |
145 | def generate_response(model, instruction):
146 | tokenizer = Tokenizer("checkpoints/lit-llama/tokenizer.model")
147 | sample = {"instruction": instruction, "input": ""}
148 | prompt = instruction
149 | if instruction_tuning:
150 | prompt = generate_prompt(sample)
151 | encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
152 |
153 | output = generate(
154 | model,
155 | idx=encoded,
156 | max_seq_length=block_size,
157 | max_new_tokens=100,
158 | )
159 | output = tokenizer.decode(output)
160 | return output # output.split("### Response:")[1].strip()
161 |
162 |
163 | @torch.no_grad()
164 | def validate(fabric: L.Fabric, model: torch.nn.Module, val_data: np.ndarray) -> torch.Tensor:
165 | fabric.print("Validating ...")
166 | model.eval()
167 | losses = torch.zeros(eval_iters)
168 | for k in range(eval_iters):
169 | input_ids, targets = get_batch(fabric, val_data)
170 | logits = model(input_ids)
171 | loss = loss_fn(logits, targets)
172 | losses[k] = loss.item()
173 | out = losses.mean()
174 |
175 | # produce an example:
176 | instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
177 |
178 | output = generate_response(model, instruction)
179 | fabric.print(instruction)
180 | fabric.print(output)
181 |
182 | model.train()
183 | return out.item()
184 |
185 |
186 | def loss_fn(logits, targets):
187 | # shift the targets such that output n predicts token n+1
188 | logits = logits[..., :-1, :].contiguous()
189 | targets = targets[..., 1:].contiguous()
190 | loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
191 | return loss
192 |
193 |
194 | def get_batch(fabric: L.Fabric, data: list):
195 | ix = torch.randint(len(data), (micro_batch_size,))
196 |
197 | input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
198 | labels = [data[i]["labels"].type(torch.int64) for i in ix]
199 |
200 | max_len = max(len(s) for s in input_ids)
201 |
202 | def pad_right(x, pad_id):
203 | # pad right based on the longest sequence
204 | n = max_len - len(x)
205 | return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
206 |
207 | x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
208 | y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
209 | x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
210 | return x, y
211 |
212 |
213 | def load_datasets(data_dir):
214 | train_data = torch.load(os.path.join(data_dir, "train.pt"))
215 | val_data = torch.load(os.path.join(data_dir, "test.pt"))
216 | return train_data, val_data
217 |
218 |
219 | if __name__ == "__main__":
220 | # Uncomment this line if you see an error: "Expected is_sm80 to be true, but got false"
221 | # torch.backends.cuda.enable_flash_sdp(False)
222 | torch.set_float32_matmul_precision("high")
223 |
224 | from jsonargparse.cli import CLI
225 |
226 | CLI(main)
227 |
--------------------------------------------------------------------------------
/finetune/lora.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | """
4 | Instruction-tuning with LoRA on the Alpaca dataset.
5 |
6 | Note: If you run into a CUDA error "Expected is_sm80 to be true, but got false", uncomment the line
7 | `torch.backends.cuda.enable_flash_sdp(False)` in the script below (see https://github.com/Lightning-AI/lit-llama/issues/101).
8 | """
9 | import sys
10 | from pathlib import Path
11 | import os
12 | import time
13 |
14 | import lightning as L
15 | import numpy as np
16 | import torch
17 |
18 | # support running without installing as a package
19 | wd = Path(__file__).parent.parent.resolve()
20 | sys.path.append(str(wd))
21 |
22 | from generate import generate
23 | from lit_llama.lora import mark_only_lora_as_trainable, lora, lora_state_dict
24 | from lit_llama.model import LLaMA, LLaMAConfig
25 | from lit_llama.tokenizer import Tokenizer
26 | from scripts.prepare_alpaca import generate_prompt
27 |
28 |
29 | instruction_tuning = True
30 | eval_interval = 100
31 | save_interval = 100
32 | eval_iters = 100
33 | log_interval = 1
34 |
35 | # Hyperparameters
36 | learning_rate = 3e-4
37 | batch_size = 128
38 | micro_batch_size = 4
39 | gradient_accumulation_iters = batch_size // micro_batch_size
40 | assert gradient_accumulation_iters > 0
41 | max_iters = 50000 * 3 // micro_batch_size
42 | weight_decay = 0.0
43 | max_seq_length = 256 # see scripts/prepare_alpaca.py
44 | lora_r = 8
45 | lora_alpha = 16
46 | lora_dropout = 0.05
47 | warmup_iters = 100
48 |
49 |
50 | def main(
51 | data_dir: str = "data/alpaca",
52 | pretrained_path: str = "checkpoints/lit-llama/7B/lit-llama.pth",
53 | tokenizer_path: str = "checkpoints/lit-llama/tokenizer.model",
54 | out_dir: str = "out/lora/alpaca",
55 | ):
56 |
57 | fabric = L.Fabric(accelerator="cuda", devices=1, precision="bf16-true")
58 | fabric.launch()
59 | fabric.seed_everything(1337 + fabric.global_rank)
60 |
61 | if fabric.global_rank == 0:
62 | os.makedirs(out_dir, exist_ok=True)
63 |
64 | train_data, val_data = load_datasets(data_dir=data_dir)
65 |
66 | config = LLaMAConfig.from_name("7B")
67 | config.block_size = max_seq_length
68 |
69 | checkpoint = torch.load(pretrained_path)
70 |
71 | with fabric.init_module(), lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
72 | model = LLaMA(config)
73 | # strict=False because missing keys due to LoRA weights not contained in checkpoint state
74 | model.load_state_dict(checkpoint, strict=False)
75 |
76 | mark_only_lora_as_trainable(model)
77 |
78 | optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
79 | model, optimizer = fabric.setup(model, optimizer)
80 | train(fabric, model, optimizer, train_data, val_data, tokenizer_path, out_dir)
81 |
82 | # Save the final LoRA checkpoint at the end of training
83 | checkpoint = lora_state_dict(model)
84 | fabric.save(os.path.join(out_dir, "lit-llama-lora-finetuned.pth"), checkpoint)
85 |
86 |
87 | def train(
88 | fabric: L.Fabric,
89 | model: torch.nn.Module,
90 | optimizer: torch.optim.Optimizer,
91 | train_data: np.ndarray,
92 | val_data: np.ndarray,
93 | tokenizer_path: str,
94 | out_dir: str,
95 | ) -> None:
96 | """The training loop.
97 |
98 | Loosely based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT.
99 | """
100 | step_count = 0
101 |
102 | for iter_num in range(max_iters):
103 |
104 | if step_count <= warmup_iters:
105 | # linear warmup
106 | lr = learning_rate * step_count / warmup_iters
107 | for param_group in optimizer.param_groups:
108 | param_group['lr'] = lr
109 |
110 | t0 = time.time()
111 |
112 | input_ids, targets = get_batch(fabric, train_data)
113 | with fabric.no_backward_sync(model, enabled=((iter_num + 1) % gradient_accumulation_iters != 0)):
114 | logits = model(input_ids)
115 | loss = loss_fn(logits, targets)
116 | fabric.backward(loss / gradient_accumulation_iters)
117 |
118 | if (iter_num + 1) % gradient_accumulation_iters == 0:
119 | optimizer.step()
120 | optimizer.zero_grad()
121 | step_count += 1
122 |
123 | if step_count % eval_interval == 0:
124 | val_loss = validate(fabric, model, val_data, tokenizer_path)
125 | fabric.print(f"step {iter_num}: val loss {val_loss:.4f}")
126 | fabric.barrier()
127 |
128 | if step_count % save_interval == 0:
129 | print(f"Saving LoRA weights to {out_dir}")
130 | # We are only saving the LoRA weights
131 | # TODO: Provide a function/script to merge the LoRA weights with pretrained weights
132 | checkpoint = lora_state_dict(model)
133 | fabric.save(os.path.join(out_dir, f"iter-{iter_num:06d}-ckpt.pth"), checkpoint)
134 |
135 | dt = time.time() - t0
136 | if iter_num % log_interval == 0:
137 | fabric.print(f"iter {iter_num}: loss {loss.item():.4f}, time: {dt*1000:.2f}ms")
138 |
139 |
140 | def generate_response(model, instruction, tokenizer_path):
141 | tokenizer = Tokenizer(tokenizer_path)
142 | sample = {"instruction": instruction, "input": ""}
143 | prompt = instruction
144 | if instruction_tuning:
145 | prompt = generate_prompt(sample)
146 | encoded = tokenizer.encode(prompt, bos=True, eos=False, device=model.device)
147 |
148 | output = generate(
149 | model,
150 | idx=encoded,
151 | max_seq_length=max_seq_length,
152 | max_new_tokens=100,
153 | )
154 | output = tokenizer.decode(output)
155 | return output # output.split("### Response:")[1].strip()
156 |
157 |
158 | @torch.no_grad()
159 | def validate(fabric: L.Fabric, model: torch.nn.Module, val_data: np.ndarray, tokenizer_path: str) -> torch.Tensor:
160 | fabric.print("Validating ...")
161 | model.eval()
162 | losses = torch.zeros(eval_iters)
163 | for k in range(eval_iters):
164 | input_ids, targets = get_batch(fabric, val_data)
165 | logits = model(input_ids)
166 | loss = loss_fn(logits, targets)
167 | losses[k] = loss.item()
168 | out = losses.mean()
169 |
170 | # produce an example:
171 | instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
172 |
173 | output = generate_response(model, instruction, tokenizer_path)
174 | fabric.print(instruction)
175 | fabric.print(output)
176 |
177 | model.train()
178 | return out.item()
179 |
180 | def loss_fn(logits, targets):
181 | # shift the targets such that output n predicts token n+1
182 | logits = logits[..., :-1, :].contiguous()
183 | targets = targets[..., 1:].contiguous()
184 | loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
185 | return loss
186 |
187 |
188 | def get_batch(fabric: L.Fabric, data: list):
189 | ix = torch.randint(len(data), (micro_batch_size,))
190 |
191 | input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
192 | labels = [data[i]["labels"].type(torch.int64) for i in ix]
193 |
194 | max_len = max(len(s) for s in input_ids)
195 |
196 | def pad_right(x, pad_id):
197 | # pad right based on the longest sequence
198 | n = max_len - len(x)
199 | return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
200 |
201 | x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
202 | y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
203 | x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
204 | return x, y
205 |
206 |
207 | def load_datasets(data_dir):
208 | train_data = torch.load(os.path.join(data_dir, "train.pt"))
209 | val_data = torch.load(os.path.join(data_dir, "test.pt"))
210 | return train_data, val_data
211 |
212 |
213 | if __name__ == "__main__":
214 | # Uncomment this line if you see an error: "Expected is_sm80 to be true, but got false"
215 | # torch.backends.cuda.enable_flash_sdp(False)
216 | torch.set_float32_matmul_precision("high")
217 |
218 | from jsonargparse.cli import CLI
219 |
220 | CLI(main)
221 |
--------------------------------------------------------------------------------
/generate.py:
--------------------------------------------------------------------------------
1 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
2 |
3 | import sys
4 | import time
5 | import warnings
6 | from pathlib import Path
7 | from typing import Optional
8 |
9 | import lightning as L
10 | import torch
11 |
12 | # support running without installing as a package
13 | wd = Path(__file__).parent.parent.resolve()
14 | sys.path.append(str(wd))
15 |
16 | from lit_llama import LLaMA, Tokenizer
17 | from lit_llama.utils import lazy_load, llama_model_lookup, quantization
18 |
19 |
20 | @torch.no_grad()
21 | def generate(
22 | model: LLaMA,
23 | idx: torch.Tensor,
24 | max_new_tokens: int,
25 | *,
26 | max_seq_length: Optional[int] = None,
27 | temperature: float = 1.0,
28 | top_k: Optional[int] = None,
29 | eos_id: Optional[int] = None,
30 | ) -> torch.Tensor:
31 | """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
32 |
33 | The implementation of this function is modified from A. Karpathy's nanoGPT.
34 |
35 | Args:
36 | model: The model to use.
37 | idx: Tensor of shape (T) with indices of the prompt sequence.
38 | max_new_tokens: The number of new tokens to generate.
39 | max_seq_length: The maximum sequence length allowed.
40 | temperature: Scales the predicted logits by 1 / temperature
41 | top_k: If specified, only sample among the tokens with the k highest probabilities
42 | eos_id: If specified, stop generating any more token once the