lora.ipynb 16.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LoRA Serving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SGLang enables the use of [LoRA adapters](https://arxiv.org/abs/2106.09685) with a base model. By incorporating techniques from [S-LoRA](https://arxiv.org/pdf/2311.03285) and [Punica](https://arxiv.org/pdf/2310.18547), SGLang can efficiently support multiple LoRA adapters for different sequences within a single batch of inputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arguments for LoRA Serving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following server arguments are relevant for multi-LoRA serving:\n",
    "\n",
30
31
    "* `enable_lora`: Enable LoRA support for the model. This argument is automatically set to True if `--lora-paths` is provided for backward compatibility.\n",
    "\n",
32
33
34
35
    "* `lora_paths`: A mapping from each adaptor's name to its path, in the form of `{name}={path} {name}={path}`.\n",
    "\n",
    "* `max_loras_per_batch`: Maximum number of adaptors used by each batch. This argument can affect the amount of GPU memory reserved for multi-LoRA serving, so it should be set to a smaller value when memory is scarce. Defaults to be 8.\n",
    "\n",
36
37
    "* `max_loaded_loras`: If specified, it limits the maximum number of LoRA adapters loaded in CPU memory at a time. The value must be greater than or equal to `max-loras-per-batch`.\n",
    "\n",
38
    "* `lora_backend`: The backend of running GEMM kernels for Lora modules. Currently we only support Triton LoRA backend. In the future, faster backend built upon Cutlass or Cuda kernels will be added.\n",
39
    "\n",
40
41
    "* `max_lora_rank`: The maximum LoRA rank that should be supported. If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of larger LoRA rank after server startup.\n",
    "\n",
42
    "* `lora_target_modules`: The union set of all target modules where LoRA should be applied (e.g., `q_proj`, `k_proj`, `gate_proj`). If not specified, it will be automatically inferred from the adapters provided in `--lora-paths`. This argument is needed when you expect to dynamically load adapters of different target modules after server startup. You can also set it to `all` to enable LoRA for all supported modules. However, enabling LoRA on additional modules introduces a minor performance overhead. If your application is performance-sensitive, we recommend only specifying the modules for which you plan to load adapters.\n",
43
    "\n",
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
    "* `tp_size`: LoRA serving along with Tensor Parallelism is supported by SGLang. `tp_size` controls the number of GPUs for tensor parallelism. More details on the tensor sharding strategy can be found in [S-Lora](https://arxiv.org/pdf/2311.03285) paper.\n",
    "\n",
    "From client side, the user needs to provide a list of strings as input batch, and a list of adaptor names that each input sequence corresponds to."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage\n",
    "\n",
    "### Serving Single Adaptor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sglang.test.test_utils import is_in_ci\n",
    "\n",
    "if is_in_ci():\n",
    "    from patch import launch_server_cmd\n",
    "else:\n",
    "    from sglang.utils import launch_server_cmd\n",
    "\n",
    "from sglang.utils import wait_for_server, terminate_process\n",
    "\n",
    "import json\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "server_process, port = launch_server_cmd(\n",
    "    \"\"\"\n",
    "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n",
86
    "    --enable-lora \\\n",
87
88
    "    --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n",
    "    --max-loras-per-batch 1 --lora-backend triton \\\n",
89
    "    --disable-radix-cache\n",
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
    "\"\"\"\n",
    ")\n",
    "\n",
    "wait_for_server(f\"http://localhost:{port}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = f\"http://127.0.0.1:{port}\"\n",
    "json_data = {\n",
    "    \"text\": [\n",
    "        \"List 3 countries and their capitals.\",\n",
106
    "        \"List 3 countries and their capitals.\",\n",
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
    "    ],\n",
    "    \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n",
    "    # The first input uses lora0, and the second input uses the base model\n",
    "    \"lora_path\": [\"lora0\", None],\n",
    "}\n",
    "response = requests.post(\n",
    "    url + \"/generate\",\n",
    "    json=json_data,\n",
    ")\n",
    "print(f\"Output 0: {response.json()[0]['text']}\")\n",
    "print(f\"Output 1: {response.json()[1]['text']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "terminate_process(server_process)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Serving Multiple Adaptors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "server_process, port = launch_server_cmd(\n",
    "    \"\"\"\n",
    "python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n",
145
    "    --enable-lora \\\n",
146
147
148
    "    --lora-paths lora0=algoprog/fact-generation-llama-3.1-8b-instruct-lora \\\n",
    "    lora1=Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16 \\\n",
    "    --max-loras-per-batch 2 --lora-backend triton \\\n",
149
    "    --disable-radix-cache\n",
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
    "\"\"\"\n",
    ")\n",
    "\n",
    "wait_for_server(f\"http://localhost:{port}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = f\"http://127.0.0.1:{port}\"\n",
    "json_data = {\n",
    "    \"text\": [\n",
    "        \"List 3 countries and their capitals.\",\n",
166
    "        \"List 3 countries and their capitals.\",\n",
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
    "    ],\n",
    "    \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n",
    "    # The first input uses lora0, and the second input uses lora1\n",
    "    \"lora_path\": [\"lora0\", \"lora1\"],\n",
    "}\n",
    "response = requests.post(\n",
    "    url + \"/generate\",\n",
    "    json=json_data,\n",
    ")\n",
    "print(f\"Output 0: {response.json()[0]['text']}\")\n",
    "print(f\"Output 1: {response.json()[1]['text']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "terminate_process(server_process)"
   ]
  },
