FourOverSix - #776
Conversation
|
Important Review skippedAuto incremental 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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughAdds a ChangesFour Over Six adaptive scaling
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Compressor as NVFP4PackedCompressor
participant ScaleAdjust as _adjust_scale_for_four_over_six
participant Forward as _quantize_dequantize
participant FourOverSix as _four_over_six_quantize_dequantize
Compressor->>Compressor: compress(weight, scale, global_scale)
alt weights.four_over_six enabled
Compressor->>ScaleAdjust: compute MSE(scale) vs MSE(scale*1.5)
ScaleAdjust-->>Compressor: adjusted scale
end
Compressor->>Compressor: quantize(weight, adjusted scale) and pack weight_packed
Forward->>Forward: _quantize_dequantize(x, scale, args)
alt args.four_over_six and num_bits==4 and type==FLOAT
Forward->>FourOverSix: dispatch(x, scale, zero_point, global_scale)
FourOverSix->>FourOverSix: quantize-dequantize with scale
FourOverSix->>FourOverSix: quantize-dequantize with scale*1.5
FourOverSix->>FourOverSix: compute per-group MSE, select via torch.where
FourOverSix-->>Forward: chosen dequantized tensor
end
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
Signed-off-by: Roderick Wu <Roderick-Wu@a100-06.nemg-001.lab.rdu2.dc.redhat.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/compressed_tensors/quantization/lifecycle/forward_helpers.py`:
- Around line 264-270: The MSE comparison in the helper currently reduces over
an empty `group_dims`, so `_process_group()` ends up making one global choice
instead of one per group. Update the logic in `forward_helpers.py` to use the
actual group dimension after `unflatten()` when computing `mse_a` and `mse_b`,
and keep `torch.where` selecting between `dq_a` and `dq_b` on a per-group basis
rather than over the full tensor.
In `@src/compressed_tensors/quantization/quant_args.py`:
- Around line 227-237: The four_over_six quantization flag is being excluded
from serialization, so it is lost in saved configs and not restored on reload.
Update the four_over_six field in QuantizationArgs to be included in
model_dump()/config serialization, and ensure the serialized quantization config
preserves this NVFP4_46 setting so input_activations reloads correctly in
forward_quantize().
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ff18ea5c-da2d-44e4-8cbc-c92c45a818b5
📒 Files selected for processing (5)
src/compressed_tensors/compressors/nvfp4/base.pysrc/compressed_tensors/quantization/lifecycle/forward_helpers.pysrc/compressed_tensors/quantization/quant_args.pysrc/compressed_tensors/quantization/quant_scheme.pysrc/compressed_tensors/quantization/utils/helpers.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
vllm-project/llm-compressor(manual) → reviewed against open PR#2903Roderick-Wu/tested-FourOverSixinstead of the default branch
| # --- Per-group MSE comparison --- | ||
| group_dims = tuple(range(scale.ndim, x.ndim)) | ||
| mse_a = ((x - dq_a) ** 2).mean(dim=group_dims, keepdim=True) | ||
| mse_b = ((x - dq_b) ** 2).mean(dim=group_dims, keepdim=True) | ||
|
|
||
| use_b = mse_b < mse_a | ||
| return torch.where(use_b, dq_b, dq_a) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect how _quantize_dequantize is invoked for group strategy and the ranks of x vs scale
rg -nP '_quantize_dequantize' src/compressed_tensors/quantization/lifecycle -C3
rg -nP 'def _process_quantization|reshape|unflatten|unsqueeze|repeat_interleave|expand' \
src/compressed_tensors/quantization/lifecycle/forward.py -C2Repository: vllm-project/compressed-tensors
Length of output: 3541
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant helper file first, then inspect the precise implementation slices.
ast-grep outline src/compressed_tensors/quantization/lifecycle/forward_helpers.py --view expanded
echo '--- forward_helpers.py lines 1-340 ---'
sed -n '1,340p' src/compressed_tensors/quantization/lifecycle/forward_helpers.py
echo '--- search for four_over_six and group-related helpers ---'
rg -n "four_over_six|group_size|group_dims|scale.ndim|reshape|expand|unsqueeze|repeat_interleave" \
src/compressed_tensors/quantization/lifecycle/forward_helpers.pyRepository: vllm-project/compressed-tensors
Length of output: 11525
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the call chain into _process_group / _apply_quantize_op and inspect shape assumptions.
rg -n "_process_group\(|_apply_quantize_op\(" src/compressed_tensors -C3
echo '--- forward.py slices around the group dispatch ---'
sed -n '1,260p' src/compressed_tensors/quantization/lifecycle/forward.py
echo '--- any tests covering four_over_six group behavior ---'
rg -n "four_over_six|_four_over_six|group_size|group quant" tests src/compressed_tensors -C2Repository: vllm-project/compressed-tensors
Length of output: 50387
🏁 Script executed (no clone):
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import torch
x = torch.randn(2, 3, 4)
scale = torch.randn(2, 3, 1)
group_dims = tuple(range(scale.ndim, x.ndim))
print("x.ndim =", x.ndim)
print("scale.ndim =", scale.ndim)
print("group_dims =", group_dims)
m = (x ** 2).mean(dim=group_dims, keepdim=True)
print("mean shape =", tuple(m.shape))
print("mean equals input:", torch.allclose(m, x ** 2))
print("mse selection shape =", tuple(torch.where(m < (m + 1), x, x).shape))
PY
Length of output: 247
Reduce the MSE over the group axis here
_process_group() passes x and scale.unsqueeze(-1) at the same rank, so group_dims is empty in this helper. mean(dim=()) collapses each candidate to a single scalar, and torch.where then applies one branch to the entire tensor instead of one decision per group. Use the actual group dimension after unflatten().
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/compressed_tensors/quantization/lifecycle/forward_helpers.py` around
lines 264 - 270, The MSE comparison in the helper currently reduces over an
empty `group_dims`, so `_process_group()` ends up making one global choice
instead of one per group. Update the logic in `forward_helpers.py` to use the
actual group dimension after `unflatten()` when computing `mse_a` and `mse_b`,
and keep `torch.where` selecting between `dq_a` and `dq_b` on a per-group basis
rather than over the full tensor.
| four_over_six: bool = Field( | ||
| default=False, | ||
| exclude=True, | ||
| description=( | ||
| "Enable Four Over Six (4/6) adaptive block scaling for NVFP4 " | ||
| "quantization. For each group of values, tries scaling to both 4 " | ||
| "and 6, selecting the scale with lower MSE. Reduces quantization " | ||
| "error for near-maximal values in FP4. Only applies to FP4 " | ||
| "quantization (num_bits=4, type=float)." | ||
| ), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether four_over_six is expected to survive serialization and be consumed at runtime for activations
rg -nP 'exclude=True' src/compressed_tensors/quantization/quant_args.py -C1
rg -nP 'four_over_six' -C2 src/
# How dynamic/LOCAL activations flow into the forward quantize-dequantize path
rg -nP '_quantize_dequantize|forward_quantize|maybe_calibrate_or_quantize' src/compressed_tensors/quantization/lifecycle -C2Repository: vllm-project/compressed-tensors
Length of output: 11776
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect how QuantizationArgs are serialized and reloaded
ast-grep outline src/compressed_tensors/quantization/quant_args.py --view expanded || true
rg -n "model_dump|dict\(|json|serialize|deserialize|exclude=.*four_over_six|exclude=True" src/compressed_tensors -C 2
# Inspect the quantization config / scheme load-save path
rg -n "QuantizationArgs|QuantizationScheme|quant_scheme|from_dict|to_dict|from_json|save.*quant|load.*quant|model_validate" src/compressed_tensors -C 2
# Inspect how forward quantization consumes the flag for activations
sed -n '1,260p' src/compressed_tensors/quantization/lifecycle/forward_helpers.pyRepository: vllm-project/compressed-tensors
Length of output: 50390
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the relevant serialization and runtime paths with line numbers
sed -n '180,260p' src/compressed_tensors/quantization/quant_args.py
sed -n '150,240p' src/compressed_tensors/quantization/quant_config.py
sed -n '1,170p' src/compressed_tensors/quantization/lifecycle/forward_helpers.py
sed -n '240,330p' src/compressed_tensors/quantization/lifecycle/forward.py
sed -n '190,240p' src/compressed_tensors/compressors/model_compressors/model_compressor.pyRepository: vllm-project/compressed-tensors
Length of output: 17632
Persist four_over_six in serialized quantization configs (src/compressed_tensors/quantization/quant_args.py:227-237)
exclude=True drops this flag from model_dump(), so saved config.json loses the NVFP4_46 activation setting. Reloaded input_activations then fall back to standard FP4 scaling in forward_quantize().
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/compressed_tensors/quantization/quant_args.py` around lines 227 - 237,
The four_over_six quantization flag is being excluded from serialization, so it
is lost in saved configs and not restored on reload. Update the four_over_six
field in QuantizationArgs to be included in model_dump()/config serialization,
and ensure the serialized quantization config preserves this NVFP4_46 setting so
input_activations reloads correctly in forward_quantize().
|
Model: Meta-Llama-3-8B-Instruct
Sanity check to make sure it is actually choosing different scales, computed by re-running the FourOverSix observer on the original unquantized weights.
|
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.
|
1 participant
Merging over from https://github.com/krishnateja95/compressed-tensors/tree/FourOverSix
Tested with this branch for llm-compressor:
https://github.com/vllm-project/llm-compressor/tree/Roderick-Wu/tested-FourOverSix
Evaluation results with Llama-3-8B, Qwen3-30B-A3B, Llama-3-70B
Tested on A100s:
FP4 activation quantization is not natively supported and falls back silently to the Marlin W4A16 kernel. W4A4 and W4A16 produce identical outputs on this hardware (slightly different results here bc sampling was not greedy). Real distinction observable on SM100+.