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?

[DN6]

@pthombre we can open up a draft branch to test different mesh configurations.

[pthombre]

Thank you so much @DN6 @sayakpaul . This is a good example device mesh we can benefit from:
device_mesh = init_device_mesh( "cuda", mesh_shape=(2, 4), mesh_dim_names=("cp", "fsdp") )

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[DN6]

@pthombre Would something like this work?

device_mesh = init_device_mesh( "cuda", mesh_shape=(2, 1, 4), mesh_dim_names=("ulysses", "ring", "fsdp") )

[pthombre]

@DN6 Yes. I think support for the following mesh types would be ideal!