Perform quantization in Chunks

[drisspg]

Summary

Previously we were seeing a huge memory spike when attempting to NF4'afy the very large tensors. This was do to quantize to nearest creating very large intermediates and realizing them in gmem. This PR instead does this in chunks of default size 1024**2. This was seen to be a pretty reasonable tradeoff between speed while drastically reducing memory usage.

[jerryzh168]

do we want to incorporate this into our quantization primitive in the future? #160

[drisspg]

@jerryzh168 like include the ability to apply quantization in chunks? I am not sure if this problem is manifest in all quantization techniques. I think this may have been a quirk of the broadcasting stuff I was doing. But if it does come again I think it would be good to standardize on. This change is a "free lunch" for qlora since we only quantize once. For dynamic quant I think that the speed up might not be worth it

[jerryzh168]

@jerryzh168 like include the ability to apply quantization in chunks? I am not sure if this problem is manifest in all quantization techniques. I think this may have been a quirk of the broadcasting stuff I was doing. But if it does come again I think it would be good to standardize on. This change is a "free lunch" for qlora since we only quantize once. For dynamic quant I think that the speed up might not be worth it

our general op is quantizing things with certain block_size, so I'm wondering if this is what you are doing here as well, if so I think it might make sense to merge