[BUG]

[butyuhao]

🐛 Describe the bug

I wrote the following code, I think the dimentions of tensors are correct. What should I do?
class LlamaMLP(nn.Module):
def init(
self,
hidden_size: int,
intermediate_size: int,
hidden_act: str,
):
super().init()
# self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.gate_proj = col_nn.Linear(hidden_size, intermediate_size, dtype=torch.float, bias=False)
# self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
self.down_proj = col_nn.Linear(intermediate_size, hidden_size, dtype=torch.float, bias=False)
# self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.up_proj = col_nn.Linear(hidden_size, intermediate_size, dtype=torch.float, bias=False)
self.act_fn = ACT2FN[hidden_act]

def forward(self, x):
    up = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
    down = self.down_proj(up)
    return down

the output I get is:
Traceback (most recent call last):
File "train.py", line 158, in
layer_outputs = decoder_layer(
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/nas-alinlp/butyuhao/GLM/colossal-ai/llama/modeling_llama.py", line 348, in forward
main()
File "train.py", line 151, in main
hidden_states, self_attn_weights, present_key_value = self.self_attn(
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/nas-alinlp/butyuhao/GLM/colossal-ai/llama/modeling_llama.py", line 218, in forward
engine.execute_schedule(data_iter, return_output_label=False)
File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/_base_engine.py", line 201, in execute_schedule
output, label, loss = self._schedule.forward_backward_step(self, data_iter, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/schedule/_non_pipeline_schedule.py", line 78, in forward_backward_step
output = self._call_engine(engine, data)
File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/schedule/_base_schedule.py", line 109, in _call_engine
return engine(inputs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/engine/_base_engine.py", line 186, in call
return self.model(*args, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/nas-alinlp/butyuhao/GLM/colossal-ai/llama/modeling_llama.py", line 661, in forward
qkv = self.query_key_value(hidden_states)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/colossalai_layer/_utils.py", line 41, in forward
return self.module(*args)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/colossalai_layer/_utils.py", line 41, in forward
return self.module(*args)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in call_impl
return forward_call(*input, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/parallel_1d/layers.py", line 697, in forward
assert input.shape[-1] == self.weight.shape[-1],
AssertionError: Invalid shapes in Linear1D_Row forward: input=torch.Size([1, 128, 4096]), weight=torch.Size([12288, 2048]). Expected last dim of input 2048.
layer_outputs = decoder_layer(
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/nas-alinlp/butyuhao/GLM/colossal-ai/llama/modeling_llama.py", line 348, in forward
hidden_states, self_attn_weights, present_key_value = self.self_attn(
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/nas-alinlp/butyuhao/GLM/colossal-ai/llama/modeling_llama.py", line 218, in forward
qkv = self.query_key_value(hidden_states)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/colossalai_layer/_utils.py", line 41, in forward
return self.module(*args)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
return forward_call(*input, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/colossalai_layer/_utils.py", line 41, in forward
return self.module(*args)
File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in call_impl
return forward_call(*input, **kwargs)
File "/opt/conda/lib/python3.8/site-packages/colossalai/nn/layer/parallel_1d/layers.py", line 697, in forward
assert input.shape[-1] == self.weight.shape[-1],
AssertionError: Invalid shapes in Linear1D_Row forward: input=torch.Size([1, 128, 4096]), weight=torch.Size([12288, 2048]). Expected last dim of input 2048.
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 60198) of binary: /opt/conda/bin/python

Environment

TENSOR_PARALLEL_SIZE = 2
TENSOR_PARALLEL_MODE = '1d'

parallel = dict(
pipeline=1,
tensor=dict(mode=TENSOR_PARALLEL_MODE, size=TENSOR_PARALLEL_SIZE),
)

batch_size=1