fix(qlinear): speed up weights transfer

By instantiating the QLinear weights directly on the device, it saves
two copies: from device to cpu and back.

bench/generation/benchmark.py

@@ -173,8 +173,9 @@ def main():
  # Very simple calibration to avoid completely off results
  with Calibration():
  generate(model, tokenizer, device, prompt=CALIBRATION_PROMPT)
+ print("Freezing")
  freeze(model)
- print(f"Finished: {time.time()-start}")
+ print(f"Finished: {time.time()-start:.2f}")
 
  memory = get_device_memory(device)
  if memory is not None:

quanto/nn/qlinear.py

@@ -24,6 +24,7 @@ def from_module(cls, module, weights=torch.int8, activations: Optional[torch.dty
  dtype=module.weight.dtype,
  weights=weights,
  activations=activations,
+ device=module.weight.device,
  )
  with torch.no_grad():
  qmodule.weight.copy_(module.weight)