Best practices to reduce memory usage in SERVE mode

[mar-muel]

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?