Arm64版のAWS LambdaでGoogleのAIモデル「Gemma」を動かす

本記事は以下の記事内にて作成したLambda関数を、arm64アーキテクチャに対応させるものとなります。合わせてご確認ください。

はじめに

先の記事にて一先ずはLambdaでGemmaを動かすことができました。
ただ、AWS Lambdaはアーキテクチャをarm64に設定することもできます。
前回は特に意識していなかったためデフォルトのx86_64になっていましたが、今回はLambdaのアーキテクチャをarm64に指定して動かしてみます。

arm64アーキテクチャへの移行について

それでは、arm64アーキテクチャへ移行させるため、Lambdaやdockerのアーキテクチャ設定に関して修正します。
具体的には以下の点を修正しています。

• dockerイメージをビルドする際、arm64用にビルドする
• Lambda関数作成時にアーキテクチャとして、arm64を指定する

修正したコード

前回の記事でイメージのビルドやLambda関数の作成に使用したスクリプトを修正します。
修正したものを以下に記載します。

install-arm.sh

#!/bin/bash
set -euo pipefail

echo -n "(Create New) Input AWS Lambda Function Name [ex. myLLMFunction]: "
read -r LAMBDANAME

REGION=$(aws configure get region)
ACCOUNTID=$(aws sts get-caller-identity --output text --query Account)
RAND=$(date +%Y%m%d%H%M%S%3N | shasum -a 512 | base64 | fold -w 16 | head -n 1 | tr 'A-Z' 'a-z')
ARCH="arm64" # amd64 arm64

# Lambda function Name
if [[ -z ${LAMBDANAME} ]]; then
  LAMBDANAME="myLLMFunction"
fi
LAMBDANAME="$LAMBDANAME-$RAND"
# Role Name
ROLENAME="$LAMBDANAME-role"
# ECR Repository - lambda container image stored
REPOSITORYNAME=$(echo "$LAMBDANAME-repo" | tr 'A-Z' 'a-z')

print_message() {
  echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1"
}

create_ecr_repository() {
  local result
  result=$(aws ecr create-repository --repository-name "${REPOSITORYNAME}")
  if grep -q 'repositoryArn' <<< "${result}"; then
    print_message "Success!!: Create Repository"
  else
    print_message "Error occurred: ${REPOSITORYNAME}"
    exit 1
  fi
}

build_and_push_container() {
  local repo_uri="${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com/${REPOSITORYNAME}"
  aws ecr get-login-password --region "${REGION}" | docker login --username AWS --password-stdin "${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com"
  docker buildx build --platform linux/${ARCH} -t "${repo_uri}:latest" --push .
}

create_role() {
  local result
  result=$(aws iam create-role --role-name ${ROLENAME} --assume-role-policy-document '{"Version": "2012-10-17","Statement": [{ "Effect": "Allow", "Principal": {"Service": "lambda.amazonaws.com"}, "Action": "sts:AssumeRole"}]}')
  aws iam attach-role-policy --role-name ${ROLENAME} --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
}

create_lambda_function() {
  if [ "${ARCH}" = "amd64" ]; then
      ARCH="x86_64"
  fi
  DIGEST=$(aws ecr list-images --repository-name ${REPOSITORYNAME} --out text --query 'imageIds[?imageTag==`latest`].imageDigest')
  aws lambda create-function --function-name ${LAMBDANAME} --out text \
      --package-type Image --code ImageUri=${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com/${REPOSITORYNAME}@${DIGEST} \
      --memory-size 10240 --timeout 900 \
      --role arn:aws:iam::${ACCOUNTID}:role/${ROLENAME} \
      --architecture ${ARCH}
}

write_uninstaller() {
  cat ./uninstall.txt | sed -e "s/BEFORELAMBDANAME/${LAMBDANAME}/" \
    -e "s/BEFOREROLENAME/${ROLENAME}/" \
    -e "s/BEFOREREPOSITORYNAME/${REPOSITORYNAME}/" \
    > "./uninstall-${LAMBDANAME}.sh"
}

print_message "Starting script..."

print_message "(1/5) Create AWS ECR Repository"
create_ecr_repository

print_message "(2/5) Create Role"
create_role

print_message "(3/5) Build and Push Container"
build_and_push_container

print_message "(4/5) Create Lambda"
create_lambda_function

print_message "(5/5) Writing Uninstaller"
write_uninstaller

print_message "******* Complete!! *******"
print_message "The following resources were created."
print_message "- Lambda function: ${LAMBDANAME}"
print_message "- Role: ${ROLENAME}"
print_message "- ECR Repository: ${REPOSITORYNAME}"

スクリプト内の次の関数が主な改修箇所になります。

• build_and_push_container()
• create_lambda_function()

このスクリプトを「install-arm.sh」などの任意の名称で、前の記事で紹介しているようにGemmaのモデルやLambda用のpythonコードなどと同一のディレクトリに格納してください。
その後、スクリプトを実行するとarm64版のLambda関数が作成されます。

CMD

$ bash ./install-arm.sh

動作確認

新たにarm64版のLambda関数が作成できましたので動作確認を行います。
プロンプトとレスポンスは以下のとおりです。

Can you recommend any hot springs (onsen) in Hokkaido?
(北海道でおすすめの温泉はありますか？)

Sure, here are some of the most popular onsen resorts and baths across different regions:
(もちろん、ここでは様々な地域で最も人気のある温泉リゾートと風呂をいくつか紹介しよう：)

East:
(東部：)

• Hakuba Onsen, Otaru : Known for its stunning views overlooking Mount Fuji.
  (* 白馬温泉、 小樽： 富士山を見下ろす絶景で知られる。)

(内容はともかく)動作することが確認できました。
ただ、アーキテクチャを変更したことによってパフォーマンスがどう変わったのかは気になるところです。
こちらは別途検証してみたいと思います。