picocreator ybelkada commited on
Commit
bb01ae9
1 Parent(s): 419bcf6

Update modeling_rwkv5.py (#5)

Browse files

- Update modeling_rwkv5.py (5d64a05fb9748b26b55162cb56162c9da83135c4)


Co-authored-by: Younes Belkada <[email protected]>

Files changed (1) hide show
  1. modeling_rwkv5.py +10 -2
modeling_rwkv5.py CHANGED
@@ -747,8 +747,16 @@ class Rwkv5Model(Rwkv5PreTrainedModel):
747
  block.attention.output.weight.mul_(2 ** int(block_id // self.config.rescale_every))
748
  block.feed_forward.value.weight.mul_(2 ** int(block_id // self.config.rescale_every))
749
  else:
750
- block.attention.output.weight.div_(2 ** int(block_id // self.config.rescale_every))
751
- block.feed_forward.value.weight.div_(2 ** int(block_id // self.config.rescale_every))
 
 
 
 
 
 
 
 
752
 
753
  self.layers_are_rescaled = not self.training
754
 
 
747
  block.attention.output.weight.mul_(2 ** int(block_id // self.config.rescale_every))
748
  block.feed_forward.value.weight.mul_(2 ** int(block_id // self.config.rescale_every))
749
  else:
750
+ # Deal with quantization statistics
751
+ if hasattr(block.attention.output.weight, "SCB"):
752
+ block.attention.output.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
753
+ block.feed_forward.value.weight.SCB.div_(2 ** int(block_id // self.config.rescale_every))
754
+ elif hasattr(block.attention.output.weight, "quant_state"):
755
+ self._bnb_4bit_dequantize_and_rescale(block.attention.output, block_id)
756
+ self._bnb_4bit_dequantize_and_rescale(block.feed_forward.value, block_id)
757
+ else:
758
+ block.attention.output.weight.div_(2 ** int(block_id // self.config.rescale_every))
759
+ block.feed_forward.value.weight.div_(2 ** int(block_id // self.config.rescale_every))
760
 
761
  self.layers_are_rescaled = not self.training
762