admin/LLM-from-scratch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update code

Browse Source
This commit is contained in:
coffee
2025-05-14 13:35:24 +08:00
parent e25a677be9
commit 9d9e858b65
4 changed files with 40 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -1,3 +1,5 @@
corpus/wikipedia-cn-20230720-filtered.json
corpus/wikipedia.txt
dolly_env_local
llm_dolly.egg-info
src_llm/__pycache__
src_llm/configuration_dolly.py
+1 -1
View File
@@ -1,7 +1,7 @@
from transformers import PretrainedConfig
from transformers import logging
logger = logging.getLogger(__name__)
logger = logging.get_logger(__name__)
class DollyConfig(PretrainedConfig):
model_type = "dolly"
src_llm/modeling_dolly.py
+22 -13
View File
@@ -1,4 +1,4 @@
from configuration_dolly import DollyConfig
from .configuration_dolly import DollyConfig
from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS
from functools import partial
from typing import Callable, Optional, Tuple, Union
@@ -23,23 +23,18 @@ class DollyRMSNorm(nn.Module):
'''
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
# 也是RMSNorm可学习的参数，参数量为hidden_size
self.weight = nn.Parameter(torch.ones(hidden_size))
#添加一个非零值，防止除零
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
#x=(x^2)/n
variance = hidden_states.pow(2).mean(-1, keepdim=True)
#归一化并缩放
hidden_states = (hidden_states * torch.rsqrt(variance + self.variance_epsilon)) * self.weight
return hidden_states.to(input_dtype)
def extra_repr(self):
@@ -74,43 +69,57 @@ class DollyRotaryEmbedding(nn.Module):
super().__init__()
#获取"rope_type"的值，如果没有使用"default"
self.rope_type = config.rope_scaling.get("rope_type", "default")
self.original_max_seq_len = config.max_position_embeddings
self.config = config
if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
else:
self.rope_type = "default"
self.original_max_seq_len = config.max_position_embeddings
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
#得到频率倒数
#得到频率倒数,基础版本公式如下：
#inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2) / dim))
#得到一个长度为attention_dim/2的小数序列，如attention_dim=128,inv_freq=64
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
#注册缓存区
self.register_buffer("inv_freq", inv_freq, persistent=False)
@torch.no_grad()
@dynamic_rope_update
def forward(self, x, position_ids):
#将inv_freq进行维度扩展首先扩展为1*64*1，然后扩展为batch_size*64*1
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
#position_ids进行扩展batch_size* 1*seq_len
position_ids_expanded = position_ids[:, None, :].float()
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
with torch.autocast(device_type=device_type, enabled=False): # Force float32
#强制float32，禁止混合精度
with torch.autocast(device_type=device_type, enabled=False):
#矩阵乘以后，形状变为batch_size*64*seq_len后，再转为batch_size*seq_len*64
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
#沿64的维度进行拼接batch_size*seq_len*128
emb = torch.cat((freqs, freqs), dim=-1)
#emb的cos()部分乘缩放系数:batch_size*seq_len*128
cos = emb.cos() * self.attention_scaling
#emb的sin()部分乘缩放系数:batch_size*seq_len*128
sin = emb.sin() * self.attention_scaling
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
def rotate_half(x):
#x分为两段，各取batch_size*heads_num*seq_len*64
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
#下部分*-1,按下+上重新组合x
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
'''
对q和k使用RoPE
'''
#由batch_size*seq_len*128增加一个维度变为batch_size*1*seq_len*128
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)
#q:batch_size*heads_num*seq_len*128和sin及con进行元素乘batch_size*1*seq_len*128
#cos部分直接乘，sin部分进行位置变化后再乘，最后相加完成位置编码的嵌入
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
test/temp_test.py
+14 -40
View File
@@ -1,42 +1,16 @@
from src_llm.configuration_dolly import DollyConfig
from transformers.modeling_rope_utils import dynamic_rope_update, ROPE_INIT_FUNCTIONS
import torch
import math
def llama_rope(x, seq_len, dim, device='cpu'):
"""
实现 Llama 中的 Rotary Positional Encoding (RoPE)
x: 输入的嵌入，大小为 (batch_size, seq_len, dim)
seq_len: 序列长度
dim: 嵌入的维度
device: 设备（默认为 'cpu'）
"""
# 生成位置索引
position = torch.arange(0, seq_len, dtype=torch.float32, device=device)
# 计算每个维度的旋转角度
freqs = torch.exp(torch.arange(0, dim, 2, dtype=torch.float32, device=device) * -(math.log(10000.0) / dim))
# 计算位置编码矩阵
angle_rates = position.unsqueeze(1) * freqs.unsqueeze(0)
# 使用 sin 和 cos 对每一对相邻维度进行旋转
sin = torch.sin(angle_rates)
cos = torch.cos(angle_rates)
# 将 sin 和 cos 分别填充到奇数和偶数维度
encoding = torch.zeros((seq_len, dim), device=device)
encoding[:, 0::2] = sin
encoding[:, 1::2] = cos
# 对输入的嵌入应用旋转
x_rot = torch.matmul(x, encoding.unsqueeze(0).transpose(0, 1))
return x_rot
batch_size = 1
seq_len = 2 # 序列长度
dim = 3 # 嵌入维度
x = torch.randn(batch_size, seq_len, dim)
encoded_x = llama_rope(x, seq_len, dim)
print(encoded_x.shape) # 输出：torch.Size([2, 10, 16])
device="cuda"
inv_freq = 1.0 / (10000 ** (torch.arange(0, 128, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / 128))
batch_size = 2
seq_len = 4
position_ids = torch.arange(seq_len).repeat(batch_size, 1).to(device=device)
inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(device)
position_ids_expanded = position_ids[:, None, :].float()
with torch.autocast(device_type="cuda"):
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
emb = torch.cat((freqs, freqs), dim=-1)
cos = emb.cos()
sin = emb.sin()