Falcon-40B 모델을 대규모 모델 추론 딥러닝 컨테이너(DLC)로 Amazon SageMaker에 배포하기

이번 게시글은 영문 게시글(Deploy Falcon-40B with large model inference DLCs on Amazon SageMaker by James Park, Abhi Shivaditya, Evandro Franco, Frank Liu, Qing Lan, and Robert Van Dusen)의 한글 번역글입니다.

2023년 6월 초에 Technology Innovation Institute (TII)는 오픈소스 기반 대규모 언어 모델(LLM)인 TII Falcon LLM을 출시했습니다. Amazon SageMaker를 통해 1조 개의 토큰으로 학습된 Falcon은 최고 수준의 성능(이 글을 쓰는 시점에 허깅페이스 리더보드에서 1위)을 자랑하며, llama-65B와 같은 다른 LLM에 비해 상대적으로 가볍고 호스팅 비용이 저렴합니다. 이번 포스팅에서는 SageMaker에서 대규모 모델 추론 딥러닝 컨테이너를 사용하여 언어 이해 및 자동화된 글쓰기 지원과 같은 애플리케이션을 위해 Falcon을 배포하는 방법을 보여드립니다.

Falcon LLM이 SageMaker에 안착했습니다.

TII는 아부다비의 첨단기술연구위원회 산하 응용 연구 조직(Advanced Technology Research Council)으로, 과학자, 연구원, 엔지니어로 구성된 팀이 혁신적인 기술을 발굴하고, 우리 사회의 미래를 보장할 과학적 돌파구를 개발하는 데 전념하고 있습니다. 올해 초, TII는 최첨단 오픈 소스 LLM을 훈련하기 시작했으며, 이 작업을 완료하기 위해 SageMaker의 인프라, 도구 및 전문 지식을 사용했습니다. (이 모델이 SageMaker에서 훈련된 방법에 대해 자세히 알아보려면 Technology Innovation Institute trains the state-of-the-art Falcon LLM 40B foundation model on Amazon SageMaker 게시물을 참조하세요.) 이러한 노력의 결과물이 바로 TII Falcon LLM입니다.

1조 개의 토큰으로 훈련된 Falcon은 Eleuther AI 언어 모델 평가 Harness와 비교하여 최고 수준의 성능을 자랑하며 현재 허깅페이스 리더보드에서 정확도(accuracy) 부문 1위를 차지하고 있습니다. 이 모델은 Falcon-40B와 Falcon-7B의 두 가지 크기로 제공되며 언어 이해, 대화 경험, 자동화된 작문 지원과 같은 애플리케이션에서 최첨단 성능을 발휘합니다. 이 게시물은 이러한 유형의 도메인에서 높은 정확도의 추론을 위해 SageMaker에 Falcon 배포를 시작하도록 도와줍니다.

SageMaker 대규모 모델 추론 DLC로 LLM 호스팅 간소화

Falcon-40B 및 Falcon-7B와 같은 LLM을 호스팅하는 것은 까다로운 작업입니다. 모델이 클수록 수 십억 개의 파라미터를 포함하기 때문에 더 정확한 경우가 많지만, 모델 크기가 크면 추론 레이턴시(latency)가 느려지거나 처리량(throughput)이 저하될 수 있습니다. LLM을 호스팅하려면 만족스러운 성능을 달성하기 위해 더 많은 GPU 메모리와 최적화된 커널이 필요할 수 있습니다. 상황을 더욱 복잡하게 만드는 것은 Falcon-7B와 같은 소규모 모델은 일반적으로 NVIDIA A10G 단일 GPU를 사용하는 AWS G5 인스턴스 유형(예: g5.xlarge)에서 호스팅할 수 있지만, Falcon-40B와 같은 대규모 모델은 그렇지 못하다는 점입니다. 이 경우 텐서 병렬 처리와 같은 전략을 사용하여 더 큰 모델을 여러 조각으로 분할하고 다중 GPU의 메모리를 활용해야 합니다. 하지만 소규모 모델에 사용되는 레거시 호스팅 솔루션은 일반적으로 이러한 유형의 기능을 제공하지 않으므로 복잡해집니다.

