enhancements: sharing mine

[thiswillbeyourgithub]

(Thank you very much for this library, it's really nice).

I am aware this is absolutely not merge-ready, I made this PR to show my changes and if you're interested I can make dedicated PRs for it.

If you're not interested and think of this as noise, then I apologize and wish you a nice day.

Context

I had to add a bunch of features to my fork that were needed to improve performance especially on avoiding OOM errors on my parakeet-v3 requantization.

More context on why I had to apply smoothquant to parakeet v3 is listed here on my hugging face repo.

This was entirely done with Claude Code Opus (and a bit of Fable), I then used it with good success and performance.

Expected Behavior & Potential Risk

It grew from frustration from OOM and limitations on alpha parameters initially.

How has this PR been tested?

These changes are regression-tested from the downstream model repo (scripts/test_quantize-int8-smoothquant.py, T1-T30), which pins the bit-identical invariants (cache on == off, sliced dump == single-pass dump, streamed top-K percentile == stacked np.percentile, variable-length auto-alpha).

I'm calling it in run.sh in my fork if you want more reproduction.

Dependency Change?

None

What this fork changes (SmoothQuant static int8 path)

These are the changes on the diy branch versus onnx/neural-compressor, written
up as candidates for upstreaming. They all sit in the SmoothQuant static int8
path and were driven by exporting a real 0.6B FastConformer ASR encoder (long
calibration windows, wide auto-alpha grids), which surfaced several upstream bugs
and scaling walls. Grouped below as correctness fixes, memory/scaling, and new
features. Commit hashes are on this fork's diy branch.

Correctness fixes

• SmoothQuant extra_options are no longer silently ignored.

  • quantize() routes a StaticQuantConfig (with extra_options["SmoothQuant"])
    into smooth_quant_entry, which called smoother.transform(**quant_config.get_model_params_dict()).
    But get_model_params_dict() only surfaces the static-quant knobs, so NONE of
    the smooth knobs (SmoothQuantAlpha, SmoothQuantOpTypes, AutoAlphaArgs, ...)
    reached transform(): the smoother always ran with its hard-coded defaults
    (alpha=0.5, the [0.3, 0.7] auto grid, Conv+Gemm+MatMul+FusedConv) regardless
    of what the caller set.
  • A new _smoothquant_transform_params() maps the documented extra_options names
    to the transform() argument names so they take effect.
  • (in quantization/algorithm_entry.py)

• The selected execution provider reaches the smoother calibration.

  • The smoother passed execution_provider= to the Calibrator, whose constructor
    parameter is providers=, so the argument fell into **kwargs and was dropped.
  • Every smoother-calibration forward therefore ran on CPU even when the caller
    asked for CUDA. The provider list is now forwarded correctly.
  • (commit ad8c01b)

• Auto-alpha search no longer runs on an exhausted dataloader.

  • Upstream's alpha="auto" consumed the calibration reader during setup, so the
    per-node search saw no data and every layer silently fell back to alpha_min.
  • The reader is now rewound and the per-node reference activations are cached, so
    the search actually evaluates the grid (this also fixes the large-model path).
  • (commit d6745e4)

• The alpha grid now includes both endpoints.

  • Upstream built the grid with np.arange, which stops one step short of the
    stop value, so the maximum alpha (alpha_max) was never evaluated.
  • The grid is now inclusive of both alpha_min and alpha_max.
  • (commit 8d88f45)

Memory and scalability

• Streaming Entropy/Percentile static calibration.

  • Static calibration buffered every dumped activation tensor of every sample until
    the end of the loop (an in-code upstream TODO), so RAM grew as samples x tensors
    and large exports OOMed.
  • Each sample is now folded straight into its per-tensor histogram and freed; peak
    RAM no longer depends on the sample count.
  • (commit 6f69d55)

• The auto-alpha QDQ sub-graph session is built once per node.

  • It was rebuilt once per (node, alpha), and originally once per (node, alpha,
    sample), so a wide alpha grid ballooned ORT arena memory and rebuild time.
  • Candidate weights are now fed as runtime inputs to a single per-node session, so
    the alpha sweep is just repeated run() calls.
  • (commits 594f6f9, f2a102f)

• Variable-length calibration windows are accepted.

  • That once-per-node QDQ sub-graph baked the FIRST sample's concrete activation
    shape into its value_infos, so any later window of a different sequence length
    was rejected (Got invalid dimensions for input ... Got: N Expected: M).
    Calibration was silently constrained to uniformly-sized windows.
  • The activation input/output dims are now left dynamic (ORT resolves the real
    shape per run(), so equal-length calibration is numerically unchanged), letting
    one calibration set mix short and long clips.
  • (commit 0a22d36)

• Per-node activation caching across the alpha grid.

  • Each alpha evaluation re-harvested the same reference activations for a node.
  • They are now harvested once per node and reused across the whole grid.
  • (commit d6745e4)

• The large-model alpha search no longer leaks proto arena memory.

  • Each (node, alpha) evaluation rewrote the node's weight initializer twice (scale
    then recover), and even numpy_helper.to_array materializes external weights into
    the proto. Under protobuf's upb backend every such write abandons the old bytes in
    the ModelProto's arena (freed only when the proto dies), so retained RAM grew with
    nodes x alphas and survived into static calibration. A 0.1-step grid OOMed an
    export where a 0.2-step grid fit.
  • Candidate weights are now scaled in numpy and fed to the QDQ-loss session directly,
    and external weights are read via a throwaway TensorProto, so the search retains
    nothing.
  • (commit 72de740)

