riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 18:22:34 +06:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

[RWKV] Rwkv fix for 8bit inference (#23468)

Browse Source 
* rwkv fix for 8bit inference

* add comment
This commit is contained in:
Younes Belkada 2023-05-19 16:12:25 +02:00 committed by GitHub
parent 1c460a5273
commit 21bd3be172
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/transformers/models/rwkv/modeling_rwkv.py
View File

@ -708,6 +708,11 @@ class RwkvModel(RwkvPreTrainedModel):
if self.training:
block.attention.output.weight.mul_(2 ** int(block_id // self.config.rescale_every))
block.feed_forward.value.weight.mul_(2 ** int(block_id // self.config.rescale_every))
else:
# Deal with quantization statistics
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))
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))