Upload configuration_phogpt.py

configuration_phogpt.py

@@ -0,0 +1,142 @@
+"""A HuggingFace-style model configuration for PhoGPT."""
+import warnings
+from typing import Any, Dict, Optional, Union
+from transformers import PretrainedConfig
+from .attention import check_alibi_support, is_flash_v1_installed, is_flash_v2_installed
+from .blocks import attn_config_defaults as phogpt_attn_defaults
+from .fc import FC_CLASS_REGISTRY
+from .norm import LPLayerNorm
+from .ffn import FFN_CLASS_REGISTRY
+from .warnings import VersionedDeprecationWarning
+
+ffn_config_defaults: Dict = {'ffn_type': 'phogpt_mlp'}
+init_config_defaults: Dict = {
+    'name': 'kaiming_normal_',
+    'fan_mode': 'fan_in',
+    'init_nonlinearity': 'relu',
+    'init_div_is_residual': True,
+    'emb_init_std': None,
+    'emb_init_uniform_lim': None,
+    'init_std': None,
+    'init_gain': 0.0
+}
+
+class PhoGPTConfig(PretrainedConfig):
+    model_type = 'phogpt'
+
+    def __init__(
+        self,
+        hidden_size: int = 4096,
+        num_attention_heads: int = 32,
+        num_hidden_layers: int = 32,
+        expansion_ratio: Union[int, float] = 4,
+        max_seq_len: int = 4096,
+        vocab_size: int = 51200,
+        resid_pdrop: float = 0.0,
+        emb_pdrop: float = 0.0,
+        learned_pos_emb: bool = True,
+        attn_config: Dict = phogpt_attn_defaults,
+        ffn_config: Dict = ffn_config_defaults,
+        init_device: str = 'cpu',
+        logit_scale: Optional[Union[float, str]] = None,
+        no_bias: bool = False,
+        embedding_fraction: float = 1.0,
+        norm_type: str = 'low_precision_layernorm',
+        use_cache: bool = False,
+        init_config: Dict = init_config_defaults,
+        fc_type: str = 'torch',
+        tie_word_embeddings: bool = True,
+        use_pad_tok_in_ffn: bool = True,
+        **kwargs: Any
+    ):
+        """PhoGPT configuration class.
+
+        Args:
+            hidden_size (int): Model hidden size (embedding dimension)
+            num_attention_heads (int): Number of attention heads
+            num_hidden_layers (int): Number of transformer layers
+            expansion_ratio (int | float): FFN expansion ratio
+            max_seq_len (int): Max sequence length
+            vocab_size (int): Vocabulary size
+            resid_pdrop (float): Dropout on residuals
+            emb_pdrop (float): Dropout on embeddings
+            learned_pos_emb (bool): Use learned positional embeddings
+            attn_config (dict): Attention configuration dictionary
+            ffn_config (dict): Feedforward network config dictionary
+            init_device (str): Device for initialization
+            logit_scale (float | str): Logit scaling
+            no_bias (bool): Disable biases
+            embedding_fraction (float): Scale embedding gradients
+            norm_type (str): LayerNorm type
+            use_cache (bool): Return past key/value
+            init_config (dict): Weight initialization config
+            fc_type (str): Fully connected layer type ('torch' or 'te')
+            tie_word_embeddings (bool): Tie input/output embeddings
+            use_pad_tok_in_ffn (bool): Forward pad tokens through FFN
+        """
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.expansion_ratio = expansion_ratio
+        self.max_seq_len = max_seq_len
+        self.vocab_size = vocab_size
+        self.resid_pdrop = resid_pdrop
+        self.emb_pdrop = emb_pdrop
+        self.learned_pos_emb = learned_pos_emb
+        self.attn_config = attn_config
+        self.ffn_config = ffn_config
+        self.init_device = init_device
+        self.logit_scale = logit_scale
+        self.no_bias = no_bias
+        self.embedding_fraction = embedding_fraction
+        self.norm_type = norm_type
+        self.use_cache = use_cache
+        self.init_config = init_config
+        self.fc_type = fc_type
+        self.use_pad_tok_in_ffn = use_pad_tok_in_ffn
+
+        if 'name' in kwargs:
+            del kwargs['name']
+        if 'loss_fn' in kwargs:
+            del kwargs['loss_fn']
+
+        if self.attn_config.get('alibi', False) or self.attn_config.get('rope', False):
+            self.learned_pos_emb = False
+            warnings.warn("alibi or rope is enabled, setting learned_pos_emb to False.")
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self._validate_config()
+
+    def _set_config_defaults(self, config: Dict[str, Any], config_defaults: Dict[str, Any]) -> Dict[str, Any]:
+        for k, v in config_defaults.items():
+            if k not in config:
+                config[k] = v
+            elif isinstance(v, dict):
+                config[k] = self._set_config_defaults(config.get(k, {}), v)
+        return config
+
+    def _validate_config(self) -> None:
+        self.attn_config = self._set_config_defaults(self.attn_config, phogpt_attn_defaults)
+        self.ffn_config = self._set_config_defaults(self.ffn_config, ffn_config_defaults)
+        self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
+
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+
+        for prob in [self.attn_config.get('attn_pdrop', 0.0), self.resid_pdrop, self.emb_pdrop]:
+            if not 0.0 <= prob <= 1.0:
+                raise ValueError("Dropout probabilities must be between 0 and 1")
+
+        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
+            raise ValueError("embedding_fraction must be in (0, 1]")
+
+        if not (self.learned_pos_emb or self.attn_config.get('alibi', False) or self.attn_config.get('rope', False)):
+            warnings.warn("No positional encoding used: learned_pos_emb, alibi, or rope should be enabled.")
+
+        if self.fc_type == 'te' or self.ffn_config.get('ffn_type') == 'te_ln_mlp':
+            try:
+                import transformer_engine.pytorch as te
+                del te
+            except ImportError:
+                raise ImportError("fc_type='te' requires TransformerEngine installed.")
+