      build_and_train_models/sm-model_trainer/model_trainer_overview.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ModelTrainer - SageMaker PySDK Training Redesign\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introductions\n",
    "\n",
    "The `ModelTrainer` class in the SageMaker Python SDK simplifies the process of launching and managing training jobs on AWS SageMaker. It provides an intuitive interface for customizing training jobs, including the ability to specify custom scripts, custom containers, distributed training configurations and executing training locally. In this notebook, we outline how to get started with the ModelTrainer class, its features, and examples to help you effectively leverage its capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Benefits of the ModelTrainer\n",
    "\n",
    "The ModelTrainer is designed to address the usability challenges associated with the Estimator class. Moving training with the SageMaker PySDK towards achieving a best-in-class developer experience.\n",
    "\n",
    "Key Improvements Include:\n",
    "1. **Improved Intuitiveness** - The ModelTrainer reduces complexity by leveraging configuration classes and minimizing the interface to only a few core parameters.\n",
    "1. **Simplified Script Mode and BYOC** - The ModelTrainer natively supports script mode and removes the coupling to the SageMaker Training Toolkit for running a job in Script mode. By removing this runtime dependency, users can bring their own image to launch a training job without a needing to adapt it for script mode on SageMaker.\n",
    "1. **Simplified Distributed Training** - The ModelTrainer provides enhanced flexibility for users to specify custom commands and distributed training configurations by specify the exect commands to execut in thier container using the `command` parameter in the `SourceCode` class or by leveraging a distributed training configuration class like `Torchrun()`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install SageMaker PySDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker \"datasets[s3]\" \"requests<2.32.0\" \"protobuf<3.20\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "print(sagemaker.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelTrainer - Basic Exuction\n",
    "\n",
    "This case example shows a minimal setup for a ModelTrainer. A user need only to provide a desired training image and the commands they wish to execute in the container using the `SourceCode` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules.train import ModelTrainer\n",
    "from sagemaker.modules.configs import SourceCode\n",
    "\n",
    "pytorch_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
    "\n",
    "source_code = SourceCode(command=\"echo 'Hello World'\")\n",
    "\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=pytorch_image,\n",
    "    source_code=source_code,\n",
    "    base_job_name=\"model-trainer-basic-execution\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelTrainer - Script Mode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This examples show cases an abstracted setup for script mode where a user can provide their training image and a `SourceCode` config with path to their `source_dir`, `enty_script`, and any additional `requirements` to install in the training container for their job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules.train import ModelTrainer\n",
    "from sagemaker.modules.configs import SourceCode\n",
    "\n",
    "pytorch_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
    "\n",
    "source_code = SourceCode(\n",
    "    source_dir=\"basic-script-mode\",\n",
    "    requirements=\"requirements.txt\",\n",
    "    entry_script=\"custom_script.py\",\n",
    ")\n",
    "\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=pytorch_image,\n",
    "    source_code=source_code,\n",
    "    base_job_name=\"model-trainer-script-mode\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelTrainer - Local Container Mode\n",
    "\n",
    "This example show cases how a user can leverage the `LOCAL_CONTAINER` mode to run their training job in their local enviornment as docker containers for local experimentation and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules.train.model_trainer import ModelTrainer, Mode\n",
    "from sagemaker.modules.configs import InputData, SourceCode\n",
    "\n",
    "pytorch_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
    "\n",
    "\n",
    "source_code = SourceCode(\n",
    "    source_dir=\"basic-script-mode\",\n",
    "    entry_script=\"local_training_script.py\",\n",
    ")\n",
    "\n",
    "train_data = InputData(\n",
    "    channel_name=\"train\",\n",
    "    data_source=\"basic-script-mode/data/train/\",\n",
    ")\n",
    "\n",
    "test_data = InputData(\n",
    "    channel_name=\"test\",\n",
    "    data_source=\"basic-script-mode/data/test/\",\n",
    ")\n",
    "\n",
    "model_trainer = ModelTrainer(\n",
    "    training_mode=Mode.LOCAL_CONTAINER,\n",
    "    training_image=pytorch_image,\n",
    "    source_code=source_code,\n",
    "    input_data_config=[train_data, test_data],\n",
    "    base_job_name=\"local-container-mode\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Training\n",
    "\n",
    "In this section, we will walk through how the ModelTrainer can be used for more complex Distributed Training jobs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"openlm-research/open_llama_7b\"\n",
    "dataset_name = \"tatsu-lab/alpaca\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# Load Tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "\n",
    "\n",
    "# Load dataset from huggingface.co\n",
    "dataset = load_dataset(dataset_name)\n",
    "\n",
    "dataset = dataset.shuffle(seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"validation\" not in dataset.keys():\n",
    "    dataset[\"validation\"] = load_dataset(dataset_name, split=\"train[:1%]\")\n",
    "\n",
    "    dataset[\"train\"] = load_dataset(dataset_name, split=\"train[1%:]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import chain\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "def group_texts(examples, block_size=2048):\n",
    "    # Concatenate all texts.\n",
    "    concatenated_examples = {k: list(chain(*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\n",
    "\n",
    "\n",
    "column_names = dataset[\"train\"].column_names\n",
    "\n",
    "lm_dataset = dataset.map(\n",
    "    lambda sample: tokenizer(sample[\"text\"], return_token_type_ids=False),\n",
    "    batched=True,\n",
    "    remove_columns=list(column_names),\n",
    ").map(\n",
    "    partial(group_texts, block_size=2048),\n",
    "    batched=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save data locally\n",
    "\n",
    "training_input_path = f\"distributed-training/processed/data/\"\n",
    "lm_dataset.save_to_disk(training_input_path)\n",
    "\n",
    "print(f\"Saved data to: {training_input_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelTrainer - Distributed Training - Explicit Commands\n",
    "\n",
    "This example shows how a user could perform a more complex setup for DistributedTraining using `torchrun` directly through the `command` parameter in the `SourceCode` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules.train import ModelTrainer\n",
    "from sagemaker.modules.configs import Compute, SourceCode, InputData\n",
    "\n",
    "env = {}\n",
    "env[\"FI_PROVIDER\"] = \"efa\"\n",
    "env[\"NCCL_PROTO\"] = \"simple\"\n",
    "env[\"NCCL_SOCKET_IFNAME\"] = \"eth0\"\n",
    "env[\"NCCL_IB_DISABLE\"] = \"1\"\n",
    "env[\"NCCL_DEBUG\"] = \"WARN\"\n",
    "\n",
    "compute = Compute(\n",
    "    instance_count=1,\n",
    "    instance_type=\"ml.p4d.24xlarge\",\n",
    "    volume_size_in_gb=96,\n",
    "    keep_alive_period_in_seconds=3600,\n",
    ")\n",
    "\n",
    "hugging_face_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04\"\n",
    "\n",
    "hyperparameters = {\n",
    "    \"dataset_path\": \"/opt/ml/input/data/dataset\",\n",
    "    \"model_dir\": \"/opt/ml/model\",\n",
    "    \"cache_dir\": None,\n",
    "    \"max_train_steps\": None,\n",
    "    \"num_train_steps\": 1,\n",
    "    \"num_warmup_steps\": 0,\n",
    "    \"num_train_epochs\": 1,\n",
    "    \"forward_prefetch\": False,\n",
    "    \"limit_all_gathers\": False,\n",
    "    \"lr_scheduler_type\": \"linear\",\n",
    "    \"weight_decay\": 0.0,\n",
    "    \"learning_rate\": 5e-5,\n",
    "    \"epochs\": 1,\n",
    "    \"max_steps\": 100,\n",
    "    \"seed\": 42,\n",
    "    \"fsdp\": \"full_shard auto_wrap\",\n",
    "    \"fsdp_transformer_layer_cls_to_wrap\": \"LlamaDecoderLayer\",\n",
    "    \"gradient_checkpointing\": True,\n",
    "    \"gradient_accumulation_steps\": 1,\n",
    "    \"optimizer\": \"adamw_torch\",\n",
    "    \"per_device_train_batch_size\": 1,\n",
    "    \"model_id\": model_id,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_code = SourceCode(\n",
    "    source_dir=\"distributed-training/scripts\",\n",
    "    requirements=\"requirements.txt\",\n",
    "    command=\"torchrun --nnodes 1 \\\n",
    "            --nproc_per_node 8 \\\n",
    "            --master_addr algo-1 \\\n",
    "            --master_port 7777 \\\n",
    "            --node_rank $SM_CURRENT_HOST_RANK \\\n",
    "            run_clm_no_trainer.py\",\n",
    ")\n",
    "\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=hugging_face_image,\n",
    "    compute=compute,\n",
    "    environment=env,\n",
    "    hyperparameters=hyperparameters,\n",
    "    source_code=source_code,\n",
    "    base_job_name=f\"model-trainer-distributed-commands\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = InputData(\n",
    "    channel_name=\"dataset\",\n",
    "    data_source=training_input_path,\n",
    ")\n",
    "model_trainer.train(input_data_config=[test_data], wait=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelTrainer - Distributed Training - Abstraction\n",
    "\n",
    "This examples shows how a user could perform distributed training using an abstracted approach provided via the `Torchrun` distributed training configuration class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules.train import ModelTrainer\n",
    "from sagemaker.modules.configs import Compute, SourceCode, InputData\n",
    "\n",
    "compute = Compute(\n",
    "    instance_count=2,\n",
    "    instance_type=\"ml.p4d.24xlarge\",\n",
    "    volume_size_in_gb=96,\n",
    "    keep_alive_period_in_seconds=3600,\n",
    ")\n",
    "\n",
    "hugging_face_image = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04\"\n",
    "\n",
    "hyperparameters = {\n",
    "    \"dataset_path\": \"/opt/ml/input/data/dataset\",\n",
    "    \"model_dir\": \"/opt/ml/model\",\n",
    "    \"cache_dir\": None,\n",
    "    \"max_train_steps\": None,\n",
    "    \"num_train_steps\": 1,\n",
    "    \"num_warmup_steps\": 0,\n",
    "    \"num_train_epochs\": 1,\n",
    "    \"forward_prefetch\": False,\n",
    "    \"limit_all_gathers\": False,\n",
    "    \"lr_scheduler_type\": \"linear\",\n",
    "    \"weight_decay\": 0.0,\n",
    "    \"learning_rate\": 5e-5,\n",
    "    \"epochs\": 1,\n",
    "    \"max_steps\": 100,\n",
    "    \"seed\": 42,\n",
    "    \"fsdp\": \"full_shard auto_wrap\",\n",
    "    \"fsdp_transformer_layer_cls_to_wrap\": \"LlamaDecoderLayer\",\n",
    "    \"gradient_checkpointing\": True,\n",
    "    \"gradient_accumulation_steps\": 1,\n",
    "    \"optimizer\": \"adamw_torch\",\n",
    "    \"per_device_train_batch_size\": 1,\n",
    "    \"model_id\": model_id,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules.distributed import Torchrun\n",
    "\n",
    "source_code = SourceCode(\n",
    "    source_dir=\"distributed-training/scripts\",\n",
    "    requirements=\"requirements.txt\",\n",
    "    entry_script=\"run_clm_no_trainer.py\",\n",
    ")\n",
    "\n",
    "# Run using Torchrun\n",
    "torchrun = Torchrun()\n",
    "\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=hugging_face_image,\n",
    "    compute=compute,\n",
    "    hyperparameters=hyperparameters,\n",
    "    source_code=source_code,\n",
    "    distributed=torchrun,\n",
    "    base_job_name=f\"model-trainer-distributed-abstraction\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = InputData(\n",
    "    channel_name=\"dataset\",\n",
    "    data_source=training_input_path,\n",
    ")\n",
    "model_trainer.train(input_data_config=[test_data], wait=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelTrainer - SageMaker HyperPod Recipes\n",
    "\n",
    "This example showcases how a user could leverage the pre-defined recipe `fine-tuning/deepseek/hf_deepseek_r1_distilled_llama_70b_seq8k_gpu_fine_tuning` for fine-tuning a DeepSeek R1 Model using synthetic data.\n",
    "\n",
    "For More Available Recipes Visit - https://github.com/aws/sagemaker-hyperpod-recipes/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.modules import Session\n",
    "from sagemaker.modules.train import ModelTrainer\n",
    "from sagemaker.modules.configs import Compute, TensorBoardOutputConfig\n",
    "\n",
    "sagemaker_session = Session()\n",
    "\n",
    "default_bucket = sagemaker_session.default_bucket()\n",
    "default_bucket_prefix = sagemaker_session.default_bucket_prefix\n",
    "default_bucket_prefix_path = \"\"\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    default_bucket_prefix_path = f\"/{default_bucket_prefix}\"\n",
    "\n",
    "recipe_overrides = {\n",
    "    \"run\": {\n",
    "        \"results_dir\": \"/opt/ml/model\",\n",
    "    },\n",
    "    \"exp_manager\": {\n",
    "        \"exp_dir\": \"/opt/ml/output/\",\n",
    "        \"explicit_log_dir\": \"/opt/ml/output/tensorboard\",\n",
    "    },\n",
    "    \"model\": {\n",
    "        \"hf_model_name_or_path\": \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\",\n",
    "        \"data\": {\n",
    "            \"use_synthetic_data\": True,\n",
    "        }\n",
    "    },\n",
    "}\n",
    "\n",
    "compute = Compute(\n",
    "    instance_type=\"ml.p5.48xlarge\",\n",
    "    keep_alive_period_in_seconds=3600,\n",
    ")\n",
    "\n",
    "tensorboad_output_config = TensorBoardOutputConfig(\n",
    "    s3_output_path=f\"s3://{default_bucket}{default_bucket_prefix_path}/output/tensorboard\",\n",
    "    local_path=\"/opt/ml/output/tensorboard\",\n",
    ")\n",
    "\n",
    "smp_image = (\n",
    "    \"658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121\"\n",
    ")\n",
    "\n",
    "model_trainer = ModelTrainer.from_recipe(\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    training_image=smp_image,\n",
    "    training_recipe=\"fine-tuning/deepseek/hf_deepseek_r1_distilled_llama_8b_seq8k_gpu_fine_tuning\",\n",
    "    recipe_overrides=recipe_overrides,\n",
    "    compute=compute,\n",
    "    base_job_name=f\"model-trainer-deepseek-p5-8b\",\n",
    ").with_tensorboard_output_config(tensorboad_output_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer.train(wait=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}