def patch_zero_grad_for_master_weight_training(optimizer):
    r"""
    Patch "zero_grad" method of optimizer to support BFloat16 master weight training
    Under master weight training case, the grad is actually on 'bf16_params' or 'fp16_params'.
    So the 'zero_grad' should work on the 'bf16_params' or 'fp16_params' too.
    """

    def zero_grad(self, set_to_none: bool = True):
        for _, v in self.params_attr.items():
            _param = v.parameter
            if _param is None:
                continue
            if _param.grad is not None:
                if set_to_none:
                    _param.grad = None
                else:
                    if _param.grad.grad_fn is not None:
                        _param.grad.detach_()
                    else:
                        _param.grad.requires_grad_(False)
                    _param.grad.zero_()
        self._original_zero_grad(set_to_none)

    if not hasattr(optimizer, "_original_zero_grad"):
        setattr(optimizer, "_original_zero_grad", optimizer.zero_grad)  # noqa: B010
        optimizer.zero_grad = types.MethodType(zero_grad, optimizer)

    Patch "zero_grad" method of optimizer to support BFloat16 master weight training
    Under master weight training case, the grad is actually on 'bf16_params' or 'fp16_params'.
    So the 'zero_grad' should work on the 'bf16_params' or 'fp16_params' too.
    
intel_extension_for_pytorch/optim/_optimizer_utils.py