8.2.2. Example: Large Language Model (LLM) Inference
LLaMa 1B の推論を行う応用例
デフォルトでは Prefill と Decode の両方を MN-Core 2 を用いてコンパイルします。CPU で実行する場合は --run_prefill_on_cpu または --run_decode_on_cpu を指定してください。
--prompt 'What is the meaning of life?' のようなプロンプトを指定することもできます。また、 --num_compiler_threads でコンパイル中のスレッド数を制限できます。 詳しくは コンパイル中のエラー を参照してください。
実行方法 (Compile Prefill and Decode on MN-Core 2)
$ cd /opt/pfn/pfcomp/codegen/MLSDK/examples/
$ ./examples/run_llm_infer.sh --prepare_attention_mask_on_cpu --device mncore2:auto
想定出力 (Compile Prefill and Decode on MN-Core 2)
Prefill のコンパイル結果を格納した
codegen_dir(/tmp/mlsdk_llm_infer/prefill)Decode のコンパイル結果を格納した
codegen_dir(/tmp/mlsdk_llm_infer/decode)推論の結果
=========== Generated with compilation ==========
</s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s><s> <|system|>
You are a friendly chatbot who is an expert on MN-Core.</s>
<|user|>
The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.</s>
<|assistant|>
Yes, that's correct. The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens using 16 A100-40G GPUs. The training has started on 2023-09
========== Generated with model.generate ==========
<s> <|system|>
You are a friendly chatbot who is an expert on MN-Core.</s>
<|user|>
The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.</s>
<|assistant|>
Yes, that's correct. The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens using 16 A100-40G GPUs. The training has started on 2023-09
Generated outputs matched.
また、上の例でコンパイルした結果を直接再利用することもできます。
実行方法 (Reuse codegen_dir for Prefill and Decode)
$ cd /opt/pfn/pfcomp/codegen/MLSDK/examples/
$ ./examples/run_llm_infer.sh --prepare_attention_mask_on_cpu --device mncore2:auto --load_codegen_prefill /tmp/mlsdk_llm_infer/prefill --load_codegen_decode /tmp/mlsdk_llm_infer/decode
想定出力 (Reuse codegen_dir for Prefill and Decode)
推論の結果 (上の例の結果と等しければ良い)
=========== Generated with compilation ==========
</s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s></s><s> <|system|>
You are a friendly chatbot who is an expert on MN-Core.</s>
<|user|>
The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.</s>
<|assistant|>
Yes, that's correct. The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens using 16 A100-40G GPUs. The training has started on 2023-09
========== Generated with model.generate ==========
<s> <|system|>
You are a friendly chatbot who is an expert on MN-Core.</s>
<|user|>
The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.</s>
<|assistant|>
Yes, that's correct. The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens using 16 A100-40G GPUs. The training has started on 2023-09
Generated outputs matched.
スクリプト
1#! /bin/bash
2# Usage:
3# ./examples/run_llm_infer.sh --prepare_attention_mask_on_cpu --device mncore2:auto --load_codegen_prefill /tmp/mlsdk_llm_infer/prefill --load_codegen_decode /tmp/mlsdk_llm_infer/decode
4
5set -eux -o pipefail
6
7EXAMPLE_NAME=run_llm_infer
8VENVDIR=/tmp/${EXAMPLE_NAME}_venv
9
10CURRENT_DIR=$(realpath $(dirname $0))
11CODEGEN_DIR=$(realpath ${CURRENT_DIR}/../../)
12BUILD_DIR=${BUILD_DIR:-${CODEGEN_DIR}/build}
13
14if [[ ! -d ${VENVDIR} ]]; then
15 python3 -m venv --system-site-packages ${VENVDIR}
16fi
17
18source ${VENVDIR}/bin/activate
19# See https://discuss.huggingface.co/t/cas-bridge-xethub-hf-co-broke/158626/8 for hf_xet.
20pip3 install transformers==4.44.0 huggingface-hub==0.34.4 hf_xet==v1.1.5
21
22source "${BUILD_DIR}/codegen_pythonpath.sh"
23
24export MNCORE_USE_LEGACY_ONNX_EXPORTER=1
25export MNCORE_USE_EXTERNAL_DATA_FORMAT=1
26
27exec python3 ${CURRENT_DIR}/llm_infer.py "$@"
1import argparse
2import copy
3import os
4from typing import Mapping, Optional
5
6import torch
7from mlsdk import CacheOptions, Context, MNDevice, set_tensor_name_in_module, storage
8from mlsdk.experimental.llm.attention_mask import (
9 prepare_4d_causal_attention_mask_with_cache_position,
10)
11from mlsdk.experimental.llm.kv_cache import (
12 kv_cache_to_legacy,
13 kv_cache_to_plamo,
14 kv_cache_to_tensor,
15)
16from transformers import AutoModelForCausalLM, AutoTokenizer
17
18
19def prepare_prompt(tokenizer, prompt, system_prompt):
20 if tokenizer.chat_template:
21 messages = [
22 {
23 "role": "system",
24 "content": system_prompt,
25 },
26 {"role": "user", "content": prompt},
27 ]
28 prompt = tokenizer.apply_chat_template(
29 messages,
30 tokenize=False,
31 add_generation_prompt=True,
32 )
33 return prompt
34
35
36def infer_with_generate(
37 prompt: str,
38 model: AutoModelForCausalLM,
39 tokenizer: AutoTokenizer,
40 max_new_tokens: int,
41) -> torch.Tensor:
42 inputs = tokenizer(prompt, return_tensors="pt")
43 # Greedy decoding for simplicity for comparing the results with the compiled version.
44 output_ids = model.generate(
45 inputs["input_ids"], do_sample=False, max_new_tokens=max_new_tokens
46 )
47 assert isinstance(output_ids, torch.Tensor)
48 return output_ids
49
50
51def infer_with_compilation( # NOQA: CFQ002, CFQ001, CFQ004
52 *,
53 prompt: str,
54 model: AutoModelForCausalLM,
55 tokenizer: AutoTokenizer,
56 max_length: int,
57 max_new_tokens: int,
58 run_prefill_on_cpu: bool,
59 run_decode_on_cpu: bool,
60 device_name: str,
61 outdir: str,
62 check_intermediate_outputs: bool,
63 prepare_attention_mask_on_cpu: bool,
64 disable_cache: bool,
65 num_compiler_threads: int,
66 load_codegen_prefill: Optional[str],
67 load_codegen_decode: Optional[str],
68) -> torch.Tensor:
69 # When we recompile a model, delete the old decompressed "model.app"
70 if load_codegen_prefill is None:
71 if os.path.exists(outdir + "/prefill/model.app"):
72 os.remove(outdir + "/prefill/model.app")
73 if load_codegen_decode is None:
74 if os.path.exists(outdir + "/decode/model.app"):
75 os.remove(outdir + "/decode/model.app")
76
77 is_plamo_model = any("plamo" in a.lower() for a in model.config.architectures)
78
79 # Sharing layouts between the prefill and decode models can cause layout
80 # decision issues in the compiler. We create a shallow copy so that each model
81 # can use a different layout. We can also use `copy.deepcopy` instead
82 # if the additional host memory usage is acceptable.
83 def shallow_copy_model_with_copied_tensors(model):
84 new_model = copy.copy(model)
85 new_model._modules = model._modules.copy()
86 new_model._parameters = model._parameters.copy()
87 new_model._buffers = model._buffers.copy()
88
89 for name, child in model._modules.items():
90 if child is not None:
91 new_model._modules[name] = shallow_copy_model_with_copied_tensors(child)
92 for name, param in model._parameters.items():
93 if param is not None:
94 new_model._parameters[name] = copy.copy(param)
95 for name, buffer in model._buffers.items():
96 if buffer is not None:
97 new_model._buffers[name] = copy.copy(buffer)
98
99 return new_model
100
101 prefill_model = model
102 decode_model = shallow_copy_model_with_copied_tensors(prefill_model)
103
104 # Check if object is different while its storage is shared
105 decode_model_named_paramters = dict(decode_model.named_parameters())
106 for name, p1 in prefill_model.named_parameters():
107 p2 = decode_model_named_paramters[name]
108 assert p1 is not p2
109 assert p1.data_ptr() == p2.data_ptr(), name
110
111 def forward_prefill(
112 inputs: Mapping[str, torch.Tensor],
113 ) -> Mapping[str, torch.Tensor]:
114 assert all(isinstance(v, torch.Tensor) for v in inputs.values()), {
115 k: type(v) for k, v in inputs.items()
116 }
117 outputs = prefill_model.forward(
118 input_ids=inputs["input_ids"],
119 attention_mask=inputs["attention_mask"],
120 position_ids=inputs["position_ids"],
121 past_key_values=None,
122 use_cache=True,
123 )
124 next_input_ids = outputs["logits"].argmax(dim=2)[:, -1:]
125 return {
126 "next_input_ids": next_input_ids,
127 "next_past_key_values": kv_cache_to_tensor(outputs.past_key_values)[
128 :, :, :, :, 1:, :
129 ], # Do every operation, including the shifting, for the KV cache on device.
130 }
131
132 def forward_decode(
133 inputs: Mapping[str, torch.Tensor],
134 ) -> Mapping[str, torch.Tensor]:
135 assert all(isinstance(v, torch.Tensor) for v in inputs.values()), {
136 k: type(v) for k, v in inputs.items()
137 }
138 if is_plamo_model:
139 kv_cache_func = kv_cache_to_plamo
140 else:
141 # @todo (hvy): Stop using the deprecated legacy KV cache format of tuples.
142 kv_cache_func = kv_cache_to_legacy
143 past_key_values = kv_cache_func(inputs["past_key_values"])
144
145 outputs = decode_model.forward(
146 input_ids=inputs["input_ids"],
147 attention_mask=inputs["attention_mask"],
148 position_ids=inputs["position_ids"],
149 past_key_values=past_key_values,
150 use_cache=True,
151 )
152 next_input_ids = outputs["logits"].argmax(dim=2)[:, -1:]
153 return {
154 "next_input_ids": next_input_ids,
155 "next_past_key_values": kv_cache_to_tensor(outputs.past_key_values)[
156 :, :, :, :, 1:, :
157 ], # Do every operation, including the shifting, for the KV cache on device.
158 }
159
160 assert tokenizer.padding_side == "left"
161 inputs = tokenizer(
162 prompt, return_tensors="pt", padding="max_length", max_length=max_length
163 )
164 # @todo (hvy): Consider subtracting 1 from the position_ids to match modeling_llama.py.
165 assert "position_ids" not in inputs
166 inputs["position_ids"] = inputs["attention_mask"].cumsum(1)
167 if prepare_attention_mask_on_cpu:
168 inputs["attention_mask"] = prepare_4d_causal_attention_mask_with_cache_position(
169 inputs["attention_mask"], inputs["position_ids"], model.dtype
170 )
171 output_ids = inputs["input_ids"]
172
173 device = MNDevice(device_name)
174 context = Context(device)
175 Context.switch_context(context)
176 set_tensor_name_in_module(prefill_model, "prefill_model")
177 for p in prefill_model.parameters():
178 context.register_param(p)
179 for b in prefill_model.buffers():
180 context.register_buffer(b)
181 set_tensor_name_in_module(decode_model, "decode_model")
182 for p in decode_model.parameters():
183 context.register_param(p)
184 for b in decode_model.buffers():
185 context.register_buffer(b)
186
187 compiled_funcs = {}
188 if load_codegen_prefill is not None:
189 compiled_funcs["prefill"] = context.load_codegen_dir(
190 storage.path(load_codegen_prefill)
191 )
192 input_proxies_allocated = compiled_funcs["prefill"].allocate_input_proxy()
193 input_proxies_allocated["input_ids"].load_from(
194 inputs["input_ids"].detach(), clone=False
195 )
196 input_proxies_allocated["attention_mask"].load_from(
197 inputs["attention_mask"].detach(), clone=False
198 )
199 input_proxies_allocated["position_ids"].load_from(
200 inputs["position_ids"].detach(), clone=False
201 )
202 for model_param in prefill_model.parameters():
203 model_param_proxy = context.get_registered_value_proxy(model_param)
204 model_param_proxy.load_from(model_param)
205
206 decode_input_proxies = None
207 if load_codegen_decode is not None:
208 compiled_funcs["decode"] = context.load_codegen_dir(
209 storage.path(load_codegen_decode)
210 )
211 decode_input_proxies = compiled_funcs["decode"].allocate_input_proxy()
212
213 for step in range(max_new_tokens):
214 if step == 0:
215 if not run_prefill_on_cpu and "prefill" not in compiled_funcs:
216 # Set codegen internal flags
217 options = dict()
218 options["envs"] = {
219 "CODEGEN_TIME_SLICE_SCATTERED_INDEXING_BCAST": "1",
220 "CODEGEN_OP_DEF": "Gather=GatherBcast",
221 }
222 compiled_funcs["prefill"] = context.compile(
223 forward_prefill,
224 inputs,
225 storage.path(outdir + "/prefill"),
226 options=options,
227 cache_options=(
228 CacheOptions(outdir + "/prefill_cache")
229 if not disable_cache
230 else None
231 ),
232 num_compiler_threads=num_compiler_threads,
233 )
234 forward_for_step = compiled_funcs.get("prefill", forward_prefill)
235 else:
236 if not run_decode_on_cpu and "decode" not in compiled_funcs:
237 compiled_funcs["decode"] = context.compile(
238 forward_decode,
239 inputs,
240 storage.path(outdir + "/decode"),
241 cache_options=(
242 CacheOptions(outdir + "/decode_cache")
243 if not disable_cache
244 else None
245 ),
246 num_compiler_threads=num_compiler_threads,
247 io_spec_constraints=(
248 {
249 "past_key_values": compiled_funcs["prefill"].output_specs[
250 "next_past_key_values"
251 ],
252 # This constraint is to share the layout between decode's
253 # `past_key_values` and `next_past_key_values`.
254 "next_past_key_values": compiled_funcs[
255 "prefill"
256 ].output_specs["next_past_key_values"],
257 }
258 if not run_prefill_on_cpu
259 else None
260 ),
261 )
262 forward_for_step = compiled_funcs.get("decode", forward_decode)
263
264 outputs = forward_for_step(inputs)
265
266 if check_intermediate_outputs:
267 # @todo (hvy): Consider using a more sophisticated check for the outputs.
268 if "mncore" in device_name:
269 atol = 1.0
270 else:
271 assert device_name == "pfvm:cpu"
272 atol = 5e-3
273 n_tokens = inputs["position_ids"].max()
274 forward = forward_prefill if step == 0 else forward_decode
275 # We treat `past_key_values` as a `TensorProxy` to avoid sending it to the
276 # host. Therefore, to perform computations on the host, we need to transfer
277 # them explicity.
278 inputs_on_host = (
279 inputs
280 if step == 0
281 else {
282 **inputs,
283 "past_key_values": inputs["past_key_values"].cpu().detach(),
284 }
285 )
286 outputs_expected = forward(inputs_on_host)
287 # Since we compute argmax inside the compiled function, we only output the
288 # last token. For a more accurate comparison, we may need to output the
289 # full `logits`.
290 next_input_ids = outputs["next_input_ids"]
291 next_input_ids_expected = outputs_expected["next_input_ids"]
292 next_past_key_values = outputs["next_past_key_values"][
293 :, :, :, :, -n_tokens:
294 ]
295 next_past_key_values_expected = outputs_expected["next_past_key_values"][
296 :, :, :, :, -n_tokens:
297 ]
298
299 assert torch.allclose(next_input_ids, next_input_ids_expected, atol=atol), (
300 step,
301 (next_input_ids - next_input_ids_expected).abs().max(),
302 )
303 assert torch.allclose(
304 next_past_key_values, next_past_key_values_expected, atol=atol
305 ), (
306 step,
307 (next_past_key_values - next_past_key_values_expected).abs().max(),
308 )
309
310 next_input_ids = outputs["next_input_ids"].cpu()
311 if prepare_attention_mask_on_cpu:
312 next_attention_mask = inputs["attention_mask"][:, :, -1:, :]
313 next_attention_mask = torch.roll(next_attention_mask, shifts=-1, dims=-1)
314 next_attention_mask[:, :, :, -1] = 0
315 else:
316 next_attention_mask = inputs["attention_mask"]
317 next_attention_mask = torch.roll(next_attention_mask, shifts=-1, dims=-1)
318 next_attention_mask[:, -1] = 1
319 next_position_ids = inputs["position_ids"][:, -1:] + 1
320 next_past_key_values = outputs["next_past_key_values"]
321 inputs = {
322 "input_ids": next_input_ids.detach(),
323 "attention_mask": next_attention_mask.detach(),
324 "position_ids": next_position_ids.detach(),
325 "past_key_values": next_past_key_values,
326 }
327
328 output_ids = torch.cat([output_ids, next_input_ids], dim=1)
329
330 if next_input_ids.item() == tokenizer.eos_token_id:
331 break
332
333 if compiled_funcs.get("decode") is not None:
334 if decode_input_proxies is None:
335 decode_input_proxies = compiled_funcs["decode"].allocate_input_proxy()
336 decode_input_proxies["input_ids"].load_from(
337 next_input_ids.detach(), clone=False
338 )
339 decode_input_proxies["attention_mask"].load_from(
340 next_attention_mask.detach(), clone=False
341 )
342 decode_input_proxies["position_ids"].load_from(
343 next_position_ids.detach(), clone=False
344 )
345
346 return output_ids[:, max_new_tokens:]
347
348
349def main(args):
350 prompt = args.prompt
351 system_prompt = args.system_prompt
352 model_name = args.model_name
353 max_length = args.max_length
354 max_new_tokens = args.max_new_tokens
355 run_prefill_on_cpu = args.run_prefill_on_cpu
356 run_decode_on_cpu = args.run_decode_on_cpu
357 device_name = args.device_name
358 outdir = args.outdir
359 check_intermediate_outputs = args.check_intermediate_outputs
360 prepare_attention_mask_on_cpu = args.prepare_attention_mask_on_cpu
361 disable_cache = args.disable_cache
362 num_compiler_threads = args.num_compiler_threads
363 load_codegen_prefill = args.load_codegen_prefill
364 load_codegen_decode = args.load_codegen_decode
365
366 model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
367 model.eval() # Some models do not return the KV cache in training mode.
368 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
369 if tokenizer.pad_token_id is None:
370 tokenizer.pad_token_id = tokenizer.eos_token_id
371 tokenizer.padding_side = "left" # For static KV caching
372 tokenizer.truncation_side = "left"
373
374 prompt = prepare_prompt(tokenizer, prompt, system_prompt)
375
376 outputs = infer_with_compilation(
377 prompt=prompt,
378 model=model,
379 tokenizer=tokenizer,
380 max_length=max_length,
381 max_new_tokens=max_new_tokens,
382 run_prefill_on_cpu=run_prefill_on_cpu,
383 run_decode_on_cpu=run_decode_on_cpu,
384 device_name=device_name,
385 outdir=outdir,
386 check_intermediate_outputs=check_intermediate_outputs,
387 prepare_attention_mask_on_cpu=prepare_attention_mask_on_cpu,
388 disable_cache=disable_cache,
389 num_compiler_threads=num_compiler_threads,
390 load_codegen_prefill=load_codegen_prefill,
391 load_codegen_decode=load_codegen_decode,
392 )
393 print(
394 "=========== Generated with compilation ==========\n",
395 tokenizer.decode(outputs[0]),
396 )
397
398 outputs_expected = infer_with_generate(prompt, model, tokenizer, max_new_tokens)
399 print(
400 "========== Generated with model.generate ==========\n",
401 tokenizer.decode(outputs_expected[0]),
402 )
403
404 # @todo (hvy): Do not rely on `max_new_tokens` tokens always being generated?
405 assert torch.equal(
406 outputs[:, -max_new_tokens:], outputs_expected[:, -max_new_tokens:]
407 ), "Outputs differed. Check generated outputs above."
408 print("Generated outputs matched.")
409
410
411if __name__ == "__main__":
412 parser = argparse.ArgumentParser()
413 parser.add_argument(
414 "--prompt",
415 type=str,
416 default='The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.', # NOQA
417 )
418 parser.add_argument(
419 "--system_prompt",
420 type=str,
421 default="You are a friendly chatbot who is an expert on MN-Core.",
422 )
423 parser.add_argument(
424 "--model_name", type=str, default="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
425 )
426 parser.add_argument("--max_length", type=int, default=256)
427 parser.add_argument("--max_new_tokens", type=int, default=64)
428 parser.add_argument("--num_compiler_threads", type=int, default=-1)
429 parser.add_argument("--run_prefill_on_cpu", action="store_true")
430 parser.add_argument("--run_decode_on_cpu", action="store_true")
431 parser.add_argument("--device_name", type=str, default="mncore2:auto")
432 parser.add_argument(
433 "--outdir",
434 type=str,
435 default="/tmp/mlsdk_llm_infer",
436 help="""
437 Output directory. Compilation results for prefill and decode are stored in the
438 `/prefill` and `/decode` subdirectories, respectively.
439 """,
440 )
441 parser.add_argument("--check_intermediate_outputs", action="store_true")
442 parser.add_argument("--prepare_attention_mask_on_cpu", action="store_true")
443 parser.add_argument("--disable_cache", action="store_true")
444 parser.add_argument("--load_codegen_prefill", type=str, default=None)
445 parser.add_argument("--load_codegen_decode", type=str, default=None)
446 args = parser.parse_args()
447 main(args)