Skip to content

Conversation

@fegin
Copy link
Contributor

@fegin fegin commented Aug 29, 2025

Summary
This PR utilizes the latest APIs provided by DeviceMesh to simplify the creation of all different meshes.

The design philosophy is as follow:

  1. Create one world mesh with the shape as [world_size,]
  2. Create all 1-D submeshes by using 1) unflattening from the world mesh, or 2) slicing and flatten from other derived meshes.
  3. ParallelDims now provides an API, get_mesh(), which accepts str or list[str]. When the argument is str, the API directly return the corresponding 1-D submesh. If the argument is list[str], the dim names will be used to concatenate to form a n-D device mesh.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 29, 2025
@fegin fegin force-pushed the chienchin/new_device_mesh branch 7 times, most recently from 12eca61 to 19e4a23 Compare October 15, 2025 20:39
@fegin fegin force-pushed the chienchin/new_device_mesh branch from 19e4a23 to 178bc11 Compare October 28, 2025 20:34
@fegin fegin marked this pull request as ready for review October 28, 2025 21:01
@fegin fegin requested a review from wwwjn as a code owner October 28, 2025 21:01
Copy link
Contributor

@tianyu-l tianyu-l left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should modify FLUX train.py as it's in core now.

@ruisizhang123 let's adapt SimpleFSDP after this PR is merged.
oh it seems being fixed in #1959

@fegin fegin force-pushed the chienchin/new_device_mesh branch from d6eae58 to 70d6c03 Compare November 18, 2025 01:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants