Multi GPU training on g5.12xlarge

[H3zi]

Hi all,
I wonder if there's a built-in solution in SageMaker to run distributed training on a g5.12xlarge (4xGPUs)?
From what I can see, SMDDP does not support g5.12xlarge instances, so specifying distributed config in the estimator will not work.

I can hack my way using pytorch\accelerate but I wonder why there's no straightforward way to do it with the SM SDK?

Thanks

[ruhanprasad]

Hello, I would suggest trying the torch_distributed distribution in the estimator. This uses torchrun under the hood and should work on any multi-gpu instance type. And you can make sure to set the distributed backend to "nccl" in your training script. This will leverage DDP training without the instance type constraints of SMDDP

distribution={
        "torch_distributed": {
            "enabled": True
        }
    }