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.

Takeaways

1. 1

   SageMaker Data Parallelism can be enabled in the HuggingFace estimator by setting the distribution parameter dictionary to enable smdistributed dataparallel.

2. 2

   The HuggingFace estimator includes built-in Git support, allowing training scripts and source directories to be pulled directly from specified repository branches.

3. 3

   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.