PyTorch Lightning with Amazon SageMaker

Francesco Pochetti 0
Viewed 39

We’re currently running using Pytorch Lightning for training outside of SageMaker. Looking to use SageMaker to leverage distributed training, checkpointing, model training optimization(training compiler) etc to accelerate training process and save costs. Whats the recommended way to migrate their PyTorch Lightning scripts to run on SageMaker?

amazon-sagemaker
3 Answers

The easiest way to run Pytorch Lightning on SageMaker is to use the SageMaker PyTorch estimator (example) to get started. Ideally you will have add a requirement.txt for installing pytorch lightning along with your source code.

Regarding distributed training Amazon SageMaker recently launched native support for running Pytorch lightning based distributed training. Please follow the below link to setup your training code

https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-modify-sdp-pt-lightning.html

https://aws.amazon.com/blogs/machine-learning/run-pytorch-lightning-and-native-pytorch-ddp-on-amazon-sagemaker-training-featuring-amazon-search/

Arun Lokanatha0

Since your question is specific to migration of already working code into Sagemaker, using the link here as reference, I can try to break the process into 3 parts :

  1. Create a Pytorch Estimator - estimator
  2. entry_point = "my_model.py" - this part should be your existing Pytorch Lightning script. In the main method you can have something like this:
if __name__ ==  '__main__':
     import pytorch_lightning as pl
     trainer = pl.Trainer(max_epochs=20, 
                         devices=-1, ## in order to utilize all GPUs
                         accelerator="gpu", 
                         strategy="ddp", 
                         enable_checkpointing=True, 
                         default_root_dir="/opt/ml/checkpoints",
                         )
  1. model=estimator.fit()
Gaurav Chawla0

There's no big difference in running PyTorch Lightning and plain PyTorch scripts with SageMaker.

One caveat, however, when running distributed training jobs with DDPPlugin, is to set properly the NODE_RANK environment variable at the beginning of the script, because PyTorch Lightning knows nothing about SageMaker environment variables and relies on generic cluster variables:

os.environ["NODE_RANK"] = str(int(os.environ.get("CURRENT_HOST", "algo-1")[5:]) - 1)

or (more robust):

rc = json.loads(os.environ.get("SM_RESOURCE_CONFIG", "{}"))
os.environ["NODE_RANK"] = str(rc["hosts"].index(rc["current_host"]))
juvchan0
Related