PyTorch Custom Operator Integration

[matthewdouglas]

Overview

This PR introduces the initial scaffolding to integrate PyTorch Custom Operators as the primary mechanism for dispatching to device-specific operator implementation.

As outlined in the related RFC #1545, the intent is that this will supersede the previous backend registration interface that was developed on the multi-backend-refactor branch. The baseline CUDA operators are established in this PR, and the implementation for additional backends is to be ported over to this new interface.

Why Custom Ops?

• 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 helps to simplify the development for additional backends, while taking an idiomatic modern PyTorch approach.

Operator Definitions

We broadly categorize operator functionality into three feature groups, though there can be some overlap.

LLM.int8()

Inference requirements

• int8_vectorwise_quant(A: Tensor, threshold: float = 0.0) -> (Tensor, Tensor, Tensor?)
  • Implements the LLM.int8() quantization algorithm with the specified threshold.
  • Returns an int8 quantized tensor, a float32 tensor containing the scaling stats, and an optional int32 tensor containing a list of column indices with outliers present.
• int8_scaled_mm(A: Tensor, B: Tensor, row_stats: Tensor, col_stats: Tensor, bias: Tensor?, dtype=torch.float16)
  • By default, this is a composition of the below two operators. The choice can be made to implement one fused operator or two separately.
    • int8_linear_matmul(A: Tensor, B: Tensor) -> Tensor
      • Performs an 8-bit integer matrix multiplication between two int8 matrices.
      • Returns an int32 matrix: A @ B.T
    • int8_mm_dequant(A: Tensor, row_stats: Tensor, col_stats: Tensor, dtype=torch.float16, bias: Tensor?) -> Tensor
      • Dequantizes the result of a quantized 8-bit matrix multiplication with an optional fused bias.
      • The result is returned in the specified dtype, which is always torch.float16 for the current CUDA implementation.

Optional

• int8_vectorwise_dequant(A: Tensor, stats: Tensor)
  • Dequantizes an int8 tensor that was quantized with int8_vectorwise_quant.
  • A default implementation in PyTorch is provided, which should work with any backend.
  • This is a utility utilized by Transformers, Diffusers, PEFT, and others.
• int8_double_quant(A: Tensor, threshold: float = 0.0)
  • Quantizes the input tensor using the LLM.int8() algorithm across both dimensions.
  • This is only useful for full int8 training (e.g. not LoRA), and as such, we only recommend implementing int8_vectorwise_quant.

NF4/FP4

Minimal requirements

• dequantize_4bit(A: Tensor, absmax: Tensor, blocksize: int, quant_type: Literal["nf4" | "fp4"], shape: int[], dtype) -> Tensor
  • Dequantizes a packed 4bit tensor into the specified floating point dtype.
  • Note: Unlike bitsandbytes.functional.dequantize_4bit, this operator does not dequantize the absmax tensor. If utilized, dequantize_blockwise must be performed first.
• quantize_4bit(A: Tensor, blocksize: int, quant_type: Literal["nf4" | "fp4"], quant_storage=torch.uint8) -> (Tensor, Tensor)
  • Quantizes a floating point tensor into a packed 4bit tensor.
  • Returns a tensor with the quantized data packed into into bytes, backed by the storage type specified. The float32 absmax scaling factors are additionally returned.
  • Note: Unlike bitsandbytes.functional.quantize_4bit, this operator does not quantize the absmax tensor. If utilized, quantize_blockwise must be performed first.

Double quantization (aka compressed_statistics or nested)

• dequantize_blockwise(A: Tensor, absmax: Tensor, code: Tensor, blocksize: int, dtype) -> Tensor
  • Dequantizes an 8bit tensor that was quantized with quantize_blockwise
  • The dequantized tensor with the specified dtype.
• quantize_blockwise(A: Tensor, code: Tensor, blocksize: int) -> (Tensor, Tensor)
  • Quantizes into an 8bit blocked data type defined by code.
  • The blocksize will typically be 256 for usage with NF4/FP4 and optimizers.
  • Returns the quantized tensor in uint8 format, along with float32 absmax.

Optional

• gemv_4bit
  • Fast path for bsz=1 inference with 4bit quantization. This operator is subject to some future revision.

Optimizers

Optimizer functionality will be implemented to support the custom operators in a future update.

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Titus-von-Koeller]

The custom_ops as well as clean up parts are looking really excellent. Thanks for this impactful work!

As we said, let's do separate PRs for

• torch.compile
  • about fixing avoidable issues with graph breaks
  • fix issues around torch.compile of quantize functions
• deprecate:
  • TestSpMMFunctional tests and mark related functions with @deprecated
  • SwitchBackLinear and communicate on issues / repo of open_clip
• fix AttributeError: module 'bitsandbytes.functional' has no attribute 'double_quant'