backends/gaudi/server/text_generation_server/models/custom_modeling/flash_dbrx_modeling.py [32:265]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - ) from text_generation_server.layers import ( FastLinear, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, SpeculativeHead, get_linear, ) from text_generation_server.layers.rotary import ( PositionRotaryEmbedding, ) from text_generation_server.layers.layernorm import ( FastLayerNorm, ) from vllm_hpu_extension.ops import DynamicFusedMOE import habana_frameworks.torch as htorch class DbrxAttentionConfig(PretrainedConfig): def __init__( self, attn_pdrop: float = 0, clip_qkv: Optional[float] = None, kv_n_heads: int = 1, rope_theta: float = 10000.0, **kwargs: Any, ): super().__init__(**kwargs) self.attn_pdrop = attn_pdrop self.clip_qkv = clip_qkv self.kv_n_heads = kv_n_heads self.rope_theta = rope_theta for k in ["model_type"]: if k in kwargs: kwargs.pop(k) if len(kwargs) != 0: raise ValueError(f"Found unknown {kwargs=}") class DbrxFFNConfig(PretrainedConfig): def __init__( self, ffn_act_fn: Optional[dict] = None, ffn_hidden_size: int = 3584, moe_num_experts: int = 4, moe_top_k: int = 1, moe_jitter_eps: Optional[float] = None, moe_loss_weight: float = 0.01, moe_normalize_expert_weights: Optional[float] = 1, uniform_expert_assignment: bool = False, **kwargs: Any, ): super().__init__() if ffn_act_fn is None: ffn_act_fn = {"name": "silu"} self.ffn_act_fn = ffn_act_fn self.ffn_hidden_size = ffn_hidden_size self.moe_num_experts = moe_num_experts self.moe_top_k = moe_top_k self.moe_jitter_eps = moe_jitter_eps self.moe_loss_weight = moe_loss_weight self.moe_normalize_expert_weights = moe_normalize_expert_weights self.uniform_expert_assignment = uniform_expert_assignment if uniform_expert_assignment: raise ValueError("`uniform_expert_assignment = True` is not supported") for k in ["model_type"]: if k in kwargs: kwargs.pop(k) if len(kwargs) != 0: raise ValueError(f"Found unknown {kwargs=}") class DbrxConfig(PretrainedConfig): attribute_map = { "hidden_size": "d_model", "num_attention_heads": "n_heads", "num_hidden_layers": "n_layers", } def __init__( self, d_model: int = 2048, n_heads: int = 16, n_layers: int = 24, max_seq_len: int = 2048, vocab_size: int = 32000, resid_pdrop: float = 0.0, emb_pdrop: float = 0.0, attn_config: Optional[DbrxAttentionConfig] = None, ffn_config: Optional[DbrxFFNConfig] = None, use_cache: bool = True, initializer_range: float = 0.02, output_router_logits: bool = False, router_aux_loss_coef: float = 0.05, **kwargs: Any, ): if attn_config is None: self.attn_config = DbrxAttentionConfig() elif isinstance(attn_config, dict): self.attn_config = DbrxAttentionConfig(**attn_config) else: self.attn_config = attn_config if ffn_config is None: self.ffn_config = DbrxFFNConfig() elif isinstance(ffn_config, dict): self.ffn_config = DbrxFFNConfig(**ffn_config) else: self.ffn_config = ffn_config self.d_model = d_model self.n_heads = n_heads self.n_layers = n_layers self.max_seq_len = max_seq_len self.vocab_size = vocab_size self.resid_pdrop = resid_pdrop self.emb_pdrop = emb_pdrop self.use_cache = use_cache self.initializer_range = initializer_range self.output_router_logits = output_router_logits self.router_aux_loss_coef = router_aux_loss_coef tie_word_embeddings = kwargs.pop("tie_word_embeddings", False) if tie_word_embeddings: raise ValueError("tie_word_embeddings is not supported for Dbrx models.") super().__init__( tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def num_key_value_heads(self): # We can't use the attribute map, since this the number of KV # heads is not top-level. return self.attn_config.kv_n_heads def promote_scalar(x: torch.Tensor) -> torch.Tensor: return x.view(1) if len(x.size()) == 0 else x def load_attention(config, prefix, weights): return TensorParallelColumnLinear.load_qkv( config, prefix=f"{prefix}.Wqkv", weights=weights, bias=False, num_heads=config.n_heads, num_key_value_heads=config.attn_config.kv_n_heads, ) def _load_experts(config, prefix, weights): world_size = weights.process_group.size() rank = weights.process_group.rank() assert ( config.ffn_config.ffn_hidden_size % world_size == 0 ), f"The chosen size {config.ffn_config.ffn_hidden_size} is not compatible with sharding on {world_size} shards" expert_size = config.ffn_config.ffn_hidden_size block_size = expert_size // world_size start = rank * block_size stop = (rank + 1) * block_size tensor = torch.empty( (config.ffn_config.moe_num_experts * block_size, config.d_model), dtype=weights.dtype, device=weights.device, ) slice_ = weights._get_slice(f"{prefix}") for i in range(config.ffn_config.moe_num_experts): offset = i * expert_size expert_slice = slice_[start + offset : stop + offset] tensor[i * block_size : (i + 1) * block_size] = expert_slice.to( dtype=weights.dtype ).to(device=weights.device) return tensor def _load_experts_quantized(config, prefix, weights, cls): world_size = weights.process_group.size() rank = weights.process_group.rank() assert ( config.ffn_config.ffn_hidden_size % world_size == 0 ), f"The chosen size {config.ffn_config.ffn_hidden_size} is not compatible with sharding on {world_size} shards" expert_size = config.ffn_config.ffn_hidden_size block_size = expert_size // world_size start = rank * block_size stop = (rank + 1) * block_size slice_ = weights._get_slice(f"{prefix}") experts = [] for i in range(config.ffn_config.moe_num_experts): if config.quantize in ["gptq", "awq"]: raise NotImplementedError( "Dbrx does not support gptq/awq quantization yet." ) else: offset = i * expert_size expert_slice = ( slice_[start + offset : stop + offset] .to(dtype=weights.dtype) .to(device=weights.device) ) if cls == TensorParallelRowLinear: expert_slice = expert_slice.t().contiguous() linear = get_linear(expert_slice, None) experts.append(cls(linear, weights.process_group)) else: linear = get_linear(expert_slice, None) experts.append(cls(linear)) return experts class DbrxAttention(torch.nn.Module): def __init__( self, prefix: str, config, weights, - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - server/text_generation_server/models/custom_modeling/flash_dbrx_modeling.py [38:269]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - ) from text_generation_server.layers import ( FastLinear, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, SpeculativeHead, get_linear, ) from text_generation_server.layers.rotary import ( PositionRotaryEmbedding, ) from text_generation_server.layers.layernorm import ( FastLayerNorm, ) class DbrxAttentionConfig(PretrainedConfig): def __init__( self, attn_pdrop: float = 0, clip_qkv: Optional[float] = None, kv_n_heads: int = 1, rope_theta: float = 10000.0, **kwargs: Any, ): super().__init__(**kwargs) self.attn_pdrop = attn_pdrop self.clip_qkv = clip_qkv self.kv_n_heads = kv_n_heads self.rope_theta = rope_theta for k in ["model_type"]: if k in kwargs: kwargs.pop(k) if len(kwargs) != 0: raise ValueError(f"Found unknown {kwargs=}") class DbrxFFNConfig(PretrainedConfig): def __init__( self, ffn_act_fn: Optional[dict] = None, ffn_hidden_size: int = 3584, moe_num_experts: int = 4, moe_top_k: int = 1, moe_jitter_eps: Optional[float] = None, moe_loss_weight: float = 0.01, moe_normalize_expert_weights: Optional[float] = 1, uniform_expert_assignment: bool = False, **kwargs: Any, ): super().__init__() if ffn_act_fn is None: ffn_act_fn = {"name": "silu"} self.ffn_act_fn = ffn_act_fn self.ffn_hidden_size = ffn_hidden_size self.moe_num_experts = moe_num_experts self.moe_top_k = moe_top_k self.moe_jitter_eps = moe_jitter_eps self.moe_loss_weight = moe_loss_weight self.moe_normalize_expert_weights = moe_normalize_expert_weights self.uniform_expert_assignment = uniform_expert_assignment if uniform_expert_assignment: raise ValueError("`uniform_expert_assignment = True` is not supported") for k in ["model_type"]: if k in kwargs: kwargs.pop(k) if len(kwargs) != 0: raise ValueError(f"Found unknown {kwargs=}") class DbrxConfig(PretrainedConfig): attribute_map = { "hidden_size": "d_model", "num_attention_heads": "n_heads", "num_hidden_layers": "n_layers", } def __init__( self, d_model: int = 2048, n_heads: int = 16, n_layers: int = 24, max_seq_len: int = 2048, vocab_size: int = 32000, resid_pdrop: float = 0.0, emb_pdrop: float = 0.0, attn_config: Optional[DbrxAttentionConfig] = None, ffn_config: Optional[DbrxFFNConfig] = None, use_cache: bool = True, initializer_range: float = 0.02, output_router_logits: bool = False, router_aux_loss_coef: float = 0.05, **kwargs: Any, ): if attn_config is None: self.attn_config = DbrxAttentionConfig() elif isinstance(attn_config, dict): self.attn_config = DbrxAttentionConfig(**attn_config) else: self.attn_config = attn_config if ffn_config is None: self.ffn_config = DbrxFFNConfig() elif isinstance(ffn_config, dict): self.ffn_config = DbrxFFNConfig(**ffn_config) else: self.ffn_config = ffn_config self.d_model = d_model self.n_heads = n_heads self.n_layers = n_layers self.max_seq_len = max_seq_len self.vocab_size = vocab_size self.resid_pdrop = resid_pdrop self.emb_pdrop = emb_pdrop self.use_cache = use_cache self.initializer_range = initializer_range self.output_router_logits = output_router_logits self.router_aux_loss_coef = router_aux_loss_coef tie_word_embeddings = kwargs.pop("tie_word_embeddings", False) if tie_word_embeddings: raise ValueError("tie_word_embeddings is not supported for Dbrx models.") super().__init__( tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def num_key_value_heads(self): # We can't use the attribute map, since this the number of KV # heads is not top-level. return self.attn_config.kv_n_heads def promote_scalar(x: torch.Tensor) -> torch.Tensor: return x.view(1) if len(x.size()) == 0 else x def load_attention(config, prefix, weights): return TensorParallelColumnLinear.load_qkv( config, prefix=f"{prefix}.Wqkv", weights=weights, bias=False, num_heads=config.n_heads, num_key_value_heads=config.attn_config.kv_n_heads, ) def _load_experts(config, prefix, weights): world_size = weights.process_group.size() rank = weights.process_group.rank() assert ( config.ffn_config.ffn_hidden_size % world_size == 0 ), f"The chosen size {config.ffn_config.ffn_hidden_size} is not compatible with sharding on {world_size} shards" expert_size = config.ffn_config.ffn_hidden_size block_size = expert_size // world_size start = rank * block_size stop = (rank + 1) * block_size tensor = torch.empty( (config.ffn_config.moe_num_experts * block_size, config.d_model), dtype=weights.dtype, device=weights.device, ) slice_ = weights._get_slice(f"{prefix}") for i in range(config.ffn_config.moe_num_experts): offset = i * expert_size expert_slice = slice_[start + offset : stop + offset] tensor[i * block_size : (i + 1) * block_size] = expert_slice.to( dtype=weights.dtype ).to(device=weights.device) return tensor def _load_experts_quantized(config, prefix, weights, cls): world_size = weights.process_group.size() rank = weights.process_group.rank() assert ( config.ffn_config.ffn_hidden_size % world_size == 0 ), f"The chosen size {config.ffn_config.ffn_hidden_size} is not compatible with sharding on {world_size} shards" expert_size = config.ffn_config.ffn_hidden_size block_size = expert_size // world_size start = rank * block_size stop = (rank + 1) * block_size slice_ = weights._get_slice(f"{prefix}") experts = [] for i in range(config.ffn_config.moe_num_experts): if config.quantize in ["gptq", "awq"]: raise NotImplementedError( "Dbrx does not support gptq/awq quantization yet." ) else: offset = i * expert_size expert_slice = ( slice_[start + offset : stop + offset] .to(dtype=weights.dtype) .to(device=weights.device) ) if cls == TensorParallelRowLinear: expert_slice = expert_slice.t().contiguous() linear = get_linear(expert_slice, None) experts.append(cls(linear, weights.process_group)) else: linear = get_linear(expert_slice, None) experts.append(cls(linear)) return experts class DbrxAttention(torch.nn.Module): def __init__( self, prefix: str, config, weights, - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -