Final push to clean our overwritten arguments #1353
Conversation
Prior to this commit, it was impossible to describe more than one inplace property. This would limit our ability to set the inplace property on instruction with more than one result. With this patch, it is now possible to define one inplace property list for each output of an instruction. To use this feature simply invoke inplaceOperand method for each output on the declaration of the instruction.
The QuantizationProfile node suffers the same problem as the save node: it has implicit memory dependencies and the scheduler could be tricker to generate invalid code because of that. Unlike the save node, the problem is less likely to happen because this node only deals with its only created variables. Yet, as soon as those variables are created, it is possible to temper with them to create invalid schedules. Unlike the save node, where we taught the scheduler how to deal with the implicit memory dependencies, this patch fixes the representation of the QuantizationProfile node (and its underlying QuantizationProfileInstr) by removing the implicit memory dependencies. After this patch, the save node is the only node with implicit memory dependencies. Test Plan: $ ninja check $ ./bin/image-classifier tests/images/imagenet/*.png -image_mode=0to1 -m=resnet50 -dump_profile=resnet50_prof_async.yaml $ ./bin/image-classifier tests/images/imagenet/*.png -image_mode=0to1 -m=resnet50 -dump_profile=resnet50_prof_async2.yaml $ diff -q resnet50_prof_async2.yaml resnet50_prof_async.yaml $ ./bin/image-classifier tests/images/imagenet/*.png -image_mode=0to1 -m=resnet50 -dump_profile=resnet50_prof_async2.yaml $ diff -q resnet50_prof_async2.yaml resnet50_prof_async.yaml Results: Everything passed, no diffs.
|
@qcolombet Yes, the quantize-profile node is the last node with overwritten attributes in our tree, but there are other nodes that we don't currently support. One important node is embedding tables. Embedding tables are (potentially huge) lookup tables. When training the lookup tables we need to update a small number of nodes in each iteration. Implementing the embedding tables using the method that nodes produce new values is impractical. Yes, the optimizer can save us in some cases, but I don't think that we should rely on the optimizer to save us. Is there a problem with enforcing a rule that operators that update some variables 'own' the variables? |
@nadavrot By own, do you mean only the owner node can read from or write to the variable? |
I agree and I don't think we should expose them as variables (at least not how it is done right now).
I think that would be the way to go, but then I believe it should be some special kind of variables. One thing that we need to enforce is that "owned" variables don't have any other reads (which is where we introduce the problematic memory dependencies). The best way to do that would be to not expose these variables for other to grab, and we miss that today. Indeed, variables are available at a module level and this cannot be changed easily because of how we use the Module at runtime #1334. |
|
@qcolombet Do you know how XLA handles variable mutability? @qcolombet @jfix71 How would you feel if we split variable into 3 categories:
What do you think? |
No, looking is still on my todolist :).
As long as those are not "gettable" to form other links that would work. The thing I want to be sure is that we don't create patterns that could lead into bugs into the scheduler. |
|
I wonder if we could have the variable being own (I mean litterally in the code) by the node itself. |
|
Note: I am happy to just trash this PR and go with the verifier check for now, until we land a better designed for the variables. I just hacked this this morning to see how difficult it would be to finally clean-up the overwritten problem and turned out it was easy and not too hacky, so I figured I shared :). |
|
It does prevent us for running into memory dependencies issues as of today though, which is nice IMHO. |
|
Abandoning this PR in favor of checks in the verifier. |
|
@nadavrot regarding mutability in XLA (hlo_instruction.h): I was incidentally going in that direction :). |
The quantization profile node was the last use (aside from save nodes) of the overwritten attribute and thus, was also sensible to the same kind of scheduling bugs that was exposed in #1283.
There are several ways we could have fixed/worked around that:
I don't like #1 because it pushes more hacks in the scheduler (eventually I'd like we model the memory dependencies period.). #2 is really just a check that the situation does not happen, but if it does, we're back to square 1.
Therefore, I gave a quick stab at #3 with this PR.
The nice thing I was pursuing and that PR achieves is that we would have the invariant that only save nodes write to variable.