azureml_multinode_lightning.py

import os
from azureml.core import ScriptRunConfig, Experiment
from azureml.core.runconfig import MpiConfiguration

cluster = ws.compute_targets[cluster_name]

src = ScriptRunConfig(
    source_directory=source_dir,
    script=script_name,
    arguments=["--max_epochs", 5, "--gpus", 2,  "--num_nodes", 2, "--accelerator", "ddp"],
    compute_target=cluster,
    environment=env,
    distributed_job_config=MpiConfiguration(process_count_per_node=1, node_count=2),
)

run = Experiment(ws, experiment_name).submit(src)
run