chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .utils import SizeTracker, quantize_model_ # NOQA
|
||||
@@ -0,0 +1,211 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from collections import Counter
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class EM:
|
||||
"""
|
||||
EM algorithm used to quantize the columns of W to minimize
|
||||
|
||||
||W - W_hat||^2
|
||||
|
||||
Args:
|
||||
- W: weight matrix of size (in_features x out_features)
|
||||
- n_iter: number of k-means iterations
|
||||
- n_centroids: number of centroids (size of codebook)
|
||||
- eps: for cluster reassignment when an empty cluster is found
|
||||
- max_tentatives for cluster reassignment when an empty cluster is found
|
||||
- verbose: print error after each iteration
|
||||
|
||||
Remarks:
|
||||
- If one cluster is empty, the most populated cluster is split into
|
||||
two clusters
|
||||
- All the relevant dimensions are specified in the code
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, W, n_centroids=256, n_iter=20, eps=1e-6, max_tentatives=30, verbose=True
|
||||
):
|
||||
self.W = W
|
||||
self.n_centroids = n_centroids
|
||||
self.n_iter = n_iter
|
||||
self.eps = eps
|
||||
self.max_tentatives = max_tentatives
|
||||
self.verbose = verbose
|
||||
self.centroids = torch.Tensor()
|
||||
self.assignments = torch.Tensor()
|
||||
self.objective = []
|
||||
|
||||
def initialize_centroids(self):
|
||||
"""
|
||||
Initializes the centroids by sampling random columns from W.
|
||||
"""
|
||||
|
||||
in_features, out_features = self.W.size()
|
||||
indices = torch.randint(
|
||||
low=0, high=out_features, size=(self.n_centroids,)
|
||||
).long()
|
||||
self.centroids = self.W[:, indices].t() # (n_centroids x in_features)
|
||||
|
||||
def step(self, i):
|
||||
"""
|
||||
There are two standard steps for each iteration: expectation (E) and
|
||||
minimization (M). The E-step (assignment) is performed with an exhaustive
|
||||
search and the M-step (centroid computation) is performed with
|
||||
the exact solution.
|
||||
|
||||
Args:
|
||||
- i: step number
|
||||
|
||||
Remarks:
|
||||
- The E-step heavily uses PyTorch broadcasting to speed up computations
|
||||
and reduce the memory overhead
|
||||
"""
|
||||
|
||||
# assignments (E-step)
|
||||
distances = self.compute_distances() # (n_centroids x out_features)
|
||||
self.assignments = torch.argmin(distances, dim=0) # (out_features)
|
||||
n_empty_clusters = self.resolve_empty_clusters()
|
||||
|
||||
# centroids (M-step)
|
||||
for k in range(self.n_centroids):
|
||||
W_k = self.W[:, self.assignments == k] # (in_features x size_of_cluster_k)
|
||||
self.centroids[k] = W_k.mean(dim=1) # (in_features)
|
||||
|
||||
# book-keeping
|
||||
obj = (self.centroids[self.assignments].t() - self.W).norm(p=2).item()
|
||||
self.objective.append(obj)
|
||||
if self.verbose:
|
||||
logging.info(
|
||||
f"Iteration: {i},\t"
|
||||
f"objective: {obj:.6f},\t"
|
||||
f"resolved empty clusters: {n_empty_clusters}"
|
||||
)
|
||||
|
||||
def resolve_empty_clusters(self):
|
||||
"""
|
||||
If one cluster is empty, the most populated cluster is split into
|
||||
two clusters by shifting the respective centroids. This is done
|
||||
iteratively for a fixed number of tentatives.
|
||||
"""
|
||||
|
||||
# empty clusters
|
||||
counts = Counter(map(lambda x: x.item(), self.assignments))
|
||||
empty_clusters = set(range(self.n_centroids)) - set(counts.keys())
|
||||
n_empty_clusters = len(empty_clusters)
|
||||
|
||||
tentatives = 0
|
||||
while len(empty_clusters) > 0:
|
||||
# given an empty cluster, find most populated cluster and split it into two
|
||||
k = random.choice(list(empty_clusters))
|
||||
m = counts.most_common(1)[0][0]
|
||||
e = torch.randn_like(self.centroids[m]) * self.eps
|
||||
self.centroids[k] = self.centroids[m].clone()
|
||||
self.centroids[k] += e
|
||||
self.centroids[m] -= e
|
||||
|
||||
# recompute assignments
|
||||
distances = self.compute_distances() # (n_centroids x out_features)
|
||||
self.assignments = torch.argmin(distances, dim=0) # (out_features)
|
||||
|
||||
# check for empty clusters
|
||||
counts = Counter(map(lambda x: x.item(), self.assignments))
|
||||
empty_clusters = set(range(self.n_centroids)) - set(counts.keys())
|
||||
|
||||
# increment tentatives
|
||||
if tentatives == self.max_tentatives:
|
||||
logging.info(
|
||||
f"Could not resolve all empty clusters, {len(empty_clusters)} remaining"
|
||||
)
|
||||
raise EmptyClusterResolveError
|
||||
tentatives += 1
|
||||
|
||||
return n_empty_clusters
|
||||
|
||||
def compute_distances(self):
|
||||
"""
|
||||
For every centroid m, computes
|
||||
|
||||
||M - m[None, :]||_2
|
||||
|
||||
Remarks:
|
||||
- We rely on PyTorch's broadcasting to speed up computations
|
||||
and reduce the memory overhead
|
||||
- Without chunking, the sizes in the broadcasting are modified as:
|
||||
(n_centroids x n_samples x out_features) -> (n_centroids x out_features)
|
||||
- The broadcasting computation is automatically chunked so that
|
||||
the tensors fit into the memory of the GPU
|
||||
"""
|
||||
|
||||
nb_centroids_chunks = 1
|
||||
|
||||
while True:
|
||||
try:
|
||||
return torch.cat(
|
||||
[
|
||||
(self.W[None, :, :] - centroids_c[:, :, None]).norm(p=2, dim=1)
|
||||
for centroids_c in self.centroids.chunk(
|
||||
nb_centroids_chunks, dim=0
|
||||
)
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
except RuntimeError:
|
||||
nb_centroids_chunks *= 2
|
||||
|
||||
def assign(self):
|
||||
"""
|
||||
Assigns each column of W to its closest centroid, thus essentially
|
||||
performing the E-step in train().
|
||||
|
||||
Remarks:
|
||||
- The function must be called after train() or after loading
|
||||
centroids using self.load(), otherwise it will return empty tensors
|
||||
"""
|
||||
|
||||
distances = self.compute_distances() # (n_centroids x out_features)
|
||||
self.assignments = torch.argmin(distances, dim=0) # (out_features)
|
||||
|
||||
def save(self, path, layer):
|
||||
"""
|
||||
Saves centroids and assignments.
|
||||
|
||||
Args:
|
||||
- path: folder used to save centroids and assignments
|
||||
"""
|
||||
|
||||
torch.save(self.centroids, os.path.join(path, "{}_centroids.pth".format(layer)))
|
||||
torch.save(
|
||||
self.assignments, os.path.join(path, "{}_assignments.pth".format(layer))
|
||||
)
|
||||
torch.save(self.objective, os.path.join(path, "{}_objective.pth".format(layer)))
|
||||
|
||||
def load(self, path, layer):
|
||||
"""
|
||||
Loads centroids and assignments from a given path
|
||||
|
||||
Args:
|
||||
- path: folder use to load centroids and assignments
|
||||
"""
|
||||
|
||||
self.centroids = torch.load(
|
||||
os.path.join(path, "{}_centroids.pth".format(layer))
|
||||
)
|
||||
self.assignments = torch.load(
|
||||
os.path.join(path, "{}_assignments.pth".format(layer))
|
||||
)
|
||||
self.objective = torch.load(
|
||||
os.path.join(path, "{}_objective.pth".format(layer))
|
||||
)
|
||||
|
||||
|
||||
class EmptyClusterResolveError(Exception):
|
||||
pass
|
||||
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .qconv import PQConv2d # NOQA
|
||||
from .qemb import PQEmbedding # NOQA
|
||||
from .qlinear import PQLinear # NOQA
|
||||
@@ -0,0 +1,115 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.modules.utils import _pair
|
||||
|
||||
|
||||
class PQConv2d(nn.Module):
|
||||
"""
|
||||
Quantized counterpart of nn.Conv2d module. Stores the centroid, the assignments
|
||||
and the non-quantized biases. The full weight is re-instantiated at each forward
|
||||
pass and autograd automatically computes the gradients with respect to the
|
||||
centroids.
|
||||
|
||||
Args:
|
||||
- centroids: centroids of size n_centroids x block_size
|
||||
- assignments: assignments of the centroids to the subvectors
|
||||
of size self.out_channels x n_blocks
|
||||
- bias: the non-quantized bias, must be either torch.Tensor or None
|
||||
|
||||
Remarks:
|
||||
- We refer the reader to the official documentation of the nn.Conv2d module
|
||||
for the other arguments and the behavior of the module.
|
||||
- Performance tests on GPU show that this implementation is 10% slower than
|
||||
the non-quantized nn.Conv2d module for a standard training loop.
|
||||
- During the backward, the gradients are averaged by cluster and not summed.
|
||||
This explains the hook registered to the centroids.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
centroids,
|
||||
assignments,
|
||||
bias,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
padding_mode="zeros",
|
||||
):
|
||||
super(PQConv2d, self).__init__()
|
||||
self.block_size = centroids.size(1)
|
||||
self.n_centroids = centroids.size(0)
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = _pair(kernel_size)
|
||||
self.stride = _pair(stride)
|
||||
self.padding = _pair(padding)
|
||||
self.dilation = _pair(dilation)
|
||||
self.groups = groups
|
||||
self.padding_mode = padding_mode
|
||||
# check compatibility
|
||||
if in_channels // groups * np.prod(self.kernel_size) % self.block_size != 0:
|
||||
raise ValueError("Wrong PQ sizes")
|
||||
if len(assignments) % out_channels != 0:
|
||||
raise ValueError("Wrong PQ sizes")
|
||||
if in_channels % groups != 0:
|
||||
raise ValueError("in_channels must be divisible by groups")
|
||||
if out_channels % groups != 0:
|
||||
raise ValueError("out_channels must be divisible by groups")
|
||||
# define parameters
|
||||
self.centroids = nn.Parameter(centroids, requires_grad=True)
|
||||
self.register_buffer("assignments", assignments)
|
||||
self.register_buffer("counts", torch.bincount(assignments).type_as(centroids))
|
||||
if bias is not None:
|
||||
self.bias = nn.Parameter(bias)
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
# register hook for averaging gradients per centroids instead of summing
|
||||
self.centroids.register_hook(lambda x: x / self.counts[:, None])
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return (
|
||||
self.centroids[self.assignments]
|
||||
.reshape(-1, self.out_channels, self.block_size)
|
||||
.permute(1, 0, 2)
|
||||
.reshape(
|
||||
self.out_channels, self.in_channels // self.groups, *self.kernel_size
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return F.conv2d(
|
||||
x,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
s = "{in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}"
|
||||
if self.padding != (0,) * len(self.padding):
|
||||
s += ", padding={padding}"
|
||||
if self.dilation != (1,) * len(self.dilation):
|
||||
s += ", dilation={dilation}"
|
||||
if self.groups != 1:
|
||||
s += ", groups={groups}"
|
||||
if self.bias is None:
|
||||
s += ", bias=False"
|
||||
if self.padding_mode != "zeros":
|
||||
s += ", padding_mode={padding_mode}"
|
||||
s += ", n_centroids={n_centroids}, block_size={block_size}"
|
||||
return s.format(**self.__dict__)
|
||||
@@ -0,0 +1,107 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class PQEmbedding(nn.Module):
|
||||
"""
|
||||
Quantized counterpart of nn.Embedding module. Stores the centroids and
|
||||
the assignments. The full weight is re-instantiated at each forward
|
||||
pass.
|
||||
|
||||
Args:
|
||||
- centroids: centroids of size n_centroids x block_size
|
||||
- assignments: assignments of the centroids to the subvectors
|
||||
of size self.out_features x n_blocks
|
||||
- bias: the non-quantized bias
|
||||
|
||||
Remarks:
|
||||
- We refer the reader to the official documentation of the nn.Embedding module
|
||||
for the other arguments and the behavior of the module
|
||||
- Performance tests on GPU show that this implementation is 10% slower than
|
||||
the non-quantized nn.Embedding module for a standard training loop.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
centroids,
|
||||
assignments,
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
padding_idx=None,
|
||||
max_norm=None,
|
||||
norm_type=2.0,
|
||||
scale_grad_by_freq=False,
|
||||
sparse=False,
|
||||
_weight=None,
|
||||
):
|
||||
super(PQEmbedding, self).__init__()
|
||||
self.block_size = centroids.size(1)
|
||||
self.n_centroids = centroids.size(0)
|
||||
self.num_embeddings = num_embeddings
|
||||
self.embedding_dim = embedding_dim
|
||||
if padding_idx is not None:
|
||||
if padding_idx > 0:
|
||||
assert (
|
||||
padding_idx < self.num_embeddings
|
||||
), "Padding_idx must be within num_embeddings"
|
||||
elif padding_idx < 0:
|
||||
assert (
|
||||
padding_idx >= -self.num_embeddings
|
||||
), "Padding_idx must be within num_embeddings"
|
||||
padding_idx = self.num_embeddings + padding_idx
|
||||
self.padding_idx = padding_idx
|
||||
self.max_norm = max_norm
|
||||
self.norm_type = norm_type
|
||||
self.scale_grad_by_freq = scale_grad_by_freq
|
||||
self.sparse = sparse
|
||||
# check compatibility
|
||||
if self.embedding_dim % self.block_size != 0:
|
||||
raise ValueError("Wrong PQ sizes")
|
||||
if len(assignments) % self.num_embeddings != 0:
|
||||
raise ValueError("Wrong PQ sizes")
|
||||
# define parameters
|
||||
self.centroids = nn.Parameter(centroids, requires_grad=True)
|
||||
self.register_buffer("assignments", assignments)
|
||||
self.register_buffer("counts", torch.bincount(assignments).type_as(centroids))
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return (
|
||||
self.centroids[self.assignments]
|
||||
.reshape(-1, self.num_embeddings, self.block_size)
|
||||
.permute(1, 0, 2)
|
||||
.flatten(1, 2)
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return F.embedding(
|
||||
input,
|
||||
self.weight,
|
||||
self.padding_idx,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.sparse,
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
s = "{num_embeddings}, {embedding_dim}"
|
||||
if self.padding_idx is not None:
|
||||
s += ", padding_idx={padding_idx}"
|
||||
if self.max_norm is not None:
|
||||
s += ", max_norm={max_norm}"
|
||||
if self.norm_type != 2:
|
||||
s += ", norm_type={norm_type}"
|
||||
if self.scale_grad_by_freq is not False:
|
||||
s += ", scale_grad_by_freq={scale_grad_by_freq}"
|
||||
if self.sparse is not False:
|
||||
s += ", sparse=True"
|
||||
s += ", n_centroids={n_centroids}, block_size={block_size}"
|
||||
|
||||
return s.format(**self.__dict__)
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class PQLinear(nn.Module):
|
||||
"""
|
||||
Quantized counterpart of nn.Linear module. Stores the centroid, the assignments
|
||||
and the non-quantized biases. The full weight is re-instantiated at each forward
|
||||
pass.
|
||||
|
||||
Args:
|
||||
- centroids: centroids of size n_centroids x block_size
|
||||
- assignments: assignments of the centroids to the subvectors
|
||||
of size self.out_features x n_blocks
|
||||
- bias: the non-quantized bias
|
||||
|
||||
Remarks:
|
||||
- We refer the reader to the official documentation of the nn.Linear module
|
||||
for the other arguments and the behavior of the module
|
||||
- Performance tests on GPU show that this implementation is 15% slower than
|
||||
the non-quantized nn.Linear module for a standard training loop.
|
||||
"""
|
||||
|
||||
def __init__(self, centroids, assignments, bias, in_features, out_features):
|
||||
super(PQLinear, self).__init__()
|
||||
self.block_size = centroids.size(1)
|
||||
self.n_centroids = centroids.size(0)
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
# check compatibility
|
||||
if self.in_features % self.block_size != 0:
|
||||
raise ValueError("Wrong PQ sizes")
|
||||
if len(assignments) % self.out_features != 0:
|
||||
raise ValueError("Wrong PQ sizes")
|
||||
# define parameters
|
||||
self.centroids = nn.Parameter(centroids, requires_grad=True)
|
||||
self.register_buffer("assignments", assignments)
|
||||
self.register_buffer("counts", torch.bincount(assignments).type_as(centroids))
|
||||
if bias is not None:
|
||||
self.bias = nn.Parameter(bias)
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return (
|
||||
self.centroids[self.assignments]
|
||||
.reshape(-1, self.out_features, self.block_size)
|
||||
.permute(1, 0, 2)
|
||||
.flatten(1, 2)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return F.linear(
|
||||
x,
|
||||
self.weight,
|
||||
self.bias,
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"in_features={self.in_features},\
|
||||
out_features={self.out_features},\
|
||||
n_centroids={self.n_centroids},\
|
||||
block_size={self.block_size},\
|
||||
bias={self.bias is not None}"
|
||||
@@ -0,0 +1,128 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .em import EM, EmptyClusterResolveError
|
||||
|
||||
|
||||
class PQ(EM):
|
||||
"""
|
||||
Quantizes the layer weights W with the standard Product Quantization
|
||||
technique. This learns a codebook of codewords or centroids of size
|
||||
block_size from W. For further reference on using PQ to quantize
|
||||
neural networks, see "And the Bit Goes Down: Revisiting the Quantization
|
||||
of Neural Networks", Stock et al., ICLR 2020.
|
||||
|
||||
PQ is performed in two steps:
|
||||
(1) The matrix W (weights or fully-connected or convolutional layer)
|
||||
is reshaped to (block_size, -1).
|
||||
- If W is fully-connected (2D), its columns are split into
|
||||
blocks of size block_size.
|
||||
- If W is convolutional (4D), its filters are split along the
|
||||
spatial dimension.
|
||||
(2) We apply the standard EM/k-means algorithm to the resulting reshaped matrix.
|
||||
|
||||
Args:
|
||||
- W: weight matrix to quantize of size (in_features x out_features)
|
||||
- block_size: size of the blocks (subvectors)
|
||||
- n_centroids: number of centroids
|
||||
- n_iter: number of k-means iterations
|
||||
- eps: for cluster reassignment when an empty cluster is found
|
||||
- max_tentatives for cluster reassignment when an empty cluster is found
|
||||
- verbose: print information after each iteration
|
||||
|
||||
Remarks:
|
||||
- block_size be compatible with the shape of W
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
W,
|
||||
block_size,
|
||||
n_centroids=256,
|
||||
n_iter=20,
|
||||
eps=1e-6,
|
||||
max_tentatives=30,
|
||||
verbose=True,
|
||||
):
|
||||
self.block_size = block_size
|
||||
W_reshaped = self._reshape(W)
|
||||
super(PQ, self).__init__(
|
||||
W_reshaped,
|
||||
n_centroids=n_centroids,
|
||||
n_iter=n_iter,
|
||||
eps=eps,
|
||||
max_tentatives=max_tentatives,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
def _reshape(self, W):
|
||||
"""
|
||||
Reshapes the matrix W as expained in step (1).
|
||||
"""
|
||||
|
||||
# fully connected: by convention the weight has size out_features x in_features
|
||||
if len(W.size()) == 2:
|
||||
self.out_features, self.in_features = W.size()
|
||||
assert (
|
||||
self.in_features % self.block_size == 0
|
||||
), "Linear: n_blocks must be a multiple of in_features"
|
||||
return (
|
||||
W.reshape(self.out_features, -1, self.block_size)
|
||||
.permute(2, 1, 0)
|
||||
.flatten(1, 2)
|
||||
)
|
||||
|
||||
# convolutional: we reshape along the spatial dimension
|
||||
elif len(W.size()) == 4:
|
||||
self.out_channels, self.in_channels, self.k_h, self.k_w = W.size()
|
||||
assert (
|
||||
self.in_channels * self.k_h * self.k_w
|
||||
) % self.block_size == 0, (
|
||||
"Conv2d: n_blocks must be a multiple of in_channels * k_h * k_w"
|
||||
)
|
||||
return (
|
||||
W.reshape(self.out_channels, -1, self.block_size)
|
||||
.permute(2, 1, 0)
|
||||
.flatten(1, 2)
|
||||
)
|
||||
# not implemented
|
||||
else:
|
||||
raise NotImplementedError(W.size())
|
||||
|
||||
def encode(self):
|
||||
"""
|
||||
Performs self.n_iter EM steps.
|
||||
"""
|
||||
|
||||
self.initialize_centroids()
|
||||
for i in range(self.n_iter):
|
||||
try:
|
||||
self.step(i)
|
||||
except EmptyClusterResolveError:
|
||||
break
|
||||
|
||||
def decode(self):
|
||||
"""
|
||||
Returns the encoded full weight matrix. Must be called after
|
||||
the encode function.
|
||||
"""
|
||||
|
||||
# fully connected case
|
||||
if "k_h" not in self.__dict__:
|
||||
return (
|
||||
self.centroids[self.assignments]
|
||||
.reshape(-1, self.out_features, self.block_size)
|
||||
.permute(1, 0, 2)
|
||||
.flatten(1, 2)
|
||||
)
|
||||
|
||||
# convolutional case
|
||||
else:
|
||||
return (
|
||||
self.centroids[self.assignments]
|
||||
.reshape(-1, self.out_channels, self.block_size)
|
||||
.permute(1, 0, 2)
|
||||
.reshape(self.out_channels, self.in_channels, self.k_h, self.k_w)
|
||||
)
|
||||
@@ -0,0 +1,337 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import re
|
||||
from operator import attrgetter, itemgetter
|
||||
|
||||
import numpy as np
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
|
||||
from .modules import PQConv2d, PQEmbedding, PQLinear
|
||||
from .pq import PQ
|
||||
|
||||
|
||||
def quantize_model_(
|
||||
model,
|
||||
size_tracker,
|
||||
layers_to_quantize,
|
||||
block_sizes_config,
|
||||
n_centroids_config,
|
||||
step=0,
|
||||
n_iter=15,
|
||||
eps=1e-6,
|
||||
max_tentatives=100,
|
||||
verbose=True,
|
||||
):
|
||||
"""
|
||||
Quantize a model in-place by stages. All the targeted
|
||||
layers are replaced by their quantized counterpart,
|
||||
and the model is ready for the finetuning of the
|
||||
centroids in a standard training loop (no modifications
|
||||
required). Note that we do not quantize biases.
|
||||
|
||||
Args:
|
||||
- model: a nn.Module
|
||||
- size_tracker: useful for tracking quatization statistics
|
||||
- layers_to_quantize: a list containing regexps for
|
||||
filtering the layers to quantize at each stage according
|
||||
to their name (as in model.named_parameters())
|
||||
- block_sizes_config: dict like
|
||||
{
|
||||
'Conv2d': ('kernel_size', {'(3, 3)': 9, '(1, 1)': 4}),
|
||||
'Linear': ('in_features', {'*': 8})
|
||||
}
|
||||
For instance, all conv2d layers with kernel size 3x3 have
|
||||
a block size of 9 and all Linear layers are quantized with
|
||||
a block size of 8, irrespective of their size.
|
||||
- n_centroids_config: dict like
|
||||
{
|
||||
'Conv2d': ('kernel_size', {'*': 256}),
|
||||
'Linear': ('in_features', {'*': 256})
|
||||
}
|
||||
For instance, all conv2d layers are quantized with 256 centroids
|
||||
- step: the layers to quantize inplace corresponding
|
||||
to layers_to_quantize[step]
|
||||
"""
|
||||
|
||||
quantized_layers = get_layers(model, layers_to_quantize[step])
|
||||
|
||||
for layer in quantized_layers:
|
||||
|
||||
# book-keeping
|
||||
is_master_process = (not dist.is_initialized()) or (
|
||||
dist.is_initialized() and dist.get_rank() == 0
|
||||
)
|
||||
verbose = verbose and is_master_process
|
||||
|
||||
# get block size and centroids
|
||||
module = attrgetter(layer)(model)
|
||||
block_size = get_param(module, layer, block_sizes_config)
|
||||
n_centroids = get_param(module, layer, n_centroids_config)
|
||||
if verbose:
|
||||
logging.info(
|
||||
f"Quantizing layer {layer} with block size {block_size} and {n_centroids} centroids"
|
||||
)
|
||||
|
||||
# quantize layer
|
||||
weight = module.weight.data.clone()
|
||||
is_bias = "bias" in [x[0] for x in module.named_parameters()]
|
||||
bias = module.bias.data.clone() if is_bias else None
|
||||
quantizer = PQ(
|
||||
weight,
|
||||
block_size,
|
||||
n_centroids=n_centroids,
|
||||
n_iter=n_iter,
|
||||
eps=eps,
|
||||
max_tentatives=max_tentatives,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
# quantization performed on all GPUs with same seed
|
||||
quantizer.encode()
|
||||
centroids = quantizer.centroids.contiguous()
|
||||
assignments = quantizer.assignments.contiguous()
|
||||
|
||||
# broadcast results to make sure weights are up-to-date
|
||||
if dist.is_initialized():
|
||||
dist.broadcast(centroids, 0)
|
||||
dist.broadcast(assignments, 0)
|
||||
|
||||
# instantiate the quantized counterpart
|
||||
if isinstance(module, nn.Linear):
|
||||
out_features, in_features = map(
|
||||
lambda k: module.__dict__[k], ["out_features", "in_features"]
|
||||
)
|
||||
quantized_module = PQLinear(
|
||||
centroids, assignments, bias, in_features, out_features
|
||||
)
|
||||
elif isinstance(module, nn.Embedding):
|
||||
num_embeddings, embedding_dim = map(
|
||||
lambda k: module.__dict__[k], ["num_embeddings", "embedding_dim"]
|
||||
)
|
||||
quantized_module = PQEmbedding(
|
||||
centroids, assignments, num_embeddings, embedding_dim
|
||||
)
|
||||
elif isinstance(module, nn.Conv2d):
|
||||
out_channels, in_channels, kernel_size = map(
|
||||
lambda k: module.__dict__[k],
|
||||
["out_channels", "in_channels", "kernel_size"],
|
||||
)
|
||||
stride, padding, dilation, groups, padding_mode = map(
|
||||
lambda k: module.__dict__[k],
|
||||
["stride", "padding", "dilation", "groups", "padding_mode"],
|
||||
)
|
||||
|
||||
quantized_module = PQConv2d(
|
||||
centroids,
|
||||
assignments,
|
||||
bias,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
padding_mode=padding_mode,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Module {module} not yet supported for quantization")
|
||||
|
||||
# replace layer by its quantized counterpart
|
||||
attrsetter(layer)(model, quantized_module)
|
||||
|
||||
# update statistics
|
||||
size_tracker.update(weight, block_size, n_centroids)
|
||||
|
||||
# return name of quantized layers
|
||||
return quantized_layers
|
||||
|
||||
|
||||
def get_layers(model, filter_regexp):
|
||||
"""
|
||||
Filters out the layers according to a regexp. Note that
|
||||
we omit biases.
|
||||
|
||||
Args:
|
||||
- model: a nn.Module
|
||||
- filter_regexp: a regexp to filter the layers to keep
|
||||
according to their name in model.named_parameters().
|
||||
For instance, the regexp:
|
||||
|
||||
down_layers\\.[123456]\\.(conv[12]|identity\\.conv))
|
||||
|
||||
is keeping blocks down_layers from 1 to 6, and inside
|
||||
each block is keeping conv1, conv2 and identity.conv.
|
||||
|
||||
Remarks:
|
||||
- We add (module\\.)? at the beginning of the regexp to
|
||||
account for the possible use of nn.parallel.DataParallel
|
||||
"""
|
||||
|
||||
# get all parameter names
|
||||
all_layers = map(itemgetter(0), model.named_parameters())
|
||||
|
||||
# remove biases
|
||||
all_layers = filter(lambda x: "bias" not in x, all_layers)
|
||||
|
||||
# remove .weight in all other names (or .weight_orig is spectral norm)
|
||||
all_layers = map(lambda x: x.replace(".weight_orig", ""), all_layers)
|
||||
all_layers = map(lambda x: x.replace(".weight", ""), all_layers)
|
||||
|
||||
# return filtered layers
|
||||
filter_regexp = "(module\\.)?" + "(" + filter_regexp + ")"
|
||||
r = re.compile(filter_regexp)
|
||||
|
||||
return list(filter(r.match, all_layers))
|
||||
|
||||
|
||||
def get_param(module, layer_name, param_config):
|
||||
"""
|
||||
Given a quantization configuration, get the right parameter
|
||||
for the module to be quantized.
|
||||
|
||||
Args:
|
||||
- module: a nn.Module
|
||||
- layer_name: the name of the layer
|
||||
- param_config: a dict like
|
||||
{
|
||||
'Conv2d': ('kernel_size', {'(3, 3)': 9, '(1, 1)': 4}),
|
||||
'Linear': ('in_features', {'*': 8})
|
||||
}
|
||||
For instance, all conv2d layers with kernel size 3x3 have
|
||||
a block size of 9 and all Linear layers are quantized with
|
||||
a block size of 8, irrespective of their size.
|
||||
|
||||
Remarks:
|
||||
- if 'fuzzy_name' is passed as a parameter, layers whose layer_name
|
||||
include 'fuzzy_name' will be assigned the given parameter.
|
||||
In the following example, conv.expand layers will have a block
|
||||
size of 9 while conv.reduce will have a block size of 4 and all
|
||||
other layers will have a block size of 2.
|
||||
{
|
||||
'Conv2d': ('fuzzy_name', {'expand': 9, 'reduce': 4, '*': 2}),
|
||||
'Linear': ('fuzzy_name', {'classifier': 8, 'projection': 4})
|
||||
}
|
||||
|
||||
"""
|
||||
|
||||
layer_type = module.__class__.__name__
|
||||
|
||||
if layer_type not in param_config:
|
||||
raise KeyError(f"Layer type {layer_type} not in config for layer {module}")
|
||||
|
||||
feature, params = param_config[module.__class__.__name__]
|
||||
|
||||
if feature != "fuzzy_name":
|
||||
feature_value = str(getattr(module, feature))
|
||||
if feature_value not in params:
|
||||
if "*" in params:
|
||||
feature_value = "*"
|
||||
else:
|
||||
raise KeyError(
|
||||
f"{feature}={feature_value} not in config for layer {module}"
|
||||
)
|
||||
else:
|
||||
feature_values = [name for name in params if name in layer_name]
|
||||
if len(feature_values) == 0:
|
||||
if "*" in params:
|
||||
feature_value = "*"
|
||||
else:
|
||||
raise KeyError(f"name={layer_name} not in config for {module}")
|
||||
else:
|
||||
feature_value = feature_values[0]
|
||||
|
||||
return params[feature_value]
|
||||
|
||||
|
||||
class SizeTracker(object):
|
||||
"""
|
||||
Class to keep track of the compressed network size with iPQ.
|
||||
|
||||
Args:
|
||||
- model: a nn.Module
|
||||
|
||||
Remarks:
|
||||
- The compressed size is the sum of three components
|
||||
for each layer in the network:
|
||||
(1) Storing the centroids given by iPQ in fp16
|
||||
(2) Storing the assignments of the blocks in int8
|
||||
(3) Storing all non-compressed elements such as biases
|
||||
- This cost in only valid if we use 256 centroids (then
|
||||
indexing can indeed by done with int8).
|
||||
"""
|
||||
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
self.size_non_compressed_model = self.compute_size()
|
||||
self.size_non_quantized = self.size_non_compressed_model
|
||||
self.size_index = 0
|
||||
self.size_centroids = 0
|
||||
self.n_quantized_layers = 0
|
||||
|
||||
def compute_size(self):
|
||||
"""
|
||||
Computes the size of the model (in MB).
|
||||
"""
|
||||
|
||||
res = 0
|
||||
for _, p in self.model.named_parameters():
|
||||
res += p.numel()
|
||||
return res * 4 / 1024 / 1024
|
||||
|
||||
def update(self, W, block_size, n_centroids):
|
||||
"""
|
||||
Updates the running statistics when quantizing a new layer.
|
||||
"""
|
||||
|
||||
# bits per weights
|
||||
bits_per_weight = np.log2(n_centroids) / block_size
|
||||
self.n_quantized_layers += 1
|
||||
|
||||
# size of indexing the subvectors of size block_size (in MB)
|
||||
size_index_layer = bits_per_weight * W.numel() / 8 / 1024 / 1024
|
||||
self.size_index += size_index_layer
|
||||
|
||||
# size of the centroids stored in float16 (in MB)
|
||||
size_centroids_layer = n_centroids * block_size * 2 / 1024 / 1024
|
||||
self.size_centroids += size_centroids_layer
|
||||
|
||||
# size of non-compressed layers, e.g. LayerNorms or biases (in MB)
|
||||
size_uncompressed_layer = W.numel() * 4 / 1024 / 1024
|
||||
self.size_non_quantized -= size_uncompressed_layer
|
||||
|
||||
def __repr__(self):
|
||||
size_compressed = (
|
||||
self.size_index + self.size_centroids + self.size_non_quantized
|
||||
)
|
||||
compression_ratio = self.size_non_compressed_model / size_compressed # NOQA
|
||||
return (
|
||||
f"Non-compressed model size: {self.size_non_compressed_model:.2f} MB. "
|
||||
f"After quantizing {self.n_quantized_layers} layers, size "
|
||||
f"(indexing + centroids + other): {self.size_index:.2f} MB + "
|
||||
f"{self.size_centroids:.2f} MB + {self.size_non_quantized:.2f} MB = "
|
||||
f"{size_compressed:.2f} MB, compression ratio: {compression_ratio:.2f}x"
|
||||
)
|
||||
|
||||
|
||||
def attrsetter(*items):
|
||||
def resolve_attr(obj, attr):
|
||||
attrs = attr.split(".")
|
||||
head = attrs[:-1]
|
||||
tail = attrs[-1]
|
||||
|
||||
for name in head:
|
||||
obj = getattr(obj, name)
|
||||
return obj, tail
|
||||
|
||||
def g(obj, val):
|
||||
for attr in items:
|
||||
resolved_obj, resolved_attr = resolve_attr(obj, attr)
|
||||
setattr(resolved_obj, resolved_attr, val)
|
||||
|
||||
return g
|
||||
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
def parse_config_yaml(yaml_data):
|
||||
# Initialize to default options.
|
||||
quantization_options = {
|
||||
"n_centroids": {
|
||||
"Linear": ["in_features", {"*": 256}],
|
||||
"Embedding": ["embedding_dim", {"*": 256}],
|
||||
},
|
||||
"block_sizes": {
|
||||
"Linear": ["fuzzy_name", {"fc": 8, "attn": 4, "emb": 4}],
|
||||
"Embedding": ["fuzzy_name", {"emb": 8}],
|
||||
},
|
||||
"layers_to_quantize": [
|
||||
"decoder\\.layers\\.\\d+\\.fc[12]",
|
||||
"decoder\\.embed_tokens\\.embeddings\\.[012]\\.[01]",
|
||||
"decoder\\.layers\\.\\d+\\.self_attn\\.(k_proj|v_proj|q_proj|out_proj)",
|
||||
],
|
||||
}
|
||||
|
||||
if "n_centroids" in yaml_data:
|
||||
quantization_options["n_centroids"] = {
|
||||
layer: convert_yaml_to_tuple(layer_data)
|
||||
for layer, layer_data in yaml_data["n_centroids"].items()
|
||||
}
|
||||
if "block_sizes" in yaml_data:
|
||||
quantization_options["block_sizes"] = {
|
||||
layer: convert_yaml_to_tuple(layer_data)
|
||||
for layer, layer_data in yaml_data["block_sizes"].items()
|
||||
}
|
||||
if "layers_to_quantize" in yaml_data:
|
||||
quantization_options["layers_to_quantize"] = yaml_data["layers_to_quantize"]
|
||||
|
||||
return quantization_options
|
||||
|
||||
|
||||
def convert_yaml_to_tuple(yaml_dictionary):
|
||||
"""Converts a yaml dictionary with two keys: `key` and `value` into a two
|
||||
argument tuple of those values."""
|
||||
return (yaml_dictionary["key"], yaml_dictionary["value"])
|
||||
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .utils import quantize_model_ # NOQA
|
||||
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .qact import ActivationQuantizer # NOQA
|
||||
from .qconv import IntConv2d # NOQA
|
||||
from .qemb import IntEmbedding # NOQA
|
||||
from .qlinear import IntLinear # NOQA
|
||||
@@ -0,0 +1,88 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
|
||||
from ..ops import emulate_int
|
||||
|
||||
|
||||
class ActivationQuantizer:
|
||||
"""
|
||||
Fake scalar quantization of the activations using a forward hook.
|
||||
|
||||
Args:
|
||||
- module. a nn.Module for which we quantize the *post-activations*
|
||||
- p: proportion of activations to quantize, set by default to 1
|
||||
- update_step: to recompute quantization parameters
|
||||
- bits: number of bits for quantization
|
||||
- method: choose among {"tensor", "histogram", "channel"}
|
||||
- clamp_threshold: to prevent gradients overflow
|
||||
|
||||
Remarks:
|
||||
- Parameters scale and zero_point are recomputed every update_step
|
||||
forward pass to reduce the overhead
|
||||
- For the list of quantization methods and number of bits, see ops.py
|
||||
- To remove the hook from the module, simply call self.handle.remove()
|
||||
- At test time, the activations are fully quantized
|
||||
- We use the straight-through estimator so that the gradients
|
||||
back-propagate nicely in the network, this is implemented with
|
||||
the detach() trick
|
||||
- The activations are hard-clamped in [-clamp_threshold, clamp_threshold]
|
||||
to prevent overflow during the backward pass
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
module,
|
||||
p=1,
|
||||
update_step=1000,
|
||||
bits=8,
|
||||
method="histogram",
|
||||
clamp_threshold=5,
|
||||
):
|
||||
self.module = module
|
||||
self.p = p
|
||||
self.update_step = update_step
|
||||
self.counter = 0
|
||||
self.bits = bits
|
||||
self.method = method
|
||||
self.clamp_threshold = clamp_threshold
|
||||
self.handle = None
|
||||
self.register_hook()
|
||||
|
||||
def register_hook(self):
|
||||
# forward hook
|
||||
def quantize_hook(module, x, y):
|
||||
|
||||
# update parameters every 1000 iterations
|
||||
if self.counter % self.update_step == 0:
|
||||
self.scale = None
|
||||
self.zero_point = None
|
||||
self.counter += 1
|
||||
|
||||
# train with QuantNoise and evaluate the fully quantized network
|
||||
p = self.p if self.module.training else 1
|
||||
|
||||
# quantize activations
|
||||
y_q, self.scale, self.zero_point = emulate_int(
|
||||
y.detach(),
|
||||
bits=self.bits,
|
||||
method=self.method,
|
||||
scale=self.scale,
|
||||
zero_point=self.zero_point,
|
||||
)
|
||||
|
||||
# mask to apply noise
|
||||
mask = torch.zeros_like(y)
|
||||
mask.bernoulli_(1 - p)
|
||||
noise = (y_q - y).masked_fill(mask.bool(), 0)
|
||||
|
||||
# using straight-through estimator (STE)
|
||||
clamp_low = -self.scale * self.zero_point
|
||||
clamp_high = self.scale * (2 ** self.bits - 1 - self.zero_point)
|
||||
return torch.clamp(y, clamp_low.item(), clamp_high.item()) + noise.detach()
|
||||
|
||||
# register hook
|
||||
self.handle = self.module.register_forward_hook(quantize_hook)
|
||||
@@ -0,0 +1,149 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.modules.conv import _ConvNd
|
||||
from torch.nn.modules.utils import _pair
|
||||
|
||||
from ..ops import emulate_int
|
||||
|
||||
|
||||
class IntConv2d(_ConvNd):
|
||||
"""
|
||||
Quantized counterpart of the nn.Conv2d module that applies QuantNoise during training.
|
||||
|
||||
Args:
|
||||
- standard nn.Conv2d parameters
|
||||
- p: amount of noise to inject (0 = no quantization, 1 = quantize all the weights)
|
||||
- bits: number of bits
|
||||
- method: choose among {"tensor", "histogram", "channel"}
|
||||
- update_step: recompute scale and zero_point every update_steps iterations
|
||||
|
||||
Remarks:
|
||||
- We use the straight-thgourh estimator so that the gradients
|
||||
back-propagate nicely in the network, this is implemented with
|
||||
the detach() trick
|
||||
- Parameters scale and zero_point are recomputed every update_step
|
||||
forward pass to reduce the overhead
|
||||
- At test time, the weights are fully quantized
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
p=0,
|
||||
bits=8,
|
||||
method="histogram",
|
||||
update_step=1000,
|
||||
):
|
||||
kernel_size = _pair(kernel_size)
|
||||
stride = _pair(stride)
|
||||
padding = _pair(padding)
|
||||
dilation = _pair(dilation)
|
||||
super(IntConv2d, self).__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
False,
|
||||
_pair(0),
|
||||
groups,
|
||||
bias,
|
||||
padding_mode,
|
||||
)
|
||||
|
||||
# quantization parameters
|
||||
self.p = p
|
||||
self.bits = bits
|
||||
self.method = method
|
||||
self.update_step = update_step
|
||||
self.counter = 0
|
||||
|
||||
def _conv_forward(self, input, weight):
|
||||
if self.padding_mode != "zeros":
|
||||
return F.conv2d(
|
||||
F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
|
||||
weight,
|
||||
self.bias,
|
||||
self.stride,
|
||||
_pair(0),
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
return F.conv2d(
|
||||
input,
|
||||
weight,
|
||||
self.bias,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
# train with QuantNoise and evaluate the fully quantized network
|
||||
p = self.p if self.training else 1
|
||||
|
||||
# update parameters every 100 iterations
|
||||
if self.counter % self.update_step == 0:
|
||||
self.scale = None
|
||||
self.zero_point = None
|
||||
self.counter += 1
|
||||
|
||||
# quantize weight
|
||||
weight_quantized, self.scale, self.zero_point = emulate_int(
|
||||
self.weight.detach(),
|
||||
bits=self.bits,
|
||||
method=self.method,
|
||||
scale=self.scale,
|
||||
zero_point=self.zero_point,
|
||||
)
|
||||
|
||||
# mask to apply noise
|
||||
mask = torch.zeros_like(self.weight)
|
||||
mask.bernoulli_(1 - p)
|
||||
noise = (weight_quantized - self.weight).masked_fill(mask.bool(), 0)
|
||||
|
||||
# using straight-through estimator (STE)
|
||||
clamp_low = -self.scale * self.zero_point
|
||||
clamp_high = self.scale * (2 ** self.bits - 1 - self.zero_point)
|
||||
weight = (
|
||||
torch.clamp(self.weight, clamp_low.item(), clamp_high.item())
|
||||
+ noise.detach()
|
||||
)
|
||||
|
||||
# return output
|
||||
output = self._conv_forward(input, weight)
|
||||
return output
|
||||
|
||||
def extra_repr(self):
|
||||
return (
|
||||
"in_channels={}, out_channels={}, kernel_size={}, stride={}, "
|
||||
"padding={}, dilation={}, groups={}, bias={}, quant_noise={}, "
|
||||
"bits={}, method={}".format(
|
||||
self.in_channels,
|
||||
self.out_channels,
|
||||
self.kernel_size,
|
||||
self.stride,
|
||||
self.padding,
|
||||
self.dilation,
|
||||
self.groups,
|
||||
self.bias is not None,
|
||||
self.p,
|
||||
self.bits,
|
||||
self.method,
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,147 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ..ops import emulate_int
|
||||
|
||||
|
||||
class IntEmbedding(nn.Module):
|
||||
"""
|
||||
Quantized counterpart of the nn.Embedding module that applies QuantNoise during training.
|
||||
|
||||
Args:
|
||||
- num_embeddings: number of tokens
|
||||
- embedding_dim: embedding dimension
|
||||
- p: amount of noise to inject (0 = no quantization, 1 = quantize all the weights)
|
||||
- bits: number of bits
|
||||
- method: choose among {"tensor", "histogram", "channel"}
|
||||
- update_step: recompute scale and zero_point every update_steps iterations
|
||||
|
||||
Remarks:
|
||||
- We use the straight-through estimator so that the gradients
|
||||
back-propagate nicely in the network, this is implemented with
|
||||
the detach() trick
|
||||
- Parameters scale and zero_point are recomputed every update_step
|
||||
forward pass to reduce the overhead
|
||||
- At test time, the weights are fully quantized
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
padding_idx=None,
|
||||
max_norm=None,
|
||||
norm_type=2.0,
|
||||
scale_grad_by_freq=False,
|
||||
sparse=False,
|
||||
_weight=None,
|
||||
p=0,
|
||||
update_step=1000,
|
||||
bits=8,
|
||||
method="histogram",
|
||||
):
|
||||
super(IntEmbedding, self).__init__()
|
||||
self.num_embeddings = num_embeddings
|
||||
self.embedding_dim = embedding_dim
|
||||
if padding_idx is not None:
|
||||
if padding_idx > 0:
|
||||
assert (
|
||||
padding_idx < self.num_embeddings
|
||||
), "Padding_idx must be within num_embeddings"
|
||||
elif padding_idx < 0:
|
||||
assert (
|
||||
padding_idx >= -self.num_embeddings
|
||||
), "Padding_idx must be within num_embeddings"
|
||||
padding_idx = self.num_embeddings + padding_idx
|
||||
self.padding_idx = padding_idx
|
||||
self.max_norm = max_norm
|
||||
self.norm_type = norm_type
|
||||
self.scale_grad_by_freq = scale_grad_by_freq
|
||||
if _weight is None:
|
||||
self.weight = nn.Parameter(torch.Tensor(num_embeddings, embedding_dim))
|
||||
self.reset_parameters()
|
||||
else:
|
||||
assert list(_weight.shape) == [
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
], "Shape of weight does not match num_embeddings and embedding_dim"
|
||||
self.weight = nn.Parameter(_weight)
|
||||
self.sparse = sparse
|
||||
|
||||
# quantization parameters
|
||||
self.p = p
|
||||
self.bits = bits
|
||||
self.method = method
|
||||
self.update_step = update_step
|
||||
self.counter = 0
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.normal_(self.weight)
|
||||
if self.padding_idx is not None:
|
||||
with torch.no_grad():
|
||||
self.weight[self.padding_idx].fill_(0)
|
||||
|
||||
def forward(self, input):
|
||||
# train with QuantNoise and evaluate the fully quantized network
|
||||
p = self.p if self.training else 1
|
||||
|
||||
# update parameters every 1000 iterations
|
||||
if self.counter % self.update_step == 0:
|
||||
self.scale = None
|
||||
self.zero_point = None
|
||||
self.counter += 1
|
||||
|
||||
# quantize weight
|
||||
weight_quantized, self.scale, self.zero_point = emulate_int(
|
||||
self.weight.detach(),
|
||||
bits=self.bits,
|
||||
method=self.method,
|
||||
scale=self.scale,
|
||||
zero_point=self.zero_point,
|
||||
)
|
||||
|
||||
# mask to apply noise
|
||||
mask = torch.zeros_like(self.weight)
|
||||
mask.bernoulli_(1 - p)
|
||||
noise = (weight_quantized - self.weight).masked_fill(mask.bool(), 0)
|
||||
|
||||
# using straight-through estimator (STE)
|
||||
clamp_low = -self.scale * self.zero_point
|
||||
clamp_high = self.scale * (2 ** self.bits - 1 - self.zero_point)
|
||||
weight = (
|
||||
torch.clamp(self.weight, clamp_low.item(), clamp_high.item())
|
||||
+ noise.detach()
|
||||
)
|
||||
|
||||
# return output
|
||||
output = F.embedding(
|
||||
input,
|
||||
weight,
|
||||
self.padding_idx,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.sparse,
|
||||
)
|
||||
return output
|
||||
|
||||
def extra_repr(self):
|
||||
s = "{num_embeddings}, {embedding_dim}"
|
||||
if self.padding_idx is not None:
|
||||
s += ", padding_idx={padding_idx}"
|
||||
if self.max_norm is not None:
|
||||
s += ", max_norm={max_norm}"
|
||||
if self.norm_type != 2:
|
||||
s += ", norm_type={norm_type}"
|
||||
if self.scale_grad_by_freq is not False:
|
||||
s += ", scale_grad_by_freq={scale_grad_by_freq}"
|
||||
if self.sparse is not False:
|
||||
s += ", sparse=True"
|
||||
s += "quant_noise={p}, bits={bits}, method={method}"
|
||||
return s.format(**self.__dict__)
|
||||
@@ -0,0 +1,113 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ..ops import emulate_int
|
||||
|
||||
|
||||
class IntLinear(nn.Module):
|
||||
"""
|
||||
Quantized counterpart of the nn.Linear module that applies QuantNoise during training.
|
||||
|
||||
Args:
|
||||
- in_features: input features
|
||||
- out_features: output features
|
||||
- bias: bias or not
|
||||
- p: amount of noise to inject (0 = no quantization, 1 = quantize all the weights)
|
||||
- bits: number of bits
|
||||
- method: choose among {"tensor", "histogram", "channel"}
|
||||
- update_step: recompute scale and zero_point every update_steps iterations
|
||||
|
||||
Remarks:
|
||||
- We use the straight-through estimator so that the gradients
|
||||
back-propagate nicely in the network, this is implemented with
|
||||
the detach() trick.
|
||||
- Parameters scale and zero_point are recomputed every update_step
|
||||
forward pass to reduce the overhead
|
||||
- At test time, the weights are fully quantized
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
bias=True,
|
||||
p=0,
|
||||
update_step=3000,
|
||||
bits=8,
|
||||
method="histogram",
|
||||
):
|
||||
super(IntLinear, self).__init__()
|
||||
self.in_features = int(in_features)
|
||||
self.out_features = int(out_features)
|
||||
self.weight = torch.nn.Parameter(torch.Tensor(out_features, in_features))
|
||||
self.chosen_bias = bias
|
||||
if self.chosen_bias:
|
||||
self.bias = torch.nn.Parameter(torch.Tensor(out_features))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.reset_parameters()
|
||||
|
||||
# quantization parameters
|
||||
self.p = p
|
||||
self.bits = bits
|
||||
self.method = method
|
||||
self.update_step = update_step
|
||||
self.counter = 0
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
if self.chosen_bias:
|
||||
nn.init.constant_(self.bias, 0.0)
|
||||
return
|
||||
|
||||
def forward(self, input):
|
||||
# train with QuantNoise and evaluate the fully quantized network
|
||||
p = self.p if self.training else 1
|
||||
|
||||
# update parameters every 100 iterations
|
||||
if self.counter % self.update_step == 0:
|
||||
self.scale = None
|
||||
self.zero_point = None
|
||||
self.counter += 1
|
||||
|
||||
# quantize weight
|
||||
weight_quantized, self.scale, self.zero_point = emulate_int(
|
||||
self.weight.detach(),
|
||||
bits=self.bits,
|
||||
method=self.method,
|
||||
scale=self.scale,
|
||||
zero_point=self.zero_point,
|
||||
)
|
||||
|
||||
# mask to apply noise
|
||||
mask = torch.zeros_like(self.weight)
|
||||
mask.bernoulli_(1 - p)
|
||||
noise = (weight_quantized - self.weight).masked_fill(mask.bool(), 0)
|
||||
|
||||
# using straight-through estimator (STE)
|
||||
clamp_low = -self.scale * self.zero_point
|
||||
clamp_high = self.scale * (2 ** self.bits - 1 - self.zero_point)
|
||||
weight = (
|
||||
torch.clamp(self.weight, clamp_low.item(), clamp_high.item())
|
||||
+ noise.detach()
|
||||
)
|
||||
|
||||
# return output
|
||||
output = F.linear(input, weight, self.bias)
|
||||
return output
|
||||
|
||||
def extra_repr(self):
|
||||
return "in_features={}, out_features={}, bias={}, quant_noise={}, bits={}, method={}".format(
|
||||
self.in_features,
|
||||
self.out_features,
|
||||
self.bias is not None,
|
||||
self.p,
|
||||
self.bits,
|
||||
self.method,
|
||||
)
|
||||
@@ -0,0 +1,49 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def emulate_int(w, bits, method, scale=None, zero_point=None):
|
||||
q = globals()[f"emulate_int{bits}_{method}"]
|
||||
return q(w, scale=scale, zero_point=zero_point)
|
||||
|
||||
|
||||
def quantize(w, scale, zero_point):
|
||||
return (
|
||||
torch.clamp(torch.round(w / scale + zero_point), 0, 255) - zero_point
|
||||
) * scale
|
||||
|
||||
|
||||
def emulate_int8_histogram(w, scale=None, zero_point=None):
|
||||
if scale is None:
|
||||
obs = torch.quantization.observer.HistogramObserver()
|
||||
_ = obs(w.float())
|
||||
scale, zero_point = obs.calculate_qparams()
|
||||
scale = scale.cuda().type_as(w)
|
||||
zero_point = zero_point.cuda().type_as(w)
|
||||
return quantize(w, scale, zero_point), scale, zero_point
|
||||
|
||||
|
||||
def emulate_int8_channel(w, scale=None, zero_point=None):
|
||||
if scale is None:
|
||||
obs = torch.quantization.observer.PerChannelMinMaxObserver(
|
||||
ch_axis=-1, qscheme=torch.per_channel_symmetric
|
||||
)
|
||||
_ = obs(w)
|
||||
scale, zero_point, ch_axis = obs.get_qparams()
|
||||
scale = scale.cuda().type_as(w)
|
||||
zero_point = zero_point.cuda().type_as(w)
|
||||
return quantize(w, scale, zero_point), scale, zero_point
|
||||
|
||||
|
||||
def emulate_int8_tensor(w, scale=None, zero_point=None):
|
||||
if scale is None:
|
||||
obs = torch.quantization.observer.MinMaxObserver()
|
||||
_ = obs(w)
|
||||
scale, zero_point = obs.calculate_qparams()
|
||||
scale = scale.cuda().type_as(w)
|
||||
zero_point = zero_point.cuda().type_as(w)
|
||||
return quantize(w, scale, zero_point), scale, zero_point
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
from operator import attrgetter
|
||||
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
|
||||
from ..pq.utils import attrsetter, get_layers
|
||||
from .modules import ActivationQuantizer, IntConv2d, IntEmbedding, IntLinear
|
||||
|
||||
|
||||
MAPPING = {nn.Linear: IntLinear, nn.Embedding: IntEmbedding, nn.Conv2d: IntConv2d}
|
||||
|
||||
|
||||
def quantize_model_(model, p=0.2, bits=8, update_step=3000):
|
||||
"""
|
||||
Replaces all modules with their scalar quantized counterpart and
|
||||
registers hooks to quantize the post-ativations of those modules.
|
||||
|
||||
Args:
|
||||
- model: a nn.Module
|
||||
- p: amount of noise (0 for no noise, 1 to quantize all the weights/activations)
|
||||
- bits: number of bits
|
||||
- update_step: update quantization parameters every update_step steps
|
||||
"""
|
||||
|
||||
# quantize all layers
|
||||
quantized_layers = get_layers(model, "(.*?)")
|
||||
|
||||
for layer in quantized_layers:
|
||||
|
||||
# book-keeping
|
||||
is_master_process = (not dist.is_initialized()) or (
|
||||
dist.is_initialized() and dist.get_rank() == 0
|
||||
)
|
||||
|
||||
# recover module
|
||||
module = attrgetter(layer)(model)
|
||||
if is_master_process:
|
||||
logging.info(
|
||||
f"Quantizing layer {layer} with bits={bits} and QuantNoise={p}"
|
||||
)
|
||||
|
||||
# quantization params
|
||||
q_params = {
|
||||
"p": p,
|
||||
"update_step": update_step,
|
||||
"bits": bits,
|
||||
"method": "histogram",
|
||||
"counter": 0,
|
||||
}
|
||||
|
||||
# instantiate the quantized counterpart
|
||||
if isinstance(module, tuple(MAPPING.keys())):
|
||||
QuantizedModule = MAPPING[module.__class__]
|
||||
quantized_module = QuantizedModule.__new__(QuantizedModule)
|
||||
params = module.__dict__
|
||||
params.update(q_params)
|
||||
quantized_module.__dict__.update(params)
|
||||
|
||||
else:
|
||||
if is_master_process:
|
||||
logging.info(f"Module {module} not yet supported for quantization")
|
||||
continue
|
||||
|
||||
# activation quantization
|
||||
a_q = ActivationQuantizer(quantized_module, p=0, bits=bits, method="histogram")
|
||||
|
||||
# replace layer by its quantized counterpart
|
||||
attrsetter(layer)(model, quantized_module)
|
||||
|
||||
# return name of quantized layers
|
||||
return quantized_layers
|
||||
Reference in New Issue
Block a user