[RWKV] Fix RWKV 4bit (huggingface#23910)

fix RWKV 4bit
marma · May 31, 2023 · c63bfc3 · c63bfc3
1 parent 55451c6
commit c63bfc3
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/transformers/models/rwkv/modeling_rwkv.py b/src/transformers/models/rwkv/modeling_rwkv.py
@@ -710,6 +710,13 @@ def _rescale_layers(self):
                         if hasattr(block.attention.output.weight, "SCB"):
                             block.attention.output.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
                             block.feed_forward.value.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
+                        elif hasattr(block.attention.output.weight, "quant_state"):
+                            block.attention.output.weight.quant_state[0].div_(
+                                2 ** int(block_id // self.config.rescale_every)
+                            )
+                            block.feed_forward.value.weight.quant_state[0].div_(
+                                2 ** int(block_id // self.config.rescale_every)
+                            )
                         else:
                             block.attention.output.weight.div_(2 ** int(block_id // self.config.rescale_every))
                             block.feed_forward.value.weight.div_(2 ** int(block_id // self.config.rescale_every))