Allow users to pass in a custom device_mesh to the enable_parallelism method in the diffusers library

[pthombre]

def enable_parallelism( self, *, config: Union[ParallelConfig, ContextParallelConfig], cp_plan: Optional[Dict[str, ContextParallelModelPlan]] = None, mesh: Optional[DeviceMesh] = None, # Add this parameter ):

If users have already initialized a device mesh for other parallelism strategies (FSDP, tensor parallelism, pipeline parallelism), they can reuse it instead of creating a new one. This prevents conflicts and redundant initialization.

[sayakpaul]

Thanks! Could you also provide as a few device_mesh values that you think would be the most reasonable for us to start working on this?

But one thing we need to consider is currently, we only support CP and supporting other parallelism strategies will probably go beyond just letting users pass a mesh through enable_parallelism. I am not sure if we will be able to ship that soon.

Cc: @DN6

[pthombre]

I think the immediate usecase on our end is to enable FSDP with CP. So basically supporting a device mesh with dp and cp values specified.

We also want to try combining tensor parallelism as well, but this can be added as a future if this complicates the ask too much.

[sayakpaul]

That sounds reasonable.

I think the immediate usecase on our end is to enable FSDP with CP. So basically supporting a device mesh with dp and cp values specified.

Do you have an example device_mesh value for us in that case?