[DeviceMesh] Simplifying internal bookkeeping with CuTe layout #161106
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
[ghstack-poisoned]
pytorch-bot
bot
added
the
oncall: distributed
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/161106
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 69 PendingAs of commit 99623a9 with merge base f638854 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
fduwjj
changed the title
…Te layout" cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta [ghstack-poisoned]
…Te layout" cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta [ghstack-poisoned]
…Te layout" cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta [ghstack-poisoned]
…Te layout" We want to implement the cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta [ghstack-poisoned]
fduwjj
changed the title
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: 1. Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout) 2. Separate backend from the mesh operations so that we eventually can let users do lots of operations without initializing any backend. Concretely, in this PR, we do the following: 1. Replaced all index/offset and its mappings with CuTe Layout and a backend class which handles all the bookkeeping and create backend if needed. Use CuTe layout for both slicing and _flatten. 2. We also started to make devicemesh more functional (first from the backend perspective). Each newly created device mesh is like a universe, all devicemesh transformed out from it (slicing, flatten, unflatten, etc) will share the same backend (PG) while creating a new device mesh will be different universe. So we changed our unit tests accordingly as well. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta [ghstack-poisoned]
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: 1. Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout) 2. Separate backend from the mesh operations so that we eventually can let users do lots of operations without initializing any backend. Concretely, in this PR, we do the following: 1. Replaced all index/offset and its mappings with CuTe Layout and a backend class which handles all the bookkeeping and create backend if needed. Use CuTe layout for both slicing and _flatten. 2. We also started to make devicemesh more functional (first from the backend perspective). Each newly created device mesh is like a universe, all devicemesh transformed out from it (slicing, flatten, unflatten, etc) will share the same backend (PG) while creating a new device mesh will be different universe. So we changed our unit tests accordingly as well. cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 18, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 18, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 18, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 19, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 19, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 19, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 19, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 24, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 24, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 25, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. 6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`. 7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig` The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 25, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. 6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`. 7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig` The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 25, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. 6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`. 7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig` The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 25, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. 6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`. 7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig` The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 26, 2025
…ing with CuTe layout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. 6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`. 7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig` The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 26, 2025
…yout" We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the _Layout to handle all index operations rather use a map to record mesh dims. 2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout). 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use a new function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. 6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`. 7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig` The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of #161106 (original one got messed with EasyCLA) cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 27, 2025
…ing with CuTe layout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`.
7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig`
Another function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 27, 2025
…yout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Separated the Devicemesh creation via mesh into a util function `_create_mesh_from_ranks`.
7. Added a type alias for `tuple[Optional[str], Optional[C10dBackend.Options]]` to be `BackendConfig`
Another function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 30, 2025
…ing with CuTe layout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Add an util function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Sep 30, 2025
…yout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Add an util function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Oct 1, 2025
…ing with CuTe layout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Add an util function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Oct 1, 2025
…yout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Add an util function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Oct 1, 2025
…ing with CuTe layout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Add an util function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Oct 1, 2025
…yout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the _Layout to handle all index operations rather use a map to record mesh dims.
2. Replaced`flatten_name_to_root_dims` with `flatten_name_to_root_layout`. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout).
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use a new function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe.
6. Add an util function (the name is open for discussion) is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
pytorchmergebot
pushed a commit
that referenced
this pull request
Oct 2, 2025
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the `_MeshLayout` to handle all index operations rather use a map to record mesh dims.
2. Removed `flatten_name_to_root_dims`, because now we can directly get layout from a flattened device mesh.
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use the newly added function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The reason is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
Verified that loss curve is very close for DeepSeekV3 on torchtitan, note that exact same match is challenging because even if we run the baseline twice, the loss curve does not exactly match.
<img width="1113" height="490" alt="image" src="https://github.com/user-attachments/assets/7877b5a4-337e-4ad8-b878-2378f4f0f38d" />
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
Pull Request resolved: #163213
Approved by: https://github.com/lw, https://github.com/fegin
fduwjj
added a commit
that referenced
this pull request
Oct 3, 2025
…ing with CuTe layout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the `_MeshLayout` to handle all index operations rather use a map to record mesh dims.
2. Removed `flatten_name_to_root_dims`, because now we can directly get layout from a flattened device mesh.
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use the newly added function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The reason is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
Verified that loss curve is very close for DeepSeekV3 on torchtitan, note that exact same match is challenging because even if we run the baseline twice, the loss curve does not exactly match.
<img width="1113" height="490" alt="image" src="https://github.com/user-attachments/assets/7877b5a4-337e-4ad8-b878-2378f4f0f38d" />
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
fduwjj
added a commit
that referenced
this pull request
Oct 3, 2025
…yout"
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the `_MeshLayout` to handle all index operations rather use a map to record mesh dims.
2. Removed `flatten_name_to_root_dims`, because now we can directly get layout from a flattened device mesh.
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use the newly added function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The reason is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
Verified that loss curve is very close for DeepSeekV3 on torchtitan, note that exact same match is challenging because even if we run the baseline twice, the loss curve does not exactly match.
<img width="1113" height="490" alt="image" src="https://github.com/user-attachments/assets/7877b5a4-337e-4ad8-b878-2378f4f0f38d" />
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
cc H-Huang awgu wanchaol fegin wz337 wconstab d4l3k pragupta ezyang msaroufim dcci
[ghstack-poisoned]
pytorchmergebot
pushed a commit
that referenced
this pull request
Oct 3, 2025
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
1. Use the `_MeshLayout` to handle all index operations rather use a map to record mesh dims.
2. Removed `flatten_name_to_root_dims`, because now we can directly get layout from a flattened device mesh.
3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`.
4. Use the newly added function `check_overlap` to check layout overlap.
5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The reason is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1).
Verified that loss curve is very close for DeepSeekV3 on torchtitan, note that exact same match is challenging because even if we run the baseline twice, the loss curve does not exactly match.
<img width="1113" height="490" alt="image" src="https://github.com/user-attachments/assets/7877b5a4-337e-4ad8-b878-2378f4f0f38d" />
The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout.
This is a continue of #161106 (original one got messed with EasyCLA)
Pull Request resolved: #163213
Approved by: https://github.com/lw, https://github.com/fegin
Chao1Han
pushed a commit
to Chao1Han/pytorch
that referenced
this pull request
Oct 21, 2025
…ch#163213) We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the `_MeshLayout` to handle all index operations rather use a map to record mesh dims. 2. Removed `flatten_name_to_root_dims`, because now we can directly get layout from a flattened device mesh. 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use the newly added function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The reason is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1). Verified that loss curve is very close for DeepSeekV3 on torchtitan, note that exact same match is challenging because even if we run the baseline twice, the loss curve does not exactly match. <img width="1113" height="490" alt="image" src="https://github.com/user-attachments/assets/7877b5a4-337e-4ad8-b878-2378f4f0f38d" /> The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of pytorch#161106 (original one got messed with EasyCLA) Pull Request resolved: pytorch#163213 Approved by: https://github.com/lw, https://github.com/fegin
Chao1Han
pushed a commit
to Chao1Han/pytorch
that referenced
this pull request
Oct 21, 2025
…ch#163213) We want to refactor the internal bookkeeping of DeviceMesh so that: Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible. Concretely, in this PR, we do the following: 1. Use the `_MeshLayout` to handle all index operations rather use a map to record mesh dims. 2. Removed `flatten_name_to_root_dims`, because now we can directly get layout from a flattened device mesh. 3. Replaced `_get_slice_mesh_dims` with `_get_slice_mesh_layout`. 4. Use the newly added function `check_overlap` to check layout overlap. 5. Use a new function `to_remapping_tensor` to use layout ranks as indices when the mesh tensor is not representable as CuTe. The reason is that layout acts as a backend of mesh tensor bookkeeping (indexing indices), it needs to be used as indices for remap back to the mesh tensor for new DeviceMesh generation and backend init. For example, in the case of 2K to 4K, the underlying layout is (2K, 1) but the actual value of the mesh tensor is [2K, 2K+1, ....,]. While flattening, slicing, we need to remap the layout back to the new mesh tensor so it maps the actual device allocation. For example, in the 2K to 4K case, if the shape is (1K, 1K) with dim_names ("dp", "tp"). Then when slicing "tp", the mesh tensor should be (2K, 2K+1, ..., 3K-1) or (3K, 3K+1, ... 4K-1). not the global ranks generated from the layout. (1K, 1). Verified that loss curve is very close for DeepSeekV3 on torchtitan, note that exact same match is challenging because even if we run the baseline twice, the loss curve does not exactly match. <img width="1113" height="490" alt="image" src="https://github.com/user-attachments/assets/7877b5a4-337e-4ad8-b878-2378f4f0f38d" /> The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor. With this refactoring we also enabled the slicing and flatten of non-contiguous dims of a device mesh which is hard to implement without cute layout. This is a continue of pytorch#161106 (original one got messed with EasyCLA) Pull Request resolved: pytorch#163213 Approved by: https://github.com/lw, https://github.com/fegin
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
We want to refactor the internal bookkeeping of DeviceMesh so that:
Simply the bookkeeping logics and make it generic enough so that it is easy to support new transformations like flatten noncontiguous dim, reshape and unflatten. (We leveraged the CuTe layout). This new layout also let us handle non-contiguous slicing, flatten, transpose possible.
Concretely, in this PR, we do the following:
flatten_name_to_root_dimswithflatten_name_to_root_layout. Basically one (size, stride) pair maps to one PG. One mesh_dim_name can only map to only layout. (More than one mesh_dim_name can map to the same layout)._get_slice_mesh_dimswith_get_slice_mesh_layout._maybe_create_backendto create PG based on mesh_dim_name and layout._from_layoutsto create a new DeviceMesh from a given layout list.The PR looks big indeed but we don't change any existing behavior of DeviceMesh, so it is a pure refactor.
cc @H-Huang @awgu @wanchaol @fegin @wz337 @wconstab @d4l3k @pragupta