wehub-resource-sync
910 lines
36 KiB
Python
910 lines
36 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import re
|
|
|
|
import paddle
|
|
from paddle import _C_ops, pir
|
|
from paddle.base import core, framework
|
|
from paddle.base.dygraph import base as imperative_base
|
|
from paddle.base.framework import Variable, in_dynamic_or_pir_mode, in_pir_mode
|
|
from paddle.base.libpaddle import DataType
|
|
from paddle.distributed import fleet
|
|
from paddle.optimizer.adamw import AdamW
|
|
from paddle.pir import Value
|
|
|
|
from paddlenlp.utils.log import logger
|
|
|
|
try:
|
|
from .adamw_triton import adamw_triton
|
|
except:
|
|
adamw_triton = None
|
|
|
|
|
|
from ..quantization.qat_utils import dequantize, quantize
|
|
|
|
|
|
class AdamWMini(AdamW):
|
|
def __init__(
|
|
self,
|
|
named_parameters=None,
|
|
learning_rate=0.001,
|
|
beta1=0.9,
|
|
beta2=0.999,
|
|
epsilon=1e-8,
|
|
weight_decay=0.0,
|
|
use_lowprecision_moment=False,
|
|
lr_ratio=None,
|
|
apply_decay_param_fun=None,
|
|
grad_clip=None,
|
|
lazy_mode=False,
|
|
multi_precision=False,
|
|
amsgrad=False,
|
|
dim=2048,
|
|
n_heads=32,
|
|
n_kv_heads=None,
|
|
verbose=True,
|
|
name=None,
|
|
):
|
|
self.dim = dim
|
|
self.n_heads = n_heads
|
|
self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
|
|
self.head_numel = self.dim * self.dim // self.n_heads
|
|
self.verbose = verbose
|
|
self.check_block_name = True
|
|
self._already_create_accumulator = set() # Initialize accumulator tracking set
|
|
|
|
# Block naming patterns
|
|
self.embd_names = {"embed", "embd", "wte"}
|
|
self.output_names = {"lm_head", "output", "final_layer"}
|
|
self.wqk_names = {"k_proj", "q_proj", "wq", "wk", "query", "key"}
|
|
self.wv_names = {"v_proj", "wv", "value"}
|
|
self.attn_proj_names = {"o_proj", "wo", "attn.proj"}
|
|
self.mlp_names = {"feed_forward", "linear", "mlp"}
|
|
self.adam_block_names = {"bias"}
|
|
|
|
# Validation
|
|
if not self.dim == int(self.dim):
|
|
raise ValueError(f"Invalid dim value: {self.dim}")
|
|
if not self.n_heads == int(self.n_heads):
|
|
raise ValueError(f"Invalid n_heads value: {self.n_heads}")
|
|
if not self.n_kv_heads == int(self.n_kv_heads):
|
|
raise ValueError(f"Invalid n_kv_heads value: {self.n_kv_heads}")
|
|
if not self.n_heads % self.n_kv_heads == 0:
|
|
raise ValueError(f"n_heads {self.n_heads} must be divisible by n_kv_heads {self.n_kv_heads}")
|
|
|
|
parameters = []
|
|
for param_name, param in named_parameters:
|
|
param_name = param_name.lower()
|
|
param.name = param_name
|
|
parameters.append(param)
|
|
|
|
super().__init__(
|
|
learning_rate=learning_rate,
|
|
beta1=beta1,
|
|
beta2=beta2,
|
|
epsilon=epsilon,
|
|
parameters=parameters,
|
|
weight_decay=weight_decay,
|
|
use_lowprecision_moment=use_lowprecision_moment,
|
|
lr_ratio=lr_ratio,
|
|
apply_decay_param_fun=apply_decay_param_fun,
|
|
grad_clip=grad_clip,
|
|
lazy_mode=lazy_mode,
|
|
multi_precision=multi_precision,
|
|
amsgrad=amsgrad,
|
|
name=name,
|
|
)
|
|
|
|
def _add_moments_pows(self, p):
|
|
"""Add moment accumulators with shapes based on block type."""
|
|
name = p.name
|
|
|
|
# Get accumulator data type
|
|
acc_dtype = p.dtype
|
|
if self._is_dtype_fp16_or_bf16(acc_dtype) and not self._use_lowprecision_moment:
|
|
acc_dtype = DataType.FLOAT32 if in_pir_mode() else core.VarDesc.VarType.FP32
|
|
|
|
# Add accumulators based on block type
|
|
if any(adam_block_name in name for adam_block_name in self.adam_block_names):
|
|
# Standard Adam for bias terms
|
|
super()._add_moments_pows(p)
|
|
elif any(wqk_name in name for wqk_name in self.wqk_names):
|
|
# One accumulator per head for Q/K blocks
|
|
total_size = paddle.numel(p)
|
|
shape_moment1 = [total_size // self.head_numel, self.head_numel]
|
|
shape_moment2 = [total_size // self.head_numel, 1]
|
|
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype, shape=shape_moment1)
|
|
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype, shape=shape_moment2)
|
|
self._add_accumulator(
|
|
name=self._beta1_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=0.9 if isinstance(self._beta1, (Variable, Value)) else self._beta1,
|
|
shape=[1],
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
device="cpu",
|
|
)
|
|
self._add_accumulator(
|
|
name=self._beta2_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=0.999 if isinstance(self._beta2, (Variable, Value)) else self._beta2,
|
|
shape=[1],
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
device="cpu",
|
|
)
|
|
elif (
|
|
any(embd_name in name for embd_name in self.embd_names)
|
|
or any(output_name in name for output_name in self.output_names)
|
|
or any(wv_name in name for wv_name in self.wv_names)
|
|
or any(mlp_name in name for mlp_name in self.mlp_names)
|
|
or any(attn_proj_name in name for attn_proj_name in self.attn_proj_names)
|
|
):
|
|
# One accumulator per neuron for other blocks
|
|
if any(embd_name in name for embd_name in self.embd_names):
|
|
shape = [p.shape[0], 1] if len(p.shape) > 1 else [1]
|
|
else:
|
|
shape = [1, p.shape[1]] if len(p.shape) > 1 else [1]
|
|
|
|
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
|
|
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype, shape=shape)
|
|
self._add_accumulator(
|
|
name=self._beta1_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=0.9 if isinstance(self._beta1, (Variable, Value)) else self._beta1,
|
|
shape=[1],
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
device="cpu",
|
|
)
|
|
self._add_accumulator(
|
|
name=self._beta2_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=0.999 if isinstance(self._beta2, (Variable, Value)) else self._beta2,
|
|
shape=[1],
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
device="cpu",
|
|
)
|
|
else:
|
|
self._add_accumulator(self._moment1_acc_str, p, dtype=acc_dtype)
|
|
self._add_accumulator(self._moment2_acc_str, p, dtype=acc_dtype, shape=[1])
|
|
self._add_accumulator(
|
|
name=self._beta1_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=0.9 if isinstance(self._beta1, (Variable, Value)) else self._beta1,
|
|
shape=[1],
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
device="cpu",
|
|
)
|
|
self._add_accumulator(
|
|
name=self._beta2_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=0.999 if isinstance(self._beta2, (Variable, Value)) else self._beta2,
|
|
shape=[1],
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
device="cpu",
|
|
)
|
|
|
|
def _append_optimize_op(self, block, param_and_grad):
|
|
"""Implement optimization operations for different block types."""
|
|
assert isinstance(block, (framework.Block, pir.Block))
|
|
if isinstance(param_and_grad, dict):
|
|
param_and_grad = self._update_param_group(param_and_grad)
|
|
|
|
param = param_and_grad[0]
|
|
name = param.name
|
|
|
|
# Whether we should do weight decay for the parameter.
|
|
with_decay = True
|
|
if self._apply_decay_param_fun is not None and not self._apply_decay_param_fun(param.name):
|
|
with_decay = False
|
|
|
|
# Get moment accumulators
|
|
moment1 = self._get_accumulator_master(self._moment1_acc_str, param)
|
|
moment2 = self._get_accumulator_master(self._moment2_acc_str, param)
|
|
beta1_pow_acc = self._get_accumulator_master(self._beta1_pow_acc_str, param)
|
|
beta2_pow_acc = self._get_accumulator_master(self._beta2_pow_acc_str, param)
|
|
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(param.dtype)
|
|
master_weight = self._master_weights[name] if find_master else None
|
|
lr = self._create_param_lr(param_and_grad)
|
|
|
|
# create the adamw optimize op
|
|
if in_dynamic_or_pir_mode():
|
|
lr_ratio_ = 1.0 if self._lr_ratio is None else self._lr_ratio(param)
|
|
|
|
_beta1 = self._beta1 if not isinstance(self._beta1, Variable) else self._beta1.item(0)
|
|
_beta2 = self._beta2 if not isinstance(self._beta2, Variable) else self._beta2.item(0)
|
|
found_inf = self._get_auxiliary_var("found_inf") if in_pir_mode() else None
|
|
|
|
self.adamw_python(
|
|
param_and_grad[0],
|
|
param_and_grad[1],
|
|
lr,
|
|
moment1,
|
|
moment2,
|
|
beta1_pow_acc,
|
|
beta2_pow_acc,
|
|
master_weight,
|
|
found_inf,
|
|
_beta1,
|
|
_beta2,
|
|
self._epsilon,
|
|
lr_ratio_,
|
|
self._weight_decay,
|
|
with_decay,
|
|
find_master,
|
|
name,
|
|
)
|
|
return None
|
|
else:
|
|
raise NotImplementedError("Not implemented yet.")
|
|
|
|
def adamw_python(
|
|
self,
|
|
param,
|
|
grad,
|
|
learning_rate,
|
|
moment1,
|
|
moment2,
|
|
beta1_pow,
|
|
beta2_pow,
|
|
master_weight,
|
|
skip_update,
|
|
beta1,
|
|
beta2,
|
|
epsilon,
|
|
lr_ratio,
|
|
coeff,
|
|
with_decay,
|
|
multi_precision,
|
|
name,
|
|
):
|
|
if skip_update:
|
|
return
|
|
if not with_decay:
|
|
coeff = 0.0
|
|
if "norm" in name or "ln" in name or "bias" in name:
|
|
coeff = 0.0
|
|
if not multi_precision:
|
|
master_weight = None
|
|
|
|
if any(adam_block_name in name for adam_block_name in self.adam_block_names):
|
|
_, _, _, _, _, _, _ = _C_ops.adamw_(
|
|
param,
|
|
grad,
|
|
learning_rate,
|
|
moment1,
|
|
moment2,
|
|
None,
|
|
beta1_pow,
|
|
beta2_pow,
|
|
master_weight,
|
|
skip_update,
|
|
beta1,
|
|
beta2,
|
|
epsilon,
|
|
lr_ratio,
|
|
coeff,
|
|
with_decay,
|
|
self._lazy_mode,
|
|
1000,
|
|
multi_precision,
|
|
False,
|
|
self._amsgrad,
|
|
)
|
|
|
|
else:
|
|
lr = learning_rate * lr_ratio
|
|
if master_weight is not None:
|
|
p = master_weight
|
|
else:
|
|
p = param
|
|
p *= 1.0 - lr * coeff
|
|
|
|
# Block-specific updates with per-block learning rates
|
|
if any(wqk_name in name for wqk_name in self.wqk_names):
|
|
# Q/K blocks: reshape and compute per-head learning rates
|
|
grad_reshaped = paddle.reshape(grad, [-1, self.head_numel])
|
|
mom1 = paddle.reshape(moment1, [-1, self.head_numel])
|
|
mom2 = moment2 # Already shaped correctly
|
|
|
|
# Compute per-head second moment
|
|
mom2_update = paddle.mean(grad_reshaped * grad_reshaped, axis=1, keepdim=True)
|
|
# Update moments with correct beta values
|
|
mom1 = mom1 * beta1 + (1.0 - beta1) * grad_reshaped
|
|
mom2 = mom2 * beta2 + (1.0 - beta2) * mom2_update
|
|
|
|
# Compute adaptive learning rate
|
|
denom = mom2.sqrt() / ((1.0 - beta2_pow).sqrt()) + epsilon
|
|
|
|
# Apply updates
|
|
update = (mom1 / denom) * (-(lr / (1.0 - beta1_pow)))
|
|
p += paddle.reshape(update, param.shape)
|
|
|
|
elif (
|
|
any(embd_name in name for embd_name in self.embd_names)
|
|
or any(output_name in name for output_name in self.output_names)
|
|
or any(wv_name in name for wv_name in self.wv_names)
|
|
or any(mlp_name in name for mlp_name in self.mlp_names)
|
|
or any(attn_proj_name in name for attn_proj_name in self.attn_proj_names)
|
|
):
|
|
mom1 = moment1
|
|
mom2 = moment2 # Already shaped correctly
|
|
|
|
mom1 = mom1 * beta1 + (1.0 - beta1) * grad
|
|
|
|
if any(embd_name in name for embd_name in self.embd_names):
|
|
mom2 = mom2 * beta2 + (1.0 - beta2) * (grad * grad).mean(axis=1, keepdim=True)
|
|
else:
|
|
mom2 = mom2 * beta2 + (1.0 - beta2) * (grad * grad).mean(axis=0, keepdim=True)
|
|
|
|
denom = mom2.sqrt() / ((1.0 - beta2_pow).sqrt()) + epsilon
|
|
p += (mom1 / denom) * (-(lr / (1.0 - beta1_pow)))
|
|
|
|
else:
|
|
# Other blocks
|
|
mom1 = moment1
|
|
mom2 = moment2 # Already shaped correctly
|
|
|
|
mom1 = mom1 * beta1 + (1.0 - beta1) * grad
|
|
mom2 = mom2 * beta2 + (1.0 - beta2) * (grad * grad).mean()
|
|
|
|
denom = mom2.sqrt() / ((1.0 - beta2_pow).sqrt()) + epsilon
|
|
p += (mom1 / denom) * (-(lr / (1.0 - beta1_pow)))
|
|
|
|
# Update param in-place
|
|
if master_weight is not None:
|
|
master_weight[:] = p
|
|
param[:] = p.astype(param.dtype)
|
|
else:
|
|
param[:] = p
|
|
|
|
# Update accumulators in-place
|
|
moment1[:] = mom1
|
|
moment2[:] = mom2
|
|
beta1_pow[:] = beta1 * beta1_pow[:]
|
|
beta2_pow[:] = beta2 * beta2_pow[:]
|
|
|
|
return None
|
|
|
|
def _count_block(self):
|
|
"""Count the number of each block type for logging."""
|
|
if not self.verbose:
|
|
return
|
|
|
|
counts = {
|
|
"embedding": 0,
|
|
"output": 0,
|
|
"query/key": 0,
|
|
"value": 0,
|
|
"attention_proj": 0,
|
|
"mlp": 0,
|
|
}
|
|
|
|
for name in self._already_create_accumulator:
|
|
if "bias" in name:
|
|
continue
|
|
if any(embd_name in name for embd_name in self.embd_names):
|
|
counts["embedding"] += 1
|
|
if any(output_name in name for output_name in self.output_names):
|
|
counts["output"] += 1
|
|
if any(wqk_name in name for wqk_name in self.wqk_names):
|
|
counts["query/key"] += 1
|
|
if any(wv_name in name for wv_name in self.wv_names):
|
|
counts["value"] += 1
|
|
if any(attn_proj_name in name for attn_proj_name in self.attn_proj_names):
|
|
counts["attention_proj"] += 1
|
|
if any(mlp_name in name for mlp_name in self.mlp_names):
|
|
counts["mlp"] += 1
|
|
|
|
logger.info("\nAdam-mini found blocks:")
|
|
logger.info(f"- {counts['embedding']} embedding layers")
|
|
logger.info(f"- {counts['output']} output layers")
|
|
logger.info(f"- {counts['query/key']} Query and Key layers")
|
|
logger.info(f"- {counts['value']} Value layers")
|
|
logger.info(f"- {counts['attention_proj']} Attention projection layers")
|
|
logger.info(f"- {counts['mlp']} MLP layers\n")
|
|
|
|
# Print warnings for missing blocks
|
|
if counts["embedding"] == 0:
|
|
logger.warning("Warning: No embedding layers found")
|
|
if counts["output"] == 0:
|
|
logger.warning("Warning: No output layers found (ignore if using weight tying)")
|
|
if counts["query/key"] == 0:
|
|
logger.warning("Warning: No Query/Key layers found")
|
|
if counts["value"] == 0:
|
|
logger.warning("Warning: No Value layers found")
|
|
if counts["attention_proj"] == 0:
|
|
logger.warning("Warning: No attention projection layers found")
|
|
if counts["mlp"] == 0:
|
|
logger.warning("Warning: No MLP layers found")
|
|
if sum(counts.values()) == 0:
|
|
logger.warning("Warning: No Transformer blocks found")
|
|
|
|
def _create_accumulators(self, block, parameters):
|
|
"""Create accumulators for parameters."""
|
|
assert isinstance(block, (framework.Block, pir.Block))
|
|
if isinstance(parameters, dict):
|
|
parameters = self._update_param_group(parameters)
|
|
|
|
for p in parameters:
|
|
if p.name in self._already_create_accumulator:
|
|
continue
|
|
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
|
|
master_p = self._create_master_weight(p)
|
|
self._add_moments_pows(master_p)
|
|
self._already_create_accumulator.add(p.name)
|
|
continue
|
|
if self._is_dtype_fp16_or_bf16(p.dtype) and not self._multi_precision:
|
|
logger.warning(
|
|
"Accumulating with FP16 or BF16 in optimizer can lead to poor accuracy or slow convergence."
|
|
"Consider using multi_precision=True option of the Adam optimizer."
|
|
)
|
|
self._add_moments_pows(p)
|
|
self._already_create_accumulator.add(p.name)
|
|
|
|
if self.check_block_name:
|
|
self._count_block()
|
|
self.check_block_name = False
|
|
|
|
|
|
class AdamWCustom(AdamW):
|
|
def __init__(self, quantization_config, tensorwise_offload_optimizer, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.quant_scale_mapping = {}
|
|
for p in self._param_groups:
|
|
if "quantization_linear" in p.name and "w_1" in p.name:
|
|
self.quant_scale_mapping[p.name.replace("w_1", "w_0")] = p
|
|
self.quantization_config = quantization_config
|
|
self._hcg = fleet.get_hybrid_communicate_group()
|
|
self.mp_group = self._hcg.get_model_parallel_group()
|
|
self.tensorwise_offload_optimizer = tensorwise_offload_optimizer
|
|
|
|
def _add_moments_pows(self, p, moment_dtype=core.VarDesc.VarType.FP32):
|
|
acc_dtype = p.dtype
|
|
|
|
self._add_accumulator(self._moment1_acc_str, p, dtype=moment_dtype)
|
|
self._add_accumulator(self._moment2_acc_str, p, dtype=moment_dtype)
|
|
try:
|
|
type = core.VarDesc.VarType.DENSE_TENSOR
|
|
except:
|
|
type = core.VarDesc.VarType.LOD_TENSOR
|
|
self._add_accumulator(
|
|
name=self._beta1_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=(0.9 if isinstance(self._beta1, (Variable, Value)) else self._beta1),
|
|
shape=[1],
|
|
type=type,
|
|
)
|
|
self._add_accumulator(
|
|
name=self._beta2_pow_acc_str,
|
|
param=p,
|
|
dtype=acc_dtype,
|
|
fill_value=(0.999 if isinstance(self._beta2, (Variable, Value)) else self._beta2),
|
|
shape=[1],
|
|
type=type,
|
|
)
|
|
|
|
def _create_accumulators(self, block, parameters):
|
|
assert isinstance(block, (framework.Block, pir.Block))
|
|
if isinstance(parameters, dict):
|
|
parameters = self._update_param_group(parameters)
|
|
|
|
# Create accumulator tensors for first and second moments
|
|
for p in parameters:
|
|
if p.name in self._already_create_accumulator:
|
|
continue
|
|
if self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype):
|
|
master_p = self._create_master_weight(p)
|
|
if self._use_lowprecision_moment:
|
|
if p.name in self.quant_scale_mapping:
|
|
p_scale = self.quant_scale_mapping[p.name]
|
|
if str(p_scale.dtype) == "paddle.float16":
|
|
moment_dtype = core.VarDesc.VarType.FP16
|
|
elif str(p_scale.dtype) == "paddle.bfloat16":
|
|
moment_dtype = core.VarDesc.VarType.BF16
|
|
else:
|
|
if str(p.dtype) == "paddle.float16":
|
|
moment_dtype = core.VarDesc.VarType.FP16
|
|
elif str(p.dtype) == "paddle.bfloat16":
|
|
moment_dtype = core.VarDesc.VarType.BF16
|
|
else:
|
|
moment_dtype = core.VarDesc.VarType.FP32
|
|
|
|
self._add_moments_pows(master_p, moment_dtype)
|
|
self._already_create_accumulator.add(p.name)
|
|
|
|
elif self._is_dtype_fp16_or_bf16(p.dtype) and not self._multi_precision:
|
|
raise NotImplementedError("AdamWCustom only support AMP training")
|
|
else:
|
|
self._add_moments_pows(p)
|
|
self._already_create_accumulator.add(p.name)
|
|
if self.tensorwise_offload_optimizer:
|
|
self.offload_optim(p)
|
|
|
|
def _create_master_weight(self, param):
|
|
if param.name in self._master_weights:
|
|
var = self._master_weights[param.name]
|
|
else:
|
|
var_name = self._gen_master_weight_var_name(param)
|
|
if param.name in self.quant_scale_mapping:
|
|
quant_scale = self.quant_scale_mapping[param.name]
|
|
if self.quantization_config.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
|
|
var = dequantize(
|
|
param,
|
|
quant_scale,
|
|
"weight",
|
|
self.quantization_config.weight_quantize_algo,
|
|
self.quantization_config,
|
|
apply_hadamard=self.quantization_config.apply_hadamard,
|
|
side="left",
|
|
).astype("float32")
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Unknown weight_quantize_algo {self.quantization_config.weight_quantize_algo}"
|
|
)
|
|
else:
|
|
var = paddle.cast(param, "float32")
|
|
var.name = var_name
|
|
self._master_weights[param.name] = var
|
|
return var
|
|
|
|
def _is_dtype_fp16_or_bf16(self, dtype):
|
|
"""
|
|
check the dtype is fp16 or the dtype is bf16
|
|
:param dtype: instance of core.VarDesc.VarType
|
|
:return: True if dtype is one of fp16 or bf16, False otherwise
|
|
"""
|
|
if dtype == paddle.int8 or dtype == paddle.float8_e4m3fn:
|
|
return True
|
|
assert isinstance(
|
|
dtype, (core.VarDesc.VarType, core.DataType)
|
|
), "The dtype should be an instance of core.VarDesc.VarType or core.DataType."
|
|
if isinstance(dtype, core.VarDesc.VarType):
|
|
return dtype == core.VarDesc.VarType.FP16 or dtype == core.VarDesc.VarType.BF16
|
|
else:
|
|
return dtype == core.DataType.FLOAT16 or dtype == core.DataType.BFLOAT16
|
|
|
|
def _append_optimize_op(self, block, param_and_grad):
|
|
assert isinstance(block, (framework.Block, pir.Block))
|
|
if isinstance(param_and_grad, dict):
|
|
param_and_grad = self._update_param_group(param_and_grad)
|
|
param, grad = param_and_grad
|
|
|
|
# Whether we should do weight decay for the parameter.
|
|
with_decay = True
|
|
if self._apply_decay_param_fun is not None and not self._apply_decay_param_fun(param.name):
|
|
with_decay = False
|
|
|
|
if self.tensorwise_offload_optimizer:
|
|
self.reload_optim(param)
|
|
|
|
moment1 = self._get_accumulator_master(self._moment1_acc_str, param_and_grad[0])
|
|
moment2 = self._get_accumulator_master(self._moment2_acc_str, param_and_grad[0])
|
|
beta1_pow_acc = self._get_accumulator_master(self._beta1_pow_acc_str, param_and_grad[0])
|
|
beta2_pow_acc = self._get_accumulator_master(self._beta2_pow_acc_str, param_and_grad[0])
|
|
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(param_and_grad[0].dtype)
|
|
master_weight = self._master_weights[param_and_grad[0].name] if find_master else None
|
|
if param.name in self.quant_scale_mapping:
|
|
quant_scale = self.quant_scale_mapping[param.name]
|
|
else:
|
|
quant_scale = None
|
|
lr = self._create_param_lr(param_and_grad)
|
|
# create the adamw optimize op
|
|
if in_dynamic_or_pir_mode():
|
|
lr_ratio_ = 1.0 if self._lr_ratio is None else self._lr_ratio(param_and_grad[0])
|
|
|
|
_beta1 = self._beta1 if not isinstance(self._beta1, Variable) else self._beta1.item(0)
|
|
_beta2 = self._beta2 if not isinstance(self._beta2, Variable) else self._beta2.item(0)
|
|
|
|
found_inf = self._get_auxiliary_var("found_inf") if in_pir_mode() else None
|
|
skip_update_param = quant_scale is not None
|
|
apply_adamw = self.adamw_custom if adamw_triton is None else adamw_triton
|
|
apply_adamw(
|
|
param_and_grad[0],
|
|
param_and_grad[1],
|
|
lr,
|
|
moment1,
|
|
moment2,
|
|
beta1_pow_acc,
|
|
beta2_pow_acc,
|
|
master_weight,
|
|
found_inf,
|
|
_beta1,
|
|
_beta2,
|
|
self._epsilon,
|
|
lr_ratio_,
|
|
self._weight_decay,
|
|
with_decay,
|
|
find_master,
|
|
skip_update_param,
|
|
)
|
|
if skip_update_param:
|
|
if param.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
|
|
if "parallel_quantization_linear" not in param.name:
|
|
group = None
|
|
elif param.weight_quantize_algo in ["a8w8linear", "a8w4linear"] and "row" in param.name:
|
|
group = None
|
|
else:
|
|
group = self.mp_group
|
|
param[:], quant_scale[:] = quantize(
|
|
x=master_weight.astype(quant_scale.dtype),
|
|
weight_quantize_algo=self.quantization_config.weight_quantize_algo,
|
|
tensor_type="weight",
|
|
quantization_config=self.quantization_config,
|
|
side="left",
|
|
apply_hadamard=self.quantization_config.apply_hadamard,
|
|
group=group,
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Please check your weight_quantize_algo {self.quantization_config.weight_quantize_algo}."
|
|
)
|
|
if self.tensorwise_offload_optimizer:
|
|
self.offload_optim(param)
|
|
|
|
return None
|
|
else:
|
|
raise NotImplementedError("Not implemented yet.")
|
|
|
|
def adamw_custom(
|
|
self,
|
|
param,
|
|
grad,
|
|
learning_rate,
|
|
moment1,
|
|
moment2,
|
|
beta1_pow,
|
|
beta2_pow,
|
|
master_weight,
|
|
skip_update,
|
|
beta1,
|
|
beta2,
|
|
epsilon,
|
|
lr_ratio,
|
|
coeff,
|
|
with_decay,
|
|
multi_precision,
|
|
skip_update_param,
|
|
):
|
|
if skip_update:
|
|
return
|
|
if not with_decay:
|
|
coeff = 0.0
|
|
if not multi_precision:
|
|
master_weight = None
|
|
lr = learning_rate * lr_ratio
|
|
if master_weight is not None:
|
|
p = master_weight
|
|
else:
|
|
p = param
|
|
|
|
p *= 1.0 - lr * coeff
|
|
moment_dtype = moment1.dtype
|
|
mom1 = moment1.astype("float32")
|
|
mom2 = moment2.astype("float32")
|
|
|
|
mom1 = beta1 * mom1 + (1.0 - beta1) * grad
|
|
mom2 = beta2 * mom2 + (1.0 - beta2) * grad * grad
|
|
denom = mom2.sqrt() / (1.0 - beta2_pow).sqrt() + epsilon
|
|
p += (mom1 / denom) * (-(lr / (1.0 - beta1_pow)))
|
|
|
|
if master_weight is not None:
|
|
master_weight[:] = p
|
|
if not skip_update_param:
|
|
param[:] = p.astype(param.dtype)
|
|
else:
|
|
param[:] = p
|
|
moment1[:] = mom1.astype(moment_dtype)
|
|
moment2[:] = mom2.astype(moment_dtype)
|
|
beta1_pow[:], beta2_pow[:] = beta1 * beta1_pow[:], beta2 * beta2_pow[:]
|
|
return
|
|
|
|
def offload_optim(self, p):
|
|
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype)
|
|
if find_master:
|
|
self._master_weights[p.name] = self._master_weights[p.name].pin_memory()
|
|
target_name = self._master_weights[p.name].name
|
|
else:
|
|
target_name = p.name
|
|
for name in [self._moment1_acc_str, self._moment2_acc_str]:
|
|
if self._name is not None:
|
|
name = self._name + "_" + name
|
|
self._accumulators[name][target_name] = self._accumulators[name][target_name].pin_memory()
|
|
|
|
def reload_optim(self, p):
|
|
find_master = self._multi_precision and self._is_dtype_fp16_or_bf16(p.dtype)
|
|
if find_master:
|
|
self._master_weights[p.name] = self._master_weights[p.name].cuda()
|
|
target_name = self._master_weights[p.name].name
|
|
else:
|
|
target_name = p.name
|
|
for name in [self._moment1_acc_str, self._moment2_acc_str]:
|
|
if self._name is not None:
|
|
name = self._name + "_" + name
|
|
self._accumulators[name][target_name] = self._accumulators[name][target_name].cuda()
|
|
|
|
|
|
class AdamWLoRAPro(AdamW):
|
|
def __init__(self, scaling_factor=2.0, x_mode="zero", *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
assert scaling_factor is not None
|
|
if x_mode not in ["zero", "sylvester", "symmetry"]:
|
|
raise ValueError(
|
|
f"Invalid x_mode value: {x_mode}, " f"mode should be in ['zero', 'sylvester', 'symmetry']"
|
|
)
|
|
self.scaling_factor = scaling_factor
|
|
self.x_mode = x_mode
|
|
|
|
def _solve_sylvester(self, A, B, C, X=None):
|
|
if A.dtype in [paddle.bfloat16, paddle.float16]:
|
|
A = A.to("float32")
|
|
B = B.to("float32")
|
|
C = C.to("float32")
|
|
B = -B
|
|
m = tuple(B.shape)[-1]
|
|
n = tuple(A.shape)[-1]
|
|
R, U = paddle.linalg.eig(x=A)
|
|
S, V = paddle.linalg.eig(x=B)
|
|
|
|
CV = C @ V
|
|
|
|
U_real, U_imag = paddle.real(U), paddle.imag(U)
|
|
CV_real, CV_imag = paddle.real(CV), paddle.imag(CV)
|
|
|
|
n_dim = U_real.shape[0]
|
|
|
|
block_top = paddle.concat([U_real, -U_imag], axis=1) # (n, 2n)
|
|
block_bot = paddle.concat([U_imag, U_real], axis=1) # (n, 2n)
|
|
A_block = paddle.concat([block_top, block_bot], axis=0) # (2n, 2n)
|
|
B_block = paddle.concat([CV_real, CV_imag], axis=0) # (2n, m)
|
|
|
|
F_block = paddle.linalg.solve(A_block, B_block) # [F_real; F_imag]
|
|
|
|
F_real = F_block[:n_dim, :]
|
|
F_imag = F_block[n_dim:, :]
|
|
F = paddle.complex(F_real, F_imag)
|
|
|
|
W = R[..., :, None] - S[..., None, :]
|
|
Y = F / W
|
|
try:
|
|
V_inv = paddle.linalg.inv(V)
|
|
except RuntimeError:
|
|
# Add regularization to handle singular matrices
|
|
epsilon = 1e-6 * paddle.mean(paddle.abs(V))
|
|
V_reg = V + epsilon * paddle.eye(V.shape[-1])
|
|
V_inv = paddle.linalg.inv(V_reg)
|
|
X = U[..., :n, :n] @ Y[..., :n, :m] @ V_inv[..., :m, :m]
|
|
|
|
if all(paddle.isreal(x.flatten()[0]) for x in [A, B, C]):
|
|
return paddle.real(X)
|
|
else:
|
|
return X
|
|
|
|
@imperative_base.no_grad
|
|
@framework.non_static_only
|
|
def step(self) -> None:
|
|
"""
|
|
Execute the optimizer and update parameters once.
|
|
|
|
Returns:
|
|
None
|
|
|
|
Examples:
|
|
.. code-block:: python
|
|
|
|
>>> import paddle
|
|
|
|
>>> a = paddle.rand([2,13], dtype="float32")
|
|
>>> linear = paddle.nn.Linear(13, 5)
|
|
>>> # This can be any optimizer supported by dygraph.
|
|
>>> opt = paddle.optimizer.AdamW(learning_rate = 0.01,
|
|
... parameters = linear.parameters())
|
|
>>> out = linear(a)
|
|
>>> out.backward()
|
|
>>> opt.step()
|
|
>>> opt.clear_grad()
|
|
"""
|
|
if paddle.base.dygraph.base.in_to_static_mode():
|
|
self._declarative_step()
|
|
return
|
|
|
|
if not isinstance(self._parameter_list[0], dict):
|
|
param_id_to_idx = {id(param): idx for idx, param in enumerate(self._parameter_list)}
|
|
|
|
lora_params = {}
|
|
for idx, param in enumerate(self._parameter_list):
|
|
name = getattr(param, "name", f"param_{idx}")
|
|
match = re.match(r"lo_ra_linear_(\d+)\.w_(\d+)", name)
|
|
if match:
|
|
layer_num = int(match.group(1))
|
|
weight_type = match.group(2)
|
|
if layer_num not in lora_params:
|
|
lora_params[layer_num] = {}
|
|
lora_params[layer_num][weight_type] = param
|
|
|
|
for layer_num, weights in lora_params.items():
|
|
if "1" in weights and "2" in weights:
|
|
param_B = weights["1"]
|
|
param_A = weights["2"]
|
|
|
|
idx_B = param_id_to_idx[id(param_B)]
|
|
idx_A = param_id_to_idx[id(param_A)]
|
|
|
|
if param_A._grad_ivar() is not None and param_B._grad_ivar() is not None:
|
|
A = param_A.detach()
|
|
B = param_B.detach()
|
|
grad_A = param_A._grad_ivar()
|
|
grad_B = param_B._grad_ivar()
|
|
|
|
delta = 1e-08
|
|
AA_T = A @ A.T
|
|
B_TB = B.T @ B
|
|
AA_T_inv = paddle.linalg.pinv(AA_T + delta * paddle.eye(num_rows=AA_T.shape[0]))
|
|
B_TB_inv = paddle.linalg.pinv(B_TB + delta * paddle.eye(num_rows=B_TB.shape[0]))
|
|
|
|
if self.x_mode == "sylvester":
|
|
X = self._solve_sylvester(
|
|
B_TB, AA_T, -(1 / self.scaling_factor**2) * B_TB_inv @ grad_A @ A.T
|
|
)
|
|
elif self.x_mode == "symmetry":
|
|
X = -0.5 * (1 / self.scaling_factor**2) * B_TB_inv @ B.T @ grad_B @ AA_T
|
|
else: # zero mode
|
|
X = paddle.zeros(shape=(B_TB_inv.shape[0], B_TB_inv.shape[0]))
|
|
|
|
X = X.clone().detach().cast(A.dtype)
|
|
|
|
new_grad_A = (1 / self.scaling_factor**2) * B_TB_inv @ grad_A + X @ A
|
|
new_grad_B = (1 / self.scaling_factor**2) * (
|
|
(paddle.eye(num_rows=B.shape[0]) - B @ B_TB_inv @ B.T) @ grad_B @ AA_T_inv
|
|
) - B @ X
|
|
|
|
self._parameter_list[idx_A]._grad_ivar()[:] = new_grad_A
|
|
self._parameter_list[idx_B]._grad_ivar()[:] = new_grad_B
|
|
|
|
params_grads = []
|
|
for param in self._parameter_list:
|
|
if param.stop_gradient:
|
|
continue
|
|
if param._grad_ivar() is not None:
|
|
grad_var = param._grad_ivar()
|
|
if framework.in_dygraph_mode():
|
|
if (
|
|
hasattr(grad_var, "is_selected_rows")
|
|
and grad_var.is_selected_rows()
|
|
and self.regularization is not None
|
|
):
|
|
raise RuntimeError(
|
|
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
|
)
|
|
else:
|
|
if (
|
|
hasattr(grad_var, "_is_sparse")
|
|
and grad_var._is_sparse()
|
|
and self.regularization is not None
|
|
):
|
|
raise RuntimeError(
|
|
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
|
)
|
|
params_grads.append((param, grad_var))
|
|
|
|
self._apply_optimize(loss=None, startup_program=None, params_grads=params_grads)
|
|
else:
|
|
raise NotImplementedError("AdamWLoRAPro does not support parameter groups")
|