Build error in codellama model's word embedding layer (32016 not divisible by 64)

[CaesarWWK]

When I use build.py to build a code llama model, I got an error when processing word-embedding layer.

The cause is that the word-embedding size of code llama is 32016 which can not be divisible by 64. The script padded its newly initialized word-embedding matrix to be divisible by 64, but it did not change the original word-embedding matrix. Therefore the script tried to fit a matrix of shape 32064,hidden_size into 32016,hidden_size.

I believe there exists a bug when building lm_head for code llama in build.py as well.

I fixed the bug by expand the code llama's word-embedding to 32064 and save it in advance, then process it in build.py.
Everything works perfectly after that.

[byshiue]

@CaesarWWK, thank you for the report. Can you create a PR to introduce how do you fix this issue? It is welcome to contribute to TensorRT-LLM.

[CaesarWWK]

@CaesarWWK, thank you for the report. Can you create a PR to introduce how do you fix this issue? It is welcome to contribute to TensorRT-LLM.

Hi,
Sorry for mistaken, I actually took a "workaround" but not "fix" the bug. I tried to fix the bug but it's more complicated than I thought.

For the workaround, I load the code llama in advance then call model.resize_token_embeddings(32064) to resize the word-embedding related layers. Then I save the model weights and use the new weights to build engines. Everything works well.

I will try to fix the bug again, any update will be reported to this thread.
Thank you

[shatealaboxiaowang]

@CaesarWWK, thank you for the report. Can you create a PR to introduce how do you fix this issue? It is welcome to contribute to TensorRT-LLM.

Hi, Sorry for mistaken, I actually took a "workaround" but not "fix" the bug. I tried to fix the bug but it's more complicated than I thought.

For the workaround, I load the code llama in advance then call model.resize_token_embeddings(32064) to resize the word-embedding related layers. Then I save the model weights and use the new weights to build engines. Everything works well.

I will try to fix the bug again, any update will be reported to this thread. Thank you

Thank you, Can you tell your specific steps?
Do you take a torch load model and resize it, then save the weights, then quantize it with AWQ, and then build it as tensorrt?

[Tracin]

@CaesarWWK @shatealaboxiaowang
Hi, You can try below:

1. Use the latest main branch.
2. Just comment out the assert https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/parameter.py#L80-L82
3. Change this line https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/llama/weight.py#L1163 to if v.shape[0] % 64 != 0:
4. Change this line https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/llama/weight.py#L1173 to torch.nn.functional.pad(v[1], [0, 0, 0, 64 * (v[0].shape[1] // group_size) - v[1].shape[0] % (64 * (v[0].shape[1] // group_size))],

I pushed the MR and may be merge in few days, you can try using instruction above, it works from my side.

[CaesarWWK]

@CaesarWWK, thank you for the report. Can you create a PR to introduce how do you fix this issue? It is welcome to contribute to TensorRT-LLM.

Hi, Sorry for mistaken, I actually took a "workaround" but not "fix" the bug. I tried to fix the bug but it's more complicated than I thought.
For the workaround, I load the code llama in advance then call model.resize_token_embeddings(32064) to resize the word-embedding related layers. Then I save the model weights and use the new weights to build engines. Everything works well.
I will try to fix the bug again, any update will be reported to this thread. Thank you

Thank you, Can you tell your specific steps? Do you take a torch load model and resize it, then save the weights, then quantize it with AWQ, and then build it as tensorrt?

Yes, exactly as you said.

[gesanqiu]

@CaesarWWK @shatealaboxiaowang Hi, You can try below:

1. Use the latest main branch.
2. Just comment out the assert https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/parameter.py#L80-L82
3. Change this line https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/llama/weight.py#L1163 to if v.shape[0] % 64 != 0:
4. Change this line https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/llama/weight.py#L1173 to torch.nn.functional.pad(v[1], [0, 0, 0, 64 * 32 - v[1].shape[0] % (64 * 32)],

I pushed the MR and may be merge in few days, you can try using instruction above, it works from my side.

@Tracin Tried this on main branch with codellama-13b-instruct, it doesn't work. On main branch, the original error message is RuntimeError: Padding length too large. With your modify, the error message just the same like release branch AssertionError: The value updated is not the same shape as the original. Updated: (32064, 5120), original: (32016, 5120), which means I need to modify the vocab_size of model.

[Tracin]

@gesanqiu Hi, The file parameter.py in step two is probably installed in the environment, We actually comment out the assert, so you will not get this AssertionError anymore.

[gesanqiu]

@gesanqiu Hi, The file parameter.py in step two is probably installed in the environment, We actually comment out the assert, so you will not get this AssertionError anymore.

Sorry for misinformation, I know that I need to comment the assertion lines, what I mean is that I didn't meet vocab_size issue on main branch, so I don't see the necessity of your modify. And the issue on main branch is RuntimeError: Padding length too large, assert on https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/llama/weight.py#L1173.