train_gpt_demo.py例子中的一些疑问

[yhcc]

hello，能麻烦帮我解答下关于在ZERO中使用tensor_parallelize的疑问吗？
（1）在

ColossalAI/examples/language/gpt/train_gpt_demo.py

Lines 169 to 174 in 65f56f4

	elif 'wte' in mn or 'wpe' in mn:
	split_param_col_tp1d(param, pg) # colmn slice
	elif 'c_attn' in mn or 'c_proj' in mn:
	split_param_col_tp1d(param, pg) # colmn slice
	else:
	param.set_dist_spec(ReplicaSpec())

这个例子中，似乎只需要对column进行切分，但这样其实是会导致在每个卡上计算出来的tensor是在最后一个维度为1/n_proc原始大小，这里是CAI会做特定的处理让它参与运算后返回的activation是已经gather过，所以大小也是expected的大小了吗？
（2）对应的

ColossalAI/examples/language/gpt/train_gpt_demo.py

Lines 157 to 168 in 65f56f4

	if 'mlp.c_fc' in mn:
	if 'weight' in pn or 'bias' in pn:
	split_param_col_tp1d(param, pg) # colmn slice
	# keep the shape of the output from c_fc
	param.compute_spec.set_output_replicate(False)
	else:
	param.set_dist_spec(ReplicaSpec())
	elif 'mlp.c_proj' in mn:
	if 'weight' in pn:
	split_param_row_tp1d(param, pg) # row slice
	else:
	param.set_dist_spec(ReplicaSpec())

这个161行意思是就是指因为后续还有一个根据row切分的，所以暂时不需要gather回未切分前tensor的大小么？

[ver217]

目前为了兼容性是都会自动gather output的

[yhcc]

感谢回答。