Stream targeted block re-quantization's calibration pass - #2064
Stream targeted block re-quantization's calibration pass#2064aquilarubra wants to merge 3 commits into
Conversation
11ba30e to
3be3fae
Compare
c6e097e to
3841bb3
Compare
Rebase of PR intel#2064 onto current main. Most of the original stack (disk-streaming-core intel#2061, resumability intel#2062, AutoScheme streaming intel#2063) is already merged verbatim or superseded by intel#2220, so this commit carries only the calibration-pass streaming logic and its tests that weren't already upstream.
e193420 to
3f26e3c
Compare
|
Rebased onto current |
| export AR_RESUME_DIR=/path/to/resume/state | ||
| ``` | ||
|
|
||
| ### AR_DISK_STREAM_MODEL |
There was a problem hiding this comment.
We already have this variable; please clean up the relevant documentation.
There was a problem hiding this comment.
Good catch — the merge with main had duplicated the AR_DISK_STREAM_MODEL and AR_RESUME_DIR sections verbatim. Removed the duplicate copy, kept the richer one (mentions the parallel-scoring interaction). Pushed.
xin3he
left a comment
xin3he
left a comment
There was a problem hiding this comment.
Nice catch, please clean up the document
| export AR_DISK_STREAM_MODEL=1 | ||
| ``` | ||
|
|
||
| ### AR_RESUME_DIR |
There was a problem hiding this comment.
We already have this variable; please clean up the relevant documentation.
The merge with main duplicated these two sections verbatim; keep the richer copy (mentions parallel-scoring interaction) and drop the dupe.
38b10ff to
8f89538
Compare
|
/azp run Unit-Test-CUDA-AutoRound. |
|
/azp run Unit-Test-CUDA-AutoRound |
|
No pipelines are associated with this pull request. |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines successfully started running 1 pipeline(s). |
3 participants
Summary
Depends on #2061 (disk-streaming core) — stacked on top of it, so the diff includes those commits until it merges; only the last commit is new here.
to_quant_block_nameslets a caller restrict tuning to a subset of decoder blocks — useful for cheaply re-quantizing just a couple of blocks in an already-produced checkpoint at higher precision instead of redoing a full multi-hour run. Combined withAR_DISK_STREAM_MODEL, this crashes: the "cache block inputs" calibration forward pass needs real weights in every block leading up to (and sometimes through) the target block(s), but nothing materializes blocks outsidequant_block_listfor this specific pass — they stay meta forever, and the forward silently propagates meta-ness until it collides with a genuinely-materialized module. Full (unrestricted) runs never hit this, sincequant_block_listalready covers every block in that case.What's in this PR
auto_round/calibration/llm.py: when disk streaming is active and any decoder block still has meta parameters at the point this calibration forward runs, wraps it with the existingstream_block_forwardprimitive (from #2061), scoped to just the still-meta blocks. Only activates when there's something left meta to fix, so it's a no-op for the normal full-quantization path.Also adds an
AR_CALIB_STREAM_DEVICEenv-gated fast path: the default keeps this forward entirely on CPU (mixing a GPU-streamed block with CPU-resident hidden states crashes with a device mismatch), but a full CPU forward through every pre-target block of a 100B+ model is unusably slow for this specific targeted-requant use case. When set, every already-real tensor is moved to that device for the pass's duration and moved back afterward.Validation
Reproduced and fixed against a tiny hybrid-MoE fixture with an MTP head, with both a single restricted block and two adjacent ones, plus a control run of the full (unrestricted) path confirming no regression.