Distributed and testing tweaks

[jlamypoirier]

✨ Description

Changes extracted from #428 that do not directly relate to pipeline RL, most implortantly:

• Generalize distributed dim creation to handle non-trivial strides, ex. tensor_and_sequence_data with pipeline_first and model_and_sequence_data.
• Improve distributed test utils to make usage much simpler. The new version handles the boilerplate for running multiple "subtests" within one distributed run, and uses torch.multiprocessing.spawn instead of a subprocess call, which eliminates the need for independent scripts.