Skip to content

refactor(config): replace grid workspace env vars with MSEConfig - #3436

Merged
Qubitium merged 2 commits into
mainfrom
devin/1792844887-mse-grid-config
Oct 24, 2026
Merged

Qubitium merged 2 commits into
mainfrom
devin/1792844887-mse-grid-config

Conversation

@Qubitium

@Qubitium Qubitium commented Oct 24, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Replaces the GPTQMODEL_MSE_GRID_BUDGET_MB / GPTQMODEL_CLIP_GRID_BUDGET_MB
environment variables from #3435 with a proper config object on
QuantizeConfig, folding the legacy float mse exponent into it: the field
is now mse, taking an MSEConfig.

QuantizeConfig(
    mse=MSEConfig(norm=2.4, grid_budget="64M", clip_grid_budget="1G"),
)

What Changed

  • quantization/config.py: MSEConfig now carries norm (the MSE exponent)
    plus grid_budget / clip_grid_budget. Budgets are byte counts (ints) or
    postfixed strings ("256M", "1G", optional B); None keeps the
    device-derived default (16M CPU / 256M GPU) and 0 forces single-candidate
    chunking (and one-token AWQ product tiles). The mse field lives on
    BaseQuantizeConfig, so GPTQConfig, QQQConfig, and AWQConfig all
    accept it.
  • Back-compat normalization: mse=2.4 → MSEConfig(norm=2.4); dicts merge
    field-wise, and per-layer dynamic overrides replace norm while keeping
    configured budgets.
  • quantizer.py: find_params/_mse_shrink_grid read the exponent and grid
    budget from qcfg.mse.
  • qqq.py: Quantizer now accepts qcfg in its constructor like the GPTQ
    Quantizer and reads the grid budget from self.qcfg.mse; the
    configure(grid_budget_mb=...) shim is gone. Internal Quantizer() call
    sites pass self.qcfg.
  • awq_processor.py / qqq_processor.py / gptq_processor.py: read budgets
    and norm from qcfg.mse; dynamic mse overrides go through
    _normalize_mse.
  • models/writer.py + config.py meta payload: checkpoint meta still
    serializes the float exponent (meta["mse"] = mse.norm) so saved configs
    round-trip identically. Budgets are calibration-scratch tuning knobs, not
    part of quantization identity.
  • README updated for the new API.

Tests

  • Unit tests have been executed and passed.
  • Existing MSE/clip grid tests were updated to the new config surface
    (mse=MSEConfig / mse={...} dicts); the tiny-budget forced-chunking
    paths still hold.
cd tests && python -m pytest test_find_params_mse_grid.py \
    test_qqq_quantizer.py test_awq_clip_consistency.py -x -q
# 134 passed, 72 skipped (CUDA unavailable on this VM)

python -m pytest test_serialization.py -k meta -q
# 1 passed (meta["mse"] == cfg.mse.norm round-trip)

Review Requirements

  • Reviewed this PR.
  • Checked that the code matches existing project structure, APIs, and conventions.
  • Avoided unnecessary monkeypatching and used the project's normal extension points where possible.
  • Minimized and compacted the impacted code surface.
  • Considered and eliminated potential regressions.

Kernel Accuracy Requirements

Notes

  • norm semantics are unchanged: 0 disables the MSE search; AWQ ignores it
    and only uses clip_grid_budget.
  • No GPU on this VM: CUDA-parametrized test cases skip locally.
  • 3 unrelated test_serialization.py cases fail with HF-download OSError
    in this environment (no network access to checkpoints).

Link to Devin session: https://app.devin.ai/sessions/2eb9e33e49b54734a60b2552d71b0aca
Open in Devin Desktop: https://app.devin.ai/desktop/session/2eb9e33e49b54734a60b2552d71b0aca?variant=devin
Requested by: @Qubitium

Remove GPTQMODEL_MSE_GRID_BUDGET_MB and GPTQMODEL_CLIP_GRID_BUDGET_MB.
The per-chunk workspace budgets for the vectorized GPTQ/QQQ MSE shrink
grid and AWQ clip grid now live on a new MSEConfig dataclass exposed as
QuantizeConfig.mse_config (grid_budget_mb / clip_grid_budget_mb).

None preserves the previous device-derived defaults (16 MiB CPU,
256 MiB GPU); 0 still forces single-candidate chunking. QQQ threads the
budget through Quantizer.configure; the AWQ processor reads
qcfg.mse_config directly. The knob is calibration-time tuning only and
is not serialized into checkpoint meta.
@devin-ai-integration

Copy link
Copy Markdown
Contributor

I'll fix CI failures and address comments from users with write access that start with 'Devin'.

  • Disable automatic comment, CI, and merge conflict monitoring

…nfig

- `mse` on QuantizeConfig now takes an MSEConfig (norm exponent +
  grid_budget/clip_grid_budget); a bare float still normalizes to
  MSEConfig(norm=...), and dicts merge field-wise for dynamic overrides.
- Budgets are byte counts (ints) or 'K'/'M'/'G' postfixed strings
  ("256M", "1G"); None keeps the device default (16M CPU / 256M GPU),
  0 forces single-candidate chunking.
- QQQ's Quantizer now accepts qcfg like the GPTQ one and reads the MSE
  grid budget from it; the configure(grid_budget_mb) shim is dropped.
- Checkpoint meta still serializes the float norm as before.
@Qubitium
Qubitium merged commit e55eff4 into main Oct 24, 2026
6 checks passed
@Qubitium
Qubitium deleted the devin/1792844887-mse-grid-config branch October 24, 2026 15:11
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