🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/151497

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (4 Unrelated Failures)

As of commit 362b6e8 with merge base ba51f48 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

With #151719 (and a few changes to DTensor's flex_attention handling see here: P1789852198), I could run the E2E ring attention tests locally!

Looks pretty good let me know when ready for another review

@pytorchbot rebase

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

Successfully rebased gh/XilunWu/133/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/151497)

@wanchaol

From API side, I don't think the CP APIs should leak the "Sharder" concept to the users, it is confusing with some other concepts (i.e. the Shard placement). I think it should not be a "Sharder" actually, it is mainly a way to achieve load balancing, the API design needs some further discussion.

Regarding the naming issue, we can discuss alternative names if Sharder is not suitable. However, it's essential to note that Sharder (or any other name) in this context refers specifically to partitioning the input batch along the sequence dimension to achieve load balancing, which is unrelated to DTensor as we don't represent the input batch as a DTensor.

To clarify, is the concern related to the naming of Sharder or the exposure of the input partition concept itself? The input partition needs to be exposed as only users know how to achieve a good load balancing given a customized BlockMask. Though we also plan to do auto load balancing based on the mask, some advanced users have expressed the need to manually control the partition.

An alternative is make configuring Sharder an opt-in API that users can call before CP. Thus, most users can just use the built-in feature without worrying the load balancing concept. In such a case, users will just pass in the global BlockMask. If users would like to perform its own partitioning, then users can pass the Sharder instead of BlockMask.

@XilunWu Can you update the summary as @wanchaol mentioned?

I very much believe we want to not rely on auto-load balancing for all users and need way for users to explicit partition

@fegin @drisspg I am not debating whether to expose a way to do load balancing, I think it could be very valuable to allow users (especially the advanced users) to control how to do load balancing instead of fully rely on auto load balancing. But we should NOT leak any sort of sharding concept (or the input partition concept you called) to the users, it would confuse users on its relationship with DTensor sharding, why there's another sharding method, etc. So yes the concern is mainly the exposure of the sharding (or you called as input_partition) concept, I think we should not mix the concept of load balancing and sharding.

Regarding the naming issue, we can discuss alternative names if Sharder is not suitable. However, it's essential to note that Sharder (or any other name) in this context refers specifically to partitioning the input batch along the sequence dimension to achieve load balancing, which is unrelated to DTensor as we don't represent the input batch as a DTensor.

It is not just the naming issue. I don't think this is "unrelated to DTensor as we don't represent the input batch as a DTensor", context parallelism is a prototype feature developed within the dtensor module and it uses quite a few concepts from DTensor. It is critical to make sure CP works in a good way with DTensor. IMO If it's not working well with DTensor and do not provide good abstraction layers, I don't know if there're materialized values for users to use our solutions compared to using other OSS solutions or even writing their own one (i.e. if we are just allgathering k/v, and rely on user to define how to shard in details and does not do anything else, doing this manually is even better for user).

To be concrete, the current interaction between DTensor and CP is very weird: Although the Sharder/input_partition does not use DTensor (this is debatable, what if user just want to pass their own created DTensor to this flex attention thing?) and we don't represent the input batch as DTensor, but the current way to dispatch to the F.scaled_dot_product_attention to run ring attention (not this flex attention impl) is through constructing DTensor as inputs that sharded on sequence dim to the SDPA function, and dispatch to the op. But when using the RoundRobinSharder for load balancing, we should NOT construct DTensor anymore as the RoundRobinSharder does not align with the Shard placement behavior, so the constructed DTensors are actually not valid. This tells that the current Sharder abstraction in CP does not work in a correct way, and this is one main reason I think we should not expose it as is in the current way :)

@wanchaol I agree current CP doesn't work well with DTensor. But I don't agree with you that load balancing should be decoupled from the input sharding. How the input is sharded and arranged will affect the load balancing. Even though we just simply perform allgather on KV, Q is not gathered and how Q is sharded is decided by the input sharding.

The use case you mentioned is definitely an important one and advanced one. However, a collective (likely to be alltoall) for Q is required for this case, even for ring attention implementation. IMO, we should not sacrifice the performance of the basic use case. I definitely agree this is a direction we should solve, but I will argue that we still need to pass the information of how the input is partitioned to CP so that CP knows how to perform load balancing (reshuffling or alltoall for Q if needed).

IMO, what @kwen2501 proposed may be one way to achieve the good DTensor semantics for the load balancing? We let DTensor to support a special shard (a sharding that can represent rank0 has the chunk0 and chunk N and rank has the chunk 1 and chunk N -1, ...), I don't know a suitable name, bit-map-annotation shard?

With this representation, users for the current CP use case can simply partition the input batch using this bit-map-annotation shard to achieve load balancing without extra alltoall. Users can also opt to not shard the input in this way if they have a special use case to do attention and are okay to do alltoall. Either way, this sharding information serves a carrier to pass the information to CP to perform load balancing.

Could you give me some high level idea what you imagine how to decouple load balancing from sharding? I may misinterpret what you mean about the sharding being decoupled from sharding.

@wanchaol As for dispatcher, we can ask users to wrap SDPA/FlexAttention inside a module and use parallelize_module as we discussed. I don't have a strong opinion about this. cc., @XilunWu

But I don't agree with you that load balancing should be decoupled from the input sharding. How the input is sharded and arranged will affect the load balancing. Even though we just simply perform allgather on KV, Q is not gathered and how Q is sharded is decided by the input sharding.

@fegin Oh IIUC even allgathering KV requires one to know how the load balancing is performed, not only Q is decided by load balancing. So if user defined their own load balancing approach, it might be easier for them to shard/allgather QKVs. I do think load balancing and shard input are two separate concepts, maybe not totally orthogonal though.

IMO, we should not sacrifice the performance of the basic use case. I definitely agree this is a direction we should solve, but I will argue that this is the gap of DTensor.
IMO, what @kwen2501 proposed is the right solution to support both with a good UX -- we let DTensor to support a special shard (a sharding that can represent rank0 has the chunk0 and chunk N and rank has the chunk 1 and chunk N -1, ...), I don't know a suitable name, bit-map-annotation shard?

I have pretty strong opinion that we should NOT add new special Shard placement that could not be generalized, as that would significantly complicates DTensor without bringing much values (i.e. the current roundrobin is just one way to help load balancing, there could be many other ways, so it could not generalize well). I also feel there is actually no gap in DTensor to implement CP today, i.e. even in today's state, I could refactor how CP works with SDPA or flex attention without touching any core parts of DTensor, and still provide good UX. It is more like how the CP and its load balancing functionality should be designed.

Could you give me some high level idea what you imagine how to decouple load balancing from sharding? I may misinterpret what you mean about the sharding being decoupled from sharding.

Sure, so the high level idea is: load balancing and sharding need to be decoupled to separate layers, the load balancing layer is responsible to reorder the input sequence so that the attention computation are evenly balanced across rows/ranks. Sharding is a separate layer after it, it simply take the input reordered by the load balancer and shard it exactly as how DTensor shard tensor sequentially today. DTensor does not need to aware of load balancing at all as load balancing is a separate layer that exposed and maybe controlled by users. I actually had a draft refactor while I was refactoring the current ring attention implementation, I'll publish the changes soon :)