Sync modeling_auristream.py from TuKoResearch/AuriStream7BDeep_7k

modeling_auristream.py

@@ -0,0 +1,596 @@
+"""
+AuriStream sequence model definition.
+"""
+
+import math
+import inspect
+import random
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import numpy as np
+from huggingface_hub import PyTorchModelHubMixin
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
+from transformers import PreTrainedModel
+from .configuration_auristream import AuriStreamConfig
+
+
+class AuriStream(PreTrainedModel):
+    config_class = AuriStreamConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        # if use_rope is in the config and false, initialize a wpe layer in transformer
+        if hasattr(config, 'use_rope') and not config.use_rope:
+            self.transformer = nn.ModuleDict(dict(
+                wte = nn.Embedding(config.vocab_size, config.n_embd),
+                wpe = nn.Embedding(config.seq_len, config.n_embd),
+                drop = nn.Dropout(config.dropout),
+                h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+                ln_f = RMSNorm(config.n_embd, bias=config.bias),
+            ))
+        else:
+            self.transformer = nn.ModuleDict(dict(
+                wte = nn.Embedding(config.vocab_size, config.n_embd),
+                drop = nn.Dropout(config.dropout),
+                h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+                ln_f = RMSNorm(config.n_embd, bias=config.bias),
+            ))
+
+        # check if n_pred_steps is defined in the config, this is the number of linear heads for prediction
+        if hasattr(config, 'n_pred_steps'):
+            self.future_heads = nn.ModuleList([nn.Linear(config.n_embd, config.vocab_size, bias=False) for _ in range(config.n_pred_steps - 1)])
+        else:
+            self.future_heads = None
+
+        self.coch_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # # init all weights
+        # self.apply(self._init_weights)
+        # # apply special scaled init to the residual projections, per GPT-2 paper
+        # for pn, p in self.named_parameters():
+        #     if pn.endswith('c_proj.weight'):
+        #         torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        self.dwa = None
+        if config.skip_connections:
+            self.dwa = DWA(config.n_layer + 1)
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+   
+    def forward(self, seq, tgt=None, output_hidden_states=False, return_dict=False, up_until_layer=None):
+        """
+        Input: coch: torch.Tensor of shape (b, t)
+               tgt_coch: torch.Tensor of shape (b, t) or None
+        """
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(seq) # token embeddings of shape (b, t, n_embd)
+
+        # if wpe exists in self.transformer apply leanred positional embedding
+        if hasattr(self.transformer, 'wpe'):
+            pos = torch.arange(0, seq.size(1), dtype=torch.long, device=seq.device)
+            pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+            x = self.transformer.drop(tok_emb + pos_emb)
+        else:
+            x = self.transformer.drop(tok_emb)
+
+        # Initialize embedding accumulators
+        if self.dwa is not None:
+            x = self.dwa.init_accumulators(x)
+
+        all_hidden_states = []
+        for block_idx, block in enumerate(self.transformer.h):
+            # Forward the block
+            all_hidden_states.append(x)
+            if up_until_layer is not None and block_idx == up_until_layer:
+                break
+            x = block(x)
+            if self.dwa is not None:
+                x = self.dwa(x)
+            
+
+        # append the last hidden state if we did not exit early
+        if up_until_layer is None or block_idx == len(self.transformer.h) - 1:
+            all_hidden_states.append(x)
+
+        if output_hidden_states:
+            model_output = BaseModelOutput(
+                last_hidden_state=x,
+                hidden_states=all_hidden_states,
+            )
+            return model_output
+
+        x = self.transformer.ln_f(x)
+        logits = self.coch_head(x)
+
+        if tgt is not None:
+            loss = F.cross_entropy(
+                logits.reshape(-1, self.config.vocab_size), tgt.reshape(-1),
+            )
+
+            # If we have more than one future head, compute the loss for each head
+            if self.future_heads is not None:
+                for i, head in enumerate(self.future_heads):
+                    future_logits = head(x[:, :-(i+1)])
+                    loss += F.cross_entropy(
+                        future_logits.reshape(-1, self.config.vocab_size), tgt[:, (i+1):].reshape(-1),
+                    )
+                # divide loss by number of future heads
+                loss = loss / (len(self.future_heads) + 1)
+
+            if return_dict:
+                model_output = CausalLMOutput(
+                    loss=loss,
+                    logits=logits,
+                )
+                return model_output
+            
+            return logits, loss
+
+        return logits, None
+
+    def sample_logits(self, logits: torch.FloatTensor, temperature: float = 0.9, 
+                      top_k: int = 500, top_p: float = 0.5) -> torch.LongTensor:
+        """
+        Samples an integer from the distribution of logits
+        Parameters:
+            logits (torch.FloatTensor): The logits of the distribution
+            temp (float): The temperature of the sampling, if 0.0, then argmax is used
+            top_k (int): The number of top k tokens to consider during sampling
+            top_p (float): The cumulative probability threshold for nucleus (top-p) sampling
+        Returns:
+            torch.LongTensor: The sampled integer
+        """
+        # If temperature is 0.0, use argmax
+        if temperature == 0.0:
+            return torch.argmax(logits, dim=-1)
+        
+        # Apply temperature
+        logits = logits / temperature
+
+        # Apply top-k filtering if specified
+        if top_k is not None:
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[..., [-1]]] = -float('Inf')
+
+        # Apply top-p (nucleus) filtering if specified
+        if top_p is not None:
+            # Sort the logits in descending order
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            # Compute the sorted softmax probabilities
+            sorted_probs = F.softmax(sorted_logits, dim=-1)
+            # Compute the cumulative probabilities
+            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+            # Create a mask for tokens to remove
+            sorted_indices_to_remove = cumulative_probs > top_p
+            # Shift the mask right to keep at least one token
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = 0
+            # Scatter the mask back to the original indices
+            indices_to_remove = sorted_indices_to_remove.scatter(dim=-1, index=sorted_indices, src=sorted_indices_to_remove)
+            logits[indices_to_remove] = -float('Inf')
+        
+        # Compute softmax probabilities
+        probs = F.softmax(logits, dim=-1)
+        # Flatten probabilities to (batch_size * sequence_length, vocab_size)
+        flat_probs = probs.view(-1, probs.size(-1))
+        # Sample from the distribution
+        sampled = torch.multinomial(flat_probs, num_samples=1)
+        # Reshape to original shape except for the last dimension
+        sampled = sampled.view(*logits.shape[:-1])
+        return sampled
+
+    @torch.no_grad()
+    def generate(self, seq: torch.Tensor, n_tokens: int = 1, temp=1.0, 
+        top_k=500, top_p=0.5, seed=None):
+        """
+        Parameters: 
+            seq: torch.Tensor of shape (b, t, n_freq_bins)
+                Input cochleagram to use for generation
+            n_tokens: int
+                Number of time bins to predict
+            temp: float
+                Temperature for sampling logits
+            seed: int
+                Random seed for sampling
+        
+        Returns:
+            pred_coch: torch.Tensor of shape (b, t, n_freq_bins)
+                The predicted cochleagram
+            all_logits: (optional if return_logits is True) torch.Tensor of shape (b, n_tokens, n_freq_bins)
+                The logits for each time step
+            all_embs: (optional if return_embs is not None) list of torch.Tensor
+                The embeddings for each transformer block
+        """
+
+        # Set seed if provided
+        if seed is not None:
+            random.seed(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+
+        # make a list of logits to return
+        all_logits = []
+        device = seq.device
+
+        # grab shape of the cochleagram
+        b, t = seq.size()
+
+        # TODO: double check this works then delete the block bellow:
+        # pass the given input through the model to get the predictions and cache
+        # the k and v values for each transformer block in the process
+        # pos = torch.arange(0, t, dtype=torch.long, device=device)
+        # tok_emb = self.transformer.wte(seq) # token embeddings of shape (b, t, n_embd)
+        # pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        # x = self.transformer.drop(tok_emb + pos_emb)
+
+        #### Embed conditioning sequence into KV cache
+
+        tok_emb = self.transformer.wte(seq) # token embeddings of shape (b, t, n_embd)
+        # if wpe exists in self.transformer apply leanred positional embedding
+        if hasattr(self.transformer, 'wpe'):
+            pos = torch.arange(0, seq.size(1), dtype=torch.long, device=seq.device)
+            pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+            x = self.transformer.drop(tok_emb + pos_emb)
+        else:
+            x = self.transformer.drop(tok_emb)
+
+        # Initialize list to store k and v for each transformer block
+        k_list = []
+        v_list = []
+        for block_idx, block in enumerate(self.transformer.h):
+            # Pass through the transformer block, and store k and v
+            x, k, v = block(x, return_kv=True)
+            k_list.append(k)
+            v_list.append(v)
+        # k_cache and v_cache have shape (n_layer, b, n_head, t, n_embd//n_head)
+        k_cache = torch.stack(k_list, dim=0)
+        v_cache = torch.stack(v_list, dim=0)
+        # Pass through the final layer norm
+        x = self.transformer.ln_f(x)
+
+        # First prediction of the model is the decoding of the last time bin
+        logits = self.coch_head(x[:, [-1]])
+        predictions = [self.sample_logits(logits, temperature=temp)]
+        all_logits.append(logits)
+
+        ### Predict future tokens
+
+        # Now we pass the last time bin through the model to predict the next time bin
+        # we subtract 1 from max_new_tokens because we already predicted the first time bin
+        # using the last embedding of the input
+        for i in range(n_tokens-1):
+
+            # TODO: double check this works then delete the block bellow:
+            # # Get the emb and pos embedding of just the last token
+            # pos = torch.arange(t+i, t+i+1, dtype=torch.long, device=device) # shape (t)
+            # tok_emb = self.transformer.wte(predictions[-1]) # token embeddings of shape (b, t, n_embd)
+            # pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+            # x = self.transformer.drop(tok_emb + pos_emb)
+
+            # Get the emb and pos embedding of just the last token
+            tok_emb = self.transformer.wte(predictions[-1]) # token embeddings of shape (b, t, n_embd)
+            # if wpe exists in self.transformer apply leanred positional embedding
+            if hasattr(self.transformer, 'wpe'):
+                pos = torch.arange(t+i, t+i+1, dtype=torch.long, device=device) # shape (t)
+                pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+                x = self.transformer.drop(tok_emb + pos_emb)
+            else:
+                x = self.transformer.drop(tok_emb)
+
+            # Pass through transformer block
+            k_list = []
+            v_list = []
+            for block_idx, block in enumerate(self.transformer.h):
+                x, k, v = block(x, k_cache=k_cache[block_idx], v_cache=v_cache[block_idx])
+                k_list.append(k)
+                v_list.append(v)
+            x = self.transformer.ln_f(x)
+            # create the cache with the new embeddings
+            k_cache = torch.stack(k_list, dim=0)
+            v_cache = torch.stack(v_list, dim=0)
+            # predict next time bin
+            logits = self.coch_head(x)
+            predictions.append(self.sample_logits(logits, temperature=temp, top_k=top_k, top_p=top_p))
+            all_logits.append(logits)
+
+        pred_coch = torch.cat(predictions, dim=1)
+        all_logits = torch.cat(all_logits, dim=1)
+
+        return pred_coch, all_logits
+        
+    
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        # filter out those that do not require grad
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        # Create AdamW optimizer and use the fused version if it is available
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, T, dt, gpu_type='A40'):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+        # first estimate the number of flops we do per iteration.
+        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+        N = self.unsharded_param_count
+        cfg = self.config
+        L, H, Q = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head
+        # L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        # express our flops throughput as ratio of A100 bfloat16 peak flops
+        flops_achieved = flops_per_iter * (1.0/dt) # per second
+
+        # grab promised flops based on GPU type
+        if gpu_type == 'A40':
+            flops_promised = 149.7e12 # A40 GPU bfloat16 peak flops is 149.7 TFLOPS
+        elif gpu_type == 'A100':
+            flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
+        elif gpu_type == 'H100':
+            flops_promised = 756e12 # H100 GPU bfloat16 peak flops is 756 TFLOPS
+        elif gpu_type == 'TPUv4':
+            flops_promised = 275e12
+        elif gpu_type == 'TPUv5e':
+            flops_promised = 197e12
+
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+
+#########################################################
+#####              Layer Definitions                #####
+#########################################################
+
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.attn = CausalSelfAttention(config)
+        self.mlp = MLP(config)
+        self.attn_scale = 1.0 # (1 / (2 * config.n_layer)**0.5)
+        self.norm1 = RMSNorm(config.n_embd, bias=config.bias)
+        self.norm2 = RMSNorm(config.n_embd, bias=config.bias)
+
+    def forward(self, x, return_kv=False, k_cache=None, v_cache=None):
+        # If we are given a key and value cache, we will use the pre-computed values to minimize
+        # the computation cost
+        if k_cache is not None and v_cache is not None:
+            # Pass the key and value cache to the attention layer, obtain new key and value caches
+            x_attn, k, v = self.attn.kv_cache_forward(self.norm1(x), k_cache, v_cache)
+            x = x + x_attn
+            x = x + self.mlp(self.norm2(x))
+            return x, k, v
+        # We might want to encode the caches of a whole block of keys and values at once using the
+        # fast flash attention impelmentation while still returning the key and value caches
+        elif return_kv:
+            # Pass the key and value cache to the attention layer, obtain new key and value caches
+            x_attn, k, v = self.attn(self.norm1(x), return_kv=True)
+            x = x + x_attn
+            x = x + self.mlp(self.norm2(x))
+            return x, k, v
+        
+        x = x + self.attn_scale * self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=False)
+        # output projection
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+        
+        rope_theta = 10000
+        if hasattr(config, 'rope_theta') and config.rope_theta is not None:
+            rope_theta = config.rope_theta
+
+        self.rotary = Rotary(self.head_dim, base=rope_theta)
+
+        if hasattr(config, 'use_rope') and not config.use_rope:
+            self.rotary = None
+
+    def forward(self, x, return_kv=False, return_attn_maps=False):
+
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        qkv = self.c_attn(x)
+        q, k, v = qkv.split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, self.head_dim)
+        q = q.view(B, T, self.n_head, self.head_dim)
+        v = v.view(B, T, self.n_head, self.head_dim)
+
+        if self.rotary is not None:
+            cos, sin = self.rotary(q)
+            q = apply_rotary_emb(q, cos, sin)
+            k = apply_rotary_emb(k, cos, sin)
+
+        if not return_kv and not return_attn_maps:
+            y = F.scaled_dot_product_attention(
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), 
+                is_causal=True)
+        else:
+            # manual implementation of attention
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+            att = torch.einsum('bnsh,bnkh->bnsk', q, k) * (1.0 / math.sqrt(k.size(-1)))
+            mask = torch.triu(torch.ones(T, T), diagonal=1).to(dtype=torch.bool).to(x.device)
+            mask = mask.view(1, 1, T, T)
+            masked_att = att.masked_fill(mask, float('-inf'))
+            # upcast to float32 for numerical stability, as per llama implementation
+            masked_att = F.softmax(masked_att, dim=-1, dtype=torch.float32).to(q.dtype)
+            # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+            y = torch.einsum('bnsk,bnkh->bnsh', masked_att, v)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.c_proj(y)
+    
+        # return attention maps if requested
+        if return_attn_maps:
+            return y, F.softmax(att, dim=-1)
+
+        # return key and value caches if requested
+        if return_kv:
+            return y, k, v
+
+        return y
+
+    def kv_cache_forward(self, x, k_cache=None, v_cache=None):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # append cached keys and values with new keys and values
+        if k_cache is not None:
+            k = torch.cat((k_cache, k), dim=2)
+        if v_cache is not None:
+            v = torch.cat((v_cache, v), dim=2)
+
+        # manual implementation of attention
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = F.softmax(att, dim=-1)
+        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.c_proj(y)
+
+        return y, k, v
+
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.SiLU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class Rotary(torch.nn.Module):
+    def __init__(self, dim, base=10000, learned=False):
+        super().__init__()
+        # Compute the base inverse frequencies as before.
+        # self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    def forward(self, x):
+        seq_len = x.shape[1]
+        # Create a tensor of positions.
+        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+        # Outer product to compute angles; this uses the (possibly learnable) frequencies.
+        freqs = torch.outer(t, self.inv_freq).to(x.device)
+        cos_cached = freqs.cos()
+        sin_cached = freqs.sin()
+        return cos_cached[None, :, None, :], sin_cached[None, :, None, :]
+
+
+def apply_rotary_emb(x, cos, sin):
+    assert x.ndim == 4  # multihead attention expected
+    d = x.shape[3] // 2
+    x1 = x[..., :d]
+    x2 = x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=3)
+
+
+class RMSNorm(nn.Module):
+    """ Root Mean Square Normalization """
+    def __init__(self, dim: int, weight: bool = True, bias: bool = False, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim)) if weight else None
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        if self.weight is not None:
+            return output * self.weight
+        return output
+
+class DWA(nn.Module):
+    """ Depth Weighted Average layer that averages representations across the layers of a transformer """
+    """ From: https://arxiv.org/pdf/2402.02622"""
+    
+    def __init__(self, n_layers: int):
+        super().__init__()
+        self.alphas = nn.Parameter(torch.zeros(n_layers, n_layers))
+        self.alphas.data = torch.eye(n_layers)
+        self.accumulators = []
+    
+    def init_accumulators(self, x):
+        self.accumulators = [x]
+        return x * self.alphas[0, 0]
+    
+    def forward(self, x):
+        self.accumulators.append(x)
+        x = 0.0
+        for i in range(len(self.accumulators)):
+            x = x + self.alphas[i, len(self.accumulators)-1] * self.accumulators[i]
+        return x
+