Skip to content

perf: avoid repeated INT8 cast-warning records - #2127

Open
yfukuda-oqc wants to merge 2 commits into
bitsandbytes-foundation:mainfrom
yfukuda-oqc:perf/int8-cast-warning-once
Open

yfukuda-oqc wants to merge 2 commits into
bitsandbytes-foundation:mainfrom
yfukuda-oqc:perf/int8-cast-warning-once

Conversation

@yfukuda-oqc

Copy link
Copy Markdown

Problem

MatMul8bitLt.forward logs a warning on every non-FP16 eager call. Repeated FP32 or BF16 inputs therefore construct repeated LogRecord objects, even when a downstream handler discards duplicate messages.

Change

Emit the cast warning once per input dtype per process. Keep the existing warning text and level, and preserve the FP16 intermediate cast and all quantization and autograd operations.

A bitsandbytes-owned set tracks the warned dtypes. A lock serializes first emission for concurrent callers; subsequent calls return before acquiring the lock or constructing a record. Calls with WARNING disabled do not consume the first warning. The existing compilation guard remains in place.

This introduces no public API. Related: #2031 addresses a different 4-bit warning; #1985 proposes native BF16 quantization, which changes the numerical path.

Validation

Tested on Windows 11, Python 3.12.16, Torch 2.8.0+cu128, CUDA 12.8, and an NVIDIA GeForce RTX 4060 Laptop GPU (8 GiB, driver 617.42).

python -m pytest tests/test_int8_cast_warning.py -q -o log_cli=false -o addopts='' -p no:cacheprovider
python -m pytest tests/test_autograd.py tests/test_linear8bitlt.py -k '(test_matmullt or test_linear) and not torch_compile' -q -o log_cli=false -o addopts='' -p no:cacheprovider
python -m pre_commit run --files bitsandbytes/autograd/_functions.py tests/test_int8_cast_warning.py
  • Focused tests: 19 passed, 4 skipped. The skipped CPU weight-gradient cases require int8_double_quant, which this backend does not implement.
  • Existing INT8 tests: 840 passed, 800 deselected.
  • All applicable pre-commit hooks passed.
  • One hundred FP32 calls plus one hundred BF16 calls construct 2 warning records, compared with 200 on baseline. FP16 does not warn. The new count and concurrent-caller tests fail on the baseline and pass with this change.
  • Sixteen simultaneous first callers produce one record. Tests also cover disabled WARNING level, unrelated messages, empty inputs and the compilation guard.
  • CPU/CUDA tests compare FP32/BF16 outputs and supported input/weight/bias gradients with the original logging behavior, including thresholds 0 and 6 and repeated calls, using rtol=0, atol=0.
  • torch.compile(backend="eager", fullgraph=True) with frozen, prequantized weights preserves repeated forward results and input gradients on CPU/CUDA without emitting cast warnings.

Baseline Python revision: a616886cacf86701514f0cb1d6be760f967481c0. An additional local comparison loaded that baseline source into a separate module and compared 12 CPU/CUDA cases over three momentum-SGD updates each. All 288 output, gradient, parameter and momentum-buffer comparisons matched exactly. This supplemental comparison is separate from the committed regression suite.

Both baseline and candidate used libbitsandbytes_cuda128.dll from the official bitsandbytes 0.50.2 Windows wheel (SHA256 391825688f575c1eeed41c84441ef0b94f5b05c4089231b7125f27ecf0a44ce7). The Python package was imported from the source checkout; no native code changed or new native build was used.

The full platform suite, inductor/CUDA graphs, multiple GPUs and bitsandbytes-specific optimizer comparisons were not exercised. No end-to-end training speedup is claimed.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant