 archived/smp-train-gpt-neox-fsdp-tp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Train GPT-NeoX (or Llama-v2) with SageMaker-PyTorch FSDP At Large Scale\n",
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.\n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "---\n",
    "\n",
    "In this notebook, you will learn how to train the Hugging Face Transformers GPT-NeoX and Llama-v2 models,\n",
    "with hybrid sharding, tensor parallelism, and activation offloading with SageMaker.\n",
    "You can either launch this notebook from an Amazon SageMaker notebook instance which handles all credentials automatically,\n",
    "or by running it locally and setting credentials manually.\n",
    "\n",
    "The notebook is accompanied by the following files:\n",
    "- `train.py`: The entry point script that'll be passed to the SageMaker PyTorch estimator later in this notebook when launching the training job.\n",
    "- `arguments.py`: This file has functions for argument parsing (i.e. hyperparameters).\n",
    "- `checkpoints.py`: This file has functions for saving and loading checkpoints.\n",
    "- `data_utils`: This file has functions for handling S3 URLs.\n",
    "- `data`: This directory has scripts for preparing and loading data.\n",
    "- `fsdp_utils.py`: This file has util functions for fully sharded data parallelism.\n",
    "- `learning_rates.py`: This file has functions for learning rate schedule.\n",
    "- `logging_utils.py`: This file has functions to handle logging.\n",
    "- `memory_tracker.py`: This file has functions to track memory usage.\n",
    "- `requirements.txt`: This file installs the dependencies, including HuggingFace transformers.\n",
    "- `train_lib.py`: This file has functions for running an end-to-end training of the GPT-NeoX or Llama-v2 model with SMP FSDP, settings for hybrid sharding applied, and implemented with code lines to save, load, and fine-tune the model.\n",
    "- `train_utils.py`: This file has utility functions for training.\n",
    "\n",
    "## Additional Resources\n",
    "- To learn more about launching a multi-node distributed PyTorch training job, see [Launching a Distributed Training Job](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#launching-a-distributed-training-job).\n",
    "- To learn more about using the SageMaker Python SDK with PyTorch, see [Using PyTorch with the SageMaker Python SDK](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html).\n",
    "- To learn more about launching a training job in Amazon SageMaker with your own training image, see [Use Your Own Training Algorithms](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html).\n",
    "\n",
    "## Prerequisites\n",
    "You need to create an `S3` bucket to store the input data for training.\n",
    "This bucket must be located in the same AWS Region that you choose to launch your training job. To learn how to create a `S3` bucket,\n",
    "see [Create your first S3 bucket](https://docs.aws.amazon.com/AmazonS3/latest/userguide/creating-bucket.html) in the Amazon S3 documentation.\n",
    "\n",
    "## Launching Environment\n",
    "\n",
    "### Amazon SageMaker Notebook\n",
    "You can run the notebook with an Amazon SageMaker notebook instance without manually setting your aws credentials.\n",
    "\n",
    "1. Create a new SageMaker notebook instance and open it.\n",
    "2. Zip the contents of this folder & upload to the instance with the `Upload` button on the top-right.\n",
    "3. Open a new terminal with `New -> Terminal`.\n",
    "4. Within the terminal, enter the correct directory and unzip the file.\n",
    "    1. `cd SageMaker && unzip <your-zip-name-here>.zip`\n",
    "\n",
    "### Locally\n",
    "You can run locally by launching a Jupyter notebook server with `jupyter notebook`.\n",
    "This requires you to set your aws credentials in the environment manually.\n",
    "See [Configure the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html) for more details.\n",
    "\n",
    "## Amazon SageMaker Initialization\n",
    "Run the following cell to import SageMaker modules and retrieve information of your current SageMaker work environment,\n",
    "such as your AWS account ID, the AWS Region, and the ARN of your Amazon SageMaker execution role.\n",
    "Upgrade SageMaker SDK to the latest version.\n",
    "\n",
    "**NOTE:** This step might require a kernel restart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade \"sagemaker>=2.224\"\n",
    "%pip install sagemaker-experiments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### AWS Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import os\n",
    "\n",
    "import boto3\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "role = get_execution_role()\n",
    "print(f\"SageMaker Execution Role: {role}\")\n",
    "\n",
    "client = boto3.client(\"sts\")\n",
    "account = client.get_caller_identity()[\"Account\"]\n",
    "print(f\"AWS account: `{account}`.\")\n",
    "\n",
    "session = boto3.session.Session()\n",
    "region = session.region_name\n",
    "print(f\"AWS region: `{region}`.\")\n",
    "\n",
    "sm_boto_client = boto3.client(\"sagemaker\")\n",
    "sagemaker_session = sagemaker.session.Session(boto_session=session)\n",
    "\n",
    "# get default bucket\n",
    "default_bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "print(f\"\\nDefault bucket for this session: `{default_bucket}`.\")\n",
    "\n",
    "# FSX setup if you plan to use it.\n",
    "# - Specify FSx Lustre file system id.\n",
    "FSX_FILE_SYSTEM_TYPE = \"FSxLustre\"\n",
    "FSX_FILE_SYSTEM_ID = \"fs-XXXXXXXXXXXXXXXXX\"  # Sample: \"fs-[a-z0-9]+\"\n",
    "\n",
    "# - Specify the SG and subnet used by the FSX, these are passed to SM Estimator so jobs use this as well.\n",
    "FSX_SECURITY_GROUP_ID = \"sg-XXXXXXXXXXXXXXXXX\"  # Sample: \"sg-[a-z0-9]+\"\n",
    "FSX_SUBNET_ID = \"subnet-XXXXXXXXXXXXXXXXX\"  # Sample: \"subnet-[a-z0-9]+\"\n",
    "\n",
    "# - Specify directory path for input data on the file system.\n",
    "# You need to provide normalized and absolute path below.\n",
    "# Your mount name can be provided by you when creating fsx, or generated automatically.\n",
    "# You can find this mount_name on the FSX page in console.\n",
    "# Example of fsx generated mount_name: \"3x5lhbmv\"\n",
    "# Example base path: \"/3x5lhbmv\"\n",
    "FSX_FILE_SYSTEM_BASE_PATH = \"/XXXXXXXX\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download and Prepare the GLUE/SST2 Dataset\n",
    "\n",
    "Here you will download, prepare the GLUE/SST2 dataset and then copy the files to S3."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install the Hugging Face Transformers and Datasets libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q datasets transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Imports.\n",
    "\n",
    "import datasets\n",
    "from datasets import load_dataset, load_from_disk, load_metric\n",
    "from sagemaker.pytorch import PyTorch\n",
    "import transformers\n",
    "import logging\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "from transformers.testing_utils import CaptureLogger\n",
    "\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose Model\n",
    "\n",
    "Choose to train either the `GPT-NeoX` or `Llama-v2` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_type = \"gpt_neox\"  # [\"gpt_neox\", \"llama_v2\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data\n",
    "\n",
    "This section loads the [GLUE/SST2](https://huggingface.co/datasets/glue/viewer/sst2/train) dataset\n",
    "and splits it to training and validation datasets.\n",
    "You can update this section to load any HuggingFace dataset you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Prepare datasets.\n",
    "\n",
    "hyperparameters = {\n",
    "    \"cache_dir\": \"tmp\",\n",
    "    \"dataset_config_name\": \"sst2\",\n",
    "    \"dataset_name\": \"glue\",\n",
    "    \"do_train\": True,\n",
    "    \"do_eval\": True,\n",
    "}\n",
    "\n",
    "raw_datasets = load_dataset(\n",
    "    hyperparameters[\"dataset_name\"],\n",
    "    hyperparameters[\"dataset_config_name\"],\n",
    ")\n",
    "\n",
    "# NOTE: Remove existing validation dataset, as it is too small to shard across all ranks.\n",
    "del raw_datasets[\"validation\"]\n",
    "\n",
    "if \"validation\" not in raw_datasets.keys():\n",
    "    validation_percentage = \"10%\"\n",
    "\n",
    "    raw_datasets[\"validation\"] = load_dataset(\n",
    "        hyperparameters[\"dataset_name\"],\n",
    "        hyperparameters[\"dataset_config_name\"],\n",
    "        split=f\"train[:{validation_percentage}]\",\n",
    "        cache_dir=hyperparameters[\"cache_dir\"],\n",
    "    )\n",
    "\n",
    "    raw_datasets[\"train\"] = load_dataset(\n",
    "        hyperparameters[\"dataset_name\"],\n",
    "        hyperparameters[\"dataset_config_name\"],\n",
    "        split=f\"train[{validation_percentage}:]\",\n",
    "        cache_dir=hyperparameters[\"cache_dir\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load tokenizer\n",
    "Nearly every NLP task begins with a tokenizer,\n",
    "which converts your text data into a format (token) that can be processed by the NLP model.\n",
    "\n",
    "The following cell loads a tokenizer for GPT-NeoX-7B using [AutoTokenizer.from_pretrained()](https://huggingface.co/docs/transformers/v4.19.4/en/autoclass_tutorial#autotokenizer)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Tokenizer.\n",
    "\n",
    "tokenizer_kwargs = {\n",
    "    \"cache_dir\": hyperparameters[\"cache_dir\"],\n",
    "}\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\", **tokenizer_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preprocess data\n",
    "\n",
    "Set up a function to run the tokenizer and group texts into chunks smaller than the block size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Util functions.\n",
    "\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    tok_logger = transformers.utils.logging.get_logger(\"transformers.tokenization_utils_base\")\n",
    "\n",
    "    with CaptureLogger(tok_logger) as cl:\n",
    "        output = tokenizer(examples[text_column_name])\n",
    "        # clm input could be much much longer than block_size\n",
    "        if \"Token indices sequence length is longer than the\" in cl.out:\n",
    "            tok_logger.warning(\n",
    "                \"^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model.\"\n",
    "            )\n",
    "    return output\n",
    "\n",
    "\n",
    "# Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.\n",
    "def group_texts(examples):\n",
    "    # Concatenate all texts.\n",
    "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n",
    "    # customize this part to your needs.\n",
    "    if total_length >= block_size:\n",
    "        total_length = (total_length // block_size) * block_size\n",
    "        # Split by chunks of max_len.\n",
    "        result = {\n",
    "            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
    "            for k, t in concatenated_examples.items()\n",
    "        }\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Tokenize datasets.\n",
    "\n",
    "column_names = raw_datasets[\"train\"].column_names\n",
    "text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n",
    "\n",
    "# since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function\n",
    "tok_logger = transformers.utils.logging.get_logger(\"transformers.tokenization_utils_base\")\n",
    "\n",
    "tokenized_datasets = raw_datasets.map(\n",
    "    tokenize_function,\n",
    "    batched=True,\n",
    "    num_proc=1,\n",
    "    remove_columns=column_names,\n",
    "    desc=\"Running tokenizer on dataset\",\n",
    ")\n",
    "\n",
    "\n",
    "block_size = tokenizer.model_max_length\n",
    "if block_size > 1024:\n",
    "    logger.warning(\n",
    "        f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n",
    "        \"Picking 1024 instead. You can change that default value by passing --block_size xxx.\"\n",
    "    )\n",
    "    block_size = 1024\n",
    "else:\n",
    "    if args.block_size > tokenizer.model_max_length:\n",
    "        logger.warning(\n",
    "            f\"The block_size passed ({block_size}) is larger than the maximum length for the model\"\n",
    "            f\"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}.\"\n",
    "        )\n",
    "    block_size = min(block_size, tokenizer.model_max_length)\n",
    "\n",
    "lm_datasets = tokenized_datasets.map(\n",
    "    group_texts,\n",
    "    batched=True,\n",
    "    #     num_proc=args.preprocessing_num_workers,\n",
    "    desc=f\"Grouping texts in chunks of {block_size}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set additional hyperparameters and S3 paths for mapping the train and validation datasets properly depending on the phase (training or validation) of the training job in each epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Prepare datasets.\n",
    "\n",
    "if hyperparameters[\"do_train\"]:\n",
    "    train_dataset = lm_datasets[\"train\"]\n",
    "    train_dataset.to_json(\"./training.json\")\n",
    "    training_dataset_location = f\"s3://{default_bucket}/dataset/train/\"\n",
    "\n",
    "    command = f\"aws s3 cp ./training.json {training_dataset_location}\"\n",
    "    os.system(command)\n",
    "else:\n",
    "    training_dataset_location = None\n",
    "\n",
    "if hyperparameters[\"do_eval\"]:\n",
    "    eval_dataset = lm_datasets[\"validation\"]\n",
    "    eval_dataset.to_json(\"./validation.json\")\n",
    "    validation_dataset_location = f\"s3://{default_bucket}/dataset/validation/\"\n",
    "\n",
    "    command = f\"aws s3 cp ./validation.json {validation_dataset_location}\"\n",
    "    os.system(command)\n",
    "else:\n",
    "    validation_dataset_location = None\n",
    "\n",
    "\n",
    "if hyperparameters[\"do_train\"]:\n",
    "    command = \"rm ./training.json\"\n",
    "    os.system(command)\n",
    "\n",
    "if hyperparameters[\"do_eval\"]:\n",
    "    command = \"rm ./validation.json\"\n",
    "    os.system(command)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store training_dataset_location\n",
    "%store validation_dataset_location\n",
    "\n",
    "%store"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify Amazon S3 Bucket Paths\n",
    "Here you need to specify the paths for training data to be used by your job. The bucket used must be in the same region as where training will run. In the cells above you downloaded the GLUE/SST2 training and validation split datasets and uploaded the json files in an S3 bucket in your account. This example will train on those json files.\n",
    "\n",
    "After you successfully run this example tensor parallel + fully sharded data parallel training job, you can modify the S3 bucket to where your own dataset is stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r training_dataset_location\n",
    "%store -r validation_dataset_location\n",
    "\n",
    "# if you're bringing your own data, uncomment the following lines and specify the locations there\n",
    "# training_dataset_location = YOUR_S3_BUCKET/training\n",
    "# validation_dataset_location = YOUR_S3_BUCKET/validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_train_bucket = training_dataset_location\n",
    "s3_test_bucket = validation_dataset_location"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following S3 bucket will store the output artifacts of the training job. You can modify this as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_output_bucket = f\"s3://sagemaker-{region}-{account}/smp-fsdp-tp/{model_type}-outputdir/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Data Channels for SageMaker Training Using Amazon S3\n",
    "In this step, define SageMaker training data channels to the S3 buckets.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Setup if you do not use fsx.\n",
    "\n",
    "use_fsx = True\n",
    "\n",
    "if not use_fsx:\n",
    "    if s3_train_bucket != None:\n",
    "        train = sagemaker.inputs.TrainingInput(\n",
    "            s3_train_bucket, distribution=\"FullyReplicated\", s3_data_type=\"S3Prefix\"\n",
    "        )\n",
    "        data_channels = {\"train\": train}\n",
    "    else:\n",
    "        data_channels = {\"train\": mock_data}\n",
    "    if s3_test_bucket != None:\n",
    "        test = sagemaker.inputs.TrainingInput(\n",
    "            s3_test_bucket, distribution=\"FullyReplicated\", s3_data_type=\"S3Prefix\"\n",
    "        )\n",
    "        data_channels[\"test\"] = test\n",
    "    else:\n",
    "        data_channels[\"test\"] = mock_data\n",
    "    print(data_channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Optional) Set Up and Use Amazon FSx for Data Channels and Checkpoints\n",
    "While the previous option of using Amazon S3 is easier to setup, using an FSx can be beneficial for performance when dealing with large input sizes and large model sizes and is more stable. In general, checkpointing should be done using FSx.\n",
    "\n",
    "Please see the instructions from [Distributed Training of Mask-RCNN in Amazon SageMaker Using FSx](https://github.com/aws/amazon-sagemaker-examples/blob/main/advanced_functionality/distributed_tensorflow_mask_rcnn/mask-rcnn-scriptmode-fsx.ipynb) to create an FSx Lustre file system and import the dataset from the S3 bucket to your FSx file system. Note that the FSx file system must be created in a private subnet with internet gateway to ensure that training job has access to the internet. For general guidance on setting an FSx Lustre file system as data input channel, see Configure Data Input Channel to Use Amazon FSx for Lustre."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Setup FSX as needed.\n",
    "\n",
    "# Instructions obtained from:\n",
    "# https://github.com/aws/amazon-sagemaker-examples/blob/main/advanced_functionality/distributed_tensorflow_mask_rcnn/mask-rcnn-scriptmode-fsx.ipynb\n",
    "\n",
    "if use_fsx:\n",
    "    from sagemaker.inputs import FileSystemInput\n",
    "\n",
    "    train = FileSystemInput(\n",
    "        directory_path=FSX_FILE_SYSTEM_BASE_PATH,\n",
    "        file_system_access_mode=\"rw\",\n",
    "        file_system_id=FSX_FILE_SYSTEM_ID,\n",
    "        file_system_type=FSX_FILE_SYSTEM_TYPE,\n",
    "    )\n",
    "\n",
    "    data_channels = {\"train\": train, \"test\": train}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Hyperparameters, Metric Definitions, and MPI Options"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tensor Parallelism\n",
    "\n",
    "Tensor parallelism is a type of model parallelism in which specific model weights, gradients, and/ or optimizer states are split across devices,\n",
    "by replacing specific submodules in the model with their distributed implementations.\n",
    "The tensor parallel degree controls the sharding level and can be set from 1 to `world_size`,\n",
    "though we only recommend setting 1 to 8,\n",
    "assuming an 8-gpu node such as `ml.p4d.24xlarge`.\n",
    "This is because inter-node tensor parallel communication is much slower than intra-node tensor parallel communication.\n",
    "\n",
    "For more information,\n",
    "see [tensor parallelism](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-tensor-parallelism.html).\n",
    "\n",
    "### Hybrid Sharding\n",
    "\n",
    "Hybrid sharding is a memory saving technique in between `FULL_SHARD` and `NO_SHARD` with `FULL_SHARD` saving the most and `NO_SHARD` not saving any. \n",
    "This technique shards parameters within the hybrid shard degree group and replicates parameters across groups.\n",
    "\n",
    "The hybrid shard degree controls sharding across GPUs and can be set to `0` to `world_size // tensor_parallel_degree`.\n",
    "\n",
    "- A value of `8` applies `FULL_SHARD` within a node and then replicates parameters across nodes since there are 8 GPUs in the nodes we are using.\n",
    "This results in reduced communication volume as expensive all-gathers and reduce-scatters are only done within a node,\n",
    "which can be more performant for medium-sized models.\n",
    "    - Generally, you want to use the smallest HSD that does not cause Out of Memory (OOM) errors. If you are hitting OOM, try increasing the hybrid shard degree to reduce memory usage on each node.\n",
    "- An value of `0` falls back to the native PyTorch implementation and API in the script.\n",
    "If `FULL_SHARD` was the strategy set, it would shard across the whole cluster of GPUs.\n",
    "If `HYBRID_SHARD` or `_HYBRID_SHARD_ZERO2` was the strategy, the default is equivalent to `8`.\n",
    "\n",
    "For more information, see [fsdp.ShardingStrategy](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp).\n",
    "\n",
    "### Activation Offloading\n",
    "\n",
    "Activation offloading is a memory saving technique which requires activation checkpointing to be enabled.\n",
    "Enabling this offloads activations onto CPU memory to save GPU memory.\n",
    "This is useful when our model is too large to fit in our nodes or when we want to train with a larger batch size.\n",
    "\n",
    "#### SageMaker Activation Offloading\n",
    "\n",
    "SageMaker activation offloading improves performance by pre-fetching activations from the CPU before they are needed,\n",
    "so that the GPU does not wait for the activations to be loaded.\n",
    "\n",
    "Setting `\"sm_activation_offloading\": True` enables our improved version.\n",
    "\n",
    "Note: we generally only need activation offloading for models >= 20B parameters or if we are getting OOM with a given batch size.\n",
    "We use it here simply to illustrate how to enable it.\n",
    "\n",
    "#### Activation Loading Horizon\n",
    "\n",
    "The activation loading horizon is the maximum number of loaded tensors that can be in the GPU memory simultaneously.\n",
    "It is a positive integer with a default value of `2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Hyperparameters.\n",
    "\n",
    "tensor_parallel_degree = 2  # An integer in [1, world_size]\n",
    "hybrid_shard_degree = (\n",
    "    4  # An integer in [0, world_size // tensor_parallel_degree] and its default value is 0.\n",
    ")\n",
    "offload_activations = True  # Enables SM activation offloading implementation.\n",
    "activation_loading_horizon = (\n",
    "    2  # Activation loading horizon, a positive integer and its default value is 2.\n",
    ")\n",
    "save_steps = 10  # Save step interval.\n",
    "max_steps = 15  # Maximum training steps.\n",
    "\n",
    "hyperparameters = {\n",
    "    \"activation_checkpointing\": 1,\n",
    "    \"auto_wrap_policy\": \"transformer_auto_wrap_policy\",\n",
    "    \"backward_fetch_policy\": \"backward_pre\",\n",
    "    \"beta1\": 0.9,\n",
    "    \"beta2\": 0.95,\n",
    "    \"bf16\": 1,\n",
    "    \"checkpoint_dir\": \"/opt/ml/checkpoints\",\n",
    "    \"checkpoint_freq\": save_steps,\n",
    "    \"clean_cache\": 0,\n",
    "    \"delayed_param\": 1,\n",
    "    \"enable_memory_profiling\": 0,\n",
    "    \"epochs\": 100,\n",
    "    \"fast_validation\": 0,\n",
    "    \"forward_prefetch\": 1,\n",
    "    \"limit_all_gathers\": 1,\n",
    "    \"logging_freq\": 1,\n",
    "    \"lr\": 0.0001,\n",
    "    \"lr_decay_iters\": 47683,\n",
    "    \"lr_decay_style\": \"cosine\",\n",
    "    \"max_steps\": max_steps,\n",
    "    \"min_lr\": 1e-05,\n",
    "    \"model_type\": model_type,\n",
    "    \"num_kept_checkpoints\": 2,\n",
    "    \"plateau\": 0.0,\n",
    "    \"seed\": 12345,\n",
    "    \"sharding_strategy\": \"hybrid_shard\",\n",
    "    \"train_batch_size\": 2,\n",
    "    \"use_smp_flash_attn\": 1,\n",
    "    \"use_smp_implementation\": 1,\n",
    "    \"val_batch_size\": 4,\n",
    "    \"validation_freq\": save_steps,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"warmup\": 0.0032,\n",
    "    \"weight_decay\": 0.2,\n",
    "    \"zipped_data\": 1,\n",
    "}\n",
    "\n",
    "if use_fsx:\n",
    "    # make sure to update paths for training_dir and test_dir based on the paths of datasets in fsx\n",
    "    # If you want to resume training, set checkpoint_dir to the same path as a previous job.\n",
    "    SM_TRAIN_DIR = \"/opt/ml/input/data/train\"\n",
    "    hyperparameters[\"checkpoint_dir\"] = f\"{SM_TRAIN_DIR}/smp-v2/{model_type}/checkpointdir\"\n",
    "    hyperparameters[\n",
    "        \"training_dir\"\n",
    "    ] = f\"{SM_TRAIN_DIR}/datasets/pytorch-gpt2-data/pytorch_gpt2/train_synthetic\"\n",
    "    hyperparameters[\n",
    "        \"test_dir\"\n",
    "    ] = f\"{SM_TRAIN_DIR}/datasets/pytorch-gpt2-data/pytorch_gpt2/val_synthetic\"\n",
    "\n",
    "# The checkpoint path (hyperparameters['checkpoint_dir'] or checkpoint_s3_uri) is not unique per job.\n",
    "# You need to modify as needed for different runs.\n",
    "# If same path is used for unrelated runs, this may increase time when downloading unnecessary checkpoints,\n",
    "# and cause conflicts when loading checkpoints.\n",
    "\n",
    "metric_definitions = [\n",
    "    {\"Name\": \"base_metric\", \"Regex\": \"<><><><><><>\"}\n",
    "]  # Add your custom metric definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Specify model config based on model name and model size.\n",
    "\n",
    "model_size = 7  # 7B or 65B parameters.\n",
    "\n",
    "model_configs = {\n",
    "    \"gpt_neox\": {\n",
    "        7: {\n",
    "            \"hidden_width\": 4096,\n",
    "            \"max_context_width\": 1024,\n",
    "            \"num_heads\": 32,\n",
    "            \"num_layers\": 32,\n",
    "        },\n",
    "        65: {\n",
    "            \"hidden_width\": 8192,\n",
    "            \"max_context_width\": 1024,\n",
    "            \"num_heads\": 64,\n",
    "            \"num_layers\": 80,\n",
    "        },\n",
    "    },\n",
    "    \"llama_v2\": {\n",
    "        7: {\n",
    "            \"hidden_width\": 4096,\n",
    "            \"llama_intermediate_size\": 11008,\n",
    "            \"max_context_width\": 1024,\n",
    "            \"num_heads\": 32,\n",
    "            \"num_layers\": 32,\n",
    "        },\n",
    "        65: {\n",
    "            \"hidden_width\": 8192,\n",
    "            \"llama_intermediate_size\": 22016,\n",
    "            \"max_context_width\": 1024,\n",
    "            \"num_heads\": 64,\n",
    "            \"num_layers\": 80,\n",
    "        },\n",
    "    },\n",
    "}\n",
    "\n",
    "model_params = model_configs.get(model_type, {}).get(model_size)\n",
    "if model_params is None:\n",
    "    raise RuntimeError(\n",
    "        f\"Unknown model config from (name, size) = ({model_type}, {model_size:02d}B).\"\n",
    "    )\n",
    "\n",
    "hyperparameters.update(model_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify Essential Parameters for a SageMaker Training Job\n",
    "\n",
    "Next, you use the `SageMaker Estimator class` to define a SageMaker Training Job,\n",
    "passing values through the following parameters for training job name,\n",
    "the number of EC2 instances, the instance type,\n",
    "and the size of the volume attached to the instances.\n",
    "\n",
    "- `instance_count`\n",
    "- `instance_type`\n",
    "- `volume_size`\n",
    "- `base_job_name`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Update the Type and Number of EC2 Instance to Use\n",
    "\n",
    "The instance type and the number of instances you specify to the `instance_type` and `instance_count` parameters,\n",
    "respectively, determine the total number of GPUs (`world_size`).\n",
    "\n",
    "$$\\text{world_size = (the number of GPUs on a single instance)} \\times \\text{(the number of instances)}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Setup instances: type and count.\n",
    "\n",
    "instance_type = \"ml.p4d.24xlarge\"\n",
    "\n",
    "# You need >= 1 p4d for 7b model.\n",
    "# You need >= 8 p4d for 65b model.\n",
    "instance_count = 1\n",
    "\n",
    "# Set to the number of GPUs on that instance.\n",
    "processes_per_host = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify a Base Job Name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Job name.\n",
    "\n",
    "instance_type_str = instance_type.split(\".\")[1] + instance_type.split(\".\")[2][:3]\n",
    "\n",
    "base_job_name = f'smp-{model_size:02d}b-{instance_type_str}-hs{hybrid_shard_degree}-tp{tensor_parallel_degree}-ao{offload_activations}-bs{hyperparameters[\"train_batch_size\"]:02d}'\n",
    "\n",
    "print(f\"Base job name: `{base_job_name}`.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not use_fsx:\n",
    "    # If you want to resume training, set checkpoint_s3_uri to the same path as a previous job.\n",
    "    # Previous checkpoint to load must have same model config.\n",
    "    checkpoint_bucket = f\"s3://sagemaker-{region}-{account}/\"\n",
    "    checkpoint_s3_uri = (\n",
    "        f\"{checkpoint_bucket}/experiments/smp-fsdp-tp-{model_type}-checkpoints/{base_job_name}/\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Estimator.\n",
    "\n",
    "if use_fsx:\n",
    "    # Use the security group and subnet that was used to create the fsx filesystem.\n",
    "    kwargs = {\n",
    "        \"security_group_ids\": [FSX_SECURITY_GROUP_ID],\n",
    "        \"subnets\": [FSX_SUBNET_ID],\n",
    "    }\n",
    "    print(kwargs)\n",
    "else:\n",
    "    kwargs = {}\n",
    "\n",
    "smp_estimator = PyTorch(\n",
    "    entry_point=\"train.py\",\n",
    "    hyperparameters=hyperparameters,\n",
    "    source_dir=os.path.join(os.getcwd(), \"../shared-scripts\"),\n",
    "    role=role,\n",
    "    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,\n",
    "    checkpoint_local_path=hyperparameters[\"checkpoint_dir\"] if use_fsx else None,\n",
    "    instance_type=instance_type,\n",
    "    volume_size=400,\n",
    "    instance_count=instance_count,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    distribution={\n",
    "        \"torch_distributed\": {\n",
    "            \"enabled\": True,\n",
    "        },\n",
    "        \"smdistributed\": {\n",
    "            \"modelparallel\": {\n",
    "                \"enabled\": True,\n",
    "                \"parameters\": {\n",
    "                    \"activation_loading_horizon\": activation_loading_horizon,\n",
    "                    \"hybrid_shard_degree\": hybrid_shard_degree,\n",
    "                    \"sm_activation_offloading\": offload_activations,\n",
    "                    \"tensor_parallel_degree\": tensor_parallel_degree,\n",
    "                },\n",
    "            }\n",
    "        },\n",
    "    },\n",
    "    py_version=\"py311\",\n",
    "    framework_version=\"2.3.1\",\n",
    "    # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
    "    output_path=s3_output_bucket,\n",
    "    max_run=86400,\n",
    "    debugger_hook_config=False,\n",
    "    base_job_name=base_job_name,\n",
    "    metric_definitions=metric_definitions,\n",
    "    **kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start Training\n",
    "\n",
    "Finally, run the `estimator.fit` method to launch the SageMaker training job of the model with hybrid sharding and activation offloading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Start training.\n",
    "\n",
    "smp_estimator.fit(inputs=data_channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resume from Checkpoint\n",
    "You can continue a previous SageMaker training job from a saved checkpoint. Set your checkpoint directory to the appropriate checkpoint you want to resume from and pass in the `checkpoint_s3_uri` from a previous job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Resume training.\n",
    "\n",
    "resume_from_checkpoint = True\n",
    "\n",
    "if resume_from_checkpoint:\n",
    "    hyperparameters[\"resume_from_checkpoint\"] = os.path.join(\n",
    "        hyperparameters[\"checkpoint_dir\"], f\"{model_type}-{save_steps}steps\"\n",
    "    )\n",
    "\n",
    "    smp_estimator = PyTorch(\n",
    "        entry_point=\"train.py\",\n",
    "        hyperparameters=hyperparameters,\n",
    "        source_dir=os.path.join(os.getcwd(), \"../shared-scripts\"),\n",
    "        role=role,\n",
    "        checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,\n",
    "        checkpoint_local_path=hyperparameters[\"checkpoint_dir\"] if use_fsx else None,\n",
    "        instance_type=instance_type,\n",
    "        volume_size=400,\n",
    "        instance_count=instance_count,\n",
    "        sagemaker_session=sagemaker_session,\n",
    "        distribution={\n",
    "            \"torch_distributed\": {\n",
    "                \"enabled\": True,\n",
    "            },\n",
    "            \"smdistributed\": {\n",
    "                \"modelparallel\": {\n",
    "                    \"enabled\": True,\n",
    "                    \"parameters\": {\n",
    "                        \"activation_loading_horizon\": activation_loading_horizon,\n",
    "                        \"hybrid_shard_degree\": hybrid_shard_degree,\n",
    "                        \"sm_activation_offloading\": offload_activations,\n",
    "                        \"tensor_parallel_degree\": tensor_parallel_degree,\n",
    "                    },\n",
    "                }\n",
    "            },\n",
    "        },\n",
    "        py_version=\"py311\",\n",
    "        framework_version=\"2.3.1\",\n",
    "        # image_uri=$IMAGE,  # Either provide `framework_version` or `image_uri`\n",
    "        output_path=s3_output_bucket,\n",
    "        max_run=86400,\n",
    "        debugger_hook_config=False,\n",
    "        base_job_name=f\"resume-{base_job_name}\",\n",
    "        metric_definitions=metric_definitions,\n",
    "        **kwargs,\n",
    "    )\n",
    "\n",
    "    smp_estimator.fit(inputs=data_channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Access the Launched SM Training Job\n",
    "You can access the launched training job from [SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/whatis.html).  \n",
    "Go to `Amazon SageMaker -> Training -> Training jobs`.  \n",
    "You can also access the training logs from here with `View Logs` which opens CloudWatch directly.\n",
    "\n",
    "## Access the Training Logs\n",
    "\n",
    "You can access the training logs from [Amazon CloudWatch](https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/WhatIsCloudWatch.html).\n",
    "\n",
    "You can use CloudWatch to track SageMaker GPU and memory utilization during training and inference. To view the metrics and logs that SageMaker writes to CloudWatch, see [SageMaker Jobs and Endpoint Metrics](https://docs.aws.amazon.com/sagemaker/latest/dg/monitoring-cloudwatch.html#cloudwatch-metrics-jobs) in the Amazon SageMaker Developer Guide.\n",
    "\n",
    "If you are a new user of CloudWatch, see [Getting Started with Amazon CloudWatch](https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/GettingStarted.html).\n",
    "\n",
    "For additional information on monitoring and analyzing Amazon SageMaker training jobs, see [Monitor and Analyze Training Jobs Using Metrics](https://docs.aws.amazon.com/sagemaker/latest/dg/training-metrics.html).\n",
    "\n",
    "## Deploy Trained Model for Inference\n",
    "\n",
    "In most cases, a trained model can be deployed on a single device for inference because inference only requires a small amount of memory.\n",
    "\n",
    "After you build and train your models, you can deploy them to get predictions in one of two ways:\n",
    "\n",
    "* To set up a persistent endpoint to get predictions from your models, use SageMaker hosting services. For an overview on deploying a single model or multiple models with SageMaker hosting services, see [Deploy a Model on SageMaker Hosting Services](https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-deployment.html#how-it-works-hosting).\n",
    "* To get predictions for an entire dataset, use SageMaker batch transform. For an overview on deploying a model with SageMaker Batch Transform, see [Get Inferences for an Entire Dataset with Batch Transform](https://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-batch.html).\n",
    "\n",
    "To learn more about deploying models for inference using SageMaker, see [Deploy Models for Inference](https://docs.aws.amazon.com/sagemaker/latest/dg/deploy-model.html).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/training|distributed_training|pytorch|model_parallel_v2|gpt-neox|smp-train-gpt-neox-fsdp-tp.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}