189
190
191
192
193
194
195
196
197
198
199
200
201
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dynamic LoRA loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of specifying all adapters during server startup via `--lora-paths`. You can also load & unload LoRA adapters dynamically via the `/load_lora_adapter` and `/unload_lora_adapter` API.\n",
    "\n",
202
    "When using dynamic LoRA loading, it's recommended to explicitly specify both `--max-lora-rank` and `--lora-target-modules` at startup. For backward compatibility, SGLang will infer these values from `--lora-paths` if they are not explicitly provided. However, in that case, you would have to ensure that all dynamically loaded adapters share the same shape (rank and target modules) as those in the initial `--lora-paths` or are strictly \"smaller\"."
203
204
205
206
207
208
209
210
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
211
212
213
214
215
216
217
    "lora0 = \"Nutanix/Meta-Llama-3.1-8B-Instruct_lora_4_alpha_16\"  # rank - 4, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj\n",
    "lora1 = \"algoprog/fact-generation-llama-3.1-8b-instruct-lora\"  # rank - 64, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj\n",
    "lora0_new = \"philschmid/code-llama-3-1-8b-text-to-sql-lora\"  # rank - 256, target modules - q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj\n",
    "\n",
    "\n",
    "# The `--target-lora-modules` param below is technically not needed, as the server will infer it from lora0 which already has all the target modules specified.\n",
    "# We are adding it here just to demonstrate usage.\n",
218
219
220
    "server_process, port = launch_server_cmd(\n",
    "    \"\"\"\n",
    "    python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n",
221
    "    --enable-lora \\\n",
222
223
224
    "    --cuda-graph-max-bs 2 \\\n",
    "    --max-loras-per-batch 2 --lora-backend triton \\\n",
    "    --disable-radix-cache\n",
225
226
    "    --max-lora-rank 256\n",
    "    --lora-target-modules all\n",
227
228
229
230
231
232
233
    "    \"\"\"\n",
    ")\n",
    "\n",
    "url = f\"http://127.0.0.1:{port}\"\n",
    "wait_for_server(url)"
   ]
  },
234
235
236
237
238
239
240
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load adapter lora0"
   ]
  },
241
242
243
244
245
246
247
248
249
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.post(\n",
    "    url + \"/load_lora_adapter\",\n",
    "    json={\n",
250
251
    "        \"lora_name\": \"lora0\",\n",
    "        \"lora_path\": lora0,\n",
252
253
254
255
256
257
258
259
260
261
    "    },\n",
    ")\n",
    "\n",
    "if response.status_code == 200:\n",
    "    print(\"LoRA adapter loaded successfully.\", response.json())\n",
    "else:\n",
    "    print(\"Failed to load LoRA adapter.\", response.json())"
   ]
  },
  {
262
   "cell_type": "markdown",
263
264
   "metadata": {},
   "source": [
265
    "Load adapter lora1:"
266
267
268
269
270
271
272
273
274
275
276
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.post(\n",
    "    url + \"/load_lora_adapter\",\n",
    "    json={\n",
277
278
    "        \"lora_name\": \"lora1\",\n",
    "        \"lora_path\": lora1,\n",
279
280
281
282
283
284
285
286
287
288
    "    },\n",
    ")\n",
    "\n",
    "if response.status_code == 200:\n",
    "    print(\"LoRA adapter loaded successfully.\", response.json())\n",
    "else:\n",
    "    print(\"Failed to load LoRA adapter.\", response.json())"
   ]
  },
  {
289
   "cell_type": "markdown",
290
291
   "metadata": {},
   "source": [
292
    "Check inference output:"
293
294
295
296
297
298
299
300
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
    "url = f\"http://127.0.0.1:{port}\"\n",
    "json_data = {\n",
    "    \"text\": [\n",
    "        \"List 3 countries and their capitals.\",\n",
    "        \"List 3 countries and their capitals.\",\n",
    "    ],\n",
    "    \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n",
    "    # The first input uses lora0, and the second input uses lora1\n",
    "    \"lora_path\": [\"lora0\", \"lora1\"],\n",
    "}\n",
    "response = requests.post(\n",
    "    url + \"/generate\",\n",
    "    json=json_data,\n",
    ")\n",
    "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n",
    "print(f\"Output from lora1 (updated): \\n{response.json()[1]['text']}\\n\")"
317
318
319
320
321
322
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
323
    "Unload lora0 and replace it with a different adapter:"
324
325
326
327
328
329
330
331
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
332
333
334
335
336
    "response = requests.post(\n",
    "    url + \"/unload_lora_adapter\",\n",
    "    json={\n",
    "        \"lora_name\": \"lora0\",\n",
    "    },\n",
337
338
339
340
341
    ")\n",
    "\n",
    "response = requests.post(\n",
    "    url + \"/load_lora_adapter\",\n",
    "    json={\n",
342
343
    "        \"lora_name\": \"lora0\",\n",
    "        \"lora_path\": lora0_new,\n",
344
345
346
347
348
349
350
351
352
    "    },\n",
    ")\n",
    "\n",
    "if response.status_code == 200:\n",
    "    print(\"LoRA adapter loaded successfully.\", response.json())\n",
    "else:\n",
    "    print(\"Failed to load LoRA adapter.\", response.json())"
   ]
  },
353
354
355
356
357
358
359
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check output again:"
   ]
  },
