A.3_sm-sup-train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "99bc5f31-1544-422c-a243-4788702b8ebc",
   "metadata": {},
   "source": [
    "# KoSimCSE training (Supervised) on SageMaker\n",
    "---\n",
    "\n",
    "## Overview \n",
    "바로 이전 모듈까지는 기존에 온프레미스에서 개발했던 환경과 동일한 환경으로 모델을 빌드하고 훈련했습니다. 하지만 아래와 같은 상황들에서도 기존 환경을 사용하는 것이 바람직할까요?\n",
    "\n",
    "- 온프레미스의 GPU가 총 1장으로 훈련 시간이 너무 오래 소요됨\n",
    "- 가용 서버 대수가 2대인데 10개의 딥러닝 모델을 동시에 훈련해야 함\n",
    "- 필요한 상황에만 GPU를 활용\n",
    "\n",
    "Amazon SageMaker는 데이터 과학자들 및 머신 러닝 엔지니어들을 위한 완전 관리형 머신 러닝 서비스로 훈련 및 추론 수행 시 인프라 설정에 대한 추가 작업이 필요하지 있기에, 단일 GPU 기반의 딥러닝 훈련을 포함한 멀티 GPU 및 멀티 인스턴스 분산 훈련을 보다 쉽고 빠르게 수행할 수 있습니다. SageMaker는 다양한 유즈케이스들에 적합한 예제들을 지속적으로 업데이트하고 있으며, 한국어 세션 및 자료들도 제공되고 있습니다.\n",
    "\n",
    "### Note\n",
    "- 이미 기본적인 Hugging Face 용법 및 자연어 처리에 익숙하신 분들은 앞 모듈을 생략하고 이 모듈부터 핸즈온을 시작하셔도 됩니다.\n",
    "- 이 노트북은 SageMaker 기본 API를 참조하므로, SageMaker Studio, SageMaker 노트북 인스턴스 또는 AWS CLI가 설정된 로컬 시스템에서 실행해야 합니다. SageMaker Studio 또는 SageMaker 노트북 인스턴스를 사용하는 경우 PyTorch 기반 커널을 선택하세요.\n",
    "- 훈련(Training) job 수행 시 최소 `ml.g4dn.xlarge` 훈련 인스턴스를 권장하며, 분산 훈련 수행 시에는 `ml.g5.12xlarge` 훈련 인스턴스를 권장합니다. 만약 인스턴스 사용에 제한이 걸려 있다면 [Request a service quota increase for SageMaker resources](https://docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.html#service-limit-increase-request-procedure)를 참조하여 인스턴스 제한을 해제해 주세요."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61873ad6-3aa2-44fd-95d6-b0b3ca3fd575",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 1. Preparation\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52bd530d-d95e-48d6-9f46-94d2eb0c903c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import boto3\n",
    "import sagemaker\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "region = boto3.Session().region_name\n",
    "bucket = None\n",
    "if bucket is None and sess is not None:\n",
    "    bucket = sess.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sess = sagemaker.Session(default_bucket=bucket)\n",
    "\n",
    "print(f\"SageMaker role arn: {role}\")\n",
    "print(f\"SageMaker bucket: {sess.default_bucket()}\")\n",
    "print(f\"SageMaker session region: {sess.boto_region_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81b057cb-9163-4f24-97e4-a3c38421b8d8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "aws configure set default.s3.max_concurrent_requests 100\n",
    "aws configure set default.s3.max_queue_size 10000\n",
    "aws configure set default.s3.multipart_threshold 1GB\n",
    "aws configure set default.s3.multipart_chunksize 64MB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d544a14e-1d50-404c-b1f8-4c5286ac2a60",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bucket_prefix = \"korsimcse\"\n",
    "model_name = \"roberta-base\"\n",
    "\n",
    "dataset_dir = \"dataset-sup-train\"\n",
    "s3_model_path = f\"s3://{bucket}/{bucket_prefix}/model/{model_name}/\"\n",
    "s3_dataset_path = f\"s3://{bucket}/{bucket_prefix}/dataset/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e591409-f3e1-4547-a24c-cc3d21258902",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!aws s3 sync {dataset_dir} {s3_dataset_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0699ad2d-594d-47dd-a1b8-f71ff34a768a",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 2. SageMaker Training\n",
    "---\n",
    "\n",
    "SageMaker에 대한 대표적인 오해가 여전히 많은 분들이 SageMaker 훈련을 위해 소스 코드를 전면적으로 수정해야 한다고 생각합니다. 하지만, 실제로는 별도의 소스 코드 수정 없이 기존 여러분이 사용했던 파이썬 스크립트에 SageMaker 훈련에 필요한 SageMaker 전용 환경 변수들만 추가하면 됩니다.\n",
    "\n",
    "SageMaker 훈련은 훈련 작업을 호출할 때, 1) 훈련 EC2 인스턴스 프로비저닝 - 2) 컨테이너 구동을 위한 도커 이미지 및 훈련 데이터 다운로드 - 3) 컨테이너 구동 - 4) 컨테이너 환경에서 훈련 수행 - 5) 컨테이너 환경에서 S3의 특정 버킷에 저장 - 6) 훈련 인스턴스 종료로 구성됩니다. 따라서, 훈련 수행 로직은 아래 예시와 같이 기존 개발 환경과 동일합니다.\n",
    "\n",
    "`/opt/conda/bin/python train_hf.py --num_epochs 5 --train_batch_size 32 ...`\n",
    "\n",
    "이 과정에서 컨테이너 환경에 필요한 환경 변수(예: 모델 경로, 훈련 데이터 경로) 들은 사전에 지정되어 있으며, 이 환경 변수들이 설정되어 있어야 훈련에 필요한 파일들의 경로를 인식할 수 있습니다. 대표적인 환경 변수들에 대한 자세한 내용은 https://github.com/aws/sagemaker-containers#important-environment-variables 을 참조하세요."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1c73292-2298-4cd6-99fa-4e27921c10c7",
   "metadata": {},
   "source": [
    "### Setup SageMaker Estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60062873-c1a8-4fcf-a38d-c8860e852b23",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "entry_point = \"sup_run.sh\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b61dee20-0628-46d3-b70d-cc915cc6fe1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "from sagemaker.utils import name_from_base\n",
    "from sagemaker.inputs import TrainingInput\n",
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "# Define Training Job Name \n",
    "job_name = f\"kosimcse-{model_name}-supervised-{time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())}\"\n",
    "print(job_name)\n",
    "\n",
    "# See https://github.com/aws/deep-learning-containers/blob/master/available_images.md\n",
    "image_uri = f\"763104351884.dkr.ecr.{region}.amazonaws.com/pytorch-training:2.0.1-gpu-py310-cu118-ubuntu20.04-sagemaker\"\n",
    "hparams = {}\n",
    "\n",
    "instance_type = \"ml.g5.12xlarge\"\n",
    "max_run = 6*60*60 # 6 hours\n",
    "use_spot_instances = False\n",
    "if use_spot_instances:\n",
    "    max_wait = 12*60*60 # 12 hours: spot instance waiting + max runtime\n",
    "else:\n",
    "    max_wait = None\n",
    "    \n",
    "# Create the Estimator\n",
    "estimator = PyTorch(\n",
    "    image_uri=image_uri,\n",
    "    entry_point=entry_point,           # train script\n",
    "    source_dir=\"src\",               # directory which includes all the files needed for training\n",
    "    instance_type=instance_type, # instances type used for the training job\n",
    "    instance_count=1,               # the number of instances used for training\n",
    "    base_job_name=job_name,         # the name of the training job\n",
    "    role=role,                      # Iam role used in training job to access AWS ressources, e.g. S3\n",
    "    volume_size=300,                # the size of the EBS volume in GB\n",
    "    py_version=\"py310\",             # the python version used in the training job\n",
    "    hyperparameters=hparams,\n",
    "    debugger_hook_config=False,\n",
    "    disable_profile=True,\n",
    "    use_spot_instances=use_spot_instances,\n",
    "    max_run=max_run,\n",
    "    max_wait=max_wait if use_spot_instances else None,\n",
    "    metric_definitions=[\n",
    "       {'Name': 'train:step', 'Regex': 'step:(.*?);'},\n",
    "       {'Name': 'train:loss', 'Regex': 'loss:(.*?);'},\n",
    "       {'Name': 'train:avg-sts-score', 'Regex': 'Avg. STS:(.*?);'}\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb5413d-c9bb-4506-b95a-0ea5894cee44",
   "metadata": {},
   "source": [
    "### Start Training job\n",
    "S3에서 훈련 인스턴스로 복사될 데이터를 지정한 후 SageMaker 훈련 job을 시작합니다. 모델 크기, 데이터 세트 크기에 따라서 몇십 분에서 몇 시간까지 소요될 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bea0a770-ab7d-4813-80ca-266ef96a1a40",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fast_file = lambda x: TrainingInput(x, input_mode=\"FastFile\")\n",
    "estimator.fit(\n",
    "    {\n",
    "        \"training\": fast_file(s3_dataset_path),\n",
    "    },\n",
    "    wait=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c03c75-9988-4349-8bab-89b1a81e70ad",
   "metadata": {
    "tags": []
   },
   "source": [
    "### View Training Job\n",
    "SageMaker 콘솔 창에서 훈련 내역을 직접 확인할 수도 있지만, 아래 코드 셀에서 생성되는 링크를 클릭하면 더 편리하게 훈련 내역을 확인할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbbe7d5a-3839-475a-bfe5-5356684477a4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "\n",
    "def make_console_link(region, train_job_name, train_task='[Training]'):\n",
    "    train_job_link = f'<b> {train_task} Review <a target=\"blank\" href=\"https://console.aws.amazon.com/sagemaker/home?region={region}#/jobs/{train_job_name}\">Training Job</a></b>'   \n",
    "    cloudwatch_link = f'<b> {train_task} Review <a target=\"blank\" href=\"https://console.aws.amazon.com/cloudwatch/home?region={region}#logStream:group=/aws/sagemaker/TrainingJobs;prefix={train_job_name};streamFilter=typeLogStreamPrefix\">CloudWatch Logs</a></b>'\n",
    "    return train_job_link, cloudwatch_link  \n",
    "        \n",
    "train_job_name = estimator.latest_training_job.job_name\n",
    "train_job_link, cloudwatch_link = make_console_link(region, train_job_name, '[Fine-tuning]')\n",
    "\n",
    "display(HTML(train_job_link))\n",
    "display(HTML(cloudwatch_link))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a551d7d-0887-41c1-b6be-176dfe526778",
   "metadata": {},
   "source": [
    "### Wait for the training jobs to complete\n",
    "훈련이 완료될 때까지 기다립니다. estimator.fit(...)에서 wait=False로 설정한 경우, 아래 코드 셀의 주석을 해제 후 실행하여 동기 방식으로 변경할 수도 있습니다. 훈련 완료까지는 수십 분이 소요됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "941ea394-ac7a-4dc5-962d-0a774b97dfad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.logs_for_job(job_name=train_job_name, wait=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fe30f4e-9365-4176-832f-b5e909f86149",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 3. Inference\n",
    "---\n",
    "#### Copy S3 model artifact to local directory\n",
    "S3에 저장된 모델 아티팩트를 로컬 경로로 복사하여 압축을 해제합니다. 필요 시 로컬 환경에서 모델을 로드하여 추론을 수행할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92fd2bf9-bf25-4f87-8df9-f0a1a8318f3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json, os\n",
    "\n",
    "local_model_dir = 'model_from_sagemaker'\n",
    "\n",
    "if not os.path.exists(local_model_dir):\n",
    "    os.makedirs(local_model_dir)\n",
    "\n",
    "!aws s3 cp {estimator.model_data} {local_model_dir}/model.tar.gz\n",
    "!tar -xzf {local_model_dir}/model.tar.gz -C {local_model_dir}\n",
    "!rm -rf {local_model_dir}/model.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2cad1c1-fc05-43cd-b2d9-0dbb17a57331",
   "metadata": {},
   "source": [
    "### Load model artifact "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a4f16c-76e4-49c4-b8d9-d5c28a6b7bc8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import torch\n",
    "import transformers\n",
    "import numpy as np\n",
    "from collections import OrderedDict\n",
    "from transformers import BertForSequenceClassification, AutoTokenizer\n",
    "from src.simcse import SimCSEModel\n",
    "from src.infer import show_embedding_score\n",
    "\n",
    "with open(f'{local_model_dir}/config.json') as f:\n",
    "    json_object = json.load(f)\n",
    "    \n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')        \n",
    "base_model = json_object[\"base_model\"]\n",
    "tokenizer = AutoTokenizer.from_pretrained(base_model)\n",
    "model = SimCSEModel(base_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39f7fe2b-bc22-4ada-9049-7529e9a2c5d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_filename = glob.glob(f'{local_model_dir}/*.pt')[0]\n",
    "state_dict = torch.load(model_filename)['model']\n",
    "\n",
    "new_state_dict = {}\n",
    "for key in state_dict:\n",
    "    new_key = key.replace('module.','')\n",
    "    new_state_dict[new_key] = state_dict[key]\n",
    "\n",
    "model.load_state_dict(new_state_dict)\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb19bed0-5981-4814-93b7-50a75ae0a145",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6349b0a2-ad86-4100-917c-9dacd1a2587c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sentences = ['이번 주 일요일에 분당 이마트 점은 문을 여나요?',\n",
    "             '일요일에 분당 이마트는 문 열어요?',\n",
    "             '분당 이마트 점은 토요일에 몇 시까지 하나요']\n",
    "show_embedding_score(tokenizer, model.cpu(), sentences)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ff93543-0915-4c99-85ea-34a57d8181c4",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "963be13d-b7f0-4b66-8e14-7990f9c8b559",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from src.sts import STSEvaluation\n",
    "from more_itertools import chunked\n",
    "batch_size = 128\n",
    "model = model.to(device)\n",
    "\n",
    "@torch.inference_mode()\n",
    "def encode(texts: List[str]) -> torch.Tensor:\n",
    "    embs = []\n",
    "    for text in chunked(texts, batch_size):\n",
    "        batch: BatchEncoding = tokenizer(\n",
    "            text,\n",
    "            padding=True,\n",
    "            truncation=True,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "        emb = model(**batch.to(device))#, use_mlp=False)\n",
    "        embs.append(emb.cpu())\n",
    "    return torch.cat(embs, dim=0)\n",
    "\n",
    "evaluation = STSEvaluation()\n",
    "sts_metrics = evaluation(encode=encode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21cef5b3-6500-417c-962f-3facccdea82c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display(sts_metrics)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}