Implement kv-caching, add more variance of extrapolation (CFG) and interpolation methods #55
Conversation
aria/model/model.py
Outdated
| super().__init__() | ||
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) | ||
| if device is None: # todo: maybe we don't need this... |
There was a problem hiding this comment.
Probably best to not explicitly do this. Could instead perhaps do device = self.device? I'm worried that this will break the distributed code (with accelerate) but I'm unable to test it at the moment.
There was a problem hiding this comment.
The problem with device=self.device is that somehow device=None defaults to CPU and I will get mismatch device error.
There was a problem hiding this comment.
@loubbrad I didn't realize these two lines can be safely removed now with hugging face codes. In fact device param itself is redundant!
In the original implementation of RotaryEmbedding (from neox repo) this device needs to be known at initialization-time because the cached cos_cached, sin_cached need to be explicitly set. But now they are registered as model parameters and moving across device is automatic.
aria/model/model.py
Outdated
| self.rotary_emb = RotaryEmbedding(self.d_head) | ||
| if use_yarn: | ||
| # todo: need more testing on this | ||
| self.rotary_emb = DynamicYaRNScaledRotaryEmbedding(self.d_head, |
There was a problem hiding this comment.
This import is missing I think?
There was a problem hiding this comment.
Oh yes you are right. Will fix it in a moment.
with torch.inference_mode():around decoding)