A bug with n_fused

[JiayiFeng]

When a attn_qkv Layer is set with n_fused>1 and reversed=False, the shape of its sliced weight is incorrect.

Seems that the root cause is here:

parallelformers/parallelformers/parallel/slicing.py

Lines 79 to 95 in 436573b

	dim = dim if not reversed or is_bias else abs(dim - 1)
	n_fused = 1 if not n_fused else n_fused
	proj_layer = proj_layer.chunk(
	n_fused * self.world_size,
	dim=dim,
	)
	if n_fused > 1:
	ranks = (len(proj_layer) + self.world_size - 1) // self.world_size
	proj_layer = [
	proj_layer[i * self.world_size : (i + 1) * self.world_size]
	for i in range(ranks)
	]
	proj_layer = list(
	map(lambda x: torch.cat([*x], dim=-1), zip(*proj_layer))
	)

For a attn_qkv weight, the arg dim is 0. So when the reversed=False and n_fused>1, the tensor is chunked on the dim 0 and then concatenated on the dim 1. Which make its shape incorrect.

[hyunwoongko]

which model did you use?

[JiayiFeng]

I used a modified GPT-NeoX model, which is not officially supported. So I written a custom policy, and find this issue.

[JiayiFeng]

Maybe the

proj_layer = list( 
         map(lambda x: torch.cat([*x], dim=-1), zip(*proj_layer)) 
     ) 

should be:

proj_layer = list( 
         map(lambda x: torch.cat([*x], dim=dim), zip(*proj_layer)) 
     ) 

I guess.

[hyunwoongko]

okay so could you test it with other models?