Skip to content

Fixed UT error for bf16 op list for QAT mode #200

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Dec 2, 2022
Merged

Fixed UT error for bf16 op list for QAT mode #200

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bb51543
Fixed UT error for bf16 op list for QAT mode
PenghuiCheng Dec 1, 2022
ed09111
Add bf16 op list for eager QAT mode
PenghuiCheng Dec 1, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 36 additions & 9 deletions neural_compressor/adaptor/pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1029,7 +1029,7 @@ def _get_quantizable_ops(self, model):


# get bf16 capability
if (CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1') and \
if self.use_bf16 and (CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1') and \
(self.version.release >= Version("1.11.0").release):
self.bf16_ops = self.query_handler.get_op_types_by_precision("bf16")
bf16_ops = []
Expand Down Expand Up @@ -1308,19 +1308,34 @@ def _pre_hook_for_qat(self, dataloader=None):
qscheme=torch.per_tensor_affine,
reduce_range=REDUCE_RANGE),
weight=torch.quantization.default_weight_fake_quant)
self.non_quant_dict = self.get_non_quant_modules(self.model.kwargs)
quantizable_ops = []
self._get_quantizable_ops_recursively(self.model._model, '', quantizable_ops)
self.bf16_ops = self.query_handler.get_op_types_by_precision("bf16")
bf16_ops = []
if self.version.release >= Version("1.11.0").release and self.use_bf16 and \
(CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1'): # pragma: no cover
self._get_bf16_ops_recursively(self.model._model, '', bf16_ops)
bf16_ops_list = [(op) for op in bf16_ops if op not in quantizable_ops]
self.model.model.training = True
torch.quantization.prepare_qat(self.model._model, inplace=True)

def _post_hook_for_qat(self):
torch.quantization.convert(self.model._model, inplace=True)
# This is a flag for reloading
self.model.q_config = {
'is_oneshot': True,
'framework': 'pytorch',
'reduce_range': REDUCE_RANGE,
'approach': 'quant_aware_training'
'approach': 'quant_aware_training',
'bf16_ops_list': bf16_ops_list,
}

def _post_hook_for_qat(self):
torch.quantization.convert(self.model._model, inplace=True)
if len(self.model.q_config['bf16_ops_list']) > 0 and \
self.version.release >= Version("1.11.0").release and self.use_bf16 and \
(CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1'): # pragma: no cover
self.model._model = torch_utils.bf16_convert.Convert(self.model._model, self.model.q_config)

def _pre_hook_for_hvd(self, dataloader=None):
# TODO: lazy init here
hvd.init()
Expand Down Expand Up @@ -2220,7 +2235,8 @@ def quantize(self, tune_cfg, model, dataloader, q_func=None):
self.model_calibration(q_model, dataloader, iterations, None,
tune_cfg.get('calib_sampling_size', 1))
q_model.save_qconf_summary(qconf_summary=self.ipex_config_path)
if self.use_bf16:
if self.use_bf16 and (CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1') and \
(self.version.release >= Version("1.11.0").release):
with torch.no_grad():
with torch.cpu.amp.autocast():
q_model = ipex.quantization.convert(q_model)
Expand Down Expand Up @@ -2487,7 +2503,7 @@ def _get_quantizable_ops_recursively(self, model, prefix, quantizable_ops):
if isinstance(self.q_dataloader, BaseDataLoader):
self.q_dataloader.batch(batch_size)
logger.info('Recovery `calibration.dataloader.batchsize` {} according \
to config.yaml'.format(batch_size))
to config.yaml' .format(batch_size))
del init_model
with open(self.ipex_config_path, 'r') as f:
self.cfgs = json.load(f)
Expand Down Expand Up @@ -2773,7 +2789,7 @@ def quantize(self, tune_cfg, model, dataloader, q_func=None):
q_model._model, prefix='')

if len(self.tune_cfg['bf16_ops_list']) > 0 and \
self.version.release >= Version("1.11.0").release and \
self.version.release >= Version("1.11.0").release and self.use_bf16 and \
(CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1'): # pragma: no cover
q_model._model = torch_utils.bf16_convert.Convert(q_model._model, self.tune_cfg)

Expand Down Expand Up @@ -2843,6 +2859,12 @@ def _pre_hook_for_qat(self, dataloader=None):
quantizable_ops = []
tmp_model = self.fuse_fx_model(self.model, is_qat=True)
self._get_quantizable_ops_recursively(tmp_model, '', quantizable_ops)
self.bf16_ops = self.query_handler.get_op_types_by_precision("bf16")
bf16_ops = []
if self.version.release >= Version("1.11.0").release and self.use_bf16 and \
(CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1'): # pragma: no cover
self._get_bf16_ops_recursively(tmp_model, '', bf16_ops)
bf16_ops_list = [(op) for op in bf16_ops if op not in quantizable_ops]
quantized_ops = OrderedDict()
for op in quantizable_ops:
if op[1] in [
Expand All @@ -2851,7 +2873,7 @@ def _pre_hook_for_qat(self, dataloader=None):
quantized_ops[op[0]] = torch.quantization.default_dynamic_qconfig
else:
quantized_ops[op[0]] = q_cfgs
# build for fetching scale and zeropoint
# build for fetching scale and zeropoint
op_config_dict = {}
for op in quantizable_ops:
op_config_dict[op] = {'weight': {'dtype': 'int8'}, 'activation': {'dtype': 'uint8'}}
Expand Down Expand Up @@ -2901,6 +2923,7 @@ def _pre_hook_for_qat(self, dataloader=None):
'framework': 'pytorch_fx',
'reduce_range': REDUCE_RANGE,
'quantizable_ops': quantizable_ops,
'bf16_ops_list': bf16_ops_list,
'op': op_config_dict,
'sub_module_list': self.sub_module_list,
'approach': 'quant_aware_training'
Expand All @@ -2926,6 +2949,10 @@ def _post_hook_for_qat(self):

if self.approach != 'post_training_dynamic_quant':
self._get_scale_zeropoint(self.model._model, self.model.q_config)
if len(self.model.q_config['bf16_ops_list']) > 0 and \
self.version.release >= Version("1.11.0").release and self.use_bf16 and \
(CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1'): # pragma: no cover
self.model._model = torch_utils.bf16_convert.Convert(self.model._model, self.model.q_config)
self._dump_model_op_stats(self.model._model, self.model.q_config, self.approach)
torch_utils.util.get_embedding_contiguous(self.model._model)

Expand Down Expand Up @@ -3102,7 +3129,7 @@ def _dump_model_op_stats(self, model, tune_cfg, approach):
res = dict()
self._get_sub_module_op_stats(model, tune_cfg, approach, res)

if (self.version.release >= Version("1.11.0").release) and \
if self.use_bf16 and (self.version.release >= Version("1.11.0").release) and \
(CpuInfo().bf16 or os.getenv('FORCE_BF16') == '1'): # pragma: no cover
bf16_ops_list = tune_cfg['bf16_ops_list']
if len(bf16_ops_list) > 0:
Expand Down