• Memory-conservative ORT sessions for the dump/calibration passes.

  • The augmented dump graphs return hundreds of activation tensors per forward, and
    ORT's BFC arenas grow ahead of demand and never release, so the calibration
    sessions carried ~1 GB of allocator slack on a 0.6B encoder and could abort with a
    BFCArena allocation failure on a loaded host.
  • These sessions run only a handful of forwards, so a new
    conservative_session_resources() disables the CPU arena and pins the CUDA arena
    to kSameAsRequested growth, applied to both the static-calibration session and the
    smoother's dump session.
  • (commit 6c6e0cf)

• Streaming per-channel percentile in the smoother calibration.

  • The smoother calibration stacked every sample's activations and called
    np.percentile once at the end, so RAM grew as samples x frames x channels
    (~6 GB per ~400 s window, hundreds of GB for a multi-window export) and thrashed.
  • Because the high percentiles SmoothQuant uses (99.999 by default) depend only on
    the largest values per channel, each sample is now folded into a per-channel
    running top-K reducer (StreamingChannelPercentile) and freed. The reducer
    reproduces np.percentile bit-for-bit (its float64 promotion and t >= 0.5 lerp
    branch included) and re-checks that K was large enough for the actual row count, so
    a result can never be silently wrong.
  • (commit e61fb24)

• Sliced static-calibration dump to bound VRAM.

  • Static int8 calibration augments the model so EVERY calibrated tensor becomes a
    graph output, and ORT keeps all graph outputs resident for the whole forward, so a
    long window peaks every activation at once (per-layer attention-score MatMuls are
    ~0.8 GB each near a FastConformer's ~400 s reach) and overflows VRAM.
  • With extra_options["CalibDumpBatch"] > 0 the tensors are dumped in slices of that
    many graph outputs, each slice its own augment + forward over the same (rewound)
    windows, so only one slice is resident at a time. The per-tensor ranges stay
    bit-identical to the single-pass dump (it only trades extra forwards for a smaller
    peak).
  • (commit a495323)

New features and usability

• Per-op-type alpha override in the auto-alpha search.

  • Pin an op type to a fixed alpha, or give it its own min:max:step grid, instead
    of one global grid for all smoothed ops.
  • (commit 280b123)

• Sensitivity-based mixed precision (extra_options["SmoothQuantExcludeWorst"]).

  • The auto-alpha search records every smoothed node's best achievable QDQ loss,
    normalized by that node's reference-output energy so it is comparable across nodes
    (Smoother.auto_alpha_losses).
  • Setting the option to an int n (or a fraction in (0, 1)) keeps the n worst-quantizing
    nodes out of quantization entirely (they stay fp32), trading a little file size for
    accuracy on the layers int8 hurts most. Only searched nodes are rankable, so with a
    fixed alpha it is a logged no-op.
  • (commit 89b4363)

• Resumable, crash-safe exports.

  • Given a checkpoint directory (extra_options["SmoothQuantCheckpointDir"] plus
    CalibParamsCheckpointFile), the smoother persists its expensive intermediates as
    they are produced and reloads whatever already exists on the next run: the smoother
    calibration, the per-node alpha search (append-only alphas.jsonl, one line per
    completed node, so an interrupted multi-hour search resumes from the last completed
    node), and the static-calibration quantization params.
  • Both per-sample calibration loops also checkpoint MID-pass, time-gated by
    CheckpointIntervalSec (default 1200 s) so a fast pass pays no IO, so an OOM- or
    time-killed pass resumes from the last sample instead of restarting.
  • All writes are atomic (temp + os.replace) so a crash mid-write never leaves a
    truncated file. Checkpoints carry no model fingerprint (the caller keys the
    directory by its full run configuration and owns invalidation) but DO record the
    slice partition behind the static-calibration partial, so a resume with a changed
    CalibDumpBatch or tensor set is discarded rather than misapplied.
  • (commits 55a3c75, 628f57f, 9eb2d4e, 0c58d7b)

• Per-layer alpha summary logged after the search.

  • A histogram per op type plus one line per smoothed node showing the alpha the
    search picked, so the export is auditable instead of opaque.
  • (commit 0a532e1)

• Richer quantization statistics table.

  • fp32 ops are split into "quantizable" (weight-bearing, the quantizer could convert
    them directly) versus "needs an int8 input" (weightless/pass-through ops that only
    convert inside an int8 region), and the Reshape/Transpose glue rows are always shown.
  • (commit e792411)

• Progress bars and quieter logs.

  • tqdm progress bars over the alpha search and both calibration passes, and the
    per-evaluation ORT warning spam is silenced.
  • (commits 82e3451, 9f0e94b)

---