[huggingface_pytorch] Training - DLC for Transformers to 4.46.1 - Accelerate 1.1.0 - PyTorch 2.3

huggingface/pytorch/training/buildspec-2-1.yml

@@ -0,0 +1,32 @@
+account_id: &ACCOUNT_ID <set-$ACCOUNT_ID-in-environment>
+region: &REGION <set-$REGION-in-environment>
+base_framework: &BASE_FRAMEWORK pytorch
+framework: &FRAMEWORK !join [ "huggingface_", *BASE_FRAMEWORK]
+version: &VERSION 2.1.0
+short_version: &SHORT_VERSION "2.1"
+contributor: huggingface
+arch_type: x86
+
+repository_info:
+  training_repository: &TRAINING_REPOSITORY
+    image_type: &TRAINING_IMAGE_TYPE training
+    root: !join [ "huggingface/", *BASE_FRAMEWORK, "/", *TRAINING_IMAGE_TYPE ]
+    repository_name: &REPOSITORY_NAME !join ["pr", "-", "huggingface", "-", *BASE_FRAMEWORK, "-", *TRAINING_IMAGE_TYPE]
+    repository: &REPOSITORY !join [ *ACCOUNT_ID, .dkr.ecr., *REGION, .amazonaws.com/, *REPOSITORY_NAME ]
+
+images:
+  BuildHuggingFacePytorchGpuPy310Cu121TrainingDockerImage:
+    <<: *TRAINING_REPOSITORY
+    build: &HUGGINGFACE_PYTORCH_GPU_TRAINING_PY3 false
+    image_size_baseline: &IMAGE_SIZE_BASELINE 21500
+    device_type: &DEVICE_TYPE gpu
+    python_version: &DOCKER_PYTHON_VERSION py3
+    tag_python_version: &TAG_PYTHON_VERSION py310
+    cuda_version: &CUDA_VERSION cu121
+    os_version: &OS_VERSION ubuntu20.04
+    transformers_version: &TRANSFORMERS_VERSION 4.36.0
+    datasets_version: &DATASETS_VERSION 2.18.0
+    tag: !join [ *VERSION, '-', 'transformers', *TRANSFORMERS_VERSION, '-', *DEVICE_TYPE, '-', *TAG_PYTHON_VERSION, '-',
+                 *CUDA_VERSION, '-', *OS_VERSION ]
+    docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /,
+                         *CUDA_VERSION, /Dockerfile., *DEVICE_TYPE ]

huggingface/pytorch/training/buildspec.yml

@@ -2,8 +2,8 @@ account_id: &ACCOUNT_ID <set-$ACCOUNT_ID-in-environment>
 region: &REGION <set-$REGION-in-environment>
 base_framework: &BASE_FRAMEWORK pytorch
 framework: &FRAMEWORK !join [ "huggingface_", *BASE_FRAMEWORK]
-version: &VERSION 2.1.0
-short_version: &SHORT_VERSION "2.1"
+version: &VERSION 2.3.0
+short_version: &SHORT_VERSION "2.3"
 contributor: huggingface
 arch_type: x86
 
@@ -15,17 +15,17 @@ repository_info:
     repository: &REPOSITORY !join [ *ACCOUNT_ID, .dkr.ecr., *REGION, .amazonaws.com/, *REPOSITORY_NAME ]
 
 images:
-  BuildHuggingFacePytorchGpuPy310Cu121TrainingDockerImage:
+  BuildHuggingFacePytorchGpuPy311Cu121TrainingDockerImage:
     <<: *TRAINING_REPOSITORY
     build: &HUGGINGFACE_PYTORCH_GPU_TRAINING_PY3 false
     image_size_baseline: &IMAGE_SIZE_BASELINE 21500
     device_type: &DEVICE_TYPE gpu
     python_version: &DOCKER_PYTHON_VERSION py3
-    tag_python_version: &TAG_PYTHON_VERSION py310
+    tag_python_version: &TAG_PYTHON_VERSION py311
     cuda_version: &CUDA_VERSION cu121
     os_version: &OS_VERSION ubuntu20.04
-    transformers_version: &TRANSFORMERS_VERSION 4.36.0
-    datasets_version: &DATASETS_VERSION 2.18.0
+    transformers_version: &TRANSFORMERS_VERSION 4.46.1
+    datasets_version: &DATASETS_VERSION 3.1.0
     tag: !join [ *VERSION, '-', 'transformers', *TRANSFORMERS_VERSION, '-', *DEVICE_TYPE, '-', *TAG_PYTHON_VERSION, '-',
                  *CUDA_VERSION, '-', *OS_VERSION ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /,

huggingface/pytorch/training/docker/2.3/py3/cu121/Dockerfile.gpu

@@ -0,0 +1,63 @@
+# https://github.com/aws/deep-learning-containers/blob/master/available_images.md
+# refer to the above page to pull latest Pytorch image
+
+# docker image region us-west-2
+FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.3.0-gpu-py311-cu121-ubuntu20.04-sagemaker
+
+LABEL maintainer="Amazon AI"
+LABEL dlc_major_version="1"
+
+# version args
+ARG TRANSFORMERS_VERSION
+ARG DATASETS_VERSION
+ARG HUGGINGFACE_HUB_VERSION=0.25.1
+ARG DIFFUSERS_VERSION=0.31.0
+ARG EVALUATE_VERSION=0.4.3
+ARG ACCELERATE_VERSION=1.1.0
+ARG TRL_VERSION=0.11.4
+ARG PEFT_VERSION=0.13.2
+ARG FLASH_ATTN_VERSION=2.6.3
+ARG PYTHON=python3
+
+# TODO: Remove when the base image is updated
+RUN pip install --upgrade "pip==24.1.2" \
+ && pip uninstall -y transformer-engine flash-attn pyarrow cryptography \
+ && pip install --no-cache-dir -U pyarrow cryptography pyopenssl Pillow \
+ && pip --no-cache-dir install --upgrade wheel setuptools \
+ && pip install --no-cache-dir -U "werkzeug==3.0.6" 
+
+# install Hugging Face libraries and its dependencies
+RUN pip install --no-cache-dir \
+	# hf_transfer will be a built-in feature, remove the extra then
+    huggingface_hub[hf_transfer]==${HUGGINGFACE_HUB_VERSION} \
+	transformers[sklearn,sentencepiece,audio,vision,pipelines]==${TRANSFORMERS_VERSION} \
+	datasets==${DATASETS_VERSION} \
+	diffusers==${DIFFUSERS_VERSION} \
+	Jinja2 \
+	tensorboard \
+	bitsandbytes \
+	evaluate==${EVALUATE_VERSION} \
+	accelerate==${ACCELERATE_VERSION} \
+	trl==${TRL_VERSION} \
+	peft==${PEFT_VERSION} \
+	flash-attn==${FLASH_ATTN_VERSION}
+
+# hf_transfer will be a built-in feature, remove the env variavle then
+ENV HF_HUB_ENABLE_HF_TRANSFER="1"
+
+RUN apt-get update \
+ # TODO: Remove upgrade statements once packages are updated in base image
+ && apt-get -y upgrade --only-upgrade systemd openssl cryptsetup libkrb5-3 \
+ && apt install -y git git-lfs \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*
+
+RUN HOME_DIR=/root \
+ && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
+ && unzip -o ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
+ && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
+ && chmod +x /usr/local/bin/testOSSCompliance \
+ && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
+ && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
+ && rm -rf ${HOME_DIR}/oss_compliance*
+

test/sagemaker_tests/huggingface_pytorch/training/integration/local/test_single_machine_training.py

@@ -34,8 +34,9 @@ def test_distilbert_base(
     # hyperparameters, which are passed into the training job
     hyperparameters = {
         "max_steps": 5,
+        "warmup_steps": 3,
         "train_batch_size": 4,
-        "model_name": "distilbert-base-uncased",
+        "model_name": "distilbert/distilbert-base-uncased",
     }
 
     estimator = HuggingFace(
@@ -67,8 +68,9 @@ def test_distilbert_base_torch_compiled(
     # hyperparameters, which are passed into the training job
     hyperparameters = {
         "max_steps": 5,
+        "warmup_steps": 3,
         "train_batch_size": 4,
-        "model_name": "distilbert-base-uncased",
+        "model_name": "distilbert/distilbert-base-uncased",
     }
 
     estimator = HuggingFace(

test/sagemaker_tests/huggingface_pytorch/training/integration/local/test_trcomp.py

@@ -39,7 +39,7 @@ def test_single_node_single_gpu_tcc_default(
     hyperparameters = {
         "max_steps": 3,
         "train_batch_size": 4,
-        "model_name": "distilbert-base-uncased",
+        "model_name": "distilbert/distilbert-base-uncased",
     }
 
     estimator = HuggingFace(

test/sagemaker_tests/huggingface_pytorch/training/integration/sagemaker/test_trcomp.py

@@ -33,7 +33,7 @@
 
 
 hyperparameters = {
-    "model_name_or_path": "bert-large-uncased-whole-word-masking",
+    "model_name_or_path": "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad",
     "dataset_name": "squad",
     "do_train": True,
     "do_eval": True,

test/sagemaker_tests/huggingface_pytorch/training/resources/scripts/train.py

@@ -20,6 +20,7 @@
     parser.add_argument("--train-batch-size", type=int, default=32)
     parser.add_argument("--eval-batch-size", type=int, default=64)
     parser.add_argument("--warmup_steps", type=int, default=500)
+    parser.add_argument("--max_steps", type=int, default=5000)
     parser.add_argument("--model_name", type=str)
     parser.add_argument("--learning_rate", type=str, default=5e-5)
 
@@ -44,15 +45,15 @@
     tokenizer = AutoTokenizer.from_pretrained(args.model_name)
 
     # load dataset
-    dataset = load_dataset("imdb")
+    dataset = load_dataset("stanfordnlp/imdb")
     metric = evaluate.load("accuracy")
 
     # tokenizer helper function
     def tokenize(batch):
-        return tokenizer(batch["text"], padding="max_length", truncation=True)
+        return tokenizer(batch["text"], padding="max_length", max_length=512, truncation=True)
 
     # load dataset
-    train_dataset, test_dataset = load_dataset("imdb", split=["train", "test"])
+    train_dataset, test_dataset = load_dataset("stanfordnlp/imdb", split=["train", "test"])
     test_dataset = test_dataset.shuffle().select(
         range(100)
     )  # smaller the size for test dataset to 10k
@@ -84,7 +85,8 @@ def compute_metrics(eval_pred):
         per_device_train_batch_size=args.train_batch_size,
         per_device_eval_batch_size=args.eval_batch_size,
         warmup_steps=args.warmup_steps,
-        evaluation_strategy="epoch",
+        max_steps=args.max_steps,
+        eval_strategy="epoch",
         logging_dir=f"{args.output_data_dir}/logs",
         learning_rate=float(args.learning_rate),
     )

test/sagemaker_tests/huggingface_pytorch/training/resources/scripts/train_torch_compiled.py

@@ -20,6 +20,7 @@
     parser.add_argument("--train-batch-size", type=int, default=32)
     parser.add_argument("--eval-batch-size", type=int, default=64)
     parser.add_argument("--warmup_steps", type=int, default=500)
+    parser.add_argument("--max_steps", type=int, default=5000)
     parser.add_argument("--model_name", type=str)
     parser.add_argument("--learning_rate", type=str, default=5e-5)
 
@@ -52,7 +53,7 @@
 
     # tokenizer helper function
     def tokenize(batch):
-        return tokenizer(batch["text"], padding="max_length", truncation=True)
+        return tokenizer(batch["text"], padding="max_length", max_length=512, truncation=True)
 
     # load dataset
     train_dataset, test_dataset = load_dataset("imdb", split=["train", "test"])
@@ -87,7 +88,8 @@ def compute_metrics(eval_pred):
         per_device_train_batch_size=args.train_batch_size,
         per_device_eval_batch_size=args.eval_batch_size,
         warmup_steps=args.warmup_steps,
-        evaluation_strategy="epoch",
+        max_steps=args.max_steps,
+        eval_strategy="epoch",
         logging_dir=f"{args.output_data_dir}/logs",
         learning_rate=float(args.learning_rate),
         torch_compile=not args.disable_torch_compile,  # optimizations