이 때, SageMaker 대규모 모델 추론(LMI; Large Model Inference) 딥러닝 컨테이너(DLC; Deep Learning Containers)가 도움이 됩니다. LMI DLC는 Falcon-40B와 같은 LLM을 호스팅하기 위한 완벽한 엔드-투-엔드 솔루션입니다. 프론트엔드에는 처리량을 높이기 위해 토큰 스트리밍 및 인스턴스 내 자동 모델 복제 등의 기능을 갖춘 대규모 모델 추론을 위해 설계된 고성능 모델 서버(DJL Serving)가 포함되어 있습니다. 백엔드에는 여러 GPU에서 모델 파라미터를 샤드하고 관리할 수 있는 DeepSpeed 및 FasterTransformer와 같은 여러 고성능 모델 병렬 엔진도 포함되어 있습니다. 이러한 엔진에는 널리 사용되는 트랜스포머 모델에 최적화된 커널도 포함되어 있어 추론 속도를 최대 3배까지 향상시킬 수 있습니다. LMI DLC를 사용하면 구성 파일을 생성하기만 하면 SageMaker에서 LLM 호스팅을 시작할 수 있습니다. SageMaker LMI DLC에 대해 자세히 알아보려면 SageMaker의 모델 병렬 처리 및 대규모 모델 추론 개발자 문서와 깃허브에 공개된 도커 이미지 목록을 참조하세요. 또한 LMI DLC를 사용하여 SageMaker에서 BLOOM-175B를 호스팅하는 방법에 대한 이전 게시물도 확인해 보세요.

솔루션 개요

이 게시물에서는 LMI DLC를 사용하여 SageMaker 상에서 DeepSpeed로 Falcon-40B를 호스팅하는 방법을 안내합니다. Falcon-40B는 여러 장의 A10 GPU를 사용해야 하는 반면, Falcon-7B는 단일 A10 GPU만 필요합니다. 또한 DeepSpeed와 Accelerate를 모두 사용하여 Falcon-40B 및 Falcon-7B를 호스팅하는 데 도움이 될 예제를 준비했습니다. 코드 예제는 깃허브에 있습니다.

이 예제는 SageMaker 노트북 인스턴스 또는 Amazon SageMaker Studio 노트북에서 실행할 수 있습니다. LMI 및 DeepSpeed를 사용하여 Falcon-40B를 호스팅하려면 ml.g5.24xlarge 인스턴스를 사용해야 합니다. 이 인스턴스는 각각 96기가바이트의 GPU 메모리를 지원하는 4개의 NVIDIA A10G GPU를 제공합니다. 또한 호스트는 96개의 vCPU와 384기가바이트의 호스트 메모리를 제공합니다. LMI 컨테이너는 모델 다운로드 및 모델 아티팩트를 파티셔닝하여 구성 파라미터를 여러 GPU에 분산할 수 있도록 하는 등 LLM 호스팅과 관련된 수많은 작업을 해결하는 데 도움이 됩니다.

SageMaker 기계 학습(ML) 인스턴스 할당량은 계정마다 다릅니다. 이 게시물의 예제를 수행하는 동안 g5.24xlarge 인스턴스에 대한 할당량을 초과했다는 오류가 표시되는 경우 Service Quotas console을 통해 한도(quota)를 늘릴 수 있습니다.

노트북 사용 방법

먼저 예제에 필요한 패키지를 설치하고 임포트하는 것으로 시작합니다. 본 예제는 Boto3 SDK와 SageMaker SDK를 사용합니다. SageMaker와 LMI가 사용하는 데 필요한 모델 아티팩트를 저장하기 위해 Amazon Simple Storage Service (Amazon S3)를 사용하므로 그에 따라 S3 접두사(prefix) 변수를 설정합니다. 다음 코드를 참조하기 바랍니다.

import sagemaker
import jinja2
from sagemaker import image_uris
import boto3
import os
import time
import json
from pathlib import Path
from sagemaker.utils import name_from_base

role = sagemaker.get_execution_role()  # execution role for the endpoint
sess = sagemaker.session.Session()  # sagemaker session for interacting with different AWS APIs
bucket = sess.default_bucket()  # bucket to house artifacts
model_bucket = sess.default_bucket()  # bucket to house artifacts
s3_code_prefix_deepspeed = "hf-large-model-djl-/code_falcon40b/deepspeed"  # folder within bucket where code artifact will go
region = sess._region_name
account_id = sess.account_id()
s3_client = boto3.client("s3")
sm_client = boto3.client("sagemaker")
smr_client = boto3.client("sagemaker-runtime")
jinja_env = jinja2.Environment()

