torch-neuronx/training/hf_sentiment_analysis/02-hf-distributed-training.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "# AWS Trainium Distributed Training - \"bert-base-cased\" for Sentiment Analysis\n",
                "This notebook shows how to fine-tune a \"bert base cased\" PyTorch model with AWS Trainium (Trn1 instances) using Neuron SDK. The original implementation is provided by HuggingFace.\n",
                "\n",
                "Our goal is building a Machine Learning model that will predict whether the tweet is offensive, neutral, or positive (<b>Sentiment Analysis</b>).\n",
                "\n",
                "The target variable is the **Sentiment**, which can be:\n",
                "* Neutral\n",
                "* Positive\n",
                "* Negative\n",
                "\n",
                "In this exercise you will do:\n",
                " - Run a Distributed Training using all the available Neuron Cores\n",
                "\n",
                "The example code referenced for this example is [trainium-distributed-training](./code/02-trainium-distributed-training/train.py)\n",
                "\n",
                "It has been tested and run on a **trn1.32xlarge**\n",
                "\n",
                "**Reference:** https://huggingface.co/bert-base-cased"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "***"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Verify that this Jupyter notebook is running the Python kernel environment that was set up according to the [PyTorch Installation Guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/torch-neuronx.html#setup-torch-neuronx). You can select the kernel from the 'Kernel -> Change Kernel' option on the top of this Jupyter notebook page."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Step 1 - Install dependencies\n",
                "\n",
                "Let's install some required dependencies for our environment."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "%env TOKENIZERS_PARALLELISM=True #Supresses tokenizer warnings making errors easier to detect\n",
                "!pip install neuronx-cc==2.* torch-neuronx torchvision datasets transformers==4.40.1"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "***"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Step 2 - Fine-Tune the model"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "Let's take a look to our train.py code"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "! pygmentize ./code/02-trainium-distributed-training/train.py"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "In order to run distributed training by using AWS Trainium Neuron Core, we have to define the number of Cores we want to use for distribution. The number of Neuron Nodes N can be 1, 2, 8, or 32."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "nproc_per_node = 32"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "collapsed": false,
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "import subprocess\n",
                "print(\"Train model\")\n",
                "RUN_CMD = f\"\"\"torchrun --nproc_per_node={nproc_per_node} ./code/02-trainium-distributed-training/train.py\"\"\"\n",
                "\n",
                "print(f'Running command: \\n{RUN_CMD}')\n",
                "if subprocess.check_call(RUN_CMD,shell=True):\n",
                "   print(\"There was an error with the fine-tune command\")\n",
                "else:\n",
                "   print(\"Fine-tune Successful!!!\")"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 2
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython2",
            "version": "3.7.16"
        },
        "vscode": {
            "interpreter": {
                "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 0
}