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

fix RWKV 4bit

[`RWKV`] Fix RWKV 4bit (#23910)
fix RWKV 4bit
c63bfc30 · Younes Belkada · GitHub · 55451c66 · c63bfc30
Unverified Commit c63bfc30 authored May 31, 2023 by Younes Belkada Committed by GitHub May 31, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 0 deletions

src/transformers/models/rwkv/modeling_rwkv.py src/transformers/models/rwkv/modeling_rwkv.py +7 -0

No files found.
--- a/src/transformers/models/rwkv/modeling_rwkv.py
+++ b/src/transformers/models/rwkv/modeling_rwkv.py
@@ -710,6 +710,13 @@ class RwkvModel(RwkvPreTrainedModel):
                        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))