이제 워크스페이스에 로컬 폴더를 생성하여 모델 아티팩트를 저장합니다.

!mkdir -p code_falcon40b_deepspeed

먼저 생성한 로컬 디렉터리에 serving.properties 구성 파일을 생성합니다. 이 serving.properties 파일은 사용하려는 모델 병렬화 및 추론 최적화 엔진을 LMI 컨테이너와 프런트엔드 DJL Serving 라이브러리에 알려줍니다. Configurations and settings 개발자 문서에서 DeepSpeed와 허깅페이스 가속 모두에 대한 구성 옵션을 찾을 수 있습니다. 여기서 option.model_id 매개변수를 설정하여 가져올 허깅페이스 모델을 정의한다는 점에 유의하세요. SageMaker는 허깅페이스 모델 작업을 간단하게 만들어주며, 이 한 줄이면 충분합니다. 또한 ml.g5.24xlarge 인스턴스에 4장의 GPU가 탑재되어 있기 때문에 option.tensor_parallel_degree를 4로 설정했습니다. 이 매개 변수는 생성 및 배포할 모델의 파티션 수를 정의합니다. ml.g5.48xlarge와 같이 GPU가 8장인 더 큰 인스턴스를 사용할 경우에 값을 4로 설정했다면 LMI는 자동으로 모델의 복제본 2개를 생성합니다.(복제본 2개가 각각 4개의 GPU에 분산됩니다.) 다음 코드를 참조하기 바랍니다.

%%writefile ./code_falcon40b_deepspeed/serving.properties
engine=Python
#to deploy falcon-40b-instruct set the model_id value to 'tiiuae/falcon-40b-instruct'
option.model_id=tiiuae/falcon-40b
option.tensor_parallel_degree=4
#option.s3url = {{s3url}}

사용자의 필요에 따라 더 적합한 경우라면, tiiuae/falcon-40b 를 tiiuae/falcon-40b-instruct로 교체할 수도 있습니다.

또한 필요한 패키지를 설치하도록 지정하는 requirements.txt 파일도 포함되어 있습니다.

%%writefile ./code_falcon40b_deepspeed/requirements.txt
einops
torch==2.0.1

마지막으로 필요한 것은 모델 추론에 사용할 model.py 파일입니다.

%%writefile ./code_falcon40b_deepspeed/model.py
from djl_python import Input, Output
import os
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from typing import Any, Dict, Tuple
import warnings

predictor = None

def get_model(properties):
    model_name = properties["model_id"]
    local_rank = int(os.getenv("LOCAL_RANK", "0"))
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    generator = pipeline(
        task="text-generation", model=model, tokenizer=tokenizer, device_map="auto"
    )
    return generator


def handle(inputs: Input) -> None:
    global predictor
    if not predictor:
        predictor = get_model(inputs.get_properties())
    if inputs.is_empty():
        # Model server makes an empty call to warmup the model on startup
        return None
    data = inputs.get_as_json()
    text = data["text"]
    text_length = data["text_length"]
    outputs = predictor(text, do_sample=True, min_length=text_length, max_length=text_length)
    result = {"outputs": outputs}
    return Output().add_as_json(result)

이제 끝입니다! 이 시점에서 DeepSpeed로 Falcon-40B를 배포하는 데 필요한 모든 아티팩트를 만들었습니다! 디렉토리를 *.tar.gz 파일로 패키징하고 Amazon S3에 업로드합니다. 실제 모델은 이 파일로 다운로드되거나 패키징되지 않았습니다. LMI 컨테이너가 허깅페이스에서 직접 모델을 다운로드합니다. 또한 다운로드 성능이 더 좋은 위치에 모델 사본을 직접 다운로드하려는 경우 S3 버킷을 대상으로 지정하는 옵션도 있습니다. LMI에는 고성능으로 Amazon S3에서 다운로드하기 위한 최적화 기능도 포함되어 있습니다. 다음 코드를 참조하기 바랍니다.

