콘텐츠로 이동

Flash Attention을 이용한 Llama3.1-8B 서빙

개요

이 페이지에서는 Ray Serve에서 RBLN SDK와 vllm-rbln을 활용하여 Llama3.1-8B 모델을 Flash Attention을 사용하여 서빙하는 방법을 소개합니다.

전체 튜토리얼의 주요 흐름은 다음과 같습니다:

  1. 환경 및 모델 컴파일 확인
  2. RBLN NPU를 활용하는 Ray Serve deployment 정의
  3. Serve CLI를 이용한 모델 서빙 애플리케이션 실행
  4. 엔드포인트 검증을 위한 추론 요청 전송

Ray Serve 환경 구성 방법에 대해서는 Ray Serve 개요를 먼저 읽어 보는 것을 권합니다. 모델 컴파일 및 배포에 대한 전체 스크립트 기반 예제는 모델 주에서 확인할 수 있습니다.

환경 설정 및 설치 확인

시작하기 전에 시스템 환경이 올바르게 구성되어 있고 필요한 패키지가 모두 설치되어 있는지 확인합니다. 필요한 항목은 다음과 같습니다:

  • 시스템 요구 사항:
    • Ubuntu 20.04 LTS (Debian bullseye) or higher
    • System with RBLN NPUs equipped (e.g., RBLN Rebel100™ or ATOM™)
  • 필수 패키지:
  • 설치 명령어: 
    1
    pip install -U ray[serve] transformers requests torch --extra-index-url https://download.pytorch.org/whl/cpu

Note

이 튜토리얼은 사용자가 RBLN SDK 기반의 모델 컴파일 및 추론에 대해 잘 이해하고 있다는 가정하에 작성되었습니다. RBLN SDK 사용법에 익숙하지 않다면 RBLN Optimum, vLLM 튜토리얼을 먼저 읽어 보는 것을 권합니다.

사전준비

Llama3.1-8B 모델 컴파일

optimum-rbln를 이용하여 Llama3.1-8B를 컴파일합니다. 해당 코드는 리벨리온 모델 주를 사용하였습니다.

Llama3.1-8B 모델을 컴파일하기 위해 사용된 파라미터는 다음과 같습니다.

  • export: 모델을 컴파일하려면 True로 설정해야 합니다.
  • rbln_batch_size: 컴파일을 위한 배치 크기를 정의합니다.
  • rbln_max_seq_len: 최대 시퀀스 길이를 지정합니다.
  • rbln_num_devices: 추론에 사용할 NPU의 수를 설정합니다.
  • rbln_kvcache_partition_len: Flash Attention을 위한 KV 캐시 파티션의 길이를 정의합니다. rbln_max_seq_len은 rbln_kvcache_partition_len의 배수여야 하며, 그보다 큰 값이어야 합니다.

Note

배치 크기는 서빙 환경에 맞게 선택할 수 있습니다. 여기에서는 1로 설정합니다.

get_model.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
import os

from optimum.rbln import RBLNLlamaForCausalLM

model_id = "meta-llama/Llama-3.1-8B-Instruct"

# Compile and export
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=model_id,
    export=True,
    rbln_batch_size=1,
    rbln_max_seq_len=131_072,
    rbln_num_devices=1,
    rbln_kvcache_partition_len=8_192,
)

# Save compiled results to disk
model.save_pretrained(os.path.basename(model_id))
get_model.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
import os

from optimum.rbln import RBLNLlamaForCausalLM

model_id = "meta-llama/Llama-3.1-8B-Instruct"

# Compile and export
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=model_id,
    export=True,
    rbln_batch_size=1,
    rbln_max_seq_len=131_072,
    rbln_num_devices=8,
    rbln_kvcache_partition_len=16_384,
)

# Save compiled results to disk
model.save_pretrained(os.path.basename(model_id))

Deployment

Deployment 개요

Step Description
1. Deployment 구현 Ray를 RBLN NPU와 함께 사용하도록 설정하고, 컴파일된 모델을 로드하여 런타임을 초기화하고 엔드포인트를 노출하는 Ray Serve deployment를 정의합니다.
2. 실행 Ray Serve CLI(serve run)를 사용하여 deployment를 실행합니다. 애플리케이션 이름, 디바이스 세트 또는 원격 Ray 클러스터를 옵션으로 구성할 수 있습니다.
3. 추론 요청 Serve 엔드포인트로 HTTP 요청을 보내고 응답을 검사하여 deployment를 검증합니다.

아래 섹션은 위의 단계를 순서대로 설명합니다.

1.1 리소스 할당

Ray Serve에서 리소스를 할당하는 방식은 각 서빙 작업(Actor, Deployment)별로 resources 파라미터를 통해 NPU 자원을 할당할 수 있습니다.

아래 Actor는 @ray.remote(resources={"RBLN": <NUM_NPUS>})로 RBLN NPU 리소스를 요청하는 방법을 보여줍니다. <NUM_NPUS>는 컴파일할 때 사용한 rbln_num_devices 값과 같게 설정합니다. RBLNActor 클래스는 할당된 RBLN NPU ID를 Serve Deployment로 전달하는데 사용합니다. 이에 대한 자세한 내용은 Ray에서 RBLN NPU 사용에서 확인할 수 있습니다.

1
2
3
4
@ray.remote(resources={"RBLN": 1})
class RBLNActor:
    def getDeviceId(self):
        return ray.get_runtime_context().get_accelerator_ids()["RBLN"]

1
2
3
4
@ray.remote(resources={"RBLN": 8})
class RBLNActor:
    def getDeviceId(self):
        return ray.get_runtime_context().get_accelerator_ids()["RBLN"]

1.2 Deployment 구현

Ray Serve deployment는 @serve.deployment 데코레이터를 활용하여 클래스를 하나의 Deployment(배포 단위)로 정의합니다. 이 데코레이터를 적용하면 해당 클래스가 Ray Serve의 서비스 엔드포인트로 등록되어, 각 Deployment의 라이프사이클(배포, 확장, 업데이트 등)과 관리를 Ray Serve가 담당하게 됩니다.

llama3_1-8b.py
  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
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 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
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
import json
import os
from unittest.mock import MagicMock

import ray
from fastapi import FastAPI, HTTPException
from ray import serve
from starlette.requests import Request
from starlette.responses import StreamingResponse
from vllm import AsyncEngineArgs, AsyncLLMEngine
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
    ChatCompletionResponse,
)
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
from vllm.renderers.online_renderer import OnlineRenderer

app = FastAPI()

ray.init(resources={"RBLN": 1})


@ray.remote(resources={"RBLN": 1})
class RBLNActor:
    def getDeviceId(self):
        return ray.get_runtime_context().get_accelerator_ids()["RBLN"]


@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 16})
@serve.ingress(app)
class Llama3_1__8B:
    def __init__(self, rbln_actor: RBLNActor):
        """
        Initialize actor.
        :return:
        """
        self.engine = None
        self.rbln_actor = rbln_actor
        self.model_name = "Llama-3.1-8B-Instruct"
        self.raw_request = None
        self.vllm_engine = None
        self.openai_serving_models = None
        self.completion_service = None
        self.chat_completion_service = None
        self.ids = ray.get(rbln_actor.getDeviceId.remote())

        self.os_environment_vars()
        self.initialize()

    def os_environment_vars(self):
        """
        Redefine the environment variables to be passed to the RBLN runtime and vLLM
        :return:
        """
        if self.ids is None or len(self.ids) <= 0:
            os.environ.pop("RBLN_DEVICES")
        os.environ["RBLN_DEVICES"] = ",".join(self.ids)
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

    def initialize(self):
        """
        Initialize model. This will be called during model loading time
        :return:
        """
        vllm_engine_args = AsyncEngineArgs(model=self.model_name)
        self.vllm_engine = AsyncLLMEngine.from_engine_args(vllm_engine_args)

        self.openai_serving_models = OpenAIServingModels(
            engine_client=self.vllm_engine,
            base_model_paths=[
                BaseModelPath(name=self.model_name, model_path=os.getcwd())
            ],
        )

        openai_serving_render = OnlineRenderer(
            model_config=self.vllm_engine.model_config,
            renderer=self.vllm_engine.renderer,
            request_logger=None,
            chat_template=None,
            chat_template_content_format="auto",
        )

        self.completion_service = OpenAIServingCompletion(
            self.vllm_engine,
            self.openai_serving_models,
            online_renderer=openai_serving_render,
            request_logger=None,
        )

        self.chat_completion_service = OpenAIServingChat(
            self.vllm_engine,
            self.openai_serving_models,
            "assistant",
            online_renderer=openai_serving_render,
            request_logger=None,
            chat_template_content_format="auto",
            chat_template=None,
        )

        async def isd():
            return False

        self.raw_request = MagicMock()
        self.raw_request.headers = {}
        self.raw_request.is_disconnected = isd

    @app.post("/v1/chat/completions")
    async def chat_completion(self, http_request: Request):
        """
        Handle chat completion request.
        :param http_request: The HTTP request object
        :return: The chat completion response
        """
        try:
            json_string: dict = await http_request.json()
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="Invalid JSON format request")
        request: ChatCompletionRequest = ChatCompletionRequest.model_validate(
            json_string
        )

        g = await self.chat_completion_service.create_chat_completion(
            request, self.raw_request
        )

        if isinstance(g, ErrorResponse):
            return [g.model_dump()]

        if request.stream:

            async def stream_generator():
                async for response in g:
                    yield response

            return StreamingResponse(stream_generator(), media_type="text/event-stream")
        else:
            return [g.model_dump()]

    @app.post("/v1/completions")
    async def completion(self, http_request: Request):
        """
        Handle completion request.
        :param http_request: The HTTP request object
        :return: The completion response
        """
        try:
            json_string: dict = await http_request.json()
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="Invalid JSON format request")
        request: CompletionRequest = CompletionRequest.model_validate(json_string)

        g = await self.completion_service.create_completion(request, self.raw_request)

        if isinstance(g, ErrorResponse):
            return [g.model_dump()]

        if request.stream:

            async def stream_generator():
                async for response in g:
                    yield response

            return StreamingResponse(stream_generator(), media_type="text/event-stream")
        else:
            return [g.model_dump()]


