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

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:

• 🤗 Transformers Documentation: Amazon SageMaker
• Example Notebooks
• Amazon SageMaker documentation for Hugging Face
• Python SDK SageMaker documentation for Hugging Face
• Deep Learning Container

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]

Tutorial

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 huggingface.co 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.

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

What are we going to do:

Model and Dataset

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