enable split sgd for embeddingbag's weight. conv/linear's bias (#127)

* enable split sgd for embeddingbag's weight. conv/linear's bias

* refine ipex.optimize code

* fix split std momentum_buffer dtype

Author: zhuhaozhe

intel_pytorch_extension_py/__init__.py

@@ -11,4 +11,6 @@
 from .ops import *
 from .utils import *
 from .weight_prepack import *
+from .optimizer_utils import *
+from .weight_cast import *
 from .optim import *

intel_pytorch_extension_py/optim/_functional.py

@@ -81,25 +81,48 @@ def sgd(params: List[Tensor],
 
     for i, param in enumerate(params):
 
+
         d_p = d_p_list[i]
+        float_d_p, float_param = None, None
+        if d_p.dtype == torch.bfloat16:
+            assert param in attr, "split sgd requires record 'trail' part of params in attr"
+            trail = attr[param]['trail']
+
+        if weight_decay != 0 or momentum != 0:
+            float_d_p = d_p.float()
+            if  d_p.dtype == torch.bfloat16:
+                float_d_p = d_p.float()
+                float_param =  torch.ops.torch_ipex.cat_bfloat16_float(param, trail)
+            else:
+                float_param =  param
+                float_d_p = d_p
+
         if weight_decay != 0:
-            d_p = d_p.add(param, alpha=weight_decay)
+            float_d_p = float_d_p.add(float_param, alpha=weight_decay)
 
         if momentum != 0:
             buf = momentum_buffer_list[i]
-
             if buf is None:
-                buf = torch.clone(d_p).detach()
+                buf = torch.clone(float_d_p).detach()
                 momentum_buffer_list[i] = buf
             else:
-                buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
+                buf.mul_(momentum).add_(float_d_p, alpha=1 - dampening)
 
             if nesterov:
-                d_p = d_p.add(buf, alpha=momentum)
+                float_d_p = d_p.add(buf, alpha=momentum)
             else:
-                d_p = buf
-        if param.dtype == torch.bfloat16 and param in attr:
-            trail = attr[param]['trail']
-            torch.ops.torch_ipex.packed_add(param, trail, d_p, alpha=-lr)
+                float_d_p = buf
+
+        if param.dtype is torch.bfloat16:
+            if float_d_p is not None and float_param is not None:
+                float_param.add_(float_d_p, alpha=-lr)
+                top_half, bot_half = torch.ops.torch_ipex.split_float_bfloat16(float_param)
+                param.copy_(top_half)
+                trail.copy_(bot_half)
+            else:
+                torch.ops.torch_ipex.packed_add(param, trail, d_p, alpha=-lr)
         else:
-            param.add_(d_p, alpha=-lr)
+            if float_d_p is not None:
+                param.add_(float_d_p, alpha=-lr)
+            else:
+                param.add_(d_p, alpha=-lr)

intel_pytorch_extension_py/optimizer_utils.py

@@ -14,75 +14,27 @@ class _ipex_optimizer(torch.optim.Optimizer):
 
     Args:
         optimizer: optimized optimizer, contains optimized model's paramerter setting.
-        weight_params_attr: the prepacked parameters' attrs, to do prepack for corresponding
-            momentum_buffer or other state according those attrs.
-        dtype: can be torch.bfloat16 or torch.float32(torch.float), determin doing bfloat16 training
-            or float training.
+        params_attr: the  parameters' attrs, to cat top_half and bottom(trail) half back to fp32
+
     """
 
-    def __init__(self, optimizer, weight_params_attr, dtype):
-        if type(optimizer) in IPEX_OPTIMIZER_MAPPING and dtype == torch.bfloat16:
-            self.optimizer = IPEX_OPTIMIZER_MAPPING[type(optimizer)] (optimizer, weight_params_attr)
+    def __init__(self, optimizer, params_attr):
+        if type(optimizer) in IPEX_OPTIMIZER_MAPPING:
+            self.optimizer = IPEX_OPTIMIZER_MAPPING[type(optimizer)] (optimizer, params_attr)
             self.master_weight_split  = True
         else:
             self.optimizer = optimizer
             self.master_weight_split  = False
-        self.weight_params_attr = weight_params_attr
+        self.params_attr = params_attr
         self.param_groups = self.optimizer.param_groups
-        self.dtype = dtype
-
-    def state_dict(self):
-        optimizer_temp = copy.deepcopy(self.optimizer)
-        weight_params_attr_ = {}
-        # For bf16 path, the optimizer's params are master weight,
-        # but self.weight_params_attr's keys are bf16 weight, it hard to
-        # query the weight's attr, so recreate a dic which using master weight
-        # as key for easily to query.
-        if self.dtype == torch.bfloat16 and not self.master_weight_split:
-           for _, values in self.weight_params_attr.items():
-                master_weight = values['master_weight']
-                weight_params_attr_[master_weight] = values
-        else:
-            weight_params_attr_ = self.weight_params_attr
-
-        for (k1, _), (_, v2) in zip(self.optimizer.state.items(), optimizer_temp.state.items()):
-            # unpack tensor state using weight's attr.
-            if k1 in weight_params_attr_:
-                weight_attr = weight_params_attr_[k1]
-                for state_key, state_value in v2.items():
-                    if isinstance(state_value, torch.Tensor):
-                        # It covers both conv and linear now. TODO: LSTM or other ops.
-                        if weight_attr['op'] is torch.nn.Conv2d:
-                            if self.master_weight_split and state_value.dtype == torch.bfloat16:
-                                state_value = torch.ops.torch_ipex.cat_bfloat16_float(state_value, weight_attr['trail'])
-                            v2[state_key] = torch.ops.torch_ipex.conv2d_weight_unpack(
-                                state_value,
-                                weight_attr['padding'],
-                                weight_attr['stride'],
-                                weight_attr['dilation'],
-                                weight_attr['kernel_size'],
-                                weight_attr['groups'],
-                                weight_attr['out_channels'],
-                                weight_attr['in_channels'],
-                                weight_attr['weight_channels_last'],
-                                weight_attr['dtype'])
-                        elif weight_attr['op'] is torch.nn.Linear:
-                            if self.master_weight_split and state_value.dtype == torch.bfloat16:
-                                state_value = torch.ops.torch_ipex.cat_bfloat16_float(state_value, weight_attr['trail'])
-                            v2[state_key] = torch.ops.torch_ipex.linear_weight_unpack(
-                                state_value,
-                                weight_attr['out_features'],
-                                weight_attr['in_features'],
-                                weight_attr['weight_transposed'],
-                                weight_attr['dtype'])
-        return optimizer_temp.state_dict()
+        self.state = self.optimizer.state
 
     def load_state_dict(self, state_dict):
         assert False, "_ipex_optimizer does not suppory load_state_dict"
 
     def zero_grad(self, set_to_none: bool = False):
-        if self.dtype == torch.bfloat16:
-            for p in self.weight_params_attr:
+        if not self.master_weight_split:
+            for p in self.params_attr:
                 if p.grad is not None:
                     if set_to_none:
                         p.grad = None
@@ -96,14 +48,14 @@ def zero_grad(self, set_to_none: bool = False):
         self.optimizer.zero_grad(set_to_none)
 
     def step(self, closure=None):
-        if self.dtype == torch.bfloat16 and not self.master_weight_split:
+        if not self.master_weight_split:
             # convert bf16 weight'grad to float.
-            for k, value in self.weight_params_attr.items():
-                value["master_weight"].grad = k.grad.detach().to(torch.float)
+            for k, value in self.params_attr.items():
+                value["master_param"].grad = k.grad.detach().to(torch.float)
         loss = self.optimizer.step(closure)
         # sync mater weight to model's paramerter
-        if self.dtype == torch.bfloat16 and not self.master_weight_split:
-            for k, value in self.weight_params_attr.items():
-                torch.ops.torch_ipex.sync_master_weight_to_bf16(value["master_weight"], k)
+        if not self.master_weight_split:
+            for k, value in self.params_attr.items():
+                torch.ops.torch_ipex.sync_master_weight_to_bf16(value["master_param"], k)
         return loss
 

intel_pytorch_extension_py/utils.py

@@ -88,12 +88,10 @@ def optimize(model, dtype=torch.bfloat16, optimizer=None, level='O1', inplace=Fa
         pass
     elif level == 'O1':
         # Do weight prepack, and convert optimizer for training case.
-        optimized_model, optimized_optimizer, weight_params_attr = _weight_prepack_with_ipex(optimized_model, optimized_optimizer, dtype)
-        if dtype == torch.bfloat16 and model.training and optimizer is not None:
-            optimized_model, optimized_optimizer, weight_params_attr = _weight_dtype_convert_with_ipex(optimized_model, optimized_optimizer, weight_params_attr)
-        if optimizer is not None:
-            assert model.training, "please call model.train() if you want to convert the optimizer to ipex optimizer."
-            optimized_optimizer = _ipex_optimizer(optimized_optimizer, weight_params_attr, dtype)
+        params_attr = {}
+        if dtype == torch.bfloat16 and model.training:
+            optimized_model, optimized_optimizer, params_attr = _weight_dtype_convert_with_ipex(optimized_model, optimized_optimizer, params_attr)
+        optimized_model, optimized_optimizer, params_attr = _weight_prepack_with_ipex(optimized_model, optimized_optimizer, params_attr)
     else:
         assert False, "Only support level O0 and O1 now for optimize"