360
361
362
363
364
365
366
367
368
369
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = f\"http://127.0.0.1:{port}\"\n",
    "json_data = {\n",
    "    \"text\": [\n",
    "        \"List 3 countries and their capitals.\",\n",
370
    "        \"List 3 countries and their capitals.\",\n",
371
372
373
374
375
376
377
378
379
    "    ],\n",
    "    \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n",
    "    # The first input uses lora0, and the second input uses lora1\n",
    "    \"lora_path\": [\"lora0\", \"lora1\"],\n",
    "}\n",
    "response = requests.post(\n",
    "    url + \"/generate\",\n",
    "    json=json_data,\n",
    ")\n",
380
381
    "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n",
    "print(f\"Output from lora1 (updated): \\n{response.json()[1]['text']}\\n\")"
382
383
   ]
  },
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LoRA GPU Pinning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another advanced option is to specify adapters as `pinned` during loading. When an adapter is pinned, it is permanently assigned to one of the available GPU pool slots (as configured by `--max-loras-per-batch`) and will not be evicted from GPU memory during runtime. Instead, it remains resident until it is explicitly unloaded.\n",
    "\n",
    "This can improve performance in scenarios where the same adapter is frequently used across requests, by avoiding repeated memory transfers and reinitialization overhead. However, since GPU pool slots are limited, pinning adapters reduces the flexibility of the system to dynamically load other adapters on demand. If too many adapters are pinned, it may lead to degraded performance, or in the most extreme case (`Number of pinned adapters == max-loras-per-batch`), halt all unpinned requests. Therefore, currently SGLang limits maximal number of pinned adapters to `max-loras-per-batch - 1` to prevent unexpected starvations. \n",
    "\n",
    "In the example below, we unload `lora1` and reload it as a `pinned` adapter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.post(\n",
    "    url + \"/unload_lora_adapter\",\n",
    "    json={\n",
    "        \"lora_name\": \"lora1\",\n",
    "    },\n",
    ")\n",
    "\n",
    "response = requests.post(\n",
    "    url + \"/load_lora_adapter\",\n",
    "    json={\n",
    "        \"lora_name\": \"lora1\",\n",
    "        \"lora_path\": lora1,\n",
    "        \"pinned\": True,  # Pin the adapter to GPU\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that the result is identical as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = f\"http://127.0.0.1:{port}\"\n",
    "json_data = {\n",
    "    \"text\": [\n",
    "        \"List 3 countries and their capitals.\",\n",
    "        \"List 3 countries and their capitals.\",\n",
    "    ],\n",
    "    \"sampling_params\": {\"max_new_tokens\": 32, \"temperature\": 0},\n",
    "    # The first input uses lora0, and the second input uses lora1\n",
    "    \"lora_path\": [\"lora0\", \"lora1\"],\n",
    "}\n",
    "response = requests.post(\n",
    "    url + \"/generate\",\n",
    "    json=json_data,\n",
    ")\n",
    "print(f\"Output from lora0: \\n{response.json()[0]['text']}\\n\")\n",
    "print(f\"Output from lora1 (pinned): \\n{response.json()[1]['text']}\\n\")"
   ]
  },
456
457
458
459
460
461
462
463
464
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "terminate_process(server_process)"
   ]
  },
465
466
467
468
469
470
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Future Works\n",
    "\n",
471
    "The development roadmap for LoRA-related features can be found in this [issue](https://github.com/sgl-project/sglang/issues/2929). Currently radix attention is incompatible with LoRA and must be manually disabled. Other features, including Unified Paging, Cutlass backend, and dynamic loading/unloadingm, are still under development."
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}