def adamw_step(self, closure=None):
    """Performs a single optimization step.

    Args:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()

    for group in self.param_groups:
        # fp32 master weight and fp32 weight(some layer no need cast)
        params_with_grad = []
        # bf16 weight(mapped to fp32 master weight) and empty tensor(empty means no need casted layer's weight)
        params2 = []
        grads = []
        exp_avgs = []
        exp_avg_sqs = []
        max_exp_avg_sqs = []
        state_steps = []
        beta1, beta2 = group["betas"]

        for p in group["params"]:
            # params_attr: {'layer.master_weight(fp32)': {'bf16_param': 'layer.weight(bf16)'}}
            grad = (
                get_bf16_grad(p, self.params_attr)
                if is_master_weight(p, self.params_attr)
                else p.grad
            )
            if grad is not None:
                params_with_grad.append(p)
                if grad.is_sparse:
                    raise RuntimeError(
                        "Adam does not support sparse gradients, please consider SparseAdam instead"
                    )
                grads.append(grad)

                state = self.state[p]
                # Lazy state initialization
                if len(state) == 0:
                    buffer_dtype = p.dtype
                    if p.dtype is not torch.float:
                        raise RuntimeError(
                            "parameter in optimizer(Adamw) is not FP32, need check"
                        )

                    state["step"] = torch.tensor(0.0)
                    # Exponential moving average of gradient values
                    state["exp_avg"] = torch.zeros_like(
                        p, memory_format=torch.preserve_format, dtype=buffer_dtype
                    )
                    # Exponential moving average of squared gradient values
                    state["exp_avg_sq"] = torch.zeros_like(
                        p, memory_format=torch.preserve_format, dtype=buffer_dtype
                    )
                    if group["amsgrad"]:
                        # Maintains max of all exp. moving avg. of sq. grad. values
                        state["max_exp_avg_sq"] = torch.zeros_like(
                            p, memory_format=torch.preserve_format, dtype=buffer_dtype
                        )

                exp_avgs.append(state["exp_avg"])
                exp_avg_sqs.append(state["exp_avg_sq"])

                if group["amsgrad"]:
                    max_exp_avg_sqs.append(state["max_exp_avg_sq"])

                state_steps.append(state["step"])

                param2 = get_param2(p, self.params_attr)
                params2.append(param2)

        adamw(
            params_with_grad,
            params2,
            grads,
            exp_avgs,
            exp_avg_sqs,
            max_exp_avg_sqs,
            state_steps,
            amsgrad=group["amsgrad"],
            beta1=beta1,
            beta2=beta2,
            lr=group["lr"],
            weight_decay=group["weight_decay"],
            eps=group["eps"],
            maximize=group["maximize"],
            foreach=group["foreach"],
        )

    return loss
Performs a single optimization step.

    Args:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    
intel_extension_for_pytorch/optim/_functional.py