Add proper support for distributed training

[tizianocitro]

All submissions

• Have you followed the guidelines in our Contributing document?
• Have you checked to ensure there are no other open Pull Requests for the same changes?
• Have you made sure you have rebased your branch to the latest commit on the target branch?

Description

• Add proper support for distributed training.

Checklist

• Does your submission pass all tests? (use make test)
• Have you written tests to cover all your changes? If not, provide a reason.
• Have you lint your code locally before submission? (use make format)
• Have you type checked your code locally before submission? (use make typecheck)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!