You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: colossalai/inference/README.md
+78-3Lines changed: 78 additions & 3 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -13,18 +13,92 @@
13
13
14
14
15
15
## 📌 Introduction
16
-
17
16
ColossalAI-Inference is a library which offers acceleration to Transformers models, especially LLMs. In ColossalAI-Inference, we leverage high-performance kernels, KV cache, paged attention, continous batching and other techniques to accelerate the inference of LLMs. We also provide a unified interface for users to easily use our library.
18
17
19
18
## 🛠 Design and Implementation
20
19
21
-
To be added.
20
+
### :book: Overview
21
+
We build ColossalAI-Inference based on **Four** core components: `engine`,`request handler`,`cache manager(block cached)`, `hand crafted modeling`. **Engine** controls inference step, it recives `requests`, calls `request handler` to schedule a decoding batch and runs `modeling` to perform a iteration and returns finished `requests`. **Cache manager** is bound with `request handler`, updates cache blocks and logical block tables during schedule.
22
+
23
+
The interaction between different components are shown below, you can also checkout detailed introduction below.:
Engine is designed as starter of inference loop. User can easily instantialize an infer engine with config and execute requests. We provids apis below in engine, you can refer to source code for more information:
31
+
-`generate`: main function, handle inputs and return outputs
32
+
-`add_request`: add request to waitting list
33
+
-`step`: perform one decoding iteration
34
+
- first, `request handler` schedules a batch to do prefill/decode
35
+
- then, invoke a model to generate a batch of token
36
+
- after that, do logit processing and sampling, check and decode finished requests
37
+
38
+
### :game_die: Design of request_handler
39
+
Request handler is responsible manage requests and schedule a proper batch from exisiting requests. According to existing work and experiments, we do believe that it is beneficial to increase the length of decoding sequences. In our design, we partition requests into three priorities depending on their lengths, the longer sequences are first considered.
We design a unified blocked type cache and cache manager to distribute memory. The physical memory is allocated before decoding and represented by a logical block table. During decoding process, cache manager administrate physical memory through `block table` and other components(i.e. engine) can focus on the light-weighted `block table`. Their details are introduced below.
47
+
-`cache block` We group physical memory into different memory blocks. A typical cache block is shaped `(num_kv_heads, head_size, block_size)`. We decide block number beforehand. The memory allocation and computation are executed with the granularity of memory block.
48
+
-`block table` Block table is the logical representation of cache blocks. Concretely, a block table of a single sequence is a 1D tensor, with each element holding a block id of allocated id or `-1` for non allocated. Each iteration we pass through a batch block table to the corresponding model. For more information, you can checkout the source code.
<figcation>Example of Batch Block Table</figcation>
55
+
</p>
56
+
</figure>
57
+
58
+
59
+
### :railway_car: Modeling
60
+
Modeling contains models and layers, which are hand-crafted for better performance easier usage. Deeply integrated with `shardformer`, we also construct policy for our models. In order to minimize users' learning costs, our models are aligned with [Transformers](https://github.com/huggingface/transformers)
22
61
23
62
## 🕹 Usage
24
63
64
+
### :arrow_right: Quick Start
65
+
You can enjoy your fast generation journey within three step
66
+
```python
67
+
# First, create a model in "transformers" way, you can provide a model config or use the default one.
68
+
model = transformers.LlamaForCausalLM(config).cuda()
prompts = {'Nice to meet you, Colossal-Inference!'}
81
+
engine.generate(prompts)
25
82
26
-
To be added.
83
+
```
27
84
85
+
### :bookmark: Customize your inference engine
86
+
Besides the basic fast-start inference, you can also customize your inference engine via modifying config or upload your own model or decoding components (logit processors or sampling strategies).
87
+
#### Inference Config
88
+
Inference Config is a unified api for generation process. You can define the value of args to control the generation, like `max_batch_size`,`max_output_len`,`dtype` to decide the how many sequences can be handled at a time, and how many tokens to output. Refer to the source code for more detail.
89
+
#### Generation Config
90
+
In colossal-inference, Generation config api is inherited from [Transformers](https://github.com/huggingface/transformers). Usage is aligned. By default, it is automatically generated by our system and you don't bother to construct one. If you have such demand, you can also create your own and send it to your engine.
91
+
92
+
#### Logit Processors
93
+
Logit Processosr receives logits and return processed ones, take the following step to make your own.
94
+
```python
95
+
@register_logit_processor("name")
96
+
defxx_logit_processor(logits, args):
97
+
logits = do_some_process(logits)
98
+
return logits
99
+
```
100
+
#### Sampling Strategies
101
+
We offer 3 main sampling strategies now (i.e. `greedy sample`, `multinomial sample`, `beam_search sample`), you can refer to [sampler](/ColossalAI/colossalai/inference/sampler.py) for more details. We would strongly appreciate if you can contribute your varities.
0 commit comments