PyTorch/Triton port?

[biobenkj]

Hey Zach,

Looking over the new additions with FactorNet, multiple distribution support, CUDA acceleration, etc. is incredible! I'm working on a couple of different ideas within the pytorch ecosystem where your work as part of RcppML would greatly accelerate things. Would you be interested in me taking a stab at porting the package to pytorch/triton as an alternative framework for python/other GPU accelerated approach? I don't want to take anything away from the incredible work you've put into this and would make sure you're involved throughout the process. Thanks for the consideration!

Ben

[zdebruine]

Hi Ben,

Yes, I am planning to build a Python version with pybind. Pytorch is an appealing idea but from a backend engineering perspective I am not sure just how much would transfer. I would be interested in researching possibilities and learning about what use cases you have for it.

Zach

[biobenkj]

That makes sense to use pybind. Switching to triton would avoid a lot of the cuda deps that would come with it and pytorch seems to be winning the day within the python ML ecosystem. Again, totally hear you on the pybind approach - was mostly thinking of the interface with existing pytorch tooling natively. Would it be helpful if I sketched out a little of what I was thinking regarding the pytorch/triton port?

[zdebruine]

Yes, that would be great if you could sketch out something, also happy to swing by VAI to chat later this week.

The most obvious Pytorch use case I see is the activation layer, providing a constrained and regularizable least squares solver (optional non-negativity) would open a lot of possibilities, along with a constrained solvable weights layer (as opposed to gradient-based updates with constrained projection). However, if used for NMF it would be far better to operate on the closed form alternating least squares optimization, rather than SGD.

The other value prop for Pytorch that I see is an NMF or SVD layer with autograd support, as existing implementations are very limited. However, this again sacrifices the closed form ALS optimization for SGD.