Skip to content

Commit

Permalink
Use torch.inference_mode() for lower memory usage during calibration (
Browse files Browse the repository at this point in the history
  • Loading branch information
mgoin committed Jul 18, 2024
1 parent 93c0d54 commit 57c31bb
Showing 1 changed file with 9 additions and 0 deletions.
9 changes: 9 additions & 0 deletions auto_fp8/quantize.py
Original file line number Diff line number Diff line change
Expand Up @@ -272,13 +272,22 @@ def quantize_activations(
cleanup_memory()

# Pass through calibration data to measure activation scales
<<<<<<< HEAD
with tqdm.tqdm(
total=calibration_tokens.shape[0], desc="Calibrating activation scales"
) as pbar:
for row_idx in range(calibration_tokens.shape[0]):
model(calibration_tokens[row_idx].reshape(1, -1))
cleanup_memory()
pbar.update(1)
=======
with torch.inference_mode():
with tqdm.tqdm(total=calibration_tokens.shape[0], desc="Calibrating activation scales") as pbar:
for row_idx in range(calibration_tokens.shape[0]):
model(calibration_tokens[row_idx].reshape(1, -1))
cleanup_memory()
pbar.update(1)
>>>>>>> b1c6ad6 (Use `torch.inference_mode()` for lower memory usage during calibration (#20))

# Replace dynamic quantizer observer with StaticLinear for export
for name, quantizer in model.named_modules():
Expand Down

0 comments on commit 57c31bb

Please sign in to comment.