Llama 2 70B: Update needed to convert.py to support 70B HF format model files

[TheBloke]

Following from discussions in the Llama 2 70B PR: #2276 :

Since that PR, converting Llama 2 70B models from Meta's original PTH format files works great.

But it is not possible to make usable Llama 2 70B models from HF format. The models convert and quantise fine, but always produce gibberish, as in this example:

 ### Human: write a story about llamas\n### Assistant:20 300202000 B00A0

@klosax reports:

It looks like the tensors gets transformed with the new permute using the GQA parameter num_local_key_value_heads and num_key_value_heads somehow:
https://github.com/huggingface/transformers/blob/b257c46a075419c09e5ce5c5aa39bc346ecdb9a5/src/transformers/models/llama/convert_llama_weights_to_hf.py#L173-L195

For reference, here are all the changes that happened in Transformers' convert_llama_weights_to_hf.py for the Llama 2 release: huggingface/transformers@07360b6#diff-110a445233a8b15a0875998eeaf75cb8607b38a5daa736291dd058766879bbdd

Would anyone be able to look into this? It's a bit beyond my experience.

I'm getting multiple requests a day for 70B fine tune quants for FreeWilly 2, Llama2-Guanaco, and the newly released Airoboros 1.4.1 70B, and would love to be able to provide them for people.

Thanks in advance.

[klosax]

It would indeed be very nice to be able to convert the 70b HF models.

I tried to look into it and figure out how it all worked but the needed skills are beyond me.

I think the permute() function in the transformer conversion script is getting reversed by the permute() in llama.cpp convert.py, but this function is not yet compatible with llama v2 which uses GQA.

[KerfuffleV2]

I think the permute() function in the transformer conversion script is getting reversed by the permute() in llama.cpp convert.py

Have you tried just temporarily disabling permute() in the llama.cpp convert.py?

Basically just

def permute(weights: NDArray, n_head: int) -> NDArray:
    return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])
                   .swapaxes(1, 2)
                   .reshape(weights.shape))

to

def permute(weights: NDArray, n_head: int) -> NDArray:
    return weights

It may not work (or some kind of reshape might still be needed) but this should be a pretty easy one to at least try. I'd test it myself but I don't have the 16bit 70B at hand.

[klosax]

Have you tried just temporarily disabling permute() in the llama.cpp convert.py?

No I dont think that that would work since all pth models, including the 70b, is transformed by the HF permute(). I guess we need a new permute() in convert.py to reverse it.

[7erminalVelociraptor]

I think the permute() function in the transformer conversion script is getting reversed by the permute() in llama.cpp convert.py

Have you tried just temporarily disabling permute() in the llama.cpp convert.py?

Basically just

def permute(weights: NDArray, n_head: int) -> NDArray:
    return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])
                   .swapaxes(1, 2)
                   .reshape(weights.shape))

to

def permute(weights: NDArray, n_head: int) -> NDArray:
    return weights

It may not work (or some kind of reshape might still be needed) but this should be a pretty easy one to at least try. I'd test it myself but I don't have the 16bit 70B at hand.

I can try it out later today, I have the airoboros 70b finetune in HF on my desktop and I think the base llama2 also.

[MrJackSpade]

What's the level of effort on porting the permutation changes?

Does it require a lot of knowledge, or is it a straight transposition?

[klosax]

What's the level of effort on porting the permutation changes?

Here is the new HF conversion script that converts the original pth models to HF format:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py

This is the old HF conversion script:

https://github.com/huggingface/transformers/blob/feb83521eca849731573dd40da89a02e4f370e5a/src/transformers/models/llama/convert_llama_weights_to_hf.py

Now, llama.cpp want the tensors in the pth layout so any transformations made to HF tensors need to be reversed in convert.py

[mj-shifu]

Hello,

I think I've managed to alter the conversion script so that the converted model does not produce gibberish any more:

https://github.com/mj-shifu/llama.cpp/blob/e15a67d6b21c10326a5cc74ab6d6ce9b8d7702bb/convert.py

I am not sure if it is correct but the converted Huggingface model produces exactly the same outputs as the converted pth model when sampling is disabled.

@TheBloke Could you probably try that out?

I'm sorry for any mistakes, this is my first time contributing here.

[TheBloke]

Wonderful, thank you! I am having dinner now but will check as soon as I am at my pc

[KerfuffleV2]

I think I've managed to alter the conversion script so that the converted model does not produce gibberish any more

Nice, I was working on this as well and what you did is very close to my approach. So theoretically if I'm not an idiot that is a good sign.

I'd suggest making the type for n_kv_head Optional[int] since it gets set to None. Also, then you don't even need a conditional to populate it, you can just do n_kv_head = config.get("num_key_value_heads")

[mj-shifu]

Thanks for the suggestions. I just updated the code accordingly. Although I also think that the similarities are a good sign, I'm sorry that I caused duplicated work.

https://github.com/mj-shifu/llama.cpp/blob/01d16e1a1efced0cfbe92ed0c94c8003d22dbe54/convert.py

[KerfuffleV2]

Absolutely no need to apologize. I'm happy to see someone else got it done! Assuming it works, you should make a pull with these changes! I can't test whether it actually works, but for whatever my opinion is worth the code looks very reasonable.

If you wanted to reduce duplication a bit you could try something like:

def permute(weights: NDArray, n_head: int, n_kv_head: Optional[int] = None) -> NDArray:
    dim1 = weights.shape[0]
    if n_kv_head is not None and n_kv_head != n_head:
        dim1 *= n_kv_head
        n_head //= n_kv_head
    return (weights.reshape(n_head, 2, dim1 // n_head // 2, *weights.shape[1:])
                .swapaxes(1, 2)
                .reshape(weights.shape))

[mj-shifu]

Thank you! I should probably wait for @TheBloke's result until I make a pull request, shouldn't I? I have only tested it with a Huggingface model that I converted locally from the 70B pth model using the official Huggingface script.

[klosax]

If you wanted to reduce duplication a bit you could try something like:

if n_kv_head is used, n_head should not be divided by n_kv_head in the third parameter.

Something like this looks correct to me:

def permute(weights: NDArray, n_head: int, n_kv_head: Optional[int] = None) -> NDArray:
    dim1 = weights.shape[0] // n_head // 2

    if n_kv_head is not None and n_kv_head != n_head:
        dim1 *= k_kv_head
        n_head //= n_kv_head

    return (weights.reshape(n_head, 2, dim1, *weights.shape[1:])
                    .swapaxes(1, 2)
                    .reshape(weights.shape))