in_dtype = x.dtype x = x.to(torch.float32) # Your Code here performing RMSNorm ... result = ... # Return the result in the original dtype return result.to(in_dtype)
import torch classRMSNorm(torch.nn.Module): def__init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None): """ Construct the RMSNorm module. Args: d_model (int): Hidden dimension of the model. eps (float): Epsilon value for numerical stability, defaults to 1e-5. device (torch.device | None): Device to store the parameters on. dtype (torch.dtype | None): Data type of the parameters. """ super().__init__()
defforward(self, x: torch.Tensor) -> torch.Tensor: """ Process an input tensor of shape (batch_size, sequence_length, d_model) and return a tensor of the same shape. """
import torch classRMSNorm(torch.nn.Module): def__init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None): """ Construct the RMSNorm module. Args: d_model (int): Hidden dimension of the model. eps (float): Epsilon value for numerical stability, defaults to 1e-5. device (torch.device | None): Device to store the parameters on. dtype (torch.dtype | None): Data type of the parameters. """ super().__init__() self.d_model = d_model self.eps = eps kwargs = {'device': device, 'dtype': dtype} # weight = 1 self.weight = torch.nn.Parameter(torch.ones(self.d_model, **kwargs))
defforward(self, x: torch.Tensor) -> torch.Tensor: """ Process an input tensor of shape (batch_size, sequence_length, d_model) and return a tensor of the same shape. """ in_dtype = x.dtype x_float = x.to(torch.float32) rms = torch.sqrt( reduce(x_float ** 2, 'b s d -> b s 1', 'mean') + self.eps ) # Apply gain (broadcast) result = x_float / rms * self.weight # Return the result in the original dtype return result.to(in_dtype)
代码都是按公式来的，按照einops书写的rms公式就是下面这句
1 2 3
rms = torch.sqrt( reduce(x_float ** 2, 'b s d -> b s 1', 'mean') + self.eps )
defrun_swiglu( d_model: int, d_ff: int, w1_weight: Float[Tensor, " d_ff d_model"], w2_weight: Float[Tensor, " d_model d_ff"], w3_weight: Float[Tensor, " d_ff d_model"], in_features: Float[Tensor, " ... d_model"], ) -> Float[Tensor, " ... d_model"]: """Given the weights of a SwiGLU network, return the output of your implementation with these weights. Args: d_model (int): Dimensionality of the feedforward input and output. d_ff (int): Dimensionality of the up-project happening internally to your swiglu. w1_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W1 w2_weight (Float[Tensor, "d_model d_ff"]): Stored weights for W2 w3_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W3 in_features (Float[Tensor, "... d_model"]): Input embeddings to the feed-forward layer. Returns: Float[Tensor, "... d_model"]: Output embeddings of the same shape as the input embeddings. """ # Example: # If your state dict keys match, you can use `load_state_dict()` # swiglu.load_state_dict(weights) # You can also manually assign the weights # swiglu.w1.weight.data = w1_weight # swiglu.w2.weight.data = w2_weight # swiglu.w3.weight.data = w3_weight module = SwiGLU(d_model, d_ff, w1_weight, w2_weight, w3_weight) return module.forward(in_features)
测试结果如下：
Relative Positional Embeddings
为了将编码的位置信息注入到模型中，我们将实现**旋转位置嵌入（Rotary Position Embeddings，RoPE）。**不同于传统的绝对位置编码（sin/cos 直接加到 token embeddings 上），RoPE不改变embedding本身，而是在注意力计算中对Q和K引入与位置相关的旋转，这样做具有推理快、可外推、数学自然、实现简单的优点。
def__**init__**(self, theta: float, d_k: int, max_seq_len: int, device=None): """ Construct the RoPE module and create buffers if needed. Args: theta (float): Θ value for the RoPE. d_k (int): Dimension of query and key vectors. max_seq_len (int): Maximum sequence length that will be inputted. device (torch.device | None): Device to store the buffer on. """
defforward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor: """ Process an input tensor of shape (..., seq_len, d_k) and return a tensor of the same shape. Note: - x can have an arbitrary number of batch dimensions. - token_positions is a tensor of shape (..., seq_len) specifying the token positions of x along the sequence dimension. - Use the token positions to slice your (possibly precomputed) cos and sin tensors along the sequence dimension. Args: x (torch.Tensor): Input tensor. token_positions (torch.Tensor): Tensor containing token positions along the sequence dimension. Returns: torch.Tensor: Output tensor of the same shape as x. """
def__init__(self, theta: float, d_k: int, max_seq_len: int, device=None): """ Construct the RoPE module and create buffers if needed. Args: theta (float): Θ value for the RoPE. d_k (int): Dimension of query and key vectors. max_seq_len (int): Maximum sequence length that will be inputted. device (torch.device | None): Device to store the buffer on. """ super().__init__() if d_k % 2 != 0: raise ValueError("d_k must be even for RoPE") self.theta = theta self.d_k = d_k self.max_seq_len = max_seq_len # 1 / *Θ^((2k - 2) / d)* frequency = 1.0 / (theta ** (torch.arange(0, d_k, 2).float() / d_k)) # position = [0,1,...,i] position = torch.arange(max_seq_len, device=device) rotate_theta = torch.outer(position, frequency)
defforward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor: """ Process an input tensor of shape (..., seq_len, d_k) and return a tensor of the same shape. Note: - x can have an arbitrary number of batch dimensions. - token_positions is a tensor of shape (..., seq_len) specifying the token positions of x along the sequence dimension. - Use the token positions to slice your (possibly precomputed) cos and sin tensors along the sequence dimension. Args: x (torch.Tensor): Input tensor. token_positions (torch.Tensor): Tensor containing token positions along the sequence dimension. Returns: torch.Tensor: Output tensor of the same shape as x. """ sin_theta = self.sin_register[token_positions] cos_theta = self.cos_register[token_positions]
