Lightning Talk: Implementing Single-Dim Strategies With Sharding Validator - Anshul Sinha, Meta

This talk focuses on implementing single dimension strategies using the strategy validator in DTensor, which enhances the efficiency of tensor operations in distributed computing. The speaker explains how current approaches to tensor placement and communication can lead to inefficiencies and error-prone strategies. By adopting single dimension strategies, the system can manage placements across multiple dimensions independently, simplifying the rule creation process and reducing unnecessary communications. The strategy validator is introduced as a tool to verify these strategies by testing them against sample inputs and comparing outputs. Through this methodology, the speaker demonstrates how to ensure both the correctness and coverage of operations while enabling contributors to easily add their own operations and verify them.

All right. Hi everyone. Uh my name is Anel and today I'll be talking about implementing single dim strategies with the strategy validator in DTensor. So for a little bit of context in the uh current DTensor ops what happens is uh you have a given input and some outputs and each input will have its own placements. Uh these can be replicate, sharded or partials. And for a specific

op what might happen is uh we might need to redistribute the placements to u via communication or we can just keep the uh placement the same based on what the op is actually doing. So in the current uh strategy what happens is we first try to uh choose one input to follow its placements. Uh this is currently done via a heristic which is not uh the most

efficient. Um after that we follow the each strategy of the followed arg and then for each placement in that input we uh basically broadcast those uh placements to uh all the other inputs and then finally we compute the redistribute cost. But this is too much to do in one just one function for each op. As a result we have uh three problems. The first is that it

is difficult to add rules very quickly. The second is it's hard to prove that you're not missing propagation rules. And uh there might be some tricky ones that you can't really reason about by yourself. And what happens then is you have these unnecessary communications uh which is inefficient. And uh finally it is just difficult to review the correctness of these strategies because as you'll see in the

next slide, it is very hard to uh reason about and these result in errorprone strategies. So here's an example of torch.mmaxima. We have uh two GPUs and on each GPU we have uh two tensors which are both partial max which is basically just an unreduced tensor that we because we want to delay the communications as far as possible. So in this example uh densor0 has a on

one rank and b on the other and then we have c and d. If we were to redistribute immediately, we would uh first reduce the tensors uh by taking the reduce op in this case max and we would take the max of A and B and the max of C and D. Those are two full tensors and then we would take the max of those to finish

the uh op that we were trying to do on these 2D tensors. However, if you look down here, you can see that uh we can first propagate then redistribute. This means that we're doing each uh we're doing the OP on each rank. So we have a the tensor of max we have DT0 which is max of AC and then then on the second rank max BD and

these are still partial maxes and then after that if we were to redistribute and uh reduce the uh still partial tensors we would get the max of uh max AC and max PD which is still the same as the result above which basically means we don't need to do a communication if all the inputs are partial max and our output placement can be partial So when we

originally discovered this, what we would have to do is we'd have to uh create a new rule and basically that means a new branch. Um as you can see just by this it's very uh clogged. Um and with each new op that you find exceptions for there are uh new rules, new branches that you have to create and that's how it becomes very difficult to actually see

if the rules you're adding are correct. So how do we solve this? We do it by two parts. The first is single dim strategies. It takes advantage that each mesh dim is independent and separates the two concerns combined in the old op strategy uh which are semantic correctness. We want all the valid input output sharding combinations for each op and we the search base pruning to avoid

combinatorial blowups uh on n dimensional meshes. So in this case what happens is single dim strategies ops will provide valid input output shardings rules that we define for a single dimension and then the infra will basically take care of expanding across the full mesh and then it will select the lowest cost strategy. So here's an example based on the torch.maxim function. Um torch.maxim belongs to a class

of ops called pointwise op strategies. So what we do first is we create infra for all the common strategies across all pointwise ops. So in this case a very easy example is if all the inputs and outputs were replicated that's the same as just having a regular tensors then if you just have all the inputs and outputs sharded sharded on the same dimension this can also be

propagated and so we created rules for the point wise ops that they all share. Then you can see uh in addition to uh the maximum function there are also series of other ops that follow the same uh rules of uh the partial max partial max partial max that we originally saw. And so we create basically extra rules just specific to these ops and voila we are done.

We have uh added the rules that we want to for the sop. Um these are extremely easy to add and review. You just see uh which specific uh rules you have defined. Um and yeah so this is an example again using uh maximum uh of how we compose across uh the multi-dim mesh. As you can see we have uh multiple combinations of uh placements that would work

for each dimension. We take these uh we take these de we take uh these rules and you can see for output input one and input two the mesh dims that would work and then based on the input that we actually have we choose the uh cheapest strategy which basically is like the least amount of communications. But how do we still verify that these rules are correct? We

created what we call the strategy validator. The strategy validator basically takes op info sample inputs um for each op and then it will uh do three things. First it just computes the regular computation of this op without having uh without having densors. The second is it enumerates all the 1D placements and then the simulates the dtensor uh execution using something we call local tensor mode and based

on the matching results of the correct outputs and the densor outputs we can see which rules are actually correct and then finally we actually query uh the registered uh densor ops to see which strategies we currently have and based off this we achieve two things one we can find missing rules and two we can find uh incorrectly implemented rules that we currently have. So this is an

example of the uh strategy validation. Um as you can see in this case I chose argmax. Um and it basically what we do is it's very simple. You run the command and then you can see a few things. One you can see the incorrect rules. In this case we had three rules uh that were incorrectly implemented back then. And so we we basically implement the fix and

then we would rerun it. see that these are now correct. In addition, we can see in this case we already have zero missing rules. In some cases, there might have been missing rules that we would have implemented. And this way we can verify that the op is now fully covered. We can also see for each uh incorrect rule the samples that are breaking this rule. So that

way it's very easy to reproduce and uh you can verify the results of the strategy validator. So if you want to contribute very easily is uh you can add ops that are currently aren't covered by detensor either by just adding it to an existing category of ops that it would fall in or creating its own uh singledom strategies. In addition, if you have your own custom ops,

what you can do is you can uh create your own sample infos, sorry, create your own op infos and then use the strategy validator to find new rules, verify them, and then implement the rules. And then again, use the strategy validator to make sure that you didn't incorrectly implement rules as a safeguard. At the bottom, you can see the folder over here, which contains all the ops

that currently have registered dutensor ops. And uh if you find new ops that belong in these files, you can just quickly add them, implement your own strategies. And uh for both again uh if you're either expanding densor coverage or creating uh creating single dim strategies for your custom ops, you can uh use the strategy validator which is also found in this folder. Um that's all I have

for you today. Thank you for listening. If anyone has any questions, I'm happy to answer them. All right. Thank you guys.