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

Summary

Hugging Face and Amazon SageMaker introduced optimized Deep Learning Containers alongside a SageMaker Python SDK estimator to streamline distributed transformer model training. This workflow demonstrates fine-tuning the 400-million-parameter facebook/bart-large-cnn model for dialogue summarization on the 16,000-conversation samsum dataset. Using SageMaker Data Parallelism enabled in the distribution configuration, training executes across two ml.p3dn.24xlarge compute instances equipped with sixteen total GPUs and a cumulative batch size of 64. The HuggingFace estimator pulls the run_summarization.py script directly from GitHub, enables mixed precision with fp16, and tracks evaluation performance through ROUGE scores. Following completion, the unzipped model artifacts and a generated model card are pushed to the Hugging Face Hub using the huggingface_hub library for hosted inference.

Context

Hugging Face and Amazon SageMaker collaborated to create optimized Deep Learning Containers and an estimator in the SageMaker Python SDK to accelerate and simplify distributed training of Transformers-based models.

Approach / What changed

The workflow configures a SageMaker HuggingFace estimator with smdistributed data parallelism enabled to fine-tune facebook/bart-large-cnn on the samsum dataset using run_summarization.py across two ml.p3dn.24xlarge instances, then exports and uploads the model artifacts and model card to the Hugging Face Hub.

Takeaways

• SageMaker Data Parallelism can be enabled in the HuggingFace estimator by setting the distribution parameter dictionary to enable smdistributed dataparallel.
• The HuggingFace estimator includes built-in Git support, allowing training scripts and source directories to be pulled directly from specified repository branches.
• Training with a per-device batch size of 4 across two ml.p3dn.24xlarge instances (16 GPUs total) yields a total effective batch size of 64.

Related reading

• huggingface.co · Mar 23, 2021

  The Partnership: Amazon SageMaker and Hugging Face

  Hugging Face and Amazon formed a strategic partnership establishing Amazon Web Services as Hugging Face's preferred cloud provider to simplify training and deploying Transformer models. To facilitate this collaboration, the teams introduced Hugging Face Deep Learning Containers optimized for PyTorch and TensorFlow in Amazon SageMaker. The integration provides an extension to the SageMaker Python SDK, enabling data science teams to configure distributed training jobs across single-node and multi-node GPU clusters. Workflows integrate directly with SageMaker Automatic Model Tuning, SageMaker Studio, and Amazon S3 for managed experiment tracking and artifact storage. The open-source containers allow trained models to be exported for inference via SageMaker or published directly to the Hugging Face Model Hub while utilizing EC2 Spot instances to reduce compute expenses.

  Philipp Schmid

• huggingface.co · Jul 8, 2021

  Deploy Hugging Face models easily with Amazon SageMaker

  Hugging Face and Amazon introduced an inference solution integrating Hugging Face Transformers with Amazon SageMaker. The SageMaker Hugging Face Inference Toolkit and specialized Deep Learning Containers enable deployment of trained models or publicly available Hub models to managed production endpoints. Users configure endpoints with minimal code using the SageMaker Python SDK, standard pipelines, or custom inference scripts that override default behaviors. The architecture supports models stored in Amazon S3 as well as direct references via HF_MODEL_ID and HF_TASK environment variables. Deployments benefit from native AWS infrastructure features, including built-in monitoring, Identity and Access Management permission controls, and Virtual Private Cloud connectivity.