Scaling vision example 2 training to 8xH100

[dougfulop]

When I try to run the vision example 2 with all 8 devices enabled I end up with a size mismatch on rng.

The issues does not immediately to be solvable through increasing the number of environments and the batch size.

Is this a known issue or just a parameter I need to tweak? Training succeeds when the device limit is set to 1.

[btaba]

I'm not sure multi-gpu with madrona-mjx is supported, @StafaH ?

[StafaH]

Currently there is no support for multi-gpu for madrona-mjx. We will keep this feature request in mind as we continue development on batch renderers for MJX.

[dougfulop]

Were all of the vision experiments in the paper then only using 1x 4090?

[StafaH]

Yes, everything from the paper was trained with a single 4090 (for vision).