Indexed batch matrix multiplication to support MoEs and FFFs
Labels: module: sparse, triaged

### 🚀 The feature, motivation and pitch

Large models are becoming increasingly modularised. Recent research on [mixture-of-expert](https://arxiv.org/abs/1701.06538) networks (MoEs) and [fast feedforward](https://arxiv.org/abs/2308.14711) networks (FFFs) proposes to handle different inputs by different sets of parameters according to their nature, where the subset of parameters to be used is determined by "gating" or "tree" routing networks. Both MoEs and FFFs have been shown to deliver significant acceleration (50-220x) of feedforward layer inference.

## Problem: PyTorch does not provide any tools to implement MoEs / FFFs on scale

Let `w1s,w2s,b1s,b2s` contain the weights and biases of the modularised feedforward layer, let `x` be the batched inputs, and suppose `modules` contains the per-sample decisions of the gating/tree network on which experts/leaves are to use which sub-network of the MoE/FFF layer.

Naively, one could do the following
```
w1s = ... # (n_modules, input_width, hidden_width)
b1s = ... # (n_modules, 1, hidden_width)
w2s = ... # (n_modules, hidden_width, output_width)
b2s = ... # (n_modules, 1, output_width)

x = ... # (batch_size, input_width)
modules = ... # (batch_size, n_modules)

# execute the MoE / FFF layer
hidden = torch.matmul(x.unsqueeze(-2), w1s[modules]) # (batch_size, 1, hidden_width)
hidden += b1s[modules]
activations = activation(hidden)
outputs = torch.matmul(activations, w2s[modules])
outputs += b2s[modules]
```

The problem is that every indexing by `modules` makes `batch_size`-many copies of the size of one module/expert/leaf before feeding them into the matrix multiplication. It may not sound like such an issue, but just to give you an idea: if we split BERT-base layers into 4 experts, this means a memory overhead of `batch_size * 768 ^2` floats before multiplying and reducing to the output of `batch_size * 768`. That's an overhead quadratic in the hidden dimension of the module/expert/leaf -- in practice, instead of doing BERT inference with batch size 512 on a single 80GB A100, you have to reduce the batch size down to a mere dozen. It's fast but the GPU memory becomes a bottleneck.

## Proposed solution #1: Indexed BMM, i.e. `ibmm`
In the spirit of the `Tensor.index_add` function which selects source elements and performs operations on the fly to produce the summation output, I propose "indexed batch matrix multiplication".
To avoid unnecessary copying before performing the matrix multiplication, one could hope for an indexed variant of batched matrix multiplication `bmm` to multiply an input `k x m` matrix with an `m x n` matrix selected by an index tensor from a larger database of `n_modules`-many `m x n` tensors.

**Note that this is only needed for inference -- no autograd support is needed.**

Proposed signature and behaviour:

```
def ibmm(a: torch.Tensor, b: torch.Tensor, indices: torch.Tensor):
    # a has shape (batch_size, k, m)
    # b has shape (n_modules, m, n)
    # indices is a (batch_size,) long tensor of numbers between 0 and n_modules - 1
    # output has shape (batch_size, k, n)

    # parallelized batch for-loop
    for i in range(batch_size):
        output[i] = torch.mm(a[i], b[indices[i]]
    return output
```

## Proposed solution #2: Finish the implementation of COO sparse `matmul`
If a hybrid `2+2` COO sparse matrix could be multiplied with a dense matrix, the above `ibmm` would not be needed at all. Currently and to the best of my knowledge, no backend supports the `matmul` of hybrid (`len(shape) > 2`) COO sparse matrices with dense matrices. If that were possible, a whole new world would open for modularisation of networks.

In a world where a hybrid COO sparse matrix can be batch-multiplied with a dense matrix, the following would do the job:

```
x_coordinates = torch.arange(batch_size, dtype=torch.long)	# (batch_size,)
y_coordinates = modules
indices = torch.stack((x_coordinates, y_coordinates), dim=0)	# (2, batch_size)

x_sparse = torch.sparse_coo_tensor(
	indices=indices,
	values=x.unsqueeze(-2),
	size=(batch_size, n_modules, 1, input_width)
)	# (batch_size, n_modules, 1, input_width)
hidden = torch.matmul(x_sparse, w1s).sum(dim=1)	# (batch_size, 1, hidden_width)
hidden = hidden.squeeze(dim=1)			# (batch_size, hidden_width)
hidden += b1s[modules]				# (batch_size, hidden_width)
activations = activation(hidden)		# (batch_size, hidden_width)
activations_sparse = torch.sparse_coo_tensor(
	indices=indices,
	values=activations.unsqueeze(-1),
	size=(batch_size, n_modules, 1, hidden_width)
)	# (batch_size, n_modules, 1, hidden_width)
output = torch.matmul(activations_sparse , w2s).squeeze(-2)
output += b2s[modules] # (batch_size, output_width)
```


### Alternatives

## Workaround: Nasty but does the job for small enough networks

While `matmul` for hybrid sparse matrices is not implemented, the `mm` is. So if one flattens and indexes their data properly, they will get a working solution (with a linear overhead but hey, it's better than nothing).

```
x_coordinates = torch.arange(batch_size, dtype=torch.long, device=x.device)	# (batch_size,)
y_coordinate_bases = (leaves.long() * input_width).unsqueeze(-1) # (batch_size, 1)
y_coordinate_offsets = torch.arange(input_width, dtype=torch.long, device=x.device).unsqueeze(0) # (1, input_width)
y_coordinates = y_coordinate_bases + y_coordinate_offsets # (batch_size, input_width)
indices = torch.stack((x_coordinates.repeat(input_width), y_coordinates.flatten()), dim=0)	# (2, batch_size * input_width)
x_sparse = torch.sparse_coo_tensor(
	indices=indices,
	values=x.flatten(),
	size=(batch_size, n_modules * input_width),
	device=x.device
)	# (batch_size, n_modules * input_width)
logits = torch.mm(x_sparse, w1s.flatten(0, 1))	# (batch_size, hidden_width)
logits += b1s[leaves] 				# (batch_size, hidden_width)
activations = activation(logits)		# (batch_size, hidden_width)


y_coordinate_bases = (leaves.long() * hidden_width).unsqueeze(-1) # (batch_size, 1)
y_coordinate_offsets = torch.arange(hidden_width, dtype=torch.long, device=x.device).unsqueeze(0) # (1, hidden_width)
y_coordinates = y_coordinate_bases + y_coordinate_offsets # (batch_size, hidden_width)
second_indices = torch.stack((x_coordinates.repeat(hidden_width), y_coordinates.flatten()), dim=0) # (2, batch_size * hidden_width)
activations_sparse = torch.sparse_coo_tensor(
	indices=second_indices,
	values=activations.flatten(),
	size=(batch_size, n_modules * hidden_width),
	device=x.device
)	# (batch_size, n_modules, 1, hidden_width)
new_logits = torch.mm(activations_sparse, w2s.flatten(0, 1))	# (batch_size, output_width)
new_logits += b2s[leaves] # (batch_size, output_width)
```

### Additional context

Ultimately, the most general solution that could be implemented to support this kind of data manipulation is introducing indexed views -- i.e. views of tensors that go beyond slicing and reshaping. Somewhat ironically, in pure C++, any of this could be usually done without copying simply by dereferencing a pointer by `*`...

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer