Megatron-DeepSpeed on AzureML

Example script for running Megatron-DeepSpeed using Azure Machine Learning.

Workspace Setup

Setup an AML workspace. Refer to: set-up doc.

Dataset Preparation

Create AML Dataset. To run remote AML job, you need to provide AML FileDataset. Refer to prepare_dataset script to upload .bin and .idx files to blob store and on how to create FileDataset.

Note: The folder bookcorpus_data used by prepare_dataset script should not be under azureml directories. It is because Azure ML does not allow to include large files (limit: 100 files or 1048576 bytes) for Docker build context.

Training

Run Megatron-DeepSpeed on Azure ML. Refer to aml_submit script.