mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
d23d62cb1e
Summary: Currently accelerators does not have the concept for fp32, it only has understandings of fp16 and int8 in terms of data input. In order to fixe the issue here, we want to make sure unaries are turned into fp16 when we have the int8 exporter turned on. Reviewed By: kennyhorror Differential Revision: D17743791 fbshipit-source-id: 7322d23eb12ac3f813b525fc0ddd066f95c8ca85
556 lines
22 KiB
Python
556 lines
22 KiB
Python
## @package sparse_lookup
|
|
# Module caffe2.python.layers.sparse_lookup
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
from caffe2.python.optimizer import FP16_ENGINES, Optimizer
|
|
from caffe2.python.helpers.arg_scope import get_current_scope
|
|
from caffe2.python import schema
|
|
from caffe2.python.layers.layers import (
|
|
get_categorical_limit,
|
|
get_key,
|
|
IdList,
|
|
IdScoreList,
|
|
IdListWithEvicted,
|
|
IdScoreListWithEvicted,
|
|
LayerPsParam,
|
|
ModelLayer,
|
|
almost_equal_schemas,
|
|
)
|
|
import collections
|
|
import functools
|
|
import logging
|
|
import math
|
|
import numpy as np
|
|
import operator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_trainer_version_based_on_optim(optim_def):
|
|
if isinstance(optim_def, Optimizer) and hasattr(optim_def, "engine"):
|
|
logger.info(
|
|
"Attempting to set trainer version for engine {}".format(optim_def.engine)
|
|
)
|
|
if optim_def.engine in FP16_ENGINES:
|
|
logger.info("Setting FP16 trainer for engine {}".format(optim_def.engine))
|
|
return "fp16"
|
|
else:
|
|
logger.info("Setting FP32 trainer for engine {}".format(optim_def.engine))
|
|
return "fp32"
|
|
else:
|
|
return "fp32"
|
|
|
|
|
|
def get_sparse_lookup_predictor_version(
|
|
version,
|
|
blob_size=None,
|
|
min_blob_size_4bits=None,
|
|
embedding_dim=None,
|
|
sparse_feature_name=None,
|
|
):
|
|
assert version in {
|
|
'fp32', 'fp16', 'uint8rowwise', 'fused_uint8rowwise', 'fused_uint4rowwise'
|
|
}, "Unexpected version of sparse_lookup layer {0}".format(version)
|
|
if version == 'fused_uint4rowwise':
|
|
if (
|
|
blob_size is not None
|
|
and min_blob_size_4bits is not None
|
|
and embedding_dim is not None
|
|
):
|
|
if blob_size < min_blob_size_4bits:
|
|
logger.info(
|
|
"{} fall back to uint8 because lookup table size {} < min_blob_size_4bits {}".format(
|
|
sparse_feature_name,
|
|
blob_size,
|
|
min_blob_size_4bits,
|
|
)
|
|
)
|
|
version = 'fused_uint8rowwise'
|
|
|
|
if embedding_dim % 2 == 1:
|
|
logger.info(
|
|
"{} fall back to uint8 because lookup table dimension {} is not divisible by 2".format(
|
|
sparse_feature_name, embedding_dim
|
|
)
|
|
)
|
|
version = 'fused_uint8rowwise'
|
|
else:
|
|
raise ValueError(
|
|
(
|
|
"When 4 bit quantization is enabled for {}, "
|
|
"(i.e., Sparse lookup predictor version:{}), "
|
|
"requires arguments blob_size:{}, "
|
|
"min_blob_size_4bits:{}, embedding_dim:{}"
|
|
).format(
|
|
sparse_feature_name,
|
|
version,
|
|
blob_size,
|
|
min_blob_size_4bits,
|
|
embedding_dim
|
|
)
|
|
)
|
|
return version
|
|
|
|
|
|
def get_sparse_lookup_trainer_version(version):
|
|
assert version in {'fp32', 'fp16'},\
|
|
"Unexpected version of sparse_lookup layer {0}".format(version)
|
|
return version
|
|
|
|
def _is_id_list(input_record):
|
|
return almost_equal_schemas(input_record, IdList)
|
|
|
|
|
|
def _is_id_score_list(input_record):
|
|
return almost_equal_schemas(input_record,
|
|
IdScoreList,
|
|
check_field_types=False)
|
|
|
|
|
|
class SparseLookup(ModelLayer):
|
|
_id_list_supported_reducers = [
|
|
'LogMeanExp', 'LogSumExp', 'Max', 'Mean', 'Sum',
|
|
'WeightedSum', 'WeightedMean', 'Sqrt', 'None']
|
|
|
|
_id_score_list_supported_reducers = [
|
|
'PositionWeighted', 'RecencyWeighted', 'Mean', 'Sum', 'WeightedSum',
|
|
'WeightedMean', 'None'
|
|
]
|
|
|
|
_fp16_compatible_init_op_types = [
|
|
'Float16UniformFill'
|
|
]
|
|
|
|
_fp16_compatible_reducers = [
|
|
'Sum', 'Mean', 'Sqrt', 'PositionWeighted', 'RecencyWeighted',
|
|
]
|
|
|
|
def __init__(self, model, input_record, inner_shape, reducer,
|
|
weight_init=None, weight_optim=None,
|
|
name='sparse_lookup', regularizer=None, use_external_weights=False, **kwargs):
|
|
|
|
super(SparseLookup, self).__init__(model, name, input_record, **kwargs)
|
|
|
|
self.sparse_key = get_key(self.input_record)()
|
|
logger.info("Setup the sparse lookup layer for " + self.sparse_key)
|
|
|
|
# TODO Add some asserts about input type
|
|
if isinstance(inner_shape, int):
|
|
inner_shape = [inner_shape]
|
|
assert isinstance(inner_shape, list) or isinstance(inner_shape, tuple),\
|
|
"Unexpected type for inner_shape, expected list or tuple, got {0} for {1}".\
|
|
format(type(inner_shape), self.sparse_key)
|
|
|
|
if reducer == "PositionWeighted":
|
|
assert _is_id_score_list(self.input_record), (
|
|
"PositionWeighted only support IdScoreList, but got {} for {}"
|
|
+ "please use PositionWeighted layer to convert IdList "
|
|
+ "to IdScoreList"
|
|
).format(repr(self.input_record), self.sparse_key)
|
|
self.external_weights = self.input_record.values()
|
|
|
|
elif reducer == "RecencyWeighted":
|
|
assert _is_id_score_list(self.input_record), (
|
|
"RecencyWeighted only supports IdScoreList, "
|
|
"while the sparse feature {} is not.".format(self.sparse_key)
|
|
)
|
|
self.external_weights = self.input_record.values()
|
|
# TODO: create a new type of reducer with external weights to wrap
|
|
# this and the above two cases since essentially their input formats
|
|
# are the same.
|
|
elif use_external_weights:
|
|
assert _is_id_score_list(self.input_record), (
|
|
"Use_external_weights only supports IdScoreList, "
|
|
"while the sparse feature {} is not.".format(self.sparse_key)
|
|
)
|
|
assert reducer in ["Sum", "WeightedSum"], (
|
|
"Use_external_weights only supports Sum reducer, "
|
|
"while the reducer is {}.".format(reducer)
|
|
)
|
|
self.external_weights = self.input_record.values()
|
|
self.reducer = reducer
|
|
self.use_external_weights = use_external_weights
|
|
|
|
input_dim = get_categorical_limit(self.input_record)
|
|
assert input_dim > 0, "{} should have categorical limit > 0, but got {}".format(
|
|
self.sparse_key, input_dim
|
|
)
|
|
|
|
self.input_dim = input_dim
|
|
self.shape = [input_dim] + inner_shape
|
|
|
|
self.trainer_version = get_trainer_version_based_on_optim(
|
|
weight_optim
|
|
)
|
|
|
|
default_init_op = self._get_default_init_op()
|
|
|
|
self.weight_init = weight_init or default_init_op
|
|
|
|
self.evicted_values = None
|
|
if schema.equal_schemas(
|
|
self.input_record, IdListWithEvicted
|
|
) or schema.equal_schemas(
|
|
self.input_record, IdScoreListWithEvicted, check_field_types=False
|
|
):
|
|
self.evicted_values = self.input_record._evicted_values
|
|
|
|
# If fp16 is used, make sure fp16 init op is used
|
|
if self.trainer_version == "fp16":
|
|
assert self.reducer in self._fp16_compatible_reducers or use_external_weights, (
|
|
"Fp16 training is enabled. The reducer specified is not supported. "
|
|
"Got {}. Supported reducers: {}. Right now, in general, sum, mean, "
|
|
"positional pooling are supported. Attention is not. Please check "
|
|
"if there is fp16 trained sparse features using advanced pooling.".format(
|
|
self.reducer, self._fp16_compatible_reducers)
|
|
)
|
|
|
|
# if init op is UniformFill, we replace it directly
|
|
if self.weight_init[0] == "UniformFill":
|
|
self.weight_init = ("Float16UniformFill", self.weight_init[1])
|
|
assert self.weight_init[0] in self._fp16_compatible_init_op_types, (
|
|
"Fp16 training is enabled. Init op for weight parameter must be fp16 "
|
|
"compatibale. Got {}. Supported ops: {}".format(
|
|
self.weight_init[0],
|
|
self._fp16_compatible_init_op_types)
|
|
)
|
|
|
|
assert regularizer is None, "Regularizer is not compatible with fp16"
|
|
|
|
if self.input_record.lengths.metadata:
|
|
avg_length = self.input_record.lengths.metadata.expected_value
|
|
else:
|
|
avg_length = None
|
|
|
|
self.w = self.create_param(
|
|
param_name='w',
|
|
shape=self.shape,
|
|
initializer=self.weight_init,
|
|
optimizer=weight_optim,
|
|
ps_param=LayerPsParam(
|
|
sparse_key=self.sparse_key,
|
|
average_length=avg_length),
|
|
regularizer=regularizer
|
|
)
|
|
if self.evicted_values:
|
|
self.reinit_vec = self.create_param(
|
|
param_name="reinit_vec",
|
|
shape=inner_shape,
|
|
initializer=self.weight_init,
|
|
optimizer=model.NoOptim,
|
|
regularizer=None,
|
|
)
|
|
|
|
self.scale_bias_init = ('ConstantFill', {'value': 0.0})
|
|
|
|
self.scale_bias = self.create_param(
|
|
param_name='scale_bias',
|
|
shape=[],
|
|
initializer=self.scale_bias_init,
|
|
optimizer=model.NoOptim,
|
|
)
|
|
|
|
self.output_schema = schema.Scalar(
|
|
(np.float32, inner_shape),
|
|
self.get_next_blob_reference('output'),
|
|
)
|
|
|
|
def get_memory_usage(self):
|
|
return functools.reduce(operator.mul, self.shape) * 4
|
|
|
|
def get_fp16_compatible_parameters(self):
|
|
return [self.w]
|
|
|
|
def support_8bit(self):
|
|
# Rowwise quantization makes sense only if shape it's 2D matrix with
|
|
# second dimension >= 8
|
|
if len(self.shape) != 2 or self.shape[1] < 8:
|
|
return False
|
|
return True
|
|
|
|
def get_8bits_compatible_parameters(self, fused=True):
|
|
if not self.support_8bit():
|
|
return []
|
|
if fused:
|
|
RowwiseQuantized8BitsWeight = collections.namedtuple(
|
|
'RowwiseQuantized8BitsWeight', 'w'
|
|
)
|
|
return [RowwiseQuantized8BitsWeight(self.w)]
|
|
else:
|
|
RowwiseQuantized8BitsWeight = collections.namedtuple(
|
|
'RowwiseQuantized8BitsWeight', 'w, scale_bias'
|
|
)
|
|
return [RowwiseQuantized8BitsWeight(self.w, self.scale_bias)]
|
|
|
|
def _get_default_init_op(self):
|
|
scale = math.sqrt(1.0 / self.input_dim)
|
|
|
|
if self.trainer_version == 'fp32':
|
|
default_weight_init = ('UniformFill', {'min': -scale, 'max': scale})
|
|
elif self.trainer_version == 'fp16':
|
|
default_weight_init = ("Float16UniformFill", {'min': -scale, 'max': scale})
|
|
else:
|
|
raise NotImplementedError(
|
|
"Train version {} is not currently supported for sparse feature {}".format(
|
|
trainer_version, self.sparse_key
|
|
)
|
|
)
|
|
|
|
return default_weight_init
|
|
|
|
def _gather_wrapper(self, net, version, in_indices, out):
|
|
# Gather can work on all kinds of input data types, and output
|
|
# data with the same type. Convert the output of Gather to float,
|
|
# because the follow-up Ops expect fp32.
|
|
if version == 'fp32':
|
|
return net.Gather([self.w, in_indices], out)
|
|
elif version == 'fp16':
|
|
gathered_w = net.Gather([self.w, in_indices], 'gathered_w')
|
|
return net.HalfToFloat(gathered_w, out)
|
|
elif version == 'uint8rowwise':
|
|
gathered_w = net.Gather([self.w, in_indices], 'gathered_w')
|
|
gathered_scale_bias = net.Gather(
|
|
[self.scale_bias, in_indices],
|
|
'gathered_scale_bias'
|
|
)
|
|
|
|
return net.Rowwise8BitQuantizedToFloat(
|
|
[gathered_w, gathered_scale_bias], out)
|
|
elif version == 'fused_uint8rowwise':
|
|
gathered_w = net.Gather([self.w, in_indices], 'gathered_w')
|
|
return net.Fused8BitRowwiseQuantizedToFloat(gathered_w, out)
|
|
elif version == 'fused_uint4rowwise':
|
|
gathered_w = net.Gather([self.w, in_indices], 'gathered_w')
|
|
return net.Fused4BitRowwiseQuantizedToFloat(gathered_w, out)
|
|
|
|
else:
|
|
raise "Unsupported version of operators in SparseLookup " +\
|
|
"layer: {0} for sparse feature {1}".format(
|
|
version, self.sparse_key
|
|
)
|
|
|
|
def _sparse_lengths_weighted_reducer(
|
|
self,
|
|
in_indices,
|
|
weights,
|
|
reducer,
|
|
net,
|
|
version,
|
|
grad_on_weights=0,
|
|
):
|
|
op_input = [
|
|
self.w,
|
|
weights,
|
|
in_indices,
|
|
self.input_record.lengths(),
|
|
]
|
|
layer_name = 'SparseLengths' + reducer
|
|
|
|
if version in ['fp32', 'fp16']:
|
|
# SparseLengths* Ops will accept either fp16 or fp32 embedding
|
|
# matrix and output fp32 pooled embedding
|
|
# A special case here is that we need FP16 engine for
|
|
# SparseLengthsWeightedSum when FP16 embeedings are used for
|
|
# correct backward updates
|
|
if reducer == "WeightedSum" and version == "fp16":
|
|
net.SparseLengthsWeightedSum(
|
|
op_input,
|
|
self.output_schema.field_blobs(),
|
|
grad_on_weights=grad_on_weights,
|
|
engine='FP16',
|
|
)
|
|
else:
|
|
net.__getattr__(layer_name)(
|
|
op_input,
|
|
self.output_schema.field_blobs(),
|
|
grad_on_weights=grad_on_weights,
|
|
)
|
|
elif version == 'uint8rowwise':
|
|
op_input.insert(len(op_input), self.scale_bias)
|
|
net.__getattr__(layer_name + '8BitsRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
elif version == 'fused_uint8rowwise':
|
|
net.__getattr__(layer_name + 'Fused8BitRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
elif version == 'fused_uint4rowwise':
|
|
net.__getattr__(layer_name + 'Fused4BitRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
else:
|
|
raise "Unsupported version of operator in SparseLookUp " +\
|
|
"layer: {0} for sparse feature {1}".format(
|
|
version, self.sparse_key
|
|
)
|
|
|
|
# deal with sparse features of id_list type
|
|
def _add_ops_id_list(self, net, version):
|
|
assert self.reducer in self._id_list_supported_reducers, (
|
|
"Unsupported reducer: {} for ID_LIST {}".format(
|
|
self.reducer, self.sparse_key
|
|
)
|
|
)
|
|
if self.reducer in ['Sum', 'Mean', 'WeightedSum', 'WeightedMean']:
|
|
op_input = [self.w,
|
|
self.input_record.items(),
|
|
self.input_record.lengths()]
|
|
|
|
# For id list features, the behaviors of 'Sum' and
|
|
# 'WeightedSum' are identical, since we can regard the weight on each
|
|
# id as 1. Similarly, for 'Mean' and 'WeightedMean'.
|
|
if self.reducer == 'WeightedSum':
|
|
self.reducer = 'Sum'
|
|
elif self.reducer == 'WeightedMean':
|
|
self.reducer = 'Mean'
|
|
|
|
layer_name = 'SparseLengths' + self.reducer
|
|
if version in ['fp32', 'fp16']:
|
|
# SparseLengths* Ops will accept either fp16 or fp32 embedding
|
|
# matrix and output fp32 pooled embedding
|
|
net.__getattr__(layer_name)(
|
|
op_input,
|
|
self.output_schema.field_blobs(),
|
|
)
|
|
elif version == 'uint8rowwise':
|
|
op_input.insert(len(op_input), self.scale_bias)
|
|
net.__getattr__(layer_name + '8BitsRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
elif version == 'fused_uint8rowwise':
|
|
net.__getattr__(layer_name + 'Fused8BitRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
elif version == 'fused_uint4rowwise':
|
|
net.__getattr__(layer_name + 'Fused4BitRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
else:
|
|
raise "Unsupported version of operator in SparseLookUp " +\
|
|
"layer: {0} for sparse feature {1}".format(
|
|
version, self.sparse_key
|
|
)
|
|
|
|
elif self.reducer == 'Sqrt':
|
|
sqrt_weight = net.LengthsToWeights(
|
|
[self.input_record.lengths()],
|
|
[net.NextScopedBlob('lengths_sqrt')],
|
|
power=0.5,
|
|
)
|
|
self._sparse_lengths_weighted_reducer(
|
|
self.input_record.items(),
|
|
sqrt_weight,
|
|
'WeightedSum', net, version)
|
|
|
|
elif self.reducer == 'None':
|
|
# Gather operator will gather the embedding for each id of
|
|
# each IdList.
|
|
self._gather_wrapper(net, version, self.input_record.items(),
|
|
self.output_schema.field_blobs())
|
|
|
|
else:
|
|
table_rows = self._gather_wrapper(
|
|
net, version, self.input_record.items(), 'table_rows')
|
|
|
|
segment_ids = net.LengthsToSegmentIds(
|
|
self.input_record.lengths(),
|
|
net.NextScopedBlob(self.input_record.lengths() + '_sid'))
|
|
net.__getattr__('SortedSegmentRange' + self.reducer)(
|
|
[table_rows, segment_ids],
|
|
self.output_schema.field_blobs(),
|
|
)
|
|
|
|
# deal with sparse features of id_score_list type
|
|
def _add_ops_id_score_list(self, net, version):
|
|
assert self.reducer in self._id_score_list_supported_reducers, (
|
|
"Unsupported reducer: {} for ID_SCORE_LIST {}".format(
|
|
self.reducer, self.sparse_key
|
|
)
|
|
)
|
|
if self.reducer in ['WeightedSum', 'WeightedMean']:
|
|
self._sparse_lengths_weighted_reducer(
|
|
self.input_record.keys(),
|
|
self.input_record.values(),
|
|
self.reducer, net, version)
|
|
|
|
elif self.reducer in ['PositionWeighted', 'RecencyWeighted'] or self.use_external_weights:
|
|
self._sparse_lengths_weighted_reducer(
|
|
self.input_record.keys(),
|
|
self.external_weights,
|
|
'WeightedSum', net, version, grad_on_weights=1)
|
|
|
|
elif self.reducer in ['Sum', 'Mean']:
|
|
op_input = [self.w,
|
|
self.input_record.keys(),
|
|
self.input_record.lengths()]
|
|
|
|
layer_name = 'SparseLengths' + self.reducer
|
|
|
|
if version in ['fp32', 'fp16']:
|
|
net.__getattr__(layer_name)(
|
|
op_input,
|
|
self.output_schema.field_blobs(),
|
|
)
|
|
elif version == 'uint8rowwise':
|
|
net.__getattr__(layer_name + '8BitsRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
elif version == 'fused_uint8rowwise':
|
|
net.__getattr__(layer_name + 'Fused8BitRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
elif version == 'fused_uint4rowwise':
|
|
net.__getattr__(layer_name + 'Fused4BitRowwise')(
|
|
op_input, self.output_schema.field_blobs())
|
|
else:
|
|
raise "Unsupported version of operator in SparseLookUp " +\
|
|
"layer: {0} for sparse feature {1}".format(
|
|
version, self.sparse_key
|
|
)
|
|
|
|
elif self.reducer == 'None':
|
|
# Gather operator will gather the embedding for each id of
|
|
# each IdList.
|
|
self._gather_wrapper(net, version, self.input_record.keys(),
|
|
self.output_schema.field_blobs())
|
|
else:
|
|
raise "Only Sum, Mean, None are supported for IdScoreList input." +\
|
|
"Trying to create with {} for sparse feature {}".format(
|
|
self.reducer, self.sparse_key
|
|
)
|
|
|
|
def _add_ops(self, net, version='fp32'):
|
|
if self.evicted_values:
|
|
net.CopyRowsToTensor(
|
|
[self.w, self.evicted_values.get(), self.reinit_vec], [self.w])
|
|
if _is_id_list(self.input_record):
|
|
self._add_ops_id_list(net, version=version)
|
|
elif _is_id_score_list(self.input_record):
|
|
self._add_ops_id_score_list(net, version=version)
|
|
else:
|
|
raise "Unsupported input type {0}".format(self.input_record)
|
|
|
|
def add_train_ops(self, net):
|
|
self._add_ops(net, self.trainer_version)
|
|
|
|
def add_ops(self, net):
|
|
version_info = get_current_scope().get(
|
|
get_sparse_lookup_predictor_version.__name__, {'version': 'fp32'}
|
|
)
|
|
lookup_table_blob_size = self.shape[0] * self.shape[1]
|
|
version = get_sparse_lookup_predictor_version(
|
|
version_info['version'],
|
|
blob_size=lookup_table_blob_size,
|
|
min_blob_size_4bits=(
|
|
version_info['min_blob_size_4bits']
|
|
if 'min_blob_size_4bits' in version_info
|
|
else None
|
|
),
|
|
embedding_dim=self.shape[1],
|
|
sparse_feature_name=self.sparse_key,
|
|
)
|
|
|
|
# TODO(amalevich): Layer should not be responsible for decision about
|
|
# quantization.
|
|
if not self.support_8bit() and version in {'uint8rowwise',
|
|
'fused_uint8rowwise',
|
|
'fused_uint4rowwise'}:
|
|
version = 'fp16'
|
|
|
|
self._add_ops(net, version)
|