rbln_actor = RBLNActor.remote()
app = Llama3_1__8B.bind(rbln_actor)
llama3_1-8b.py
  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
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 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
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
import json
import os
from unittest.mock import MagicMock

import ray
from fastapi import FastAPI, HTTPException
from ray import serve
from starlette.requests import Request
from starlette.responses import StreamingResponse
from vllm import AsyncEngineArgs, AsyncLLMEngine
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
    ChatCompletionResponse,
)
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.engine.protocol import ErrorResponse
from vllm.renderers.online_renderer import OnlineRenderer

app = FastAPI()

ray.init(resources={"RBLN": 8})


@ray.remote(resources={"RBLN": 8})
class RBLNActor:
    def getDeviceId(self):
        return ray.get_runtime_context().get_accelerator_ids()["RBLN"]


@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 16})
@serve.ingress(app)
class Llama3_1__8B:
    def __init__(self, rbln_actor: RBLNActor):
        """
        Initialize actor.
        :return:
        """
        self.engine = None
        self.rbln_actor = rbln_actor
        self.model_name = "Llama-3.1-8B-Instruct"
        self.raw_request = None
        self.vllm_engine = None
        self.openai_serving_models = None
        self.completion_service = None
        self.chat_completion_service = None
        self.ids = ray.get(rbln_actor.getDeviceId.remote())

        self.os_environment_vars()
        self.initialize()

    def os_environment_vars(self):
        """
        Redefine the environment variables to be passed to the RBLN runtime and vLLM
        :return:
        """
        if self.ids is None or len(self.ids) <= 0:
            os.environ.pop("RBLN_DEVICES")
        os.environ["RBLN_DEVICES"] = ",".join(self.ids)
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

    def initialize(self):
        """
        Initialize model. This will be called during model loading time
        :return:
        """
        vllm_engine_args = AsyncEngineArgs(model=self.model_name)
        self.vllm_engine = AsyncLLMEngine.from_engine_args(vllm_engine_args)

        self.openai_serving_models = OpenAIServingModels(
            engine_client=self.vllm_engine,
            base_model_paths=[
                BaseModelPath(name=self.model_name, model_path=os.getcwd())
            ],
        )

        openai_serving_render = OnlineRenderer(
            model_config=self.vllm_engine.model_config,
            renderer=self.vllm_engine.renderer,
            request_logger=None,
            chat_template=None,
            chat_template_content_format="auto",
        )

        self.completion_service = OpenAIServingCompletion(
            self.vllm_engine,
            self.openai_serving_models,
            online_renderer=openai_serving_render,
            request_logger=None,
        )

        self.chat_completion_service = OpenAIServingChat(
            self.vllm_engine,
            self.openai_serving_models,
            "assistant",
            online_renderer=openai_serving_render,
            request_logger=None,
            chat_template_content_format="auto",
            chat_template=None,
        )

        async def isd():
            return False

        self.raw_request = MagicMock()
        self.raw_request.headers = {}
        self.raw_request.is_disconnected = isd

    @app.post("/v1/chat/completions")
    async def chat_completion(self, http_request: Request):
        """
        Handle chat completion request.
        :param http_request: The HTTP request object
        :return: The chat completion response
        """
        try:
            json_string: dict = await http_request.json()
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="Invalid JSON format request")
        request: ChatCompletionRequest = ChatCompletionRequest.model_validate(
            json_string
        )

        g = await self.chat_completion_service.create_chat_completion(
            request, self.raw_request
        )

        if isinstance(g, ErrorResponse):
            return [g.model_dump()]

        if request.stream:

            async def stream_generator():
                async for response in g:
                    yield response

            return StreamingResponse(stream_generator(), media_type="text/event-stream")
        else:
            return [g.model_dump()]

    @app.post("/v1/completions")
    async def completion(self, http_request: Request):
        """
        Handle completion request.
        :param http_request: The HTTP request object
        :return: The completion response
        """
        try:
            json_string: dict = await http_request.json()
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="Invalid JSON format request")
        request: CompletionRequest = CompletionRequest.model_validate(json_string)

        g = await self.completion_service.create_completion(request, self.raw_request)

        if isinstance(g, ErrorResponse):
            return [g.model_dump()]

        if request.stream:

            async def stream_generator():
                async for response in g:
                    yield response

            return StreamingResponse(stream_generator(), media_type="text/event-stream")
        else:
            return [g.model_dump()]


rbln_actor = RBLNActor.remote()
app = Llama3_1__8B.bind(rbln_actor)

2. 실행

Ray Serve CLI(serve run)를 사용하여 애플리케이션을 실행합니다. 파라미터는 module:application 형식으로 지정되며, 여기서 module은 Python 파일명(확장자 .py 제외), application은 Serve 엔트리포인트 객체명입니다.

이 예시에서는 llama3_1-8b.py에서 app 객체를 정의하므로 아래와 같은 명령어로 deployment를 시작할 수 있습니다. 원격 Ray 클러스터로 연결하거나 RBLN_DEVICES를 이용하여 특정 카드를 지정하는 등 추가 옵션을 함께 사용할 수 있습니다.

1
$ serve run llama3_1-8b:app --name "llama3.1-8b"

예시 출력: 

1
Application 'llama3.1-8b' is ready at http://127.0.0.1:8000/.

3.1 추론 요청 예시(Completion API)

1
2
3
4
5
6
echo '{
"model": "Llama-3.1-8B-Instruct",
"prompt": "A robot may not injure a human being",
"stream": false
}' | curl -sN -H "Content-Type: application/json" -X POST --data-binary @- \
http://localhost:8000/v1/completions | jq .

예시 출력:

OUTPUT 
 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
[
{
    "id": "cmpl-790b3169a1c1428a8912140848b42753",
    "object": "text_completion",
    "created": 1762993969,
    "model": "Llama-3.1-8B-Instruct",
    "choices": [
    {
        "index": 0,
        "text": " or, through inaction, allow a human being to come to harm.\nA",
        "logprobs": null,
        "finish_reason": "length",
        "stop_reason": null,
        "token_ids": null,
        "prompt_logprobs": null,
        "prompt_token_ids": null
    }
    ],
    "service_tier": null,
    "system_fingerprint": null,
    "usage": {
    "prompt_tokens": 10,
    "total_tokens": 26,
    "completion_tokens": 16,
    "prompt_tokens_details": null
    },
    "kv_transfer_params": null
}
]

3.2 추론 요청 예시(Chat Completions API)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
echo '{
"model": "Llama-3.1-8B-Instruct",
"messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Who won the world series in 2020?"},
    {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
    {"role": "user", "content": "Where was it played?"}
],
"temperature": 0.0,
"max_tokens": 50,
"stream": false
}' | curl -sN -H "Content-Type: application/json" -X POST --data-binary @- \
http://localhost:8000/v1/chat/completions | jq .

예시 출력:

OUTPUT 
 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
30
31
32
33
34
35
36
37
38
[
{
    "id": "chatcmpl-bad227512c58497eb40dcc1287f29b50",
    "object": "chat.completion",
    "created": 1762994293,
    "model": "Llama-3.1-8B-Instruct",
    "choices": [
    {
        "index": 0,
        "message": {
        "role": "assistant",
        "content": "The 2020 World Series was played at Globe Life Field in Arlington, Texas. It was a neutral site due to the COVID-19 pandemic, and the Dodgers played the Tampa Bay Rays in the series.",
        "refusal": null,
        "annotations": null,
        "audio": null,
        "function_call": null,
        "tool_calls": [],
        "reasoning_content": null
        },
        "logprobs": null,
        "finish_reason": "stop",
        "stop_reason": null,
        "token_ids": null
    }
    ],
    "service_tier": null,
    "system_fingerprint": null,
    "usage": {
    "prompt_tokens": 79,
    "total_tokens": 122,
    "completion_tokens": 43,
    "prompt_tokens_details": null
    },
    "prompt_logprobs": null,
    "prompt_token_ids": null,
    "kv_transfer_params": null
}
]

참고