google / aqt

Apache License 2.0
248 stars 25 forks source link

Best practices to reduce memory usage in SERVE mode #706

Open mar-muel opened 1 week ago

mar-muel commented 1 week ago

Hi there - Thanks for the great work, this library is disparately needed in the jax ecosystem.

Am I correct in my understanding that when switching to SERVE mode, the original weights are bypassed in the forward pass and instead the quantized weights are being used?

If so, loading the original weights would be unnecessary and just adds to memory usage. What would be the easiest way to avoid having to use the original weights, while re-using the existing implementation that was used for training?