work in progress about fusion layer #2
Conversation
|
Thanks for the commit. The framework is correct and should work after some modifications. |
| qkv_layer, | ||
| o_proj, | ||
| dev=dev, | ||
| max_seq_len=max_seq_len, |
There was a problem hiding this comment.
Its RoPE implementation is different from Llama's. Require some modifications in attn.py to get everything working properly.
There was a problem hiding this comment.
ok others edit is done.
I will edit attn.py!
There was a problem hiding this comment.
I think this is too difficult task for me.
First, I checked the difference in the implementation of rotary posembedding between cohere and llama.
llama is just torch.cat after frequs operation, and cohere is torch.repeat_interleave.
Also, when llama proceeds with rotary_half
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
to proceed,
cohere is
x1 = x[...,::2]
x2 = x[..., 1::2]
There is a difference in that.
However, I'm not sure how this should work with RoPE in awq.
Can I get any ideas or hints?
There was a problem hiding this comment.
Aha, I think using a staggered rope should be fine. Now I have device to modify it.
awq/models/cohere.py
Outdated
| module.self_attn.k_proj, | ||
| module.self_attn.v_proj, | ||
| ) | ||
| norm_1 = FasterTransformerRMSNorm( |
There was a problem hiding this comment.
It use normal layernorm
norm_1 = module.input_layernorm
awq/modules/fused/block.py
Outdated
| ) | ||
|
|
||
| h = hidden_states.to(attn_output.device) + attn_output | ||
| out = h + self.mlp.forward(h) |
There was a problem hiding this comment.
out = h + self.mlp.forward(norm_out)
|
Thanks! |
I attempted to implement the fusion layer. Although it seems to have been implemented plausibly, the result only generates meaningless tokens, and the logits values from model(**inputs) contain only NaN values.
As far as I know, memory usage should increase a little after fusing, but I haven't observed any increase in memory usage either.
It seems like something is wrong, but it might be a minor issue like when I implemented quantize before. Although it's still incomplete, I thought it might be helpful, so I'm submitting a PR.
I'm not sure if this process is possible, but I have enough hardware to conduct tests. If you modify the code, I can test it.
I would also like to contribute to this AutoAWQ project.