s3_code_artifact_deepspeed= sess.upload_data("model.tar.gz", bucket, s3_code_prefix_deepspeed)
print(f"S3 Code or Model tar for deepspeed uploaded to --- > {s3_code_artifact_deepspeed}")

이 시점에서 남은 작업은 사용하려는 추론 컨테이너를 정의하고 모델 객체(model object)를 생성하는 것입니다.

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0.22.1-deepspeed0.8.3-cu118"
)
model_name_acc = name_from_base(f"falcon40b-model-ds")
create_model_response = sm_client.create_model(
    ModelName=model_name_acc,
    ExecutionRoleArn=role,
    PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact_deepspeed},
)
model_arn = create_model_response["ModelArn"]

이후 엔드포인트 구성(endpoint configuration)을 만들고 엔드포인트(endpoint)를 생성합니다.

endpoint_config_name = f"{model_name}-config"
endpoint_name = f"{model_name}-endpoint"
endpoint_config_response = sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.g5.24xlarge",
            "InitialInstanceCount": 1,
            "ModelDataDownloadTimeoutInSeconds": 3600,
            "ContainerStartupHealthCheckTimeoutInSeconds": 3600,
            # "VolumeSizeInGB": 512
        },
    ],
)
endpoint_config_response

create_endpoint_response = sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)
print(f"Created Endpoint: {create_endpoint_response['EndpointArn']}")

성공적인 호스팅을 위해 염두에 두어야 할 구성 항목

대규모 모델을 호스팅할 때 중요한 고려 사항은 허깅페이스에서 모델을 다운로드하는 데 충분한 시간을 확보하는 것입니다. 테스트 결과, Falcon-40B를 인스턴스에 다운로드하는 데 약 90분이 걸렸습니다. 이를 허용하는 주요 구성 세트는 ContainerStartupHealthCheckTimeoutInSeconds 및 ModelDataDownloadTimeoutInSeconds입니다. SageMaker 엔드포인트 구성에서 각각에 대해 3600의 값이 설정되었는지 확인하세요. 또한, 모델 다운로드 시간을 약 10분으로 단축하는 S5cmd 유틸리티를 사용하는 LLMS용으로 특별히 설계된 LMI 컨테이너를 사용하여 원래의 model zoo 대신 Amazon S3에서 다운로드하면 훨씬 더 쉽게 다운로드할 수 있습니다.

모든 것이 완료되면 DescribeEndpoint를 호출하여 엔드포인트의 상태를 모니터링할 수 있습니다. 이제 엔드포인트가 추론 요청에 응답할 준비가 되었습니다! LMI가 모델 파티셔닝과 오케스트레이션을 처리하기 때문에 각 요청은 ml.g5.12xlarge 인스턴스에서 사용 가능한 4개의 GPU를 모두 사용하여 처리됩니다. 이를 통해 LLM을 호스팅하고 GPU 가속기를 수평으로 확장할 경우 성능을 향상시킬 수 있습니다. 다음 코드를 참조하기 바랍니다.

response_model = smr_client.invoke_endpoint(
    EndpointName=endpoint_name,
    Body=json.dumps({"text": "What is the purpose of life?", "text_length": 150}),
    ContentType="application/json",
)

response_model["Body"].read().decode("utf8")

이 게시물에서 사용한 코드는 깃허브에 공개되어 있습니다.

결론

SageMaker 호스팅과 LMI DLC를 사용하면 Falcon-40B와 같은 LLM을 쉽게 호스팅할 수 있습니다. 여러 GPU에서 모델을 호스팅하는 데 필요한 사항을 오케스트레이션하는 데 있어 차별화되지 않은 무거운 작업을 처리하고 필요에 맞게 구성 가능한 옵션을 제공합니다. 또한 허깅페이스 모델에 대한 기본 지원을 통해 매우 간단하게 사용할 수 있습니다.

이 게시물에서는 SageMaker 상에서 DeepSpeed를 사용하여 Falcon-40B 모델을 호스팅하는 방법을 소개했습니다. 또한, Accelerate를 사용하여 Falcon-40B를 호스팅하는 방법과 더 작은 Falcon-7B 모델을 호스팅하는 예제를 깃허브에 공개했습니다. LMI가 포함된 SageMaker에서 이 기능을 사용해 보시고 현재까지 가장 성능이 뛰어난 공개 LLM을 직접 체험해 보시기 바랍니다!