Codecov Report

Attention: Patch coverage is 20.72727% with 218 lines in your changes missing coverage. Please review.

Project coverage is 84.25%. Comparing base (c30ffcf) to head (ddc8791).
Report is 4 commits behind head on main.

@@            Coverage Diff             @@
##             main    #2000      +/-   ##
==========================================
- Coverage   84.45%   84.25%   -0.20%     
==========================================
  Files         840      845       +5     
  Lines      104346   105439    +1093     
==========================================
+ Hits        88125    88838     +713     
- Misses      16221    16601     +380     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Turns out the weights were still being automatically dequantized to a TensorPrimitive::Float because Param sets tensor.require_grad() when loading, which currently always returns a float primitive (i.e., dequantizes QFloat tensors).

In my tests with TinyLlama, running inference in f16 is slower when we have to dequantize before every op instead of already having the dequantized weights loaded.

So until we introduce layers with ops supported in quantized types (e.g., int8), we should dequantize the weights. Once I figure out how this should be handled (as cleanly and explicitly as possible), I'll re-open for review.

/edit: I think the best way to go about this is documentation. I feel like adding a dequantize() or float() method on modules is not required, so instead I added subsection to describe how this can be achieved (and why someone might want to do this) with a ModuleMapper. This is the solution I am currently using for loading quantized Llama weights.