SQ and QM: Remove torch.cuda.empty_cache, use calibration_forward_context

[kylesayrs]

Purpose

• Fixes SmoothQuant failing on multi-gpus #1081
• Fixes Error when quantizing LLama 3.3 70b to FP8 #963
• There's really no explanation online as to why the torch.cuda.empty_cache() kernel sometimes fails to launch. Given that empty_cache does not actually free memory that wouldn't have already been freed by the python garbage collector + pytorch caching allocator, it should be safe to remove this call.

Changes

• Remove torch.cuda.empty_cache() in run_calibration_forward, which only affects smoothquant and quantization modifier (sparsegpt and wanda will soon use sequential pipelines instead)
• Use calibration_forward_context in smoothquant and quantization modifier
• Remove use of torch.cuda.empty_cache() by smoothquant modiifier

Testing

• Performed memory analysis with and without torch.cuda.empty_cache and calibration_forward_context independently

Smooth Quant

Quantization Modifier

It was also found that removing the empty_cache calls in between each operation reduced the runtime of Quantization Modifier on llama3-8B by 78%

Before

512/512 [03:18<00:00,  2.58it/s]
Duration: 199.38174653053284

After

512/512 [00:42<00:00, 11.91it/s]
Duration: 44.374401807785034

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[rahul-tuli]

Did we run smoothquant with this change? I have certainly come across cases where we run into OOM without this line (even though I know this shouldn't alleviate the issue), I also saw that error go away when CUDA_LAUNCH_BLOCKING env variable was set. I'm good with this change as long as you've verified a smoothquant run! Thanks for investigation

[kylesayrs]

@rahul-tuli That's a good enough reason to wait until some regression tests are finished. We should figure out why OOM occurs and potentially add that to the device map/ fix memory leaks

[dsikka]

This method is used by more than just smoothquant so we would need to do regression testing for modifiers outside of just smoothquant before making this change, if we’re seeing the same effects as what Rahul has described. I would update the PR title/description to reflect this.

smootbquant has its own empty cache call which is not targeted by this PR

[kylesayrs]

Yep, the title is a bit of a misnomer is more in reference to "fixing smoothquant" rather than smoothquant-specific implementation

[kylesayrs]

Here's the results of memory profiling SmoothQuant under various implementations. Constants are the model used (llama3.2-1b-instruct) and the batch_size (16x2048).

graph


Standard w/ calib_context: 10047127552
Smoothquant: 10584129536
SQ w/o empty_cache: 10584129536
SQ w/o empty_cache, w/ calib_context: 10047258624

Standard w/ calib_context

# standard forward pass
with calibration_forward_context(model):
    for batch in tqdm.tqdm(dataloader):
        model(**batch)

[kylesayrs]

Here's the results of memory profiling Quantization Modifier under various implementations. Constants are the model used (llama3.2-1b-instruct) and the batch_size (16x2048).

graph


Standard w/ calib_context: 10047127552
Quantization Modifier: 11243553901
QM w/o empty_cache: 11243553901
QM w/o empty_cache, w/ calib_context: 10047128685

Standard w/ calib_context

# standard forward pass
with calibration_forward_context(model):
    for batch in tqdm.tqdm(dataloader):
        model(**batch)

[kylesayrs]

From this brief analysis, I believe it's safe to conclude that

1. Removing torch.cuda.empty_cache does not lead to increased peak memory usage
2. Adding calibration_forward_context does decrease peak memory usage (specifically due to having disabled the HF KV cache)

[dsikka]

can you resolve conflicts

[kylesayrs]

Not sure why ruff is acting up on that test case, but this PR now fixes it