[Feature] Add layerwise hidden-state QAD for quantized block reconstruction - #3029
[Feature] Add layerwise hidden-state QAD for quantized block reconstruction#3029BeichenHuang wants to merge 2 commits into
Conversation
Implement GPTQ-initialized block-local hidden-state distillation with streaming teacher targets, validation-best restoration, and sequential error propagation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e29d330c-b320-4c9e-8e2b-17e2f355fe0e
Keep the upstream change focused on implementation, tests, and the user-facing example; design context will be included in the pull request description. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e29d330c-b320-4c9e-8e2b-17e2f355fe0e
|
👋 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. |
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
The quality checks have failed. Please run |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the LayerwiseQADModifier to perform block-local quantization-aware distillation (QAD) against a full-precision teacher model using hidden-state mean squared error. It adds support for loading and initializing a teacher model (distill_teacher) via the oneshot entrypoint, configures sequential pipeline propagation rules, and includes comprehensive unit tests. The review feedback highlights several robustness improvements for the new modifier, including handling fully-padded microbatches gracefully to avoid division-by-zero crashes, preserving namedtuple types during recursive tensor mapping, and implementing a fallback training mechanism when there are fewer than two calibration batches.
| denominator = mask.sum() | ||
| if denominator.item() == 0: | ||
| raise ValueError("Loss mask does not contain any valid tokens") | ||
| return (per_token_loss * mask).sum() / denominator |
There was a problem hiding this comment.
If a microbatch consists entirely of padding tokens (which is highly likely when using small batch sizes like 1 or 2 on padded datasets), denominator.item() == 0 will be true. Raising a ValueError here will crash the entire calibration and quantization pipeline, which could be hours into a run.
Instead of crashing, we can gracefully return a zero loss that is still connected to the computation graph (e.g., prediction.sum() * 0.0). This ensures that the fully-masked batch contributes zero gradients and training continues safely.
| denominator = mask.sum() | |
| if denominator.item() == 0: | |
| raise ValueError("Loss mask does not contain any valid tokens") | |
| return (per_token_loss * mask).sum() / denominator | |
| denominator = mask.sum() | |
| if denominator.item() == 0: | |
| return prediction.sum() * 0.0 | |
| return (per_token_loss * mask).sum() / denominator |
| def _map_tensors(value: Any, transform): | ||
| if isinstance(value, torch.Tensor): | ||
| return transform(value) | ||
| if isinstance(value, tuple): | ||
| return tuple(_map_tensors(item, transform) for item in value) | ||
| if isinstance(value, list): | ||
| return [_map_tensors(item, transform) for item in value] | ||
| if isinstance(value, dict): | ||
| return {key: _map_tensors(item, transform) for key, item in value.items()} | ||
| return value |
There was a problem hiding this comment.
The recursive _map_tensors helper converts any tuple subclass (including namedtuple) into a plain tuple. If a model's block forward pass expects a namedtuple or custom tuple subclass as an input argument, this conversion will strip the named fields and custom attributes, potentially causing AttributeError or other runtime failures.
We should check if the tuple is a namedtuple (or custom subclass) and preserve its type during mapping.
| def _map_tensors(value: Any, transform): | |
| if isinstance(value, torch.Tensor): | |
| return transform(value) | |
| if isinstance(value, tuple): | |
| return tuple(_map_tensors(item, transform) for item in value) | |
| if isinstance(value, list): | |
| return [_map_tensors(item, transform) for item in value] | |
| if isinstance(value, dict): | |
| return {key: _map_tensors(item, transform) for key, item in value.items()} | |
| return value | |
| def _map_tensors(value: Any, transform): | |
| if isinstance(value, torch.Tensor): | |
| return transform(value) | |
| if isinstance(value, tuple): | |
| if hasattr(value, "_fields"): # Preserve namedtuples | |
| return type(value)(*(_map_tensors(item, transform) for item in value)) | |
| return tuple(_map_tensors(item, transform) for item in value) | |
| if isinstance(value, list): | |
| return [_map_tensors(item, transform) for item in value] | |
| if isinstance(value, dict): | |
| return {key: _map_tensors(item, transform) for key, item in value.items()} | |
| return value |
| train_indices, validation_indices = self._split_batch_indices(len(batches)) | ||
| initial_train_loss = self._evaluate_indices( | ||
| student_block, | ||
| teacher_block, | ||
| batches, | ||
| train_indices, | ||
| device, | ||
| teacher_dtype, | ||
| ) | ||
| initial_validation_loss = self._evaluate_indices( | ||
| student_block, | ||
| teacher_block, | ||
| batches, | ||
| validation_indices, | ||
| device, | ||
| teacher_dtype, | ||
| ) | ||
| self._release_device_cache() | ||
| self._log_device_memory(name, "after initial evaluation") | ||
| optimizer = torch.optim.AdamW( | ||
| trainable, | ||
| lr=self.learning_rate, | ||
| weight_decay=self.weight_decay, | ||
| ) | ||
| optimizer_steps, epochs_completed, best_validation_loss = ( | ||
| self._train_with_validation( | ||
| student_block, | ||
| teacher_block, | ||
| batches, | ||
| optimizer, | ||
| trainable, | ||
| train_indices, | ||
| validation_indices, | ||
| device, | ||
| teacher_dtype, | ||
| ) | ||
| ) | ||
| final_train_loss = self._evaluate_indices( | ||
| student_block, | ||
| teacher_block, | ||
| batches, | ||
| train_indices, | ||
| device, | ||
| teacher_dtype, | ||
| ) | ||
| final_validation_loss = self._evaluate_indices( | ||
| student_block, | ||
| teacher_block, | ||
| batches, | ||
| validation_indices, | ||
| device, | ||
| teacher_dtype, | ||
| ) |
There was a problem hiding this comment.
If the calibration dataset is very small or the batch size is large, the total number of captured batches (len(batches)) might be less than 2. In this case, _split_batch_indices will raise a ValueError and crash the run.
Since LayerwiseQADModifier already implements a non-validation _train method, we can gracefully fallback to training on all batches without validation/early stopping when len(batches) < 2, rather than crashing.
try:
if len(batches) < 2:
logger.warning(
f"Block `{name}` has only {len(batches)} calibration batch(es). "
"Skipping validation and training on all batches."
)
optimizer = torch.optim.AdamW(
trainable,
lr=self.learning_rate,
weight_decay=self.weight_decay,
)
optimizer_steps = self._train(
student_block,
teacher_block,
batches,
optimizer,
trainable,
device,
teacher_dtype,
)
epochs_completed = self.num_epochs
best_validation_loss = float("nan")
initial_train_loss = float("nan")
final_train_loss = float("nan")
initial_validation_loss = float("nan")
final_validation_loss = float("nan")
else:
train_indices, validation_indices = self._split_batch_indices(len(batches))
initial_train_loss = self._evaluate_indices(
student_block,
teacher_block,
batches,
train_indices,
device,
teacher_dtype,
)
initial_validation_loss = self._evaluate_indices(
student_block,
teacher_block,
batches,
validation_indices,
device,
teacher_dtype,
)
self._release_device_cache()
self._log_device_memory(name, "after initial evaluation")
optimizer = torch.optim.AdamW(
trainable,
lr=self.learning_rate,
weight_decay=self.weight_decay,
)
optimizer_steps, epochs_completed, best_validation_loss = (
self._train_with_validation(
student_block,
teacher_block,
batches,
optimizer,
trainable,
train_indices,
validation_indices,
device,
teacher_dtype,
)
)
final_train_loss = self._evaluate_indices(
student_block,
teacher_block,
batches,
train_indices,
device,
teacher_dtype,
)
final_validation_loss = self._evaluate_indices(
student_block,
teacher_block,
batches,
validation_indices,
device,
teacher_dtype,
)
|
|
||
| model = oneshot( | ||
| model="student-model-or-path", | ||
| distill_teacher="full-precision-teacher-model-or-path", |
There was a problem hiding this comment.
don't think we need this
| __all__ = ["SequentialPipeline"] | ||
|
|
||
|
|
||
| def _configure_modifier_pipeline(modifiers, dataset_args): |
There was a problem hiding this comment.
i'm unsure why these pipeline changes are necessary can you explain?
| @@ -0,0 +1,53 @@ | |||
| # Layerwise MSE QAD | |||
|
|
|||
| `LayerwiseQADModifier` performs block-local quantization-aware distillation. | |||
There was a problem hiding this comment.
maybe just QAD modifier, this makes it sound like it works on one linear layer at a time, (ideally this would work on a whole subgraph forward/backward)
2 participants
SUMMARY:
This PR introduces
LayerwiseQADModifier, an experimental block-localquantization-aware distillation method for decoder-only language models.
Motivation
GPTQ provides a strong quantized initialization, but independently reconstructed
blocks may not remain optimal under accumulated quantization error. Conventional
end-to-end QAD requires complete teacher and student forwards and a full student
backward graph.
Layerwise QAD instead optimizes one decoder block at a time inside LLM
Compressor's sequential calibration pipeline.
Method
For each decoder block:
same input activation.
parameters remain frozen.
This is layerwise hidden-state distillation / quantization-aware block
reconstruction, not end-to-end final-logit KL distillation.
Implementation
The PR adds:
LayerwiseQADModifier;oneshot(distill_teacher=...)teacher-model wiring;Layerwise QAD requires:
Memory behavior
Teacher targets are generated on demand and released after each microbatch, so
target memory does not scale with the total calibration dataset.
The current implementation still caches block inputs for the calibration
dataset. Chunked or disk-backed input replay is left as future work.
Validation
Validated on Llama 3.1 8B Instruct with NVFP4 W4A4 and a 512x2048 public-six
calibration dataset:
Three-seed downstream results:
Known limitations
every downstream task.
TEST PLAN:
pytest -q tests/llmcompressor/modifiers/layerwise_qad/test_base.pysplitting, early stopping, and best-weight restoration.