From 0490ce36ae4fe1fceb20ae22f2aa1af14abc04dd Mon Sep 17 00:00:00 2001 From: Saurabh Dash Date: Wed, 3 Apr 2024 19:20:57 +0000 Subject: [PATCH 1/3] changes --- .../models/cohere/configuration_cohere.py | 4 ++ .../models/cohere/modeling_cohere.py | 63 ++++++++++++------- 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/src/transformers/models/cohere/configuration_cohere.py b/src/transformers/models/cohere/configuration_cohere.py index a310ad54302ad..7ceca2b887af7 100644 --- a/src/transformers/models/cohere/configuration_cohere.py +++ b/src/transformers/models/cohere/configuration_cohere.py @@ -85,6 +85,8 @@ class CohereConfig(PretrainedConfig): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. + use_qk_norm (`bool`, *optional*, defaults to `False`): + Whether to use query-key normalization in the attention ```python >>> from transformers import CohereModel, CohereConfig @@ -123,6 +125,7 @@ def __init__( rope_theta=10000.0, attention_bias=False, attention_dropout=0.0, + use_qk_norm=False, **kwargs, ): self.vocab_size = vocab_size @@ -145,6 +148,7 @@ def __init__( self.rope_theta = rope_theta self.attention_bias = attention_bias self.attention_dropout = attention_dropout + self.use_qk_norm = use_qk_norm super().__init__( pad_token_id=pad_token_id, diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index e949bc14482e7..5015d8cb702b9 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -76,10 +76,9 @@ def _get_unpad_data(attention_mask): class CohereLayerNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-5, bias=False): + def __init__(self, param_shape=None, eps=1e-5, bias=False): super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.bias = nn.Parameter(torch.zeros(hidden_size)) if bias else None + self.weight = nn.Parameter(torch.ones(param_shape)) self.variance_epsilon = eps def forward(self, hidden_states): @@ -89,8 +88,6 @@ def forward(self, hidden_states): variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True) hidden_states = (hidden_states - mean) * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight.to(torch.float32) * hidden_states - if self.bias is not None: - hidden_states = hidden_states + self.bias.to(torch.float32) return hidden_states.to(input_dtype) @@ -122,7 +119,7 @@ def forward(self, x, position_ids): emb = torch.repeat_interleave(freqs, 2, dim=-1) cos = emb.cos() sin = emb.sin() - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + return cos, sin def rotate_half(x): @@ -133,7 +130,6 @@ def rotate_half(x): return rot_x -# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. @@ -154,11 +150,14 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ + dtype = q.dtype + q = q.float() + k = k.float() cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed + return q_embed.to(dtype=dtype), k_embed.to(dtype=dtype) # Copied from transformers.models.llama.modeling_llama.LlamaMLP Llama->Cohere @@ -192,7 +191,6 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) -# Copied from transformers.models.llama.modeling_llama.LlamaAttention Llama->Cohere class CohereAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -216,6 +214,7 @@ def __init__(self, config: CohereConfig, layer_idx: Optional[int] = None): self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.is_causal = True + self.use_qk_norm = config.use_qk_norm if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError( @@ -223,6 +222,13 @@ def __init__(self, config: CohereConfig, layer_idx: Optional[int] = None): f" and `num_heads`: {self.num_heads})." ) + if self.use_qk_norm: + # When sharding the model using Tensor Parallelism, need to be careful to use n_local_heads + self.q_norm = CohereLayerNorm(param_shape=(self.num_heads, self.head_dim), eps=config.layer_norm_eps) + self.k_norm = CohereLayerNorm( + param_shape=(self.num_key_value_heads, self.head_dim), eps=config.layer_norm_eps + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) @@ -255,8 +261,14 @@ def forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim) + if self.use_qk_norm: + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) past_key_value = getattr(self, "past_key_value", past_key_value) @@ -335,11 +347,14 @@ def forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - # Flash attention requires the input to have the shape - # batch_size x seq_length x head_dim x hidden_dim - # therefore we just need to keep the original shape - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim) + if self.use_qk_norm: + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(value_states, position_ids) @@ -505,7 +520,7 @@ class CohereSdpaAttention(CohereAttention): SDPA API. """ - # Adapted from CohereAttention.forward + # Ignore copy def forward( self, hidden_states: torch.Tensor, @@ -538,8 +553,14 @@ def forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim) + if self.use_qk_norm: + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(value_states, position_ids) @@ -599,7 +620,7 @@ def __init__(self, config: CohereConfig, layer_idx: int): self.self_attn = COHERE_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) self.mlp = CohereMLP(config) - self.input_layernorm = CohereLayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.input_layernorm = CohereLayerNorm(param_shape=(config.hidden_size), eps=config.layer_norm_eps) def forward( self, @@ -822,7 +843,7 @@ def __init__(self, config: CohereConfig): self.layers = nn.ModuleList( [CohereDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) - self.norm = CohereLayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm = CohereLayerNorm(param_shape=(config.hidden_size), eps=config.layer_norm_eps) self.gradient_checkpointing = False # Initialize weights and apply final processing From 907dcb457cc2ea65eb069a9fe45c7d4cfe90ec32 Mon Sep 17 00:00:00 2001 From: Saurabh Dash Date: Thu, 4 Apr 2024 09:22:40 +0000 Subject: [PATCH 2/3] addressing comments --- src/transformers/models/cohere/modeling_cohere.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index 5015d8cb702b9..a835e79d52d4e 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -76,9 +76,10 @@ def _get_unpad_data(attention_mask): class CohereLayerNorm(nn.Module): - def __init__(self, param_shape=None, eps=1e-5, bias=False): + def __init__(self, hidden_size=None, eps=1e-5, bias=False): + """ The hidden size can be a tuple or an int. The tuple is used for QKNorm to normalize across head_dim """ super().__init__() - self.weight = nn.Parameter(torch.ones(param_shape)) + self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): @@ -224,9 +225,9 @@ def __init__(self, config: CohereConfig, layer_idx: Optional[int] = None): if self.use_qk_norm: # When sharding the model using Tensor Parallelism, need to be careful to use n_local_heads - self.q_norm = CohereLayerNorm(param_shape=(self.num_heads, self.head_dim), eps=config.layer_norm_eps) + self.q_norm = CohereLayerNorm(hidden_size=(self.num_heads, self.head_dim), eps=config.layer_norm_eps) self.k_norm = CohereLayerNorm( - param_shape=(self.num_key_value_heads, self.head_dim), eps=config.layer_norm_eps + hidden_size=(self.num_key_value_heads, self.head_dim), eps=config.layer_norm_eps ) self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) @@ -620,7 +621,7 @@ def __init__(self, config: CohereConfig, layer_idx: int): self.self_attn = COHERE_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) self.mlp = CohereMLP(config) - self.input_layernorm = CohereLayerNorm(param_shape=(config.hidden_size), eps=config.layer_norm_eps) + self.input_layernorm = CohereLayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps) def forward( self, @@ -843,7 +844,7 @@ def __init__(self, config: CohereConfig): self.layers = nn.ModuleList( [CohereDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) - self.norm = CohereLayerNorm(param_shape=(config.hidden_size), eps=config.layer_norm_eps) + self.norm = CohereLayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps) self.gradient_checkpointing = False # Initialize weights and apply final processing From 9326091083dd67e104c78b5c29d371f7a11ae9fc Mon Sep 17 00:00:00 2001 From: Saurabh Dash Date: Thu, 4 Apr 2024 09:53:15 +0000 Subject: [PATCH 3/3] smol fix --- src/transformers/models/cohere/modeling_cohere.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/cohere/modeling_cohere.py b/src/transformers/models/cohere/modeling_cohere.py index a835e79d52d4e..41bae6db65e15 100644 --- a/src/transformers/models/cohere/modeling_cohere.py +++ b/src/transformers/models/cohere/modeling_cohere.py @@ -77,7 +77,7 @@ def _get_unpad_data(attention_mask): class CohereLayerNorm(nn.Module): def __init__(self, hidden_size=None, eps=1e-5, bias=False): - """ The hidden size can be a tuple or an int. The tuple is used for QKNorm to normalize across head_dim """ + """The hidden size can be a tuple or an int. The tuple is used for QKNorm to normalize across head_dim""" super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps