philschmid blog

Distributed Training: Train BART/T5 for Summarization using 🤗 Transformers and Amazon SageMaker

#AWS #HuggingFace #BART
, April 09, 2021 · 10 min read

Photo by Michael Bradley on Unsplash

In case you missed it: on March 25th we announced a collaboration with Amazon SageMaker to make it easier to create State-of-the-Art Machine Learning models, and ship cutting-edge NLP features faster.

Together with the SageMaker team, we built 🤗 Transformers optimized Deep Learning Containers to accelerate training of Transformers-based models. Thanks AWS friends!🤗 🚀

With the new HuggingFace estimator in the SageMaker Python SDK, you can start training with a single line of code.


The announcement blog post provides all the information you need to know about the integration, including a “Getting Started” example and links to documentation, examples, and features.

listed again here:

If you’re not familiar with Amazon SageMaker: “Amazon SageMaker is a fully managed service that provides every developer and data scientist with the ability to build, train, and deploy machine learning (ML) models quickly. SageMaker removes the heavy lifting from each step of the machine learning process to make it easier to develop high quality models.” [REF]


We will use the new Hugging Face DLCs and Amazon SageMaker extension to train a distributed Seq2Seq-transformer model on the summarization task using the transformers and datasets libraries, and then upload the model to and test it.

As distributed training strategy we are going to use SageMaker Data Parallelism, which has been built into the Trainer API. To use data-parallelism we only have to define the distribution parameter in our HuggingFace estimator.

1 # configuration for running training on smdistributed Data Parallel
2 distribution = {'smdistributed':{'dataparallel':{ 'enabled': True }}}

In this tutorial, we will use an Amazon SageMaker Notebook Instance for running our training job. You can learn here how to set up a Notebook Instance.

What are we going to do:

  • Set up a development environment and install sagemaker
  • Choose 🤗 Transformers examples/ script
  • Configure distributed training and hyperparameters
  • Create a HuggingFace estimator and start training
  • Upload the fine-tuned model to
  • Test inference

Model and Dataset

We are going to fine-tune facebook/bart-large-cnn on the samsum dataset. “BART is sequence-to-sequence model trained with denoising as pretraining objective.” [REF]

The samsum dataset contains about 16k messenger-like conversations with summaries.

1 {"id": "13818513",
2 "summary": "Amanda baked cookies and will bring Jerry some tomorrow.",
3 "dialogue": "Amanda: I baked cookies. Do you want some?\r\nJerry: Sure!\r\nAmanda: I'll bring you tomorrow :-)"}

Set up a development environment and install sagemaker

After our SageMaker Notebook Instance is running we can select either Jupyer Notebook or JupyterLab and create a new Notebook with the conda_pytorch_p36 kernel.

Note: The use of Jupyter is optional: We could also launch SageMaker Training jobs from anywhere we have an SDK installed, connectivity to the cloud and appropriate permissions, such as a Laptop, another IDE or a task scheduler like Airflow or AWS Step Functions.

After that we can install the required dependencies

1 !pip install transformers "datasets[s3]" sagemaker --upgrade

install git-lfs for model upload.

1 !curl -s | sudo bash
2 !sudo yum install git-lfs -y
3 !git lfs install

To run training on SageMaker we need to create a sagemaker Session and provide an IAM role with the right permission. This IAM role will be later attached to the TrainingJob enabling it to download data, e.g. from Amazon S3.

1 import sagemaker
3 sess = sagemaker.Session()
4 role = sagemaker.get_execution_role()
6 print(f"IAM role arn used for running training: {role}")
7 print(f"S3 bucket used for storing artifacts: {sess.default_bucket()}")

Choose 🤗 Transformers examples/ script

The 🤗 Transformers repository contains several examples/scripts for fine-tuning models on tasks from language-modeling to token-classification. In our case, we are using the from the seq2seq/ examples.

*Note: you can use this tutorial as-is to train your model on a different examples script.*

Since the HuggingFace Estimator has git support built-in, we can specify a training script stored in a GitHub repository as entry_point and source_dir.

We are going to use the transformers 4.4.2 DLC which means we need to configure the v4.4.2 as the branch to pull the compatible example scripts.

1 #git_config = {'repo': '','branch': 'v4.4.2'} # v4.4.2 is referring to the `transformers_version you use in the estimator.
2 # used due an missing package in v4.4.2
3 git_config = {'repo': '','branch': 'master'} # v4.4.2 is referring to the `transformers_version you use in the estimator.

Configure distributed training and hyperparameters

Next, we will define our hyperparameters and configure our distributed training strategy. As hyperparameter, we can define any Seq2SeqTrainingArguments and the ones defined in

