modelparallel/sm-deepspeed-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d1941c3c-ea62-41a3-b9d5-8e78330dddc7",
   "metadata": {},
   "source": [
    "# DeepSpeed on Amazon SageMaker\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "#### Note\n",
    "\n",
    "- 이미 기본적인 Hugging Face 용법 및 자연어 처리에 익숙하신 분들은 앞 모듈을 생략하고 이 모듈부터 핸즈온을 시작하셔도 됩니다.\n",
    "- 이 노트북은 SageMaker 기본 API를 참조하므로, SageMaker Studio, SageMaker 노트북 인스턴스 또는 AWS CLI가 설정된 로컬 시스템에서 실행해야 합니다. SageMaker Studio 또는 SageMaker 노트북 인스턴스를 사용하는 경우 PyTorch 기반 커널을 선택하세요.\n",
    "- 훈련(Training) job 수행 시 최소 ml.p3.2xlarge 이상의 훈련 인스턴스가 필요하며, 분산 훈련 핸즈온은 `ml.p3.16xlarge` 인스턴스를 권장합니다. 만약 인스턴스 사용에 제한이 걸려 있다면 [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": "63ce1190-1d1d-45bf-aae6-98f9ab009032",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 1. Preparation\n",
    "---\n",
    "\n",
    "SageMaker 훈련을 위해 전처리된 데이터셋을 S3에 업로드합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0e4feff2-43eb-4aa2-93d7-a727925cbf1a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com\n",
      "Requirement already satisfied: transformers==4.30.2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from -r scripts/requirements.txt (line 1)) (4.30.2)\n",
      "Requirement already satisfied: datasets==2.13.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from -r scripts/requirements.txt (line 2)) (2.13.1)\n",
      "Requirement already satisfied: evaluate==0.4.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from -r scripts/requirements.txt (line 3)) (0.4.0)\n",
      "Collecting deepspeed (from -r scripts/requirements.txt (line 4))\n",
      "  Downloading deepspeed-0.10.0.tar.gz (836 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m836.6/836.6 kB\u001b[0m \u001b[31m29.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: filelock in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (3.12.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.14.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (0.16.4)\n",
      "Requirement already satisfied: numpy>=1.17 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (1.24.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (21.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (5.4.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (2023.6.3)\n",
      "Requirement already satisfied: requests in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (2.29.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (0.13.3)\n",
      "Requirement already satisfied: safetensors>=0.3.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (0.3.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from transformers==4.30.2->-r scripts/requirements.txt (line 1)) (4.65.0)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (12.0.0)\n",
      "Requirement already satisfied: dill<0.3.7,>=0.3.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (0.3.6)\n",
      "Requirement already satisfied: pandas in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (1.5.3)\n",
      "Requirement already satisfied: xxhash in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (3.2.0)\n",
      "Requirement already satisfied: multiprocess in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (0.70.14)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (2023.5.0)\n",
      "Requirement already satisfied: aiohttp in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from datasets==2.13.1->-r scripts/requirements.txt (line 2)) (3.8.4)\n",
      "Requirement already satisfied: responses<0.19 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from evaluate==0.4.0->-r scripts/requirements.txt (line 3)) (0.18.0)\n",
      "Collecting hjson (from deepspeed->-r scripts/requirements.txt (line 4))\n",
      "  Downloading hjson-3.1.0-py3-none-any.whl (54 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.0/54.0 kB\u001b[0m \u001b[31m16.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting ninja (from deepspeed->-r scripts/requirements.txt (line 4))\n",
      "  Downloading ninja-1.11.1-py2.py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (145 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m146.0/146.0 kB\u001b[0m \u001b[31m24.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: psutil in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from deepspeed->-r scripts/requirements.txt (line 4)) (5.9.5)\n",
      "Collecting py-cpuinfo (from deepspeed->-r scripts/requirements.txt (line 4))\n",
      "  Downloading py_cpuinfo-9.0.0-py3-none-any.whl (22 kB)\n",
      "Requirement already satisfied: pydantic<2.0.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from deepspeed->-r scripts/requirements.txt (line 4)) (1.10.7)\n",
      "Requirement already satisfied: torch in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from deepspeed->-r scripts/requirements.txt (line 4)) (2.0.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (23.1.0)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (3.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (6.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (1.9.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (1.3.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from aiohttp->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (1.3.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.14.1->transformers==4.30.2->-r scripts/requirements.txt (line 1)) (4.5.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from packaging>=20.0->transformers==4.30.2->-r scripts/requirements.txt (line 1)) (3.0.9)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->transformers==4.30.2->-r scripts/requirements.txt (line 1)) (3.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->transformers==4.30.2->-r scripts/requirements.txt (line 1)) (1.26.14)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->transformers==4.30.2->-r scripts/requirements.txt (line 1)) (2023.5.7)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from pandas->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from pandas->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (2023.3)\n",
      "Requirement already satisfied: sympy in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torch->deepspeed->-r scripts/requirements.txt (line 4)) (1.11.1)\n",
      "Requirement already satisfied: networkx in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torch->deepspeed->-r scripts/requirements.txt (line 4)) (3.1)\n",
      "Requirement already satisfied: jinja2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torch->deepspeed->-r scripts/requirements.txt (line 4)) (3.1.2)\n",
      "Requirement already satisfied: six>=1.5 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from python-dateutil>=2.8.1->pandas->datasets==2.13.1->-r scripts/requirements.txt (line 2)) (1.16.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from jinja2->torch->deepspeed->-r scripts/requirements.txt (line 4)) (2.1.2)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sympy->torch->deepspeed->-r scripts/requirements.txt (line 4)) (1.3.0)\n",
      "Building wheels for collected packages: deepspeed\n",
      "  Building wheel for deepspeed (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for deepspeed: filename=deepspeed-0.10.0-py3-none-any.whl size=877480 sha256=e46a52fadb01a5c093d6692eaa89ee4b07b35a4a7d54050c860dae2fa06a81a0\n",
      "  Stored in directory: /home/ec2-user/.cache/pip/wheels/20/7b/3f/2807682bad2fba40ed888e6309597a5fda545ab30964c835aa\n",
      "Successfully built deepspeed\n",
      "Installing collected packages: py-cpuinfo, ninja, hjson, deepspeed\n",
      "Successfully installed deepspeed-0.10.0 hjson-3.1.0 ninja-1.11.1 py-cpuinfo-9.0.0\n"
     ]
    }
   ],
   "source": [
    "!pip install -r scripts/requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "62c06ed4-05af-43e6-ad7e-d24307f297e5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{2147392357.py:28} INFO - sagemaker role arn: arn:aws:iam::143656149352:role/service-role/AmazonSageMaker-ExecutionRole-20220317T150353\n",
      "[{2147392357.py:29} INFO - sagemaker bucket: sagemaker-us-east-1-143656149352\n",
      "[{2147392357.py:30} INFO - sagemaker session region: us-east-1\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import logging\n",
    "import boto3\n",
    "import botocore\n",
    "import sagemaker\n",
    "import time\n",
    "sess = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it not exists\n",
    "sagemaker_session_bucket=None\n",
    "if sagemaker_session_bucket is None and sess is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sess.default_bucket()\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "region = boto3.Session().region_name\n",
    "sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO, \n",
    "    format='[{%(filename)s:%(lineno)d} %(levelname)s - %(message)s',\n",
    "    handlers=[\n",
    "        logging.StreamHandler(sys.stdout)\n",
    "    ]\n",
    ")\n",
    "\n",
    "logging.info(f\"sagemaker role arn: {role}\")\n",
    "logging.info(f\"sagemaker bucket: {sess.default_bucket()}\")\n",
    "logging.info(f\"sagemaker session region: {sess.boto_region_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5da250f4-93a8-4ff1-9516-78801fd0b40a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# Define the model repo\n",
    "model_id = 'bert-base-multilingual-cased'\n",
    "\n",
    "# dataset used\n",
    "dataset_name = 'nsmc'\n",
    "\n",
    "# s3 key prefix for the data\n",
    "s3_prefix = 'datasets/nsmc'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f4ba528e-a30e-4f17-a93e-f74f188a78ea",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{builder.py:835} WARNING - Found cached dataset nsmc (/home/ec2-user/.cache/huggingface/datasets/nsmc/default/1.1.0/bfd4729bf1a67114e5267e6916b9e4807010aeb238e4a3c2b95fbfa3a014b5f3)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b7f2ac6d4cd4a2f972e71455ed9a723",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{arrow_dataset.py:4230} WARNING - Loading cached shuffled indices for dataset at /home/ec2-user/.cache/huggingface/datasets/nsmc/default/1.1.0/bfd4729bf1a67114e5267e6916b9e4807010aeb238e4a3c2b95fbfa3a014b5f3/cache-dcd892fe1da31a64.arrow\n",
      "[{arrow_dataset.py:4230} WARNING - Loading cached shuffled indices for dataset at /home/ec2-user/.cache/huggingface/datasets/nsmc/default/1.1.0/bfd4729bf1a67114e5267e6916b9e4807010aeb238e4a3c2b95fbfa3a014b5f3/cache-3ee302597e2fad0e.arrow\n",
      "[{3407595512.py:8} INFO -  loaded train_dataset length is: 2000\n",
      "[{3407595512.py:9} INFO -  loaded eval_dataset length is: 2000\n",
      "[{3407595512.py:10} INFO - {'id': '10020916', 'document': 'For Carl.칼 세이건으로 시작해서 칼 세이건으로 끝난다.', 'label': 1}\n"
     ]
    }
   ],
   "source": [
    "# load dataset\n",
    "train_dataset, eval_dataset = load_dataset(dataset_name, split=['train', 'test'])\n",
    "\n",
    "num_samples_for_debug = 2000\n",
    "train_dataset = train_dataset.shuffle(seed=42).select(range(num_samples_for_debug))\n",
    "eval_dataset = eval_dataset.shuffle(seed=42).select(range(num_samples_for_debug))\n",
    "\n",
    "logging.info(f\" loaded train_dataset length is: {len(train_dataset)}\")\n",
    "logging.info(f\" loaded eval_dataset length is: {len(eval_dataset)}\")\n",
    "logging.info(train_dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d9464d37-021e-4980-8b8a-313d90022fdd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "\n",
    "def tokenize(batch):\n",
    "    return tokenizer(batch['document'], padding='max_length', max_length=128, truncation=True)\n",
    "\n",
    "# tokenize dataset\n",
    "train_dataset = train_dataset.map(tokenize, batched=True, remove_columns=['id', 'document'])\n",
    "eval_dataset = eval_dataset.map(tokenize, batched=True, remove_columns=['id', 'document'])\n",
    "\n",
    "# set format for pytorch\n",
    "train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])\n",
    "eval_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])\n",
    "\n",
    "train_dataset = train_dataset.rename_column(\"label\", \"labels\")\n",
    "eval_dataset = eval_dataset.rename_column(\"label\", \"labels\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f01498c2-ddc8-425d-932b-42348f401941",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_dir = 'train'\n",
    "eval_dir = 'eval'\n",
    "!rm -rf {train_dir} {eval_dir}\n",
    "\n",
    "os.makedirs(train_dir, exist_ok=True)\n",
    "os.makedirs(eval_dir, exist_ok=True) \n",
    "\n",
    "if not os.listdir(train_dir):\n",
    "    train_dataset.save_to_disk(train_dir)\n",
    "if not os.listdir(eval_dir):\n",
    "    eval_dataset.save_to_disk(eval_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "74e4cd58-30bb-4dc6-9bc4-3e4d609232d6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{credentials.py:1051} INFO - Found credentials from IAM Role: BaseNotebookInstanceEc2InstanceRole\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/2000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# save train_dataset to s3\n",
    "train_input_path = f's3://{sess.default_bucket()}/{s3_prefix}/{train_dir}'\n",
    "train_dataset.save_to_disk(train_input_path)\n",
    "\n",
    "# save eval_dataset to s3\n",
    "eval_input_path = f's3://{sess.default_bucket()}/{s3_prefix}/{eval_dir}'\n",
    "eval_dataset.save_to_disk(eval_input_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92e377d3-51a9-4503-8d0a-88a9b8373cdc",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 2. Debugging (Development Stage)\n",
    "---\n",
    "\n",
    "SageMaker에서 훈련을 수행하기 전에 먼저 로컬 개발 환경에서 모델 훈련 코드를 개발하고 디버깅해야 합니다. SageMaker 노트북 인스턴스에서 작업하는 경우 GPU가 탑재된 인스턴스(p-family, g-family)를 사용하셔야 합니다.\n",
    "\n",
    "### DeepSpeed 디버깅"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "711b43bd-6b96-4270-af1f-372261c16982",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running command: \n",
      "cd scripts && deepspeed train_deepspeed.py \n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "[2023-07-18 00:10:55,883] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
      "CUDA SETUP: CUDA runtime path found: /home/ec2-user/anaconda3/envs/pytorch_p310/lib/libcudart.so.11.0\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 118\n",
      "CUDA SETUP: Loading binary /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n",
      "[2023-07-18 00:10:57,065] [WARNING] [runner.py:196:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.\n",
      "[2023-07-18 00:10:57,065] [INFO] [runner.py:555:main] cmd = /home/ec2-user/anaconda3/envs/pytorch_p310/bin/python3.10 -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMCwgMSwgMiwgM119 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None train_deepspeed.py\n",
      "[2023-07-18 00:10:58,781] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
      "CUDA SETUP: CUDA runtime path found: /home/ec2-user/anaconda3/envs/pytorch_p310/lib/libcudart.so.11.0\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 118\n",
      "CUDA SETUP: Loading binary /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n",
      "[2023-07-18 00:10:59,922] [INFO] [launch.py:145:main] WORLD INFO DICT: {'localhost': [0, 1, 2, 3]}\n",
      "[2023-07-18 00:10:59,922] [INFO] [launch.py:151:main] nnodes=1, num_local_procs=4, node_rank=0\n",
      "[2023-07-18 00:10:59,922] [INFO] [launch.py:162:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0, 1, 2, 3]})\n",
      "[2023-07-18 00:10:59,922] [INFO] [launch.py:163:main] dist_world_size=4\n",
      "[2023-07-18 00:10:59,922] [INFO] [launch.py:165:main] Setting CUDA_VISIBLE_DEVICES=0,1,2,3\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
      "CUDA SETUP: CUDA runtime path found: /home/ec2-user/anaconda3/envs/pytorch_p310/lib/libcudart.so.11.0\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 118\n",
      "CUDA SETUP: Loading binary /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
      "CUDA SETUP: CUDA runtime path found: /home/ec2-user/anaconda3/envs/pytorch_p310/lib/libcudart.so.11.0\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 118\n",
      "CUDA SETUP: Loading binary /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "CUDA SETUP: CUDA runtime path found: /home/ec2-user/anaconda3/envs/pytorch_p310/lib/libcudart.so.11.0\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 118\n",
      "CUDA SETUP: Loading binary /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n",
      "bin /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so\n",
      "CUDA SETUP: CUDA runtime path found: /home/ec2-user/anaconda3/envs/pytorch_p310/lib/libcudart.so.11.0\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 118\n",
      "CUDA SETUP: Loading binary /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...\n",
      "[2023-07-18 00:11:10,866] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2023-07-18 00:11:10,983] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2023-07-18 00:11:10,990] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2023-07-18 00:11:11,019] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "07/18/2023 00:11:14 - INFO - root -   Initialized the distributed environment. world_size=4, rank=0, local_rank=0\n",
      "07/18/2023 00:11:14 - INFO - root -    loaded train_dataset length is: 2000\n",
      "07/18/2023 00:11:14 - INFO - root -    loaded test_dataset length is: 2000\n",
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForSequenceClassification: ['cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-multilingual-cased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "[2023-07-18 00:11:17,979] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.10.0, git-hash=unknown, git-branch=unknown\n",
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForSequenceClassification: ['cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.bias', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-multilingual-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForSequenceClassification: ['cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-multilingual-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertForSequenceClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-multilingual-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "07/18/2023 00:11:18 - INFO - torch.distributed.distributed_c10d -   Added key: store_based_barrier_key:2 to store for rank: 0\n",
      "07/18/2023 00:11:18 - INFO - torch.distributed.distributed_c10d -   Rank 0: Completed store-based barrier for key:store_based_barrier_key:2 with 4 nodes.\n",
      "[2023-07-18 00:11:19,793] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False\n",
      "Using /home/ec2-user/.cache/torch_extensions/py310_cu118 as PyTorch extensions root...\n",
      "Creating extension directory /home/ec2-user/.cache/torch_extensions/py310_cu118/cpu_adam...\n",
      "Using /home/ec2-user/.cache/torch_extensions/py310_cu118 as PyTorch extensions root...\n",
      "Using /home/ec2-user/.cache/torch_extensions/py310_cu118 as PyTorch extensions root...\n",
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file /home/ec2-user/.cache/torch_extensions/py310_cu118/cpu_adam/build.ninja...\n",
      "Building extension module cpu_adam...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
      "Using /home/ec2-user/.cache/torch_extensions/py310_cu118 as PyTorch extensions root...\n",
      "[1/3] /usr/local/cuda-11.8/bin/nvcc  -DTORCH_EXTENSION_NAME=cpu_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/usr/local/cuda-11.8/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include/TH -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include/THC -isystem /usr/local/cuda-11.8/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++17 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86 -DBF16_AVAILABLE -c /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/deepspeed/ops/csrc/common/custom_cuda_kernel.cu -o custom_cuda_kernel.cuda.o \n",
      "[2/3] c++ -MMD -MF cpu_adam.o.d -DTORCH_EXTENSION_NAME=cpu_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/usr/local/cuda-11.8/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include/TH -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/include/THC -isystem /usr/local/cuda-11.8/include -isystem /home/ec2-user/anaconda3/envs/pytorch_p310/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -O3 -std=c++17 -g -Wno-reorder -L/usr/local/cuda-11.8/lib64 -lcudart -lcublas -g -march=native -fopenmp -D__AVX256__ -D__ENABLE_CUDA__ -DBF16_AVAILABLE -c /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp -o cpu_adam.o \n",
      "[3/3] c++ cpu_adam.o custom_cuda_kernel.cuda.o -shared -lcurand -L/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda-11.8/lib64 -lcudart -o cpu_adam.so\n",
      "Loading extension module cpu_adam...\n",
      "Time to load cpu_adam op: 30.7181875705719 seconds\n",
      "Loading extension module cpu_adam...\n",
      "Time to load cpu_adam op: 30.76508355140686 seconds\n",
      "Loading extension module cpu_adam...\n",
      "Time to load cpu_adam op: 30.734156847000122 seconds\n",
      "Loading extension module cpu_adam...\n",
      "Time to load cpu_adam op: 30.691272735595703 seconds\n",
      "Adam Optimizer #0 is created with AVX2 arithmetic capability.\n",
      "Config: alpha=0.000030, betas=(0.900000, 0.999000), weight_decay=0.000000, adam_w=1\n",
      "Adam Optimizer #0 is created with AVX2 arithmetic capability.\n",
      "Config: alpha=0.000030, betas=(0.900000, 0.999000), weight_decay=0.000000, adam_w=1\n",
      "Adam Optimizer #0 is created with AVX2 arithmetic capability.\n",
      "Config: alpha=0.000030, betas=(0.900000, 0.999000), weight_decay=0.000000, adam_w=1\n",
      "[2023-07-18 00:11:52,814] [INFO] [logging.py:96:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adamw as basic optimizer\n",
      "Adam Optimizer #0 is created with AVX2 arithmetic capability.\n",
      "Config: alpha=0.000030, betas=(0.900000, 0.999000), weight_decay=0.000000, adam_w=1\n",
      "[2023-07-18 00:11:52,842] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = DeepSpeedCPUAdam\n",
      "[2023-07-18 00:11:52,842] [INFO] [utils.py:54:is_zero_supported_optimizer] Checking ZeRO support for optimizer=DeepSpeedCPUAdam type=<class 'deepspeed.ops.adam.cpu_adam.DeepSpeedCPUAdam'>\n",
      "[2023-07-18 00:11:52,842] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.bfloat16 ZeRO stage 2 optimizer\n",
      "[2023-07-18 00:11:52,842] [INFO] [stage_1_and_2.py:133:__init__] Reduce bucket size 500,000,000\n",
      "[2023-07-18 00:11:52,842] [INFO] [stage_1_and_2.py:134:__init__] Allgather bucket size 500,000,000\n",
      "[2023-07-18 00:11:52,842] [INFO] [stage_1_and_2.py:135:__init__] CPU Offload: True\n",
      "[2023-07-18 00:11:52,842] [INFO] [stage_1_and_2.py:136:__init__] Round robin gradient partitioning: False\n",
      "Rank: 2 partition count [4] and sizes[(44463746, False)] \n",
      "Rank: 0 partition count [4] and sizes[(44463746, False)] \n",
      "Rank: 1 partition count [4] and sizes[(44463746, False)] \n",
      "Rank: 3 partition count [4] and sizes[(44463746, False)] \n",
      "[2023-07-18 00:11:55,036] [INFO] [utils.py:785:see_memory_usage] Before initializing optimizer states\n",
      "[2023-07-18 00:11:55,037] [INFO] [utils.py:786:see_memory_usage] MA 0.5 GB         Max_MA 0.5 GB         CA 0.51 GB         Max_CA 1 GB \n",
      "[2023-07-18 00:11:55,038] [INFO] [utils.py:793:see_memory_usage] CPU Virtual Memory:  used = 18.14 GB, percent = 4.9%\n",
      "[2023-07-18 00:11:56,014] [INFO] [utils.py:785:see_memory_usage] After initializing optimizer states\n",
      "[2023-07-18 00:11:56,015] [INFO] [utils.py:786:see_memory_usage] MA 0.5 GB         Max_MA 0.5 GB         CA 0.51 GB         Max_CA 1 GB \n",
      "[2023-07-18 00:11:56,015] [INFO] [utils.py:793:see_memory_usage] CPU Virtual Memory:  used = 19.4 GB, percent = 5.2%\n",
      "[2023-07-18 00:11:56,015] [INFO] [stage_1_and_2.py:493:__init__] optimizer state initialized\n",
      "[2023-07-18 00:11:56,245] [INFO] [utils.py:785:see_memory_usage] After initializing ZeRO optimizer\n",
      "[2023-07-18 00:11:56,246] [INFO] [utils.py:786:see_memory_usage] MA 0.5 GB         Max_MA 0.5 GB         CA 0.51 GB         Max_CA 1 GB \n",
      "[2023-07-18 00:11:56,246] [INFO] [utils.py:793:see_memory_usage] CPU Virtual Memory:  used = 19.4 GB, percent = 5.2%\n",
      "[2023-07-18 00:11:56,247] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Final Optimizer = adamw\n",
      "[2023-07-18 00:11:56,247] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed using client LR scheduler\n",
      "[2023-07-18 00:11:56,247] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed LR Scheduler = None\n",
      "[2023-07-18 00:11:56,247] [INFO] [logging.py:96:log_dist] [Rank 0] step=0, skipped=0, lr=[3e-05], mom=[[0.9, 0.999]]\n",
      "[2023-07-18 00:11:56,248] [INFO] [config.py:960:print] DeepSpeedEngine configuration:\n",
      "[2023-07-18 00:11:56,248] [INFO] [config.py:964:print]   activation_checkpointing_config  {\n",
      "    \"partition_activations\": false, \n",
      "    \"contiguous_memory_optimization\": false, \n",
      "    \"cpu_checkpointing\": false, \n",
      "    \"number_checkpoints\": null, \n",
      "    \"synchronize_checkpoint_boundary\": false, \n",
      "    \"profile\": false\n",
      "}\n",
      "[2023-07-18 00:11:56,248] [INFO] [config.py:964:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}\n",
      "[2023-07-18 00:11:56,248] [INFO] [config.py:964:print]   amp_enabled .................. False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   amp_params ................... False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   autotuning_config ............ {\n",
      "    \"enabled\": false, \n",
      "    \"start_step\": null, \n",
      "    \"end_step\": null, \n",
      "    \"metric_path\": null, \n",
      "    \"arg_mappings\": null, \n",
      "    \"metric\": \"throughput\", \n",
      "    \"model_info\": null, \n",
      "    \"results_dir\": \"autotuning_results\", \n",
      "    \"exps_dir\": \"autotuning_exps\", \n",
      "    \"overwrite\": true, \n",
      "    \"fast\": true, \n",
      "    \"start_profile_step\": 3, \n",
      "    \"end_profile_step\": 5, \n",
      "    \"tuner_type\": \"gridsearch\", \n",
      "    \"tuner_early_stopping\": 5, \n",
      "    \"tuner_num_trials\": 50, \n",
      "    \"model_info_path\": null, \n",
      "    \"mp_size\": 1, \n",
      "    \"max_train_batch_size\": null, \n",
      "    \"min_train_batch_size\": 1, \n",
      "    \"max_train_micro_batch_size_per_gpu\": 1.024000e+03, \n",
      "    \"min_train_micro_batch_size_per_gpu\": 1, \n",
      "    \"num_tuning_micro_batch_sizes\": 3\n",
      "}\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   bfloat16_enabled ............. True\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   checkpoint_parallel_write_pipeline  False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   checkpoint_tag_validation_enabled  True\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   checkpoint_tag_validation_fail  False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7f737451cc40>\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   communication_data_type ...... None\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   curriculum_enabled_legacy .... False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   curriculum_params_legacy ..... False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   data_efficiency_enabled ...... False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   dataloader_drop_last ......... False\n",
      "[2023-07-18 00:11:56,249] [INFO] [config.py:964:print]   disable_allgather ............ False\n",
      "[2023-07-18 00:11:56,263] [INFO] [config.py:964:print]   dump_state ................... False\n",
      "[2023-07-18 00:11:56,263] [INFO] [config.py:964:print]   dynamic_loss_scale_args ...... None\n",
      "[2023-07-18 00:11:56,263] [INFO] [config.py:964:print]   eigenvalue_enabled ........... False\n",
      "[2023-07-18 00:11:56,263] [INFO] [config.py:964:print]   eigenvalue_gas_boundary_resolution  1\n",
      "[2023-07-18 00:11:56,263] [INFO] [config.py:964:print]   eigenvalue_layer_name ........ bert.encoder.layer\n",
      "[2023-07-18 00:11:56,263] [INFO] [config.py:964:print]   eigenvalue_layer_num ......... 0\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   eigenvalue_max_iter .......... 100\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   eigenvalue_stability ......... 1e-06\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   eigenvalue_tol ............... 0.01\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   eigenvalue_verbose ........... False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   elasticity_enabled ........... False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   flops_profiler_config ........ {\n",
      "    \"enabled\": false, \n",
      "    \"recompute_fwd_factor\": 0.0, \n",
      "    \"profile_step\": 1, \n",
      "    \"module_depth\": -1, \n",
      "    \"top_modules\": 1, \n",
      "    \"detailed\": true, \n",
      "    \"output_file\": null\n",
      "}\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   fp16_auto_cast ............... None\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   fp16_enabled ................. False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   fp16_master_weights_and_gradients  False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   global_rank .................. 0\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   grad_accum_dtype ............. None\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   gradient_accumulation_steps .. 1\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   gradient_clipping ............ 0.0\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   gradient_predivide_factor .... 1.0\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   initial_dynamic_scale ........ 1\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   load_universal_checkpoint .... False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   loss_scale ................... 1.0\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   memory_breakdown ............. False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   mics_hierarchial_params_gather  False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   mics_shard_size .............. -1\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') enabled=False\n",
      "[2023-07-18 00:11:56,264] [INFO] [config.py:964:print]   nebula_config ................ {\n",
      "    \"enabled\": false, \n",
      "    \"persistent_storage_path\": null, \n",
      "    \"persistent_time_interval\": 100, \n",
      "    \"num_of_version_in_retention\": 2, \n",
      "    \"enable_nebula_load\": true, \n",
      "    \"load_path\": null\n",
      "}\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   optimizer_legacy_fusion ...... False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   optimizer_name ............... adamw\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   optimizer_params ............. {'lr': 3e-05, 'betas': [0.9, 0.999], 'eps': 1e-08}\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   pld_enabled .................. False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   pld_params ................... False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   prescale_gradients ........... False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   scheduler_name ............... None\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   scheduler_params ............. None\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   sparse_attention ............. None\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   sparse_gradients_enabled ..... False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   steps_per_print .............. 10\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   train_batch_size ............. 128\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   train_micro_batch_size_per_gpu  32\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   use_node_local_storage ....... False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   wall_clock_breakdown ......... False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   world_size ................... 4\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   zero_allow_untested_optimizer  False\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   zero_config .................. stage=2 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=500,000,000 allgather_partitions=True allgather_bucket_size=500,000,000 overlap_comm=True load_from_fp32_weights=True elastic_checkpoint=False offload_param=None offload_optimizer=DeepSpeedZeroOffloadOptimizerConfig(device='cpu', nvme_path=None, buffer_count=4, pin_memory=True, pipeline=False, pipeline_read=False, pipeline_write=False, fast_init=True) sub_group_size=1,000,000,000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50,000,000 param_persistence_threshold=100,000 model_persistence_threshold=sys.maxsize max_live_parameters=1,000,000,000 max_reuse_distance=1,000,000,000 gather_16bit_weights_on_model_save=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_gradients=False mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   zero_enabled ................. True\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   zero_force_ds_cpu_optimizer .. True\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:964:print]   zero_optimization_stage ...... 2\n",
      "[2023-07-18 00:11:56,265] [INFO] [config.py:950:print_user_config]   json = {\n",
      "    \"train_micro_batch_size_per_gpu\": 32, \n",
      "    \"gradient_accumulation\": 1, \n",
      "    \"steps_per_print\": 10, \n",
      "    \"bf16\": {\n",
      "        \"enabled\": true\n",
      "    }, \n",
      "    \"zero_optimization\": {\n",
      "        \"stage\": 2, \n",
      "        \"contiguous_gradients\": true, \n",
      "        \"overlap_comm\": true, \n",
      "        \"offload_optimizer\": {\n",
      "            \"device\": \"cpu\", \n",
      "            \"pin_memory\": true, \n",
      "            \"fast_init\": true\n",
      "        }\n",
      "    }, \n",
      "    \"optimizer\": {\n",
      "        \"type\": \"AdamW\", \n",
      "        \"params\": {\n",
      "            \"lr\": 3e-05, \n",
      "            \"betas\": [0.9, 0.999], \n",
      "            \"eps\": 1e-08\n",
      "        }\n",
      "    }\n",
      "}\n",
      "07/18/2023 00:11:56 - INFO - root -   num_training_steps: 32\n",
      "07/18/2023 00:11:56 - INFO - root -   ==== Epoch 1 start ====\n",
      "07/18/2023 00:11:57 - INFO - root -   Train loss: 0.71484375                    1.11it/s]\u001b[0m\n",
      "Training epoch 1:  56%|\u001b[34m██████████████▋           \u001b[0m| 9/16 [00:08<00:06,  1.09it/s]\u001b[0m[2023-07-18 00:12:05,290] [INFO] [logging.py:96:log_dist] [Rank 0] step=10, skipped=0, lr=[3e-05], mom=[[0.9, 0.999]]\n",
      "[2023-07-18 00:12:05,506] [INFO] [timer.py:215:stop] epoch=0/micro_step=10/global_step=10, RunningAvgSamplesPerSec=143.11783719344018, CurrSamplesPerSec=143.97873000607967, MemAllocated=0.52GB, MaxMemAllocated=3.99GB\n",
      "07/18/2023 00:12:06 - INFO - root -   Train loss: 0.69921875                    1.10it/s]\u001b[0m\n",
      "Training epoch 1: 100%|\u001b[34m█████████████████████████\u001b[0m| 16/16 [00:14<00:00,  1.09it/s]\u001b[0m\n",
      "07/18/2023 00:12:13 - INFO - root -   Eval. loss: 0.69140625\n",
      "07/18/2023 00:12:13 - INFO - root -   {'accuracy': 0.568,\n",
      " 'f1': 0.19999999999999998,\n",
      " 'precision': 0.9642857142857143,\n",
      " 'recall': 0.1115702479338843}\n",
      "07/18/2023 00:12:13 - INFO - root -   ==== Epoch 2 start ====\n",
      "07/18/2023 00:12:13 - INFO - root -   Train loss: 0.70703125                    1.20it/s]\u001b[0m\n",
      "Training epoch 2:  19%|\u001b[34m████▉                     \u001b[0m| 3/16 [00:02<00:11,  1.12it/s]\u001b[0m[2023-07-18 00:12:16,459] [INFO] [logging.py:96:log_dist] [Rank 0] step=20, skipped=0, lr=[3e-05], mom=[[0.9, 0.999]]\n",
      "[2023-07-18 00:12:16,696] [INFO] [timer.py:215:stop] epoch=0/micro_step=20/global_step=20, RunningAvgSamplesPerSec=143.38232812981283, CurrSamplesPerSec=144.5013080916504, MemAllocated=0.52GB, MaxMemAllocated=3.99GB\n",
      "07/18/2023 00:12:22 - INFO - root -   Train loss: 0.65625                       1.13it/s]\u001b[0m\n",
      "Training epoch 2:  81%|\u001b[34m████████████████████▎    \u001b[0m| 13/16 [00:11<00:02,  1.12it/s]\u001b[0m[2023-07-18 00:12:25,374] [INFO] [logging.py:96:log_dist] [Rank 0] step=30, skipped=0, lr=[3e-05], mom=[[0.9, 0.999]]\n",
      "[2023-07-18 00:12:25,636] [INFO] [timer.py:215:stop] epoch=0/micro_step=30/global_step=30, RunningAvgSamplesPerSec=143.6556516322058, CurrSamplesPerSec=145.20660051010336, MemAllocated=0.52GB, MaxMemAllocated=3.99GB\n",
      "Training epoch 2: 100%|\u001b[34m█████████████████████████\u001b[0m| 16/16 [00:14<00:00,  1.12it/s]\u001b[0m\n",
      "07/18/2023 00:12:29 - INFO - root -   Eval. loss: 0.60546875\n",
      "07/18/2023 00:12:29 - INFO - root -   {'accuracy': 0.7,\n",
      " 'f1': 0.6462264150943395,\n",
      " 'precision': 0.7611111111111111,\n",
      " 'recall': 0.5614754098360656}\n",
      "07/18/2023 00:12:29 - INFO - root -   ==== Save Model ====\n",
      "07/18/2023 00:12:30 - INFO - root -   Elapsed time: 0:01:16.299369\n",
      "[2023-07-18 00:12:30,823] [INFO] [launch.py:347:main] Process 21031 exits successfully.\n",
      "[2023-07-18 00:12:30,824] [INFO] [launch.py:347:main] Process 21032 exits successfully.\n",
      "[2023-07-18 00:12:31,825] [INFO] [launch.py:347:main] Process 21028 exits successfully.\n",
      "[2023-07-18 00:12:31,825] [INFO] [launch.py:347:main] Process 21030 exits successfully.\n",
      "CPU times: user 487 ms, sys: 236 ms, total: 723 ms\n",
      "Wall time: 1min 39s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "TRAIN_DEEPSPEED_CMD = f\"\"\"cd scripts && deepspeed train_deepspeed.py \\\n",
    "\"\"\"\n",
    "\n",
    "print(f'Running command: \\n{TRAIN_DEEPSPEED_CMD}')\n",
    "! {TRAIN_DEEPSPEED_CMD}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b29c869-2154-4673-9ec2-8894b1a3cbe0",
   "metadata": {
    "tags": []
   },
   "source": [
    "<br>\n",
    "\n",
    "## 3. SageMaker Training (Development Stage)\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 을 참조하세요.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bfb2f6b-08c3-4c5d-8665-dbd1c71ee8f0",
   "metadata": {},
   "source": [
    "### DeepSpeed migration\n",
    "\n",
    "\n",
    "SageMaker에서 DeepSpeed를 사용하려면, `deepspeed` 커맨드가 아닌 `mpirun` 커맨드를 사용해야 합니다. `mpirun`에 대한 세부 파라미터는 SageMaker Estimator 호출 시 `distribution = {\"mpi\": mpi_options}`로 설정하시면 되며, `deepspeed.init_distributed(...)`는 호출할 필요가 없습니다.\n",
    "\n",
    "```python\n",
    "if 'WORLD_SIZE' in os.environ:\n",
    "    # Environment variables set by torch.distributed.launch or torchrun\n",
    "    world_size = int(os.environ['WORLD_SIZE'])\n",
    "    rank = int(os.environ['RANK'])\n",
    "    local_rank = int(os.environ['LOCAL_RANK'])\n",
    "elif 'OMPI_COMM_WORLD_SIZE' in os.environ:\n",
    "    # Environment variables set by mpirun \n",
    "    world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])\n",
    "    rank = int(os.environ['OMPI_COMM_WORLD_RANK'])\n",
    "    local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])\n",
    "else:\n",
    "    sys.exit(\"Can't find the evironment variables for local rank\")\n",
    "\n",
    "dist.init_process_group(backend=\"nccl\", rank=rank, world_size=world_size)\n",
    "torch.cuda.set_device(local_rank)\n",
    "\n",
    "# SageMaker Training does not need to call deepspeed.init_distributed()\n",
    "if not 'SM_CHANNEL' in os.environ:\n",
    "    deepspeed.init_distributed(\"nccl\")\n",
    "  \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ef61fd36-be73-4a69-bf1d-6b2c549873f4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{credentials.py:1051} INFO - Found credentials from IAM Role: BaseNotebookInstanceEc2InstanceRole\n"
     ]
    }
   ],
   "source": [
    "instance_type = 'ml.p3.16xlarge'\n",
    "#instance_type = 'local_gpu'\n",
    "num_gpus = 8\n",
    "instance_count = 1\n",
    "batch_size = 32\n",
    "\n",
    "if instance_type in ['local', 'local_gpu']:\n",
    "    from sagemaker.local import LocalSession\n",
    "    sagemaker_session = LocalSession()\n",
    "    sagemaker_session.config = {'local': {'local_code': True}}\n",
    "else:\n",
    "    sagemaker_session = sagemaker.session.Session()\n",
    "    \n",
    "# hyperparameters, which are passed into the training job\n",
    "hyperparameters = {\n",
    "    'num_epochs': 3,                    # number of training epochs\n",
    "    'seed': 42,                         # seed\n",
    "    'train_batch_size': batch_size,     # batch size for training\n",
    "    'eval_batch_size': batch_size*2,    # batch size for evaluation\n",
    "    'warmup_steps': 0,                  # warmup steps\n",
    "    'learning_rate': 3e-5,              # learning rate used during training\n",
    "    'model_id': model_id                # pre-trained model\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d27fc875-0d23-401a-b636-80d08ace6bf3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mpi_options = {\n",
    "    \"enabled\" : True,            # Required\n",
    "    \"processes_per_host\" : 8,    # Required\n",
    "    # \"custom_mpi_options\" : \"--mca btl_vader_single_copy_mechanism none\"\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "41af34c0-a0f5-48ba-b763-2e61b546691f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "#image_uri = f\"763104351884.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04\"\n",
    "image_uri = '763104351884.dkr.ecr.{}.amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker'.format(region)\n",
    "\n",
    "# define Training Job Name \n",
    "job_name = f'deepspeed-nsmc-{time.strftime(\"%Y-%m-%d-%H-%M-%S\", time.localtime())}'\n",
    "chkpt_s3_path = f's3://{sess.default_bucket()}/{s3_prefix}/native/checkpoints'\n",
    "\n",
    "# create the Estimator\n",
    "sm_estimator = PyTorch(\n",
    "    entry_point           = 'train_deepspeed.py',  # fine-tuning script used in training jon\n",
    "    source_dir            = './scripts',        # directory where fine-tuning script is stored\n",
    "      image_uri = image_uri,\n",
    "    instance_type         = instance_type,      # instances type used for the training job\n",
    "    instance_count        = instance_count,     # 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",
    "    sagemaker_session=sagemaker_session,\n",
    "    py_version            = 'py38',             # the python version used in the training job\n",
    "    hyperparameters       = hyperparameters,    # the hyperparameter used for running the training job\n",
    "    distribution          = {\"mpi\": mpi_options},\n",
    "    disable_profiler     = True,\n",
    "    debugger_hook_config  = False,\n",
    "    #keep_alive_period_in_seconds = 20*60     # warm pool    \n",
    "    #checkpoint_s3_uri     = chkpt_s3_path,\n",
    "    #checkpoint_local_path ='/opt/ml/checkpoints',  \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d6e150c9-dbaa-49b1-95b0-19801315e67e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using provided s3_resource\n",
      "[{session.py:873} INFO - Creating training-job with name: deepspeed-nsmc-2023-07-18-00-39-17-2023-07-18-00-39-18-276\n"
     ]
    }
   ],
   "source": [
    "# define a data input dictonary with our uploaded s3 uris\n",
    "data = {\n",
    "    'train': train_input_path,\n",
    "    'eval': eval_input_path\n",
    "}\n",
    "\n",
    "# starting the train job with our uploaded datasets as input\n",
    "sm_estimator.fit(data, wait=False)\n",
    "train_job_name = sm_estimator.latest_training_job.job_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "78f7d720-0c55-44ae-8fa8-0b553eff29d5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<b> [PyTorch DeepSpeed Training] Review <a target=\"blank\" href=\"https://console.aws.amazon.com/sagemaker/home?region=us-east-1#/jobs/deepspeed-nsmc-2023-07-18-00-39-17-2023-07-18-00-39-18-276\">Training Job</a></b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<b> [PyTorch DeepSpeed Training] Review <a target=\"blank\" href=\"https://console.aws.amazon.com/cloudwatch/home?region=us-east-1#logStream:group=/aws/sagemaker/TrainingJobs;prefix=deepspeed-nsmc-2023-07-18-00-39-17-2023-07-18-00-39-18-276;streamFilter=typeLogStreamPrefix\">CloudWatch Logs</a></b>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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_link, cloudwatch_link = make_console_link(region, train_job_name, '[PyTorch DeepSpeed Training]')\n",
    "\n",
    "display(HTML(train_job_link))\n",
    "display(HTML(cloudwatch_link))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2264b3ea-4f2f-4183-a8cf-270f211606ff",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 4. (Optional) Inference\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74e32290-3431-4c99-b1f5-41f157e2e2b0",
   "metadata": {},
   "source": [
    "### Copy S3 model artifact to local directory\n",
    "\n",
    "S3에 저장된 모델 아티팩트를 로컬 경로로 복사하여 압축을 해제합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "239d818b-30e2-4883-a7a4-08e1430a0dd1",
   "metadata": {},
   "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 {sm_estimator.model_data} {local_model_dir}/model.tar.gz\n",
    "!tar -xzf {local_model_dir}/model.tar.gz -C {local_model_dir}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67024925-3bb5-4baa-95fb-8cce14411a8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "import numpy as np\n",
    "from collections import OrderedDict\n",
    "from transformers import BertForSequenceClassification, AutoTokenizer\n",
    "        \n",
    "MODEL_NAME = \"bert-base-multilingual-cased\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)#.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00aba09c-c3af-4480-a338-0ba5f2738c91",
   "metadata": {},
   "source": [
    "### Load DDP model to a non-DDP model\n",
    "데이터 병렬화를 적용하여 모델을 훈련하면 모델의 weight의 key값에 `module`이 붙게 되어 모델 로딩 시 오류가 발생합니다. 따라서, 이를 제거해 주는 후처리 과정이 필요합니다. 후처리가 번거롭다면, DDP로 훈련 후 저장할 때 명시적으로 `module`를 제외하고 저장하는 방법도 있습니다.\n",
    "\n",
    "참조: https://discuss.pytorch.org/t/how-to-switch-model-trained-on-2-gpus-to-1-gpu/20039\n",
    "\n",
    "```python\n",
    "model_to_save = model.module if hasattr(model, 'module') else model\n",
    "...\n",
    "model_to_save.state_dict()\n",
    "torch.save({'model': model_to_save.state_dict())\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c97f735f-6fd8-4bde-bd96-d13060902d70",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "model_filename = glob.glob(f'{local_model_dir}/*.pt')[0]\n",
    "state_dict = torch.load(model_filename)\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4de79bf8-2ee3-449d-bc28-f935deab8aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"이 영화 너무 재미있어요\"\n",
    "encode_plus_token = tokenizer.encode_plus(\n",
    "    text,\n",
    "    max_length=128,\n",
    "    add_special_tokens=True,\n",
    "    return_token_type_ids=False,\n",
    "    padding=\"max_length\",\n",
    "    return_attention_mask=True,\n",
    "    return_tensors=\"pt\",\n",
    "    truncation=True\n",
    ")\n",
    "\n",
    "output = model(**encode_plus_token)\n",
    "softmax_fn = torch.nn.Softmax(dim=1)\n",
    "softmax_output = softmax_fn(output[0])\n",
    "_, prediction = torch.max(softmax_output, dim=1)\n",
    "\n",
    "predicted_class_idx = prediction.item()\n",
    "score = softmax_output[0][predicted_class_idx]\n",
    "print(f\"predicted_class: {predicted_class_idx}, score={score}\")"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}