update: llama3

[Lusfie]

No description provided.

[xiezipeng-ML]

这个文件不用改动

[xiezipeng-ML]

注释可以删掉

[xiezipeng-ML]

这个实现有问题导致的之后的k，v相关计算对不上，GQA里k，v的num_head数量和q不一样，参考chatglm的方式实现：

libai/projects/ChatGLM/chatglm.py

Lines 245 to 259 in 9dcbe3b

	self.qkv_hidden_size = 3 * self.projection_size
	if self.multi_query_attention:
	self.num_multi_query_groups_per_partition = cfg.multi_query_group_num
	self.qkv_hidden_size = (
	self.projection_size
	+ 2 * self.hidden_size_per_attention_head * cfg.multi_query_group_num
	)
	self.query_key_value = Linear(
	cfg.hidden_size,
	self.qkv_hidden_size,
	bias=cfg.add_bias_linear or cfg.add_qkv_bias,
	parallel="col",
	layer_idx=self.layer_number - 1,
	)

[xiezipeng-ML]

这里的k和v的num_head排列顺序有问题，GQA里是多个q_head对应单个k_head,v_head，
举个例子

# k的h1对应q的h1，h2
q: [h1, h2, h3, h4]
k: [  h1,    h2]
v: [  h1,    h2]

repeat k,v的head后：

# k的h1对应q的h1，h2
q: [h1, h2, h3, h4]
k: [h1, h1, h2, h2]
v: [h1, h1, h2, h2]

但是这里的实现结果k,v head排列顺序有问题，所以之后计算是有问题的

# k的h1对应q的h1，h2
q: [h1, h2, h3, h4]
k: [h1, h2, h1, h2]
v: [h1, h2, h1, h2]

参考chaglm的实现：

libai/projects/ChatGLM/chatglm.py

Lines 332 to 356 in 9dcbe3b

	if self.multi_query_attention:
	key_layer = key_layer.unsqueeze(-2)
	key_layer = key_layer.expand(
	-1,
	-1,
	-1,
	self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition,
	-1,
	)
	key_layer = key_layer.contiguous().view(
	key_layer.size()[:2]
	+ (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
	)
	value_layer = value_layer.unsqueeze(-2)
	value_layer = value_layer.expand(
	-1,
	-1,
	-1,
	self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition,
	-1,
	)
	value_layer = value_layer.contiguous().view(
	value_layer.size()[:2]
	+ (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
	)