Initializes the RMSNormalization layer. Args: fd_config (FDConfig): Arguments related to inference, containing attributes such as weight_dtype, act_dtype, mp_size, hidden_size, head_dim, num_attention_heads, and ffn_hidden_size.
(
self,
fd_config: FDConfig,
hidden_size: int,
eps: float = 1e-5,
prefix: str = "",
bias: paddle.Tensor = None,
quant_scale: float = None,
begin_norm_axis: int = 1,
dtype: str = None,
layer_id: int = -1,
)
| 40 | """ |
| 41 | |
| 42 | def __init__( |
| 43 | self, |
| 44 | fd_config: FDConfig, |
| 45 | hidden_size: int, |
| 46 | eps: float = 1e-5, |
| 47 | prefix: str = "", |
| 48 | bias: paddle.Tensor = None, |
| 49 | quant_scale: float = None, |
| 50 | begin_norm_axis: int = 1, |
| 51 | dtype: str = None, |
| 52 | layer_id: int = -1, |
| 53 | ) -> None: |
| 54 | """ |
| 55 | Initializes the RMSNormalization layer. |
| 56 | |
| 57 | Args: |
| 58 | fd_config (FDConfig): Arguments related to inference, containing |
| 59 | attributes such as weight_dtype, act_dtype, mp_size, hidden_size, head_dim, |
| 60 | num_attention_heads, and ffn_hidden_size. |
| 61 | hidden_size (int) : size of hidden state. |
| 62 | eps:(float, optional): Small value added to the variance to avoid division by zero. Defaults to 1e-5. |
| 63 | prefix(str,optional):The name of current layer. Defaults to "". |
| 64 | bias (paddle.Tensor,optional): Initial bias value for the linear layer (if used). Defaults to None. |
| 65 | quant_scale(float,optional):Quantization scale, used in quantization scenarios. Defaults to -1, indicating no quantization. |
| 66 | begin_norm_axis (int, optional): The axis along which to perform normalization. Defaults to 1. |
| 67 | |
| 68 | Raises: |
| 69 | NotImplementedError: If the specified norm_type is not supported. |
| 70 | """ |
| 71 | super().__init__() |
| 72 | self.fd_config = fd_config |
| 73 | self.prefix: str = prefix |
| 74 | self.hidden_size: int = hidden_size |
| 75 | if len(prefix) == 0: |
| 76 | self.weight_key: Optional[str] = None |
| 77 | else: |
| 78 | self.weight_key: Optional[str] = f"{prefix}.weight" |
| 79 | self.with_weight: bool = self.weight_key is not None |
| 80 | self.eps: float = eps |
| 81 | if current_platform.is_gcu(): |
| 82 | self.norm_func: Callable = fused_add_rms_norm |
| 83 | else: |
| 84 | self.norm_func: Callable = fused_rms_norm |
| 85 | self.bias: Optional[paddle.Tensor] = bias |
| 86 | self.quant_scale: Optional[float] = quant_scale |
| 87 | |
| 88 | self._norm_weight_dtype = dtype |
| 89 | if self._norm_weight_dtype is None: |
| 90 | self._norm_weight_dtype = self._helper.get_default_dtype() |
| 91 | else: |
| 92 | assert dtype in [ |
| 93 | "float32", |
| 94 | "bfloat16", |
| 95 | "float16", |
| 96 | ], f"Unsupported dtype: {dtype}. Must be one of: float32, bfloat16, float16" |
| 97 | |
| 98 | self.quant_round_type: int = ( |
| 99 | self.fd_config.quant_config.quant_round_type |
no test coverage detected