How to use create_device_mesh to create two meshes of different sizes

[sash-a]

Hi 👋

I have a program with 2 separate process each using different slices of a TPU. I would like to create each slice such that the layout is as good as it can be for both slices.

Lets say on a v3-8 I want 1 slice of 6 and one of 2. Should it be done like this:

slice1 = create_device_mesh((6,), jax.devices[:6])
slice2 = create_device_mesh((2,), jax.devices[6:])

Or should I create it like this:

device_mesh = create_device_mesh((8,), jax.devices)
slice1 = device_mesh[:6]
slice2 = device_mesh[6:]

I assume option 2 is better because create_device_mesh should create a layout where items close in the list are physically close together, but at the same time the docs for create device mesh say it returns a mesh shape "with good collective performance" so option 1 might be better for the bigger slice, but the smaller slice may be stuck with 2 core which have bad collective performance? Or does it just not make too much difference?

If it helps this question is specifically for Sebulba where a separate slice of the TPU is used for the actor and the learner.