1 # hyperparameters, which are passed into the training job
2 hyperparameters={
3 'per_device_train_batch_size': 4,
4 'per_device_eval_batch_size': 4,
5 'model_name_or_path':'facebook/bart-large-cnn',
6 'dataset_name':'samsum',
7 'do_train':True,
8 'do_predict': True,
9 'predict_with_generate': True,
10 'output_dir':'/opt/ml/model',
11 'num_train_epochs': 3,
12 'learning_rate': 5e-5,
13 'seed': 7,
14 'fp16': True,
15 }
17 # configuration for running training on smdistributed Data Parallel
18 distribution = {'smdistributed':{'dataparallel':{ 'enabled': True }}}

Since, we are using SageMaker Data Parallelism our total_batch_size will be per_device_train_batch_size * n_gpus.

Create a HuggingFace estimator and start training

The last step before training is creating a HuggingFace estimator. The Estimator handles the end-to-end Amazon SageMaker training. We define which fine-tuning script should be used as entry_point, which instance_type should be used, and which hyperparameters are passed in.

1 from sagemaker.huggingface import HuggingFace
3 # create the Estimator
4 huggingface_estimator = HuggingFace(
5 entry_point='', # script
6 source_dir='./examples/seq2seq', # relative path to example
7 git_config=git_config,
8 instance_type='ml.p3dn.24xlarge',
9 instance_count=2,
10 transformers_version='4.4.2',
11 pytorch_version='1.6.0',
12 py_version='py36',
13 role=role,
14 hyperparameters = hyperparameters,
15 distribution = distribution
16 )

As instance_type we are using ml.p3dn.24xlarge, which contains 8x NVIDIA A100 with an instance_count of 2. This means we are going to run training on 16 GPUs and a total_batch_size of 16*4=64. We are going to train a 400 Million Parameter model with a total_batch_size of 64, which is just wow. To start our training we call the .fit() method.

1 # starting the training job
1 2021-04-01 13:00:35 Starting - Starting the training job...
2 2021-04-01 13:01:03 Starting - Launching requested ML instancesProfilerReport-1617282031: InProgress
3 2021-04-01 13:02:23 Starting - Preparing the instances for training......
4 2021-04-01 13:03:25 Downloading - Downloading input data...
5 2021-04-01 13:04:04 Training - Downloading the training image...............
6 2021-04-01 13:06:33 Training - Training image download completed. Training in progress
7 ....
8 ....
9 2021-04-01 13:16:47 Uploading - Uploading generated training model
10 2021-04-01 13:27:49 Completed - Training job completed
11 Training seconds: 2882
12 Billable seconds: 2882

The training seconds are 2882 because they are multiplied by the number of instances. If we calculate 2882/2=1441 is it the duration from “Downloading the training image” to “Training job completed”. Converted to real money, our training on 16 NVIDIA Tesla V100-GPU for a State-of-the-Art summarization model comes down to ~28$.

Upload the fine-tuned model to

Since our model achieved a pretty good score we are going to upload it to, create a model_card and test it with the Hosted Inference widget. To upload a model you need to create an account here.

We can download our model from Amazon S3 and unzip it using the following snippet.

1 import os
2 import tarfile
3 from sagemaker.s3 import S3Downloader
5 local_path = 'my_bart_model'
7 os.makedirs(local_path, exist_ok = True)
9 # download model from S3
11 s3_uri=huggingface_estimator.model_data, # s3 uri where the trained model is located
12 local_path=local_path, # local path where *.tar.gz will be saved
13 sagemaker_session=sess # sagemaker session used for training the model
14 )
16 # unzip model
17 tar ="{local_path}/model.tar.gz", "r:gz")
18 tar.extractall(path=local_path)
19 tar.close()
20 os.remove(f"{local_path}/model.tar.gz")

Before we are going to upload our model to we need to create a model_card. The model_card describes the model and includes hyperparameters, results, and specifies which dataset was used for training. To create a model_card we create a in our local_path

1 # read eval and test results
2 with open(f"{local_path}/eval_results.json") as f:
3 eval_results_raw = json.load(f)
4 eval_results={}
5 eval_results["eval_rouge1"] = eval_results_raw["eval_rouge1"]
6 eval_results["eval_rouge2"] = eval_results_raw["eval_rouge2"]
7 eval_results["eval_rougeL"] = eval_results_raw["eval_rougeL"]
8 eval_results["eval_rougeLsum"] = eval_results_raw["eval_rougeLsum"]
10 with open(f"{local_path}/test_results.json") as f:
11 test_results_raw = json.load(f)
12 test_results={}
13 test_results["test_rouge1"] = test_results_raw["test_rouge1"]
14 test_results["test_rouge2"] = test_results_raw["test_rouge2"]
15 test_results["test_rougeL"] = test_results_raw["test_rougeL"]
16 test_results["test_rougeLsum"] = test_results_raw["test_rougeLsum"]

