If they're seeing these kinds of gains from relatively minor changes to their Python code, I can't help but wonder how much faster the model would run in a compiled language or a language with a good JIT (way more optimization work's gone into the mainstream Javascript runtimes than CPython).
I'd assumed that overall performance in Stable Diffusion was limited by the code running on the GPU, with Python performance being a fairly minor factor-- but I guess that's not the case?
This is PyTorch code, so the Python is setting up a bunch of kernels that are executed on the GPU. The switch from + to += might allow two of those kernels to be fused together or something, and that could lead to the large performance gain.
The Python part only runs a handful of times so JIT vs. non-JIT doesn't really make a difference.
Nah it’s because PyTorch has a different implementation for __iadd__. It’s saving a copy by mutating the LHS in-place, and possibly more divergent as comments report broken code.
I haven't played much with torch, but the game is generally that you have a graph of computations which gets JIT compiled into GPU ops. The compiler may have more or less competence at finding modifications (eg, 'fusions') to reduce the number of GPU ops required to perform the computation.
Depends on the field, python is almost never the bottleneck in CV - but it almost always a bottleneck in NLP or Search research.
Regardless, the code in the github PR should be running strictly in pytorch for the heavy lifting. Clearly something is touching python that shouldn't be, or must be for frustrating reasons.
In PyTorch `x = y + x` is actually semantically different from `x += y`, so you can't easily make the switch with a compiler.
The difference is that `x += y` modifies `x` inplace, where `x = x + y` creates a new object. In other words, if anybody had a reference to `x` before the update, the "optimized" code would break things.
Not to turn this into an evangelism subthread, but I think GP is right. The performance optimization could apply, but the bug as described should be prevented by the borrow checker. In fact you can do neither x = … nor x += … when anything else still references x, because assignment requires an exclusive reference (&mut), which is statistically checked.
I'm not sure what you mean by "same behavior". I don't think you can externally distinguish between "x = x + x" and "x += x" in Rust, because while in the first case a new temporary value will be created, it will overwrite the left side value: x will be the "same object" as it's in the same address in both cases. In Python objects are stored indirectly, so after an assignment x will then point to a new object instead of the old one.
But yes, not making copies is more efficient so the same optimization applies.
Nevertheless, in Rust (or C++) the parameter x would not have been passed as a mutable reference into the function, so the value provided as input would not have been mutated.
I don't know anything about stable diffusion, but I've been optimizing a lot of prime-field arithmetic in Rust lately, and we experienced a similar speedup going from `+ x` to `+= x` (for scalars and especially for composite structures like vectors and polynomials).
For composite structures that isn't too surprising, but for scalars, I would have expected llvm to optimize the addition and assignment into a single in place addition.
The long answer is that it’s not so clear what an “in place addition” even means at the level of CPU instructions after you consider register allocation. For example, if you have
v = x;
v += y;
f(v);
and the never mention v again, then the whole operation is performed directly in the register that is specified to receive the first argument in a function call, not in whatever register might have been allocated for v.
That’s because, with some complications I don’t want to go into, compilers look at the dependency graph of values rather than at the variable names.
way more optimization work's gone into the mainstream Javascript runtimes than CPython
Even so, there are absolutely silly things which can hint JS JITs to optimize (or to not deoptimize). Like defining and instantiating a class rather than just creating POJOs with the same values, or assigning NaN instead of null to uninitialized numeric variables/properties. Conditional control flow can deopt, but generally performs better around different function calls than within a single function. Even creating and throwing errors for control flow (which is generally expensive, and terrible for maintenance) can be optimal if your try/catch is the whole body of the function it resides in. And all of those might vary between JITs.
This change isn't a matter of Python being slower than a compiled language, it's changing the meaning of the code. The line
x = x + y
creates a copy of the array x, adds y to it, and then sets the variable x to that new array. In contrast, the line
x += y
adds the array y in-place into the array x (and so hopefully no other piece of code is relying on x being immutable). This kind of trade-off occurs in pretty much all programming, for instance you see it whenever big-integer libraries are used in C++ or Rust.
I don't think this is necessarily a minor change -- += and + are operators. I have no familiarity with this library, but I think _forward(...) works on tensors of something like that, they are probably big chunky data structures. += probably saves a copy or whatever.
Comments
If they're seeing these kinds of gains from relatively minor changes to their Python code, I can't help but wonder how much faster the model would run in a compiled language or a language with a good JIT (way more optimization work's gone into the mainstream Javascript runtimes than CPython).
I'd assumed that overall performance in Stable Diffusion was limited by the code running on the GPU, with Python performance being a fairly minor factor-- but I guess that's not the case?
This is PyTorch code, so the Python is setting up a bunch of kernels that are executed on the GPU. The switch from + to += might allow two of those kernels to be fused together or something, and that could lead to the large performance gain.
The Python part only runs a handful of times so JIT vs. non-JIT doesn't really make a difference.
Nah it’s because PyTorch has a different implementation for __iadd__. It’s saving a copy by mutating the LHS in-place, and possibly more divergent as comments report broken code.
I haven't played much with torch, but the game is generally that you have a graph of computations which gets JIT compiled into GPU ops. The compiler may have more or less competence at finding modifications (eg, 'fusions') to reduce the number of GPU ops required to perform the computation.
See, for example, XLA: https://www.tensorflow.org/xla
It looks like maybe nvFuser is an equivalent library for pytorch? https://pytorch.org/blog/introducing-nvfuser-a-deep-learning...
The Python code is run every time.
Yes but not nearly as much as the GPU code, which I think is what the parent is saying; it’s typically not the bottleneck.
Depends on the field, python is almost never the bottleneck in CV - but it almost always a bottleneck in NLP or Search research.
Regardless, the code in the github PR should be running strictly in pytorch for the heavy lifting. Clearly something is touching python that shouldn't be, or must be for frustrating reasons.
It can run much faster. For example, using the PyTorch nvFuser JIT gives a 50% speedup:
https://old.reddit.com/r/MachineLearning/comments/xa75km/p_p...
In PyTorch `x = y + x` is actually semantically different from `x += y`, so you can't easily make the switch with a compiler.
The difference is that `x += y` modifies `x` inplace, where `x = x + y` creates a new object. In other words, if anybody had a reference to `x` before the update, the "optimized" code would break things.
Compiler could use a pointer to pointer.
I guess this is the kind of this stuff that drew me to Rust. This kind of behavior gives me the creeps. Just like Ruby’s conventions.
Rust has the same behavior. https://news.ycombinator.com/item?id=32805756
Not to turn this into an evangelism subthread, but I think GP is right. The performance optimization could apply, but the bug as described should be prevented by the borrow checker. In fact you can do neither x = … nor x += … when anything else still references x, because assignment requires an exclusive reference (&mut), which is statistically checked.
I'm not sure what you mean by "same behavior". I don't think you can externally distinguish between "x = x + x" and "x += x" in Rust, because while in the first case a new temporary value will be created, it will overwrite the left side value: x will be the "same object" as it's in the same address in both cases. In Python objects are stored indirectly, so after an assignment x will then point to a new object instead of the old one.
But yes, not making copies is more efficient so the same optimization applies.
Nevertheless, in Rust (or C++) the parameter x would not have been passed as a mutable reference into the function, so the value provided as input would not have been mutated.
I don't know anything about stable diffusion, but I've been optimizing a lot of prime-field arithmetic in Rust lately, and we experienced a similar speedup going from `+ x` to `+= x` (for scalars and especially for composite structures like vectors and polynomials).
For composite structures that isn't too surprising, but for scalars, I would have expected llvm to optimize the addition and assignment into a single in place addition.
The short answer is yes.
The long answer is that it’s not so clear what an “in place addition” even means at the level of CPU instructions after you consider register allocation. For example, if you have
and the never mention v again, then the whole operation is performed directly in the register that is specified to receive the first argument in a function call, not in whatever register might have been allocated for v.That’s because, with some complications I don’t want to go into, compilers look at the dependency graph of values rather than at the variable names.
Even so, there are absolutely silly things which can hint JS JITs to optimize (or to not deoptimize). Like defining and instantiating a class rather than just creating POJOs with the same values, or assigning NaN instead of null to uninitialized numeric variables/properties. Conditional control flow can deopt, but generally performs better around different function calls than within a single function. Even creating and throwing errors for control flow (which is generally expensive, and terrible for maintenance) can be optimal if your try/catch is the whole body of the function it resides in. And all of those might vary between JITs.
I've always assumed Python was interpreted until I heard Nuitka [1].
It would be interesting to get a benchmark using CPython vs Nuitka related to this change.
[1] https://github.com/Nuitka/Nuitka
This change isn't a matter of Python being slower than a compiled language, it's changing the meaning of the code. The line
creates a copy of the array x, adds y to it, and then sets the variable x to that new array. In contrast, the line adds the array y in-place into the array x (and so hopefully no other piece of code is relying on x being immutable). This kind of trade-off occurs in pretty much all programming, for instance you see it whenever big-integer libraries are used in C++ or Rust.I don't think this is necessarily a minor change -- += and + are operators. I have no familiarity with this library, but I think _forward(...) works on tensors of something like that, they are probably big chunky data structures. += probably saves a copy or whatever.
This is like saying that passing a struct vs a pointer to a struct is a minor change for C code. I mean it's just one extra * !