distributed-pytorch-bert-fine-tuning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# BERT で始める分散機械学習"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 環境ごとの値設定\n",
        "group_name = \"\" # 共有でリソース使用時に重複を防止するための値\n",
        "vnet_resourcegroup_name = \"\"\n",
        "vnet_name = \"\"\n",
        "subnet_name = \"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Workspace の初期化\n",
        "\n",
        "Azure Machine Learning Workspace リソースに接続し、Workspace オブジェクトを作成する。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from azureml.core.workspace import Workspace\n",
        "\n",
        "ws = Workspace.from_config()\n",
        "print('Workspace name: ' + ws.name, \n",
        "      'Azure region: ' + ws.location, \n",
        "      'Subscription id: ' + ws.subscription_id, \n",
        "      'Resource group: ' + ws.resource_group, sep='\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute Cluster の作成\n",
        "Azure Machine Learning の管理下のクラスターである Computer Cluster を作成する。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from azureml.core.compute import ComputeTarget, AmlCompute\n",
        "from azureml.core.compute_target import ComputeTargetException\n",
        "\n",
        "cluster_name = f\"gpucluster\"\n",
        "\n",
        "try:\n",
        "    compute_target = ComputeTarget(workspace=ws, name=cluster_name)\n",
        "    print('Found existing compute target.')\n",
        "except ComputeTargetException:\n",
        "    # クラスター作成\n",
        "    print('Creating a new compute target...')\n",
        "    compute_config = AmlCompute.provisioning_configuration(\n",
        "        vm_size='STANDARD_NC6S_V3',\n",
        "        max_nodes=2,\n",
        "        idle_seconds_before_scaledown=300,\n",
        "        vnet_resourcegroup_name=vnet_resourcegroup_name, # VNet 閉域化時のみ\n",
        "        vnet_name=vnet_name, # VNet 閉域化時のみ\n",
        "        subnet_name=subnet_name, # VNet 閉域化時のみ\n",
        "    )\n",
        "    compute_target = ComputeTarget.create(ws, cluster_name, compute_config)\n",
        "\n",
        "    compute_target.wait_for_completion(show_output=True)\n",
        "\n",
        "print(compute_target.get_status().serialize())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 学習\n",
        "Compute Cluster にジョブを投入し、分散機械学習を実行する。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Experiment の作成\n",
        "実験の各試行をまとめて管理する [Experiment](https://docs.microsoft.com/azure/machine-learning/service/concept-azure-machine-learning-architecture#experiment) を作成する。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from azureml.core import Experiment\n",
        "\n",
        "experiment_name = f\"pytorch-distributed-bert-fine-tuning-{group_name}\"\n",
        "experiment = Experiment(ws, name=experiment_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Environment の作成\n",
        "\n",
        "Environment の実態は Docker コンテナであり、下記サンプルの場合は environment.yml に記載されたパッケージがプリセットされたコンテナイメージが作成される。\n",
        "\n",
        "初回実行時はイメージビルドが挟まるため非常に時間がかかるが、2回目以降はキャッシュされたコンテナイメージが使いまわされるためあまり時間がかからない。\n",
        "\n",
        "実行時は Environment として保存されたコンテナイメージが Compute Cluster の各ノードに配られ、コンテナが立ち上がって Python コードを実行する。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from azureml.core import Environment\n",
        "\n",
        "environment_name = f\"pytorch-distributed-env-{group_name}\"\n",
        "environment_file = \"src/environment.yml\"\n",
        "\n",
        "pytorch_env = Environment.from_conda_specification(name=environment_name, file_path=environment_file)\n",
        "\n",
        "## ベースとなる Docker イメージを指定\n",
        "## 「Environment」にプリセットされた environment の詳細画面から確認可能\n",
        "pytorch_env.docker.base_image = (\n",
        "    \"mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.1-cudnn8-ubuntu18.04:20211221.v1\"\n",
        ")\n",
        "pytorch_env.register(workspace=ws)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 学習ジョブの設定\n",
        "\n",
        "Python スクリプトを実行する諸条件を定義する設定用オブジェクトである ScriptRunConfig に各種実行条件をセットする。分散機械学習の場合、```distributed_job_config```に分散機械学習用の設定ファイルを渡すことで、Compute Cluster が適切に初期化されて分散機械学習がスムーズに実行される。\n",
        "\n",
        "今回は PyTorch DDP を使用しているため```PyTorchConfiguration```を渡しているが、OpenMPI 等ほかの分散機械学習の仕組みを使うことも可能。\n",
        "\n",
        "```distributed_job_config```に渡す Configuration の種類に応じて、```NODE_RANK```のような依存する環境変数が自動でセットされたり、通信設定が裏で行われる。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from azureml.core import ScriptRunConfig\n",
        "from azureml.core.runconfig import PyTorchConfiguration\n",
        "\n",
        "node_count = 2\n",
        "arg = [\n",
        "    \"--batch_size\", 1024,\n",
        "    \"--epochs\", 40,\n",
        "    \"--lr\", 1e-4,\n",
        "    \"--num_nodes\", node_count,\n",
        "    \"--base_model\", \"cl-tohoku/bert-base-japanese-char-v2\"\n",
        "]\n",
        "\n",
        "## cl-tohoku/bert-* から好きなものを選択\n",
        "## https://huggingface.co/cl-tohoku\n",
        "\n",
        "### ex) bert-large-japanese-char の場合\n",
        "#arg = [\n",
        "#    \"--batch_size\", 1024,\n",
        "#    \"--epochs\", 40,\n",
        "#    \"--lr\", 1e-5,\n",
        "#    \"--num_nodes\", node_count,\n",
        "#    \"--base_model\", \"cl-tohoku/bert-large-japanese-char\"\n",
        "#]\n",
        "\n",
        "\n",
        "dist_config = PyTorchConfiguration(node_count=node_count)\n",
        "\n",
        "src = ScriptRunConfig(source_directory=\"src\",\n",
        "                      script='trainer_pytorch_distributed.py',\n",
        "                      compute_target=compute_target,\n",
        "                      environment=pytorch_env,\n",
        "                      arguments=arg,\n",
        "                      distributed_job_config=dist_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### ジョブの実行\n",
        "Experiment に対してジョブを渡すことでジョブが実際に実行される。\n",
        "\n",
        "Experiment 配下に Run が作成され、メトリックやログ等が自動で記録される。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "run = experiment.submit(src)\n",
        "print(run)\n",
        "run.wait_for_completion(show_output=True) # this provides a verbose log"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "authors": [
      {
        "name": "ninhu"
      }
    ],
    "category": "training",
    "compute": [
      "AML Compute"
    ],
    "datasets": [
      "MNIST"
    ],
    "deployment": [
      "None"
    ],
    "exclude_from_index": false,
    "framework": [
      "PyTorch"
    ],
    "friendly_name": "Distributed PyTorch",
    "index_order": 1,
    "interpreter": {
      "hash": "6d65a8c07f5b6469e0fc613f182488c0dccce05038bbda39e5ac9075c0454d11"
    },
    "kernelspec": {
      "display_name": "Python 3.6",
      "language": "python",
      "name": "python3"
    },
    "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.8.1"
    },
    "tags": [
      "None"
    ],
    "task": "Train a model using the distributed training via Horovod"
  },
  "nbformat": 4,
  "nbformat_minor": 2
}