Teola Towards End-to-End Optimization of LLM-based Applications

논문 : https://arxiv.org/abs/2407.00326

개인의견 : 결국 프리미티브라고 하는 작은 단위로 나눈다는 개념을 도입하는거 같고 이게 LLM이 아닌 경우도 고려해서 한번에 serving을 한다는것이라고 이해를 했는데 vllm의 continuous batching + sequence를 unit한 단위로 나누어서 스케줄링하는거에 대비 얼마만큼의 이득이 있을까 궁금하긴하다.

아래글은 Scholar GPT와 대화를 나눈 내용입니다.

Q : 논문의 핵심 내용과 강점, 알고리즘 설명, 그리고 한계점

Answer

강점과 독창적인 지점

이 논문은 Teola라는 새로운 오케스트레이션 프레임워크를 제안하여 LLM(대형 언어 모델) 기반 애플리케이션의 엔드 투 엔드 성능 최적화를 목표로 하고 있습니다. 기존의 LLM 최적화는 대부분 개별 모듈에 한정되어 있거나 단일 엔진 최적화에 집중했으나, 이 논문은 모든 작업을 더 작은 작업 단위인 “프리미티브”로 나누고 이를 데이터 흐름 그래프로 표현함으로써 최적화를 극대화하는 혁신적인 접근 방식을 제시합니다. 이는 LLM과 비-LLM 모듈 간 병렬화 및 파이프라이닝을 통해 더욱 효율적인 스케줄링을 가능하게 하며, 이를 통해 애플리케이션 전체의 성능을 높입니다.

핵심 알고리즘과 예시

Teola의 핵심 알고리즘은 프리미티브 수준 데이터 흐름 그래프를 생성하여 각 쿼리의 작업을 더 세분화된 단위로 표현하고, 이를 통해 병렬 실행과 파이프라이닝 최적화를 수행합니다. 다음은 예시 입력과 함께 이 과정에 대한 설명입니다.

  1. 예시 입력: 사용자가 “문서 QA” 기능에서 질문(“What is LLM?”)과 문서(chapter1.txt)를 입력합니다. 문서는 300자로 분할되며, 질의 확장 및 답변 생성을 위한 워크플로가 정의됩니다.
  2. 데이터 흐름 그래프 생성: Teola는 입력 데이터를 기반으로 각 프리미티브 작업을 연결한 그래프를 생성합니다. 예를 들어, 문서에서 임베딩 벡터를 생성하는 임베딩 작업과 질문을 LLM에 전달하여 초기 응답을 생성하는 질의 확장 작업 등이 포함됩니다.
  3. 프리미티브 병렬화: 문서의 각 청크에 대한 임베딩 생성은 병렬로 실행될 수 있으므로 Teola는 이를 병렬화하여 시간 단축을 시도합니다.
  4. 최종 결과 생성: 답변 생성을 위한 LLM 프리미티브와 재랭킹 작업을 통해 최적의 답변을 산출합니다.

이 과정을 통해 Teola는 작업의 효율성을 높이고 쿼리 처리 속도를 크게 향상시킵니다.

한계점

Teola는 다양한 최적화를 통해 엔드 투 엔드 성능을 크게 개선하지만, 몇 가지 한계도 존재합니다. 첫째, 동적 그래프 처리에 제한적입니다. 일부 복잡한 질의 응답 작업에서는 실시간 데이터나 피드백 루프가 필요한데, Teola는 이러한 동적 요구에 적응하기 어렵습니다. 둘째, 프리미티브 수준의 오케스트레이션을 위해서는 특정 엔진 수정이 필요하여, 기존 LLM 또는 비-LLM 엔진과의 호환성 문제가 발생할 수 있습니다.

Q : LLM 프리미티브에 대해서 예시를 들어서 자세하게 설명해 주겠니?

Answer:

LLM 프리미티브는 Teola에서 작업을 더 세분화하여 처리하기 위한 가장 작은 단위 작업입니다. 이러한 프리미티브는 기존에 하나의 단일 작업으로 처리되던 LLM 작업을 더 작은 단위로 쪼개어 각 작업의 병렬 처리 및 최적화 가능성을 극대화합니다. LLM 프리미티브의 예시로는 프리필링(Prefilling)디코딩(Decoding)이 있으며, 각각의 작업 단계를 효율적으로 관리하기 위해 추가적인 세부 프리미티브를 포함할 수도 있습니다.

1. 프리필링(Prefilling)

