Repository: microsoft/DeepSpeed
File: deepspeed/inference/v2/model_implementations/exaone4/model.py (Python)

Code before:
    @property
    def positional_embedding_type(self) -> PositionalEmbeddingType:
        return PositionalEmbeddingType.rotate_half

    @property
    def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
        rope_theta = getattr(self._config, "rope_theta", 1000000.0)
        return RotateHalfConfig(theta_base=rope_theta)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._qk_norm = CUDARMSNorm(
            channels=self.head_size,
            fp_dtype=torch.float16 if self.activation_dtype == DtypeEnum.fp16 else torch.bfloat16,
            epsilon=getattr(self._config, "rms_norm_eps", 1e-5),
        )

    def _apply_qk_norm(self, hidden_states: torch.Tensor, q_norm_gamma: torch.Tensor,
                       k_norm_gamma: torch.Tensor) -> torch.Tensor:
        """
        Apply RMSNorm to Q and K projections independently per head.
        hidden_states shape: [tokens, (n_q + n_kv + n_kv) * head_size]
        """
        tokens = hidden_states.shape[0]
        local_n_heads = self.n_heads // max(self.tp_size, 1)
        local_n_heads_kv = self.n_heads_kv // max(self.tp_size, 1)
        q_len = local_n_heads * self.head_size
        kv_len = local_n_heads_kv * self.head_size

        q = hidden_states[:, :q_len].contiguous()
        k = hidden_states[:, q_len:q_len + kv_len].contiguous()
        v = hidden_states[:, q_len + kv_len:]

        # Reshape to [tokens * n_heads, head_size] for per-head RMSNorm
        q = q.view(-1, self.head_size)
        self._qk_norm(q, q, q_norm_gamma)
        q = q.view(tokens, q_len)

        k = k.view(-1, self.head_size)
        self._qk_norm(k, k, k_norm_gamma)
        k = k.view(tokens, kv_len)

        hidden_states[:, :q_len] = q
        hidden_states[:, q_len:q_len + kv_len] = k

        return hidden_states

    def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
        embed = self.embed(ragged_batch, self._non_transformer.word_emb)
        if embed.shape[-1] != self.model_dim:
            raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")

Review comment (bug):
As EXAONE4 has uneven Q/KV heads (GQA), I think this can produce incorrect results. Shouldn't we use these?
- `self.n_heads_q_local` instead of `self.n_heads // self.tp_size`
- `self.n_heads_kv_local` instead of `self.n_heads_kv // self.tp_size`

Code after:
    @property
    def positional_embedding_type(self) -> PositionalEmbeddingType:
        return PositionalEmbeddingType.rotate_half

    @property
    def positional_embedding_config(self) -> Optional[RotateHalfConfig]:
        rope_theta = getattr(self._config, "rope_theta", 1000000.0)
        return RotateHalfConfig(theta_base=rope_theta)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._qk_norm = CUDARMSNorm(
            channels=self.head_size,
            fp_dtype=torch.float16 if self.activation_dtype == DtypeEnum.fp16 else torch.bfloat16,
            epsilon=getattr(self._config, "rms_norm_eps", 1e-5),
        )

    def _apply_qk_norm(self, hidden_states: torch.Tensor, q_norm_gamma: torch.Tensor,
                       k_norm_gamma: torch.Tensor) -> torch.Tensor:
        """
        Apply RMSNorm to Q and K projections independently per head.
        hidden_states shape: [tokens, (n_q + n_kv + n_kv) * head_size]
        """
        tokens = hidden_states.shape[0]
        local_n_heads = self.n_heads_q_local
        local_n_heads_kv = self.n_heads_kv_local
        q_len = local_n_heads * self.head_size
        kv_len = local_n_heads_kv * self.head_size

        q = hidden_states[:, :q_len].contiguous()
        k = hidden_states[:, q_len:q_len + kv_len].contiguous()
        v = hidden_states[:, q_len + kv_len:]

        # Reshape to [tokens * n_heads, head_size] for per-head RMSNorm
        q = q.view(-1, self.head_size)
        self._qk_norm(q, q, q_norm_gamma)
        q = q.view(tokens, q_len)

        k = k.view(-1, self.head_size)
        self._qk_norm(k, k, k_norm_gamma)
        k = k.view(tokens, kv_len)

        hidden_states[:, :q_len] = q
        hidden_states[:, q_len:q_len + kv_len] = k

        return hidden_states

    def _forward_embed(self, ragged_batch: RaggedBatchWrapper) -> torch.Tensor:
        embed = self.embed(ragged_batch, self._non_transformer.word_emb)
        if embed.shape[-1] != self.model_dim:
            raise ValueError(f"Embedding output shape {embed.shape} does not match model_dim {self.model_dim}")