代码语言：javascript
复制# 完整的 GLM 模型，包括嵌入层、编码器、输出层
class ChatGLMModel(ChatGLMPreTrainedModel):
def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
super().__init__(config)
# 如果设置了`empty_init`，创建任何 PyTorch 模块时，不初始化参数
if empty_init:
init_method = skip_init
else:
init_method = default_init
init_kwargs = {}
if device is not None:
init_kwargs["device"] = device
# 单词嵌入层
self.embedding = init_method(Embedding, config, **init_kwargs)
# LC
self.num_layers = config.num_layers
# GC
self.multi_query_group_num = config.multi_query_group_num
# HS
self.kv_channels = config.kv_channels
# SL
self.seq_length = config.seq_length
rotary_dim = (
config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
)
# 位置嵌入（PE）
self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
dtype=config.torch_dtype)
# GLM 编码器
self.encoder = init_method(GLMTransformer, config, **init_kwargs)
# 输出层
self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
dtype=config.torch_dtype, **init_kwargs)
self.pre_seq_len = config.pre_seq_len
self.prefix_projection = config.prefix_projection
if self.pre_seq_len is not None:
# 如果设置了前缀序列长度（PSL）
# 关闭所有参数的自动梯度
for param in self.parameters():
param.requires_grad = False
# [0, 1, ..., PSL - 1]
self.prefix_tokens = torch.arange(self.pre_seq_len).long()
# 初始化前缀编码层和 Dropout
self.prefix_encoder = PrefixEncoder(config)
self.dropout = torch.nn.Dropout(0.1)
def get_input_embeddings(self):
return self.embedding.word_embeddings
def get_prompt(self, batch_size, device, dtype=torch.half):
# prefix_tokens = [0, 1, ..., PSL - 1]
# [PSL] => [1, PSL] => [BS, PSL]
prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
# [BS, PSL, KVS=NL * HS * 2GC]
past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
# [BS, PSL, KVS=NL * HS * 2GC] => [BS, PSL, 2NL, GC, HS]
past_key_values = past_key_values.view(
batch_size,
self.pre_seq_len,
self.num_layers * 2,
self.multi_query_group_num,
self.kv_channels
)
past_key_values = self.dropout(past_key_values)
# [BS, PSL, 2NL, GC, HS] => [2NL, PSL, BS, GC, HS] => NL * [2, PSL, BS, GC, HS]
past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
return past_key_values
def forward(
self,
input_ids,
position_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.BoolTensor] = None,
full_attention_mask: Optional[torch.BoolTensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
inputs_embeds: Optional[torch.Tensor] = None,
use_cache: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
):
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
use_cache = use_cache if use_cache is not None else self.config.use_cache
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
# 输入是单词 ID，的形状为 [BS, SL]
batch_size, seq_length = input_ids.shape
# 将单词 ID 传递给词嵌入层得到嵌入向量
if inputs_embeds is None:
inputs_embeds = self.embedding(input_ids)
# 如果设置了 PSL
if self.pre_seq_len is not None:
# 如果没有提供 KV 缓存，初始化为前 PSL 个前缀的词嵌入
if past_key_values is None:
past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
dtype=inputs_embeds.dtype)
if attention_mask is not None:
attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
attention_mask], dim=-1)
if full_attention_mask is None:
if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
# 计算 PE
# 初始化位置编码层
rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
# 如果提供了位置 ID 就是用它检索位置嵌入矩阵
# 如果没有，就返回嵌入矩阵的前 SL 个向量
if position_ids is not None:
rotary_pos_emb = rotary_pos_emb[position_ids]
else:
rotary_pos_emb = rotary_pos_emb[None, :seq_length]
# [BS, SL, ES] => [SL, BS, ES]
rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
# 将词嵌入和位置嵌入传给编码器得到编码器输出
hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
)
# 返回 GLM 输出，每层的 KV 缓存和每层的输出
if not return_dict:
return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
return BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=presents,
hidden_states=all_hidden_states,
attentions=all_self_attentions,
)
def quantize(self, weight_bit_width: int):
from .quantization import quantize
quantize(self.encoder, weight_bit_width)
return self
