hmm I actually feel it's quite weird to let unflatten (or even the existing flatten behavior) to implicitly add a flattened/unflattend dim to the original mesh, it was for convenience purpose but I think it's sth quite weird (i.e. the original mesh is not having that dimension but can somehow access it). I think we probably don't want to make this to be an feature in unflatten, and maybe in the future deprecate this semantic in flatten. User can use the returned mesh directly, but not from the original mesh

This actually goes against the UX you have agreed on above, where we are trying to slicing the unflattened dim_name from the root mesh.

# Create a dummy 1D device mesh first, so no PG will be created under the hood. This is important for large-scale because no global PG will be used in that case.
device_mesh = init_device_mesh(
    'cuda', (32,), mesh_dim_names=["world"],
)
device_mesh._unflatten(0, (4,2,4), mesh_dim_names=["dp", "cp", "tp"])
device_mesh._unflatten(0, (2,8,2), mesh_dim_names=["dp_ep", "ep", "tp_ep"])
device_mesh["cp", "tp"]["tp"]

because instead of letting users to use same mesh we are actually ask user to do the bookkeeping themsevles. So users will need to keep track of flattened or unflattened submesh themselves.. which is really does not improve UX for EP here. And if you don't mind me asking, when the logic of _flatten is designed, reviewed and approved, it looks like you are also part of it as well. So what was the assumption then and what has changed now to make you change your mind? Or more broadly how did we make the call in the past? I have not heard any negative feedback regarding the slicing of flattened mesh_dim from the root mesh.

The reason I am asking here is that per @ezyang 's suggestion, we are trying to understand your design principle and mental models and so that we can keep the same design philosophy as much as possible. This does not say this bookkeep logic is per or against that design principle, but I find it confusing why we decide to make a shift there. Furthermore, IIUC, here is what I heard or observed so far:

1. We are trying to mimic all the behaviors of either torch.tensor and c10d as much as possible. Most of mesh API is trying to have same name and similar API signature as torch.tensor. This is because DeviceMesh is essentially a device representation with a mesh tensor. This makes sense, but it somehow also justify the slicing of flattened/unflattened dim_names since even with flatten and unflatten they are still one mesh(tensor)?

2. We also borrowed concept from XLA/JAX's devicemesh as well for its design. And IIUC, under the hood it is an N-D array so that aligns with why we mimic tensor behavior in point 1 as well.

3. We believe the abstraction of the Devicemesh inside Monarch and PyTorch are kind of two different levels, which is mentioned here: [RFC] Support slicing of submesh and more flexible operations to one device mesh such as reshape/split, etc. #159013 (comment). So we won't try to converge it with Monarch as much as possible.

Anything more you want to add? @wanchaol

The unflatten API itself is the correct UX I agreed on, but the way to access the unflattened device meshes from the root mesh is not, so basically the last line is the discrepency: device_mesh["cp", "tp"]["tp"], instead user should use the unflattened result to access the device mesh I think.

The reason I go against this behavior is that accessing from the root mesh is quite an implicit behavior, this is not really my design principle, but the Python/PyTorch design principle: Explicit is better than implicit!

To answer your question, when the flatten API was designed, I was also NOT a fan of this behavior, but given flatten only add one dimension to the root mesh, it did not quite contaminate this principle. Concerns about this was also the main reason I wanted to keep flatten private first. But now with the design of this unflatten API, I think the contamination become much worse. Why? Ideally user shouldn't aware of the root mesh concept as that is mainly an implementation detail, so the "root mesh" and the unflattened mesh should simply both be device meshes that represents a n-d array of devices. But now with the unflatten + root mesh behavior, a 1-D root device mesh suddenly have a bunch of additional dimensions! Taking the example you gave:

• device_mesh is a 1-D device mesh, but suddenly after the unflatten, it got 6 more dimensions! 6-more dimensions on a 1-D device mesh does not make sense as an 1-D device array abstraction
• What's worse is that it might not only have 6, it could have 8, 10 or even more dimensions because of calling the unflatten twice or more!

The implicit addition of device mesh dimensions make the abstraction not clean, and IMO we should avoid this implicit behavior.

I see, thanks for sharing your thought finally. I do agree that explicit is better than implicit so it makes sense that ideally we can have a more functional-like behavior. Since we already have the principle, we should have rejected the contamination earlier and I was under the impression that we did have a long discussion around flatten api as well. But somehow we still end up having a somehow contaminated api which is kind of suboptimal in the first place. OK let me also ask users like @tianyu-l 's input on this as well.

can you only bookkeep the mesh_dim instead of the actual device mesh?

If we only bookkeep the mesh_dim, it is not helpful because this is unflatten which means we are creating new dims and it is not like flatten which we kind of consolidate existing dims inside root mesh. that's why we need to swap the mesh for slicing and has this map. But I am open to better suggestions if any.

hmmm if unflatten to the same shape, I guess it would just return the "cached" unflattened results? Could you explain why it needs all those checks here and only reuse the dp dimension pg, instead of reusing all dimension pgs?

ok, I guess we need to define the behavior here.

1. If unflatten into same shape and same dim_names, yes we just return the cached unflatten results. This is clear.
2. If unflatten into same shape but different dim_names (sometimes the value of shape is dynamic so it can happen incidentally), my current thinking is that we should treat it as a new PG for the dim where dim_name is different because users might need a different stream for that dim? (Same PG means same stream). But like I mentioned in the comment, I am open to discussions for this as well, no strong opinion on this. This is also where I am curious about your comment as well.

If unflatten into same shape but different dim_names (sometimes the value of shape is dynamic so it can happen incidentally), my current thinking is that we should treat it as a new PG for the dim where dim_name is different because users might need a different stream for that dim?

I think if the set of ranks are the same, they should be the same communicator, maybe they could accidentally have the same shape because of dynamism, but IMO if they are the same shape, they should just reuse the same communicators, if user really want a different stream/group that have exact same ranks, I guess they should do it by hand to make it explicit. I think this is also the behavior of process group API.

I see, based on your comment above, since we like to hold the principle of being explicit, I agree we can just reuse it to make the behavior simple and predictable.

Looking through the logic of unflatten, I think one thing I observed is that the logic is quite complicated, I wonder if we need to cut the assumptions or features of this PR, and make it as few as possible.

I agree the logic is complicated but we definitely can try to simply it. But before we do that, we need to at least ensure and agree on what UX we are providing. The current logic/check is to support:

1. unflatten from a root mesh and we can do slicing from a submesh/root mesh
2. unflatten from a submesh and we can do slicing from a submesh/root mesh

we can only allow the unflatten from root mesh but this will simplify the logic a bit but not too much. Unless we want to further cut the UX.