diff --git a/src_llm/modeling_dolly.py b/src_llm/modeling_dolly.py index db3da29..ceeae1f 100755 --- a/src_llm/modeling_dolly.py +++ b/src_llm/modeling_dolly.py @@ -72,27 +72,25 @@ class DollyMLP(nn.Module): class DollyRotaryEmbedding(nn.Module): def __init__(self, config: DollyConfig, device=None): super().__init__() - # BC: "rope_type" was originally "type" - 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.max_seq_len_cached = config.max_position_embeddings + + #获取"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 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - + #得到频率倒数 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + #注册缓存区 self.register_buffer("inv_freq", inv_freq, persistent=False) - self.original_inv_freq = self.inv_freq + @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): 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() - 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 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) @@ -103,13 +101,14 @@ class DollyRotaryEmbedding(nn.Module): def rotate_half(x): - """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] 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) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin)