The token embedding matrix de-quantization process is responsible for high peak memory usage. Instead of loading and de-quantizing the entire matrix during model init, we can de-quantize only the necessary tokens during model inference. This reduces peak memory usage by a huge margin, albeit with a small increase in latency. For our intended use-case (semantic history search) this balance of memory vs latency is optimal.
Vandana Malayil (a.k.a Vandana Rajan)
997556e555