def patch_step_for_master_weight_training(optimizer):
    r"""
    Patch "step" method of optimizer to support master weight training
    1.Convert BF16 or FP16 grad to FP32
    2.Call original "step" to update parameters
    3.Sync FP32 master weight back to BF16 or FP16 weight
    """

    def master_param_non_fused_step(self, closure=None):
        # convert bf16 or fp16 weight'grad to float.
        for k, v in self.params_attr.items():
            _param = v.parameter
            if _param is None or _param is k:
                continue
            if _param.requires_grad and _param.grad is not None:
                k.grad = _param.grad.detach().float()

        loss = self._original_step(closure)
        # sync mater weight to model's paramerter
        for k, v in self.params_attr.items():
            _param = v.parameter
            if _param is None or _param is k:
                continue
            if k.device.type == "cpu":
                if _param.dtype == torch.bfloat16:
                    torch.ops.torch_ipex.sync_master_weight_to_bf16(k, _param)
                else:
                    assert _param.dtype == torch.float16
                    torch.ops.torch_ipex.sync_master_weight_to_fp16(k, _param)
            elif k.device.type == "xpu":
                _param.data = k.data.to(dtype=torch.bfloat16)
            else:
                pass
        return loss

    # Split master_param_non_fused_step into 2 steps:
    # 1.Sync_grad: Convert grad to FP32
    # 2.step_sync_weight: Call original "step" to update parameters and
    #   Sync FP32 master weight back to weight
    # This is because gradscaler will unscale grad and
    # it needs to sync grad to the FP32's grad first. After that gradscaler
    # will update weight and it also needs to sync FP32 master weight back to weight.
    def sync_grad(self):
        for k, v in self.params_attr.items():
            _param = v.parameter
            if _param is None or _param is k:
                continue
            assert (
                _param.dtype != torch.bfloat16
            ), "GradScaler is not recommended for bf16 training"
            if _param.requires_grad and _param.grad is not None:
                k.grad = _param.grad.detach().float()

    def step_sync_weight(self, closure=None):
        loss = self._original_step(closure)
        # sync mater weight to model's paramerter
        for k, v in self.params_attr.items():
            _param = v.parameter
            if _param is None or _param is k:
                continue
            assert (
                _param.dtype != torch.bfloat16
            ), "GradScaler is not recommended for bf16 training"
            torch.ops.torch_ipex.sync_master_weight_to_fp16(k, _param)
        return loss

    if not hasattr(optimizer, "_original_step"):
        setattr(optimizer, "_original_step", optimizer.step)  # noqa: B010
        optimizer.step = types.MethodType(master_param_non_fused_step, optimizer)
        optimizer.sync_grad = types.MethodType(sync_grad, optimizer)
        optimizer.step_sync_weight = types.MethodType(step_sync_weight, optimizer)

    Patch "step" method of optimizer to support master weight training
    1.Convert BF16 or FP16 grad to FP32
    2.Call original "step" to update parameters
    3.Sync FP32 master weight back to BF16 or FP16 weight
    
intel_extension_for_pytorch/optim/_optimizer_utils.py