프리필링은 LLM이 최초의 출력을 생성하기 위해 필요한 모든 입력을 처리하는 단계입니다. 예를 들어, 입력으로 질문(“What is LLM?”)과 함께 사용자 지침(“정확하고 간결하게 답변하시오.”)이 주어졌다고 가정해 보겠습니다. 이 경우 프리필링 작업은 이 전체 입력을 한 번에 처리하여 첫 번째 출력을 생성하는 데 중점을 둡니다.

  • 부분 프리필링(Partial Prefilling): 전체 입력을 기다리지 않고 일부 입력이 준비되면 이를 먼저 처리하는 방식입니다. 예를 들어, 사용자 지침이 준비되었지만 데이터베이스에서 추가 맥락이 아직 반환되지 않은 상태라면, Teola는 사용자 지침만을 먼저 프리필링하여 결과를 준비합니다. 이후 나머지 맥락이 준비되면 이를 추가하여 완전 프리필링(Full Prefilling)을 진행할 수 있습니다.

2. 디코딩(Decoding)

디코딩 단계는 프리필링 이후에 반복적으로 실행되며, LLM이 최종 출력을 생성할 때까지 각 토큰을 하나씩 예측합니다. Teola에서는 디코딩 단계를 부분 디코딩(Partial Decoding)으로 나누어 병렬 처리를 가능하게 합니다.

  • 부분 디코딩(Partial Decoding): LLM이 전체 답변을 완료하기 전이라도 의미 있는 일부 문장을 생성하면 이를 즉시 다음 프리미티브 작업(예: 후처리 작업이나 추가 분석)으로 전달할 수 있습니다. 예를 들어, LLM이 “LLM is a machine learning model…“이라는 초기 응답을 생성했다면, 이를 바로 후속 작업에 넘겨 사용자에게 빠르게 일부 응답을 제공하거나 다른 프리미티브에서 추가 작업을 시작할 수 있습니다.

  • 파이프라인 디코딩(Pipelined Decoding): 여러 질문에 대한 응답을 동시에 생성해야 하는 경우, Teola는 각 질문에 대해 생성된 부분 답변을 바로바로 후속 작업에 연결하여 전체 응답이 완료되기 전에도 다른 프리미티브를 시작할 수 있습니다.

예시로 본 전체 LLM 프리미티브 프로세스

  1. 입력 설정: 사용자는 “What is LLM?”이라는 질문을 하고, 추가적으로 “상세한 설명 제공”이라는 지침을 추가합니다.
  2. 프리필링 실행:
    • 사용자의 지침을 기반으로 부분 프리필링을 시작하여 지침에 대한 초기 처리를 수행합니다.
    • 이후 추가적인 맥락 정보가 데이터베이스에서 반환되면, 이를 통해 완전 프리필링을 수행하여 첫 번째 토큰 출력을 준비합니다.
  3. 디코딩 실행:
    • LLM이 첫 번째 출력 토큰을 생성하면, 부분 디코딩을 통해 중간 결과를 생성하고, 이 결과를 다른 후속 프리미티브 작업(예: 검색 결과 통합)에 전달합니다.
    • LLM이 전체 답변을 생성할 때까지 파이프라인 디코딩으로 각 단계의 결과를 다른 프리미티브에 연속적으로 전달하여 작업을 빠르게 완료합니다.

이와 같이, 프리미티브 수준에서의 세분화된 작업 처리는 LLM 기반 애플리케이션의 속도를 높이고, 전체적인 워크플로를 최적화하는 데 중요한 역할을 합니다.

Q : 혹시 vllm과 같은 다른 LLM serving tool과 비교해서 어떤 차이점이 있는지 설명해줄 수 있어? 그리고 이논문에서 제시하는 방법이 기존 다른 방법에 비해서 어떤 독창성과 장점, 단점을 가지고 있는지 궁금해

Answer :

Teola와 vLLM 같은 다른 LLM 서빙 도구 비교

vLLM은 LLM 서빙을 위한 도구로서, 특히 대규모 모델의 인퍼런스 최적화에 중점을 둡니다. 주로 GPU 메모리와 처리 속도를 효율적으로 사용하도록 설계되었으며, 트랜스포머 기반의 모델에 대한 효율적 서빙을 위해 다양한 키-값(KV) 캐시 관리부분 디코딩 지원을 통해 성능을 높이는 방식입니다.

반면에 Teola는 LLM뿐만 아니라 다른 비-LLM 구성 요소도 함께 사용하는 엔드 투 엔드 애플리케이션 워크플로의 전체 최적화를 목표로 합니다. LLM 인퍼런스에만 집중하는 vLLM과 달리, Teola는 LLM을 포함한 다양한 프레임워크(데이터베이스, 검색 엔진 등) 간의 병렬 실행과 파이프라이닝을 지원하여 전체 애플리케이션 성능을 향상시키는 점에서 차별화됩니다.

주요 차이점 요약:

  • 작업 단위: vLLM은 주로 모델 단위의 요청을 다루고 인퍼런스 최적화에 집중하며, Teola는 더 작은 작업 단위(프리미티브)로 세분화하여 LLM과 비-LLM 작업을 포함한 전체 워크플로를 최적화합니다.
  • 병렬화 및 파이프라이닝: Teola는 프리미티브 간 병렬화와 파이프라이닝을 통해, 단순히 한 번에 한 요청을 최적화하는 것이 아니라 전체 작업의 흐름을 최적화합니다.
  • 스케줄링 방식: vLLM은 주로 GPU 메모리 관리와 토큰 단위의 효율적 디코딩에 초점을 맞추며, Teola는 프리미티브 간 의존성 관리 및 우선 순위 기반 스케줄링을 통해 최적의 워크플로를 제공합니다.

Teola의 독창성과 장점

  1. 프리미티브 기반 오케스트레이션: Teola는 기존의 모듈 단위가 아닌 더 세분화된 프리미티브 수준에서 작업을 구성하여, 모든 작업 단위를 데이터 흐름 그래프로 표현합니다. 이를 통해 개별 작업 간의 의존성을 더 세밀하게 제어하고 병렬화를 극대화할 수 있습니다. 이는 기존 LLM 서빙 도구가 모듈 단위로 처리하여 병렬화와 최적화 기회가 제한적이었던 점과 차별화되는 독창성입니다.

  2. 애플리케이션 전체 최적화: Teola는 단순히 LLM 인퍼런스 성능뿐만 아니라, LLM이 포함된 전체 애플리케이션 성능을 최적화합니다. LLM과 비-LLM 작업 간에 발생하는 의존성을 최적화하고, 요청의 우선 순위를 고려한 스케줄링을 통해, 애플리케이션이 사용자 요청에 대해 최단 시간에 응답할 수 있도록 설계되었습니다.

  3. 탑재 가능한 최적화 패스: Teola는 다양한 최적화 패스를 제공하여, 특정 프리미티브의 데이터 흐름을 최적화합니다. 예를 들어, 중복된 의존성을 제거하거나, 배치 크기를 최적화하는 등 정적 규칙 기반 최적화가 적용되어, 특히 대규모 애플리케이션에 대해 복잡한 워크플로를 효율적으로 관리할 수 있습니다.

  4. 토폴로지 기반 배칭: Teola는 프리미티브의 위상 정보를 고려한 배칭 전략을 사용하여, 단순히 요청을 묶어 처리하는 기존 방식과 달리, 프리미티브의 상호 관계를 고려하여 최적의 배칭과 스케줄링을 수행합니다. 이는 대기 시간을 줄이고 시스템의 전반적인 성능을 높이는 데 효과적입니다.

단점 및 한계점

  1. 다이나믹한 워크플로 관리의 어려움: Teola의 최적화 방식은 사전 정의된 그래프에 기반하기 때문에, 실행 중 실시간으로 변동하는 요구 사항이나 피드백 루프를 수용하는 데 한계가 있습니다. 예를 들어, 사용자의 입력에 따라 반복적인 요청이 필요하거나, 실시간 피드백을 반영하여 결과를 수정해야 하는 작업에는 다소 유연성이 부족할 수 있습니다.

  2. 백엔드 엔진과의 결합도 문제: Teola는 더 정교한 프리미티브 수준의 최적화를 위해 각 엔진(예: LLM, 데이터베이스 등)을 조정해야 하는 경우가 많습니다. 이는 사용자가 기존의 서빙 엔진을 사용하려 할 때 호환성 문제를 발생시킬 수 있으며, 특정 엔진에 최적화된 기능을 구현해야 하는 추가적인 공수가 필요할 수 있습니다.

  3. 설정 및 관리 복잡성: Teola는 다양한 프리미티브와 그래프 기반의 스케줄링을 지원하므로, 시스템 설정 및 관리가 복잡할 수 있습니다. 특히 복잡한 애플리케이션이나 다수의 엔진이 연결된 환경에서는 설정 및 최적화 패스의 선택이 다소 까다로울 수 있습니다.

요약

Teola는 기존 LLM 서빙 도구와 달리 LLM 및 비-LLM 구성 요소를 포함한 전체 애플리케이션을 최적화하며, 세분화된 프리미티브와 그래프 기반 최적화를 통해 독창적인 방식으로 병렬화를 극대화합니다. 다만, 다이나믹한 요구 사항 대응이나 기존 엔진과의 결합 문제에서 제한적일 수 있습니다.