ops.swiglu: Tests, baseline, benchmarks

[danthe3rd]

Stack from ghstack (oldest at bottom):

• Address nits from previous PR stack + bump version #518
• Include git hash and install date in version #511
• MemEff: Raise if wrong bias #510
• Raise if kernel not supported on device #509
• Benchmarks for entire models #505
• Update mypy to 0.982 #508
• SwiGLU: Support for no bias #504
• MemEff: Set CUDA stream properly #491
• SwiGLU further optimization in MLP bw #502
• SwiGLU optimized fw/bw #490
• unbind: Add 'stack_or_none' #498
• Should fix OOMs in CI #488
• -> ops.swiglu: Tests, baseline, benchmarks #487
• Split xformers.ops #486

[fmassa]

Thanks for the PR Daniel!

I have a few comments. In particular, I wouldn't add this nn.Module nor the torch.autograd.Function in xformers/ops, and would instead just implement a vanilla function in tests / benchmarks like the following:

def swiglu(...):
    x1 = F.linear(...)
    x2 = F.linear(...)
    x3 = F.silu(x1)
    x4 = x3 * x3
    return F.linear(...)

And for the tests, I would have a wrapper nn.Module that just calls into this swiglu.

Also, I wouldn't write the backwards by hand myself for the reference implementation, but instead use the autograd implementation as reference. This way we are sure we didn't introduce a bug in our autograd.

From a second look, you might have implemented the backwards by hand with PyTorch to have an idea of the numerics which are expected to differ for the operation, is that right?

Thoughts?

[fmassa]

nit: snake_case and CamelCase mix?

[fmassa]

note that you can also call torch.nn.functional.linear so that you can measure the performance from PyTorch, making the time measurements be compatible with what users would get

[fmassa]

I wonder why provide this implementation? Can't we just test the numerics with a functional API?

[fmassa]

Thanks!

Good to merge, I have a few suggestions for the future

[fmassa]

Suggested change

	def functional_swiglu(
	def _swiglu_reference(

?

[fmassa]

Maybe reuse the function _swiglu_reference in here?

Suggested change

	x1 = self.w1(x)
	x2 = self.w2(x)
	hidden = F.silu(x1) * x2
	return self.w3(hidden)
	return _swiglu_reference(x, self.w1.weight, self.w1.bias, self.w2.weight, self.w2.bias, self.w3.weight, self.w3.bias

[fmassa]

If we want to expose this to the user through other ops, we can make this public as well, and expose an op argument.

Suggested change

	class _SwiGLUModule(nn.Module):
	class _SwiGLU(nn.Module):

[codecov-commenter]

Codecov Report

Base: 90.59% // Head: 90.59% // Increases project coverage by +0.00% 🎉

Coverage data is based on head (0f0a602) compared to base (261b164).
Patch coverage: 90.62% of modified lines in pull request are covered.

Additional details and impacted files

@@                  Coverage Diff                  @@
##           gh/danthe3rd/50/base     #487   +/-   ##
=====================================================
  Coverage                 90.59%   90.59%           
=====================================================
  Files                        78       79    +1     
  Lines                      4572     4636   +64     
=====================================================
+ Hits                       4142     4200   +58     
- Misses                      430      436    +6     

Flag	Coverage Δ
Python	90.59% <90.62%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
xformers/ops/swiglu.py	90.47% <90.47%> (ø)
xformers/ops/__init__.py	81.25% <100.00%> (+1.25%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.