|
|
|
|
|
import numpy as np |
|
import torch |
|
from torch.nn import functional as F |
|
from einops import rearrange |
|
|
|
|
|
def blockdiag_weight_to_dense_weight(weight): |
|
""" |
|
Argumments: |
|
weight: (nblocks, out / nblocks, in / blocks) |
|
Return: |
|
dense_weight: (out / in) |
|
""" |
|
return torch.block_diag(*torch.unbind(weight, dim=0)) |
|
|
|
|
|
def blockdiag_multiply_reference(x, weight): |
|
""" |
|
This implementation is slow but more likely to be correct. |
|
Arguments: |
|
x: (..., n) |
|
weight: (nblocks, q, n / nblocks) |
|
Outputs: |
|
out: (..., nblocks * q) |
|
""" |
|
n = x.shape[-1] |
|
nblocks, q, p = weight.shape |
|
assert nblocks * p == n |
|
|
|
x_reshaped = rearrange(x, '... (nblocks p) -> ... nblocks p', nblocks=nblocks) |
|
return rearrange(torch.einsum('...kp, kqp -> ...kq', x_reshaped, weight), |
|
'... nblocks q -> ... (nblocks q)') |
|
|
|
|
|
class BlockdiagMultiply(torch.autograd.Function): |
|
|
|
"""This is a faster implementation, with careful memory copies for the fastest |
|
bmm performance. |
|
The backward pass is also written manually with careful memory copies. |
|
Arguments: |
|
x: (..., n) |
|
weight: (nblocks, q, n / nblocks) |
|
Outputs: |
|
out: (..., nblocks * q) |
|
""" |
|
|
|
@staticmethod |
|
@torch.cuda.amp.custom_fwd(cast_inputs=torch.bfloat16) |
|
def forward(ctx, x, weight): |
|
ctx.save_for_backward(x, weight) |
|
batch_shape, n = x.shape[:-1], x.shape[-1] |
|
batch_dim = np.prod(batch_shape) |
|
nblocks, q, p = weight.shape |
|
assert nblocks * p == n |
|
x_reshaped = x.reshape(batch_dim, nblocks, p).transpose(0, 1) |
|
out = torch.empty(batch_dim, nblocks, q, device=x.device, dtype=x.dtype).transpose(0, 1) |
|
out = torch.bmm(x_reshaped, weight.transpose(-1, -2), out=out).transpose(0, 1) |
|
return out.reshape(*batch_shape, nblocks * q) |
|
|
|
@staticmethod |
|
@torch.cuda.amp.custom_bwd |
|
def backward(ctx, dout): |
|
x, weight = ctx.saved_tensors |
|
batch_shape, n = x.shape[:-1], x.shape[-1] |
|
batch_dim = np.prod(batch_shape) |
|
nblocks, q, p = weight.shape |
|
assert nblocks * p == n |
|
dx, dweight = None, None |
|
dout_reshaped = dout.reshape(batch_dim, nblocks, q).transpose(0, 1) |
|
if ctx.needs_input_grad[0]: |
|
dx = torch.empty(batch_dim, nblocks, p, device=x.device, dtype=x.dtype) |
|
dx = torch.bmm(dout_reshaped, weight.conj(), |
|
out=dx.transpose(0, 1)).transpose(0, 1).reshape(*batch_shape, n) |
|
if ctx.needs_input_grad[1]: |
|
x_reshaped = x.reshape(batch_dim, nblocks, p).transpose(0, 1) |
|
dweight = torch.bmm(dout_reshaped.transpose(-1, -2), x_reshaped.conj()) |
|
return dx, dweight |
|
|
|
|
|
blockdiag_multiply = BlockdiagMultiply.apply |
