콘텐츠로 이동

Llama3-8B (텍스트 생성)

Llama3-8B는 Meta에서 개발한 오픈소스 LLM(대규모 언어 모델)로, 다양한 자연어 처리 작업에서 인간을 뛰어넘는 성능을 보여줍니다. 하지만 이 모델은 방대한 매개변수를 가지고 있고 그에 따른 필요 메모리 요구량도 매우 크기 때문에, 단일 기기에서 이 모델을 지원하기란 매우 어렵습니다.

이러한 문제를 해결하기 위해 리벨리온에서는 RSD(Rebellions Scalable Design)라고 부르는 텐서 병렬화 기법을 통해 대규모 언어 모델을 지원합니다. ATOM™ 여러 대를 동시에 사용하는 상황에서는 ATOM™ 한대로 연산을 처리할 때와는 다른 패턴이 나타납니다. 이 페이지에서는 RSD를 통해 대규모 언어 모델을 지원할 때 나타나는 몇몇 특징적인 상황에 대해 소개합니다.

사전 준비

설치

1
2
$ cd RBLN_MODEL_ZOO_PATH/huggingface/text2text-generation/llama/llama3-8b
$ pip install -r requirements.txt

Compile the Model and Extract Profiled Data

1
2
3
# 예시 텍스트: "Hey, are you conscious? Can you talk to me?”
$ python3 compile.py
$ RBLN_PROFILER=1 python3 inference.py

환경변수 대신 RBLN 런타임 API를 사용할 때에는 rbln_activate_profiler=True를 설정하는 것으로 RBLN 프로파일러를 활성화할 수 있습니다.

1
2
3
4
5
6
# Load compiled model
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=os.path.basename(model_id),
    export=False,
    rbln_activate_profiler=True,
)

Analysis of Profiled data from Llama3-8B with RBLN Profiler

이 절에서는 LlamaDecoderLayer의 계층 정규화(layer normalization) 연산과, 컴파일러 샤딩(sharding)에 의해 나뉘어진 어텐션(attention) 계층의 q_proj와 o_proj 연산이 어떤 방식으로 동기화와 입/출력을 처리하는지에 대해 설명합니다. 또한, 어텐션 연산의 처리 과정의 이해를 돕기 위해 일부 명령어의 이름에 새로이 추가되는 attention 에 대해 설명합니다.

LlamaDecoderLayer
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
# https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/src/transformers/models/llama/modeling_llama.py#L332
class LlamaDecoderLayer(nn.Module):
    ...
    def forward(...) -> ...:
        ...

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            ...
        )
        ...
LlamaAttation
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
# https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/src/transformers/models/llama/modeling_llama.py#L269
class LlamaAttention(nn.Module):
    ...
    def forward(
        self,
        hidden_states: torch.Tensor,
        ...
    ) -> ...:
        ...
        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        ...
        attn_output = self.o_proj(attn_output)
        ...

루트 노드에서의 계층 정규화 실행

0_causal_lm.model.0.input_layernorm은 그 크기가 작아 단일 ATOM™에서 충분히 처리할 수 있습니다. 따라서, 이 계층은 ATOM 0에서 처리한 후 그 결과를 리프(leaf) 노드들로 전송합니다. 루트(root) ATOM™에서 input_layernorm을 처리하는 동안 리프 노드들은 Neural DMA를 통해 나중에 사용할 가중치 매개변수들을 불러옵니다. input_layernorm 연산이 끝나면, 출력 텐서는 Device HDMA를 통해 리프 ATOM™들로 전송됩니다.

출력 데이터 전송 및 동기화

input_layernorm 연산이 끝나면, 해당 결과값은 우선 Task DMA를 통해 ATOM 0의 DRAM으로 전송됩니다. 또한, 이 값을 리프 ATOM™들도 사용해야 하기 때문에, 뒤이어 Device HDMA를 통해 동일한 텐서를 각 리프 ATOM™들로 전송합니다. 이 과정이 거의 시간 지연 없이 이루어지면서 리프 노드들이 빠르게 다음 계층을 처리 할 수 있습니다. 데이터가 전송된 후, 해당 데이터를 받는 각 리프 ATOM™에서 Device Sync 명령어를 실행합니다. 이 명령어를 처리한 후, 각 ATOM™은 External HDMA를 통해 받은 데이터를 불러오는 것처럼 해당 데이터를 불러옵니다. Device HDMA 명령어는 데이터를 받는 쪽이 아니라 보내는 쪽에서 수행한다는 점을 유의하시기 바랍니다.

루트 노드로의 데이터 모으기(gather)

Device HDMA 명령어는 루트 ATOM™에서 처리한 출력 텐서를 리프 ATOM™으로 분산할 때 뿐만 아니라 각 리프 ATOM™에서 처리한 출력 텐서들을 루트 ATOM™으로 모을 때에도 사용합니다. 위 예시는 리프 ATOM™에서 계산한 출력 텐서가 어떻게 루트 ATOM™으로 전송되는지를 보여줍니다. 각 리프 ATOM™에서 Device HDMA를 통해 전송한 데이터는 루트 ATOM™에서 Task DMA를 통해 내부 Shared Memory로 전송하는데, 마지막으로 전송된 데이터를 가져오는 Task DMA 명령어에서 각 출력 텐서를 모아 연산을 수행하는 Neural Engine Clusters에 의존성을 갖습니다.

attention 명명 규칙

Transformer 기반 네트워크에서는 어텐션 연산이 어떤 방식으로 이루어지는지를 이해하는 것이 매우 중요합니다. 하지만, 주로 모듈 단위로 모델을 구성하는 기존 프레임워크에서는 어텐션 연산이 언제 어떤 식으로 이루어지는지 알아보기가 매우 어렵습니다. RBLN 프로파일러는 어텐션 연산에 해당하는 부분에 attention을 붙여 사용자의 이해를 돕습니다.