| |
|
|
| 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 |