After we extract all the metrics we want to include we are going to create our Additionally to the automated generation of the results table we add the metrics manually to the metadata of our model card under model-index

1 import json
4 ---
5 language: en
6 tags:
7 - sagemaker
8 - bart
9 - summarization
10 license: apache-2.0
11 datasets:
12 - samsum
13 model-index:
14 - name: {model_name}
15 results:
16 - task:
17 name: Abstractive Text Summarization
18 type: abstractive-text-summarization
19 dataset:
20 name: "SAMSum Corpus: A Human-annotated Dialogue Dataset for Abstractive Summarization"
21 type: samsum
22 metrics:
23 - name: Validation ROGUE-1
24 type: rogue-1
25 value: 42.621
26 - name: Validation ROGUE-2
27 type: rogue-2
28 value: 21.9825
29 - name: Validation ROGUE-L
30 type: rogue-l
31 value: 33.034
32 - name: Test ROGUE-1
33 type: rogue-1
34 value: 41.3174
35 - name: Test ROGUE-2
36 type: rogue-2
37 value: 20.8716
38 - name: Test ROGUE-L
39 type: rogue-l
40 value: 32.1337
41 widget:
42 - text: |
43 Jeff: Can I train a 🤗 Transformers model on Amazon SageMaker?
44 Philipp: Sure you can use the new Hugging Face Deep Learning Container.
45 Jeff: ok.
46 Jeff: and how can I get started?
47 Jeff: where can I find documentation?
48 Philipp: ok, ok you can find everything here.
49 ---
51 ## `{model_name}`
53 This model was trained using Amazon SageMaker and the new Hugging Face Deep Learning container.
55 For more information look at:
56 - [🤗 Transformers Documentation: Amazon SageMaker](
57 - [Example Notebooks](
58 - [Amazon SageMaker documentation for Hugging Face](
59 - [Python SDK SageMaker documentation for Hugging Face](
60 - [Deep Learning Container](
62 ## Hyperparameters
64 {hyperparameters}
67 ## Usage
68 from transformers import pipeline
69 summarizer = pipeline("summarization", model="philschmid/{model_name}")
71 conversation = '''Jeff: Can I train a 🤗 Transformers model on Amazon SageMaker?
72 Philipp: Sure you can use the new Hugging Face Deep Learning Container.
73 Jeff: ok.
74 Jeff: and how can I get started?
75 Jeff: where can I find documentation?
76 Philipp: ok, ok you can find everything here.
77 '''
78 nlp(conversation)
80 ## Results
82 | key | value |
83 | --- | ----- |
84 {eval_table}
85 {test_table}
89 """
91 # Generate model card (todo: add more data from Trainer)
92 model_card = MODEL_CARD_TEMPLATE.format(
93 model_name=f"{hyperparameters['model_name_or_path'].split('/')[1]}-{hyperparameters['dataset_name']}",
94 hyperparameters=json.dumps(hyperparameters, indent=4, sort_keys=True),
95 eval_table="\n".join(f"| {k} | {v} |" for k, v in eval_results.items()),
96 test_table="\n".join(f"| {k} | {v} |" for k, v in test_results.items()),
97 )
99 with open(f"{local_path}/", "w") as f:
100 f.write(model_card)

After we have our unzipped model and model card located in my_bart_model we can use the either huggingface_hub SDK to create a repository and upload it to – or just to an create a new repository and upload it.

1 from getpass import getpass
2 from huggingface_hub import HfApi, Repository
4 hf_username = "philschmid" # your username on
5 hf_email = "" # email used for commit
6 repository_name = f"{hyperparameters['model_name_or_path'].split('/')[1]}-{hyperparameters['dataset_name']}" # repository name on
7 password = getpass("Enter your password:") # creates a prompt for entering password
9 # get hf token
10 token = HfApi().login(username=hf_username, password=password)
12 # create repository
13 repo_url = HfApi().create_repo(token=token, name=repository_name, exist_ok=True)
15 # create a Repository instance
16 model_repo = Repository(use_auth_token=token,
17 clone_from=repo_url,
18 local_dir=local_path,
19 git_user=hf_username,
20 git_email=hf_email)
22 # push model to the hub
23 model_repo.push_to_hub()

Test inference

After we uploaded our model we can access it at{hf_username}/{repository_name}

1 print(f"{hf_username}/{repository_name}")

And use the “Hosted Inference API” widget to test it.


You can find the code here. Feel free to contact us or the forum.

Thanks for reading. If you have any questions, feel free to contact me, through Github, or on the forum. You can also connect with me on Twitter or LinkedIn.