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

update modeling_dolly code

Browse Source
This commit is contained in:
coffee
2025-05-13 17:33:33 +08:00
parent c537905e9b
commit 575dcce05b
1 changed files with 12 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src_llm/modeling_dolly.py
+12 -13
View File
@@ -72,27 +72,25 @@ class DollyMLP(nn.Module):
class DollyRotaryEmbedding(nn.Module): class DollyRotaryEmbedding(nn.Module):
def __init__(self, config: DollyConfig, device=None): def __init__(self, config: DollyConfig, device=None):
super().__init__() super().__init__()
# BC: "rope_type" was originally "type"
if hasattr(config, "rope_scaling") and config.rope_scaling is not None: #获取"rope_type"的值，如果没有使用"default"
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) self.rope_type = config.rope_scaling.get("rope_type", "default")
else:
self.rope_type = "default"
self.max_seq_len_cached = config.max_position_embeddings
self.original_max_seq_len = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings
self.config = config self.config = config
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
#得到频率倒数
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
#注册缓存区
self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("inv_freq", inv_freq, persistent=False)
self.original_inv_freq = self.inv_freq
@torch.no_grad() @torch.no_grad()
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) @dynamic_rope_update
def forward(self, x, position_ids): def forward(self, x, position_ids):
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
position_ids_expanded = position_ids[:, None, :].float() 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" 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 with torch.autocast(device_type=device_type, enabled=False): # Force float32
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
@@ -103,13 +101,14 @@ class DollyRotaryEmbedding(nn.Module):
def rotate_half(x): def rotate_half(x):
"""Rotates half the hidden dims of the input."""
x1 = x[..., : x.shape[-1] // 2] x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :] x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1) return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
'''
对q和k使用RoPE
'''
cos = cos.unsqueeze(unsqueeze_dim) cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim)
q_embed = (q * cos) + (rotate_half(q) * sin) q_embed = (q * cos) + (rotate_half(q) * sin)