Add tests to cover multi-gpu training #1273
Description
This is an internal note, but I don't think any of our tests cover multi-gpu, i've noticed some instability. This cannot run on github actions so it'll have to be something we invoke before a release.