test_trainer_distributed.py
test_deepspeed.py

To jump right into the execution point, search for the execute_subprocess_async call in those tests.
You will need at least 2 GPUs to see these tests in action: