[RFC] PyTorch Custom Operators & Multi-Backend Support

[matthewdouglas]

Purpose

We intend to integrate PyTorch Custom Operators as the primary mechanism for dispatching to device-specific operator implementations. An initial scaffolding of this is presented in PR #1544. This RFC will serve as a guideline to collect community feedback and refine our development plans moving forward.

Why?

• Registering operators with torch.library allows us to take advantage of the existing device dispatch mechanisms in PyTorch.
• We can treat calls to functionality in our CUDA kernels, or other low-level backend implementations, as opaque for improved torch.compile support.
• We can provide naive implementations of operators with only PyTorch code as a fallback option. This may additionally serve as a secondary CPU baseline, as per [RFC] Cross-Platform Refactor: CPU-only implementation #1021.
• This helps to simplify the development for additional backends, while taking an idiomatic modern PyTorch approach.

What about the multi-backend-refactor branch?

We are planning to deprecate further development on that branch upon the merging of #1544. After that point, the expectation is that we will implement backends using the new custom operator registration mechanisms. We expect to be able to reuse much of the existing implementations in the refactoring process.

Our goal is to aggressively mainline our in-tree backends, while additionally enabling out-of-tree backends. We will expand on this topic in the near future.

Supersedure

This RFC is intended to supersede topics which were covered in previously related RFCs which remained open as of this writing:

• [RFC] Extend bitsandbytes to support Intel hardware platforms #894
• [RFC] Cross-Platform Refactor: Mac M1 support #1020
• [RFC] Cross-Platform Refactor: CPU-only implementation #1021

Related Issues

Related issues and discussions include:

• ROCM Support #47
• Feature Request: ROCm support (AMD GPU) #107
• M1.M2 MacOS Users #485
• [RFC] Cross-Platform Refactor: Testing and CI/CD Strategy #1031
• ROCm Backend Status Tracker #1271
• Multi-backend refactor: Alpha release ( INTEL ONLY ) #1338
• Multi-backend refactor: Alpha release ( AMD ROCm ONLY ) #1339
• Multi-backend support: Apple Silicon / Mac (call for contributors + help fleshing out the details) #1340
• Support running on CPU  #1402
• aarch64 whl in PyPi #1437
• bitsandbytes for macos M1,M2,M3 chips #1460
• Building on NVidia GH-200 - Is this a supported platform? #1526

Additionally, this relates to the following issues and discussions which have been closed:

• Enabling Cross-platform Support #990
• [RFC] Cross-Platform Refactor: Overview + Link Hub #997
• [RFC] Cross-Platform Refactor: Build System and Binary Distribution #1032
• Release v44 not available for Mac #1378

Relevant contributors

The following contributors may have particular interest and feedback on this topic:

@Titus-von-Koeller
@christoph-koehncke
@jiqing-feng
@pnunna93
@akx
@rickardp
@ji-huazhong
@SlightwindSec

[jiqing-feng]

Hi @matthewdouglas . That sounds pretty cool. One thing I want to make sure that we don't need to build any C++ codes in CPU/XPU, so we do not need cmake in the CPU/XPU backend. I assume it won't break your design right?

[matthewdouglas]

@jiqing-feng That's perfectly fine. We won't expect that every backend would need to build a library, or that any library it builds would necessarily need to expose the same API.

For the CPU/XPU case we'll guard around what the cextension module does. Right now there is still CPU version of the C library that exposes a couple of functions: cquantize_blockwise_cpu_fp32 and cdequantize_blockwise_cpu_fp32. Those correspond to the operator definitions in this RFC:

torch.library.define("bitsandbytes::quantize_blockwise", "(Tensor A, Tensor code, int blocksize) -> (Tensor, Tensor)")

torch.library.define(
    "bitsandbytes::dequantize_blockwise",
    "(Tensor A, Tensor absmax, Tensor code, int blocksize, ScalarType dtype) -> Tensor",
)

I have not benchmarked it, but I expect that we'll be able to implement it with torch, like what was already done here:

bitsandbytes/bitsandbytes/backends/cpu_xpu_common.py

Lines 407 to 418 in 2640753

	def dequant_8bit(A, offset, quant_state):
	assert A.dtype == torch.uint8
	absmax = quant_state.code[A.reshape(-1).int()]
	blocks = absmax.shape[-1] // 256
	res = absmax.shape[-1] % 256
	if res != 0:
	absmax = F.pad(absmax, (0, 256 - res), mode="constant", value=0)
	absmax = (absmax.view(-1, 256) * quant_state.absmax.view(-1, 1)).to(quant_state.dtype).reshape(-1)
	absmax = absmax[: blocks * 256 + res]
	absmax = absmax.reshape(A.shape)
	absmax += offset
	return absmax

[matthewdouglas]

I'm going to close this out as this is implemented and merged with #1544.

General note: the custom operators in bitsandbytes are still subject to revision, and optimizers have not been implemented yet, but it is stable for internal use and initial backend integrations.