Add support for CPU

amplify.py

@@ -4,7 +4,7 @@
 
 import torch
 from torch import nn
-
+from torch.nn.functional import scaled_dot_product_attention
 from xformers.ops import SwiGLU, memory_efficient_attention
 
 from .rmsnorm import RMSNorm
@@ -13,6 +13,7 @@ from .rotary import precompute_freqs_cis, apply_rotary_emb
 from transformers import PreTrainedModel, PretrainedConfig
 from transformers.modeling_outputs import MaskedLMOutput
 
+
 class DotDict(dict):
     """Dictionary that supports the dot notation to access attributes (similarly to HuggingFace)."""
 
@@ -20,8 +21,10 @@ class DotDict(dict):
     __setattr__ = dict.__setitem__
     __delattr__ = dict.__delitem__
 
+
 class AMPLIFYConfig(PretrainedConfig):
     model_type = "AMPLIFY"
+
     # All config parameters must have a default value.
     def __init__(
         self,
@@ -45,7 +48,7 @@ class AMPLIFYConfig(PretrainedConfig):
         **kwargs,
     ):
         super().__init__(**kwargs)
-        
+
         self.hidden_size = hidden_size
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
@@ -63,7 +66,7 @@ class AMPLIFYConfig(PretrainedConfig):
         self.att_bias = att_bias
         self.pad_token_id = pad_token_id
         self.max_length = max_length
-        
+
 
 class EncoderBlock(nn.Module):
     """Transformer encoder block."""
@@ -119,8 +122,16 @@ class EncoderBlock(nn.Module):
                     nn.Linear(config.intermediate_size, config.hidden_size, bias=config.ffn_bias),
                 )
 
-        self.attention_norm = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
-        self.ffn_norm = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
+        self.attention_norm = (
+            RMSNorm(config.hidden_size, config.norm_eps)
+            if config.rms_norm
+            else nn.LayerNorm(config.hidden_size, config.norm_eps)
+        )
+        self.ffn_norm = (
+            RMSNorm(config.hidden_size, config.norm_eps)
+            if config.rms_norm
+            else nn.LayerNorm(config.hidden_size, config.norm_eps)
+        )
 
         self.ffn_dropout = nn.Dropout(config.dropout_prob)
 
@@ -130,7 +141,9 @@ class EncoderBlock(nn.Module):
         x = x + self._ff_block(self.ffn_norm(x))
         return x, contact
 
-    def _att_block(self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool):
+    def _att_block(
+        self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool
+    ):
         batch_size, seq_len, _ = x.shape
         xq, xk, xv = self.q(x), self.k(x), self.v(x)
 
@@ -140,22 +153,37 @@ class EncoderBlock(nn.Module):
         xv = xv.view(batch_size, seq_len, self.config.num_attention_heads, self.d_head)
         xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
 
-        attn = memory_efficient_attention(
-            query=xq,
-            key=xk,
-            value=xv,
-            attn_bias=attention_mask,
-            p=self.config.dropout_prob if self.training else 0,
-        )
-
-        _attn = None
+        # Compute the attention weight
+        attn_weights = None
         if output_attentions:
-            _attn = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
+            attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
             if attention_mask is not None:
-                _attn = _attn + attention_mask
-            _attn = _attn.softmax(-1)
-
-        return self.resid_dropout(self.wo(attn.view(batch_size, seq_len, self.config.num_attention_heads * self.d_head))), _attn
+                attn_weights = attn_weights + attention_mask
+            attn_weights = attn_weights.softmax(-1)
+
+        # Compute the attention using xformers if the tensors are on GPU
+        if x.is_cuda:
+            # Input and output are of dimension (B, M, H, K) where B is the batch size, M the sequence length,
+            # H the number of heads, and K the embeding size per head
+            attn = memory_efficient_attention(
+                query=xq,
+                key=xk,
+                value=xv,
+                attn_bias=attention_mask,
+                p=self.config.dropout_prob if self.training else 0,
+            )
+        else:
+            # Input and output are of dimension (B, H, M, K)
+            attn = scaled_dot_product_attention(
+                query=xq.transpose(1, 2),
+                key=xk.transpose(1, 2),
+                value=xv.transpose(1, 2),
+                attn_mask=attention_mask,
+                dropout_p=self.config.dropout_prob if self.training else 0,
+            ).transpose(1, 2)
+
+        attn_scores = self.wo(attn.view(batch_size, seq_len, self.config.num_attention_heads * self.d_head))
+        return (self.resid_dropout(attn_scores), attn_weights)
 
     def _ff_block(self, x: torch.Tensor):
         return self.ffn_dropout(self.ffn(x))
@@ -176,9 +204,10 @@ class AMPLIFYPreTrainedModel(PreTrainedModel):
 class AMPLIFY(AMPLIFYPreTrainedModel):
     """The main model class.
 
-       Args:
-          config (amplify.model.amplify.AMPLIFYConfig): model configuration, usually defined from the Hydra configuration.
+    Args:
+       config (amplify.model.amplify.AMPLIFYConfig): model configuration, usually defined from the Hydra configuration.
     """
+
     def __init__(self, config: AMPLIFYConfig, **kwargs):
         super().__init__(config)
 
@@ -187,19 +216,27 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
 
         if config.layer_norm_after_embedding:
-            self.layer_norm_1 = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
+            self.layer_norm_1 = (
+                RMSNorm(config.hidden_size, config.norm_eps)
+                if config.rms_norm
+                else nn.LayerNorm(config.hidden_size, config.norm_eps)
+            )
 
         self.transformer_encoder = nn.ModuleList()
         for _ in range(config.num_hidden_layers):
             self.transformer_encoder.append(EncoderBlock(config))
 
         if config.layer_norm_before_last_layer:
-            self.layer_norm_2 = RMSNorm(config.hidden_size, config.norm_eps) if config.rms_norm else nn.LayerNorm(config.hidden_size, config.norm_eps)
+            self.layer_norm_2 = (
+                RMSNorm(config.hidden_size, config.norm_eps)
+                if config.rms_norm
+                else nn.LayerNorm(config.hidden_size, config.norm_eps)
+            )
 
         self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
 
         self.freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)
-        
+
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -209,7 +246,11 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
 
         # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
         if attention_mask is not None and not torch.all(attention_mask == 0):
-            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
+            attention_mask = (
+                attention_mask.unsqueeze(1)
+                .unsqueeze(1)
+                .repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
+            )
         else:
             attention_mask = None
 
@@ -234,4 +275,4 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         logits = self.decoder(self.layer_norm_2(x) if self.config.layer_norm_before_last_layer else x)
 
         # Return logits or the output of the last hidden layer
-        return MaskedLMOutput(logits=logits, hidden_states=hidden_states, attentions=attentions)
+        return MaskedLMOutput(logits=logits, hidden_states=hidden_states, attentions=attentions)