ML Katas

Set up a distributed data parallel training script using torch.nn.parallel.DistributedDataParallel and torch.distributed. You'll need to use torch.multiprocessing.spawn to launch multiple processes, with each process handling a replica of the model on a different GPU (or CPU for testing). This improves training speed by distributing the workload.

Verification: Ensure that the training script runs on multiple devices. The simplest way to check is to print the rank of each process. The total time for training should be noticeably faster than single-GPU training, and the final model weights should be identical across all replicas at the end of each training step due to the gradient synchronization.