fix rwkv v5 fp16 (#10474)

2024-03-20 13:15:08 +08:00 · 2024-03-20 13:15:08 +08:00 · 749bedaf1e
commit 749bedaf1e
parent 72bcc27da9
1 changed files with 1 additions and 0 deletions
--- a/python/llm/src/bigdl/llm/transformers/convert.py
+++ b/python/llm/src/bigdl/llm/transformers/convert.py
@ -538,6 +538,7 @@ def _optimize_pre(model):
    # for rwkv models (verified RWKV/rwkv-4-world-7b)
    if model.config.model_type == "rwkv":
        model.rwkv._rescale_layers()
+        model.rwkv.layers_are_rescaled = True
    # process NormHead module in Baichuan2 7B and 13B
    if model.config.model_type == "baichuan" and model.config.vocab_size == 125696:
        # NormHead do normalization on the weights just once at inference time.