fix: handle NVFP4 scale-only checkpoint shards - #17188
Open
zcxGGmu wants to merge 1 commit into
Open
Conversation
Signed-off-by: zq <zhouquan1511@163.com>
Contributor
WalkthroughThe NVFP4 load hook now processes unified Hugging Face scale tensors before packed weights are available. It computes and clamps ChangesNVFP4 sharded checkpoint loading
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py`:
- Around line 135-161: The
test_nvfp4_load_hook_maps_scale_only_shard_with_prefix test only verifies alpha
and input_scale, so it does not validate scale conversion. Annotate the test’s
monkeypatch parameter with pytest.MonkeyPatch and its unused third parameter
with None, make block_scale_interleave return distinguishable uint8 data, and
assert the resulting stored weight scale’s shape and byte values after
nvfp4_imp.load_hook.
🪄 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: Enterprise
Run ID: 365b03f3-2e16-4e10-b3e8-ccded0cea512
📒 Files selected for processing (2)
tensorrt_llm/_torch/auto_deploy/transform/library/quantization.pytests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py
Comment on lines
+135
to
+161
| def test_nvfp4_load_hook_maps_scale_only_shard_with_prefix(monkeypatch): | ||
| from tensorrt_llm._torch.auto_deploy.transform.library.quantization import ( | ||
| NVFP4LinearQuantizationFromConfig, | ||
| ) | ||
|
|
||
| monkeypatch.setattr( | ||
| torch.ops.trtllm, | ||
| "block_scale_interleave", | ||
| lambda weight_scale: weight_scale, | ||
| raising=False, | ||
| ) | ||
|
|
||
| config = TransformConfig(stage="pattern_matcher") | ||
| nvfp4_imp = NVFP4LinearQuantizationFromConfig(config) | ||
|
|
||
| weight_name = "layer.proj.weight" | ||
| prefix = "nested." | ||
| mock_state_dict = { | ||
| prefix + weight_name + "_scale_2": torch.tensor(2.0, dtype=torch.float32), | ||
| prefix + weight_name + "_scale": torch.ones(128, 4, dtype=torch.float8_e4m3fn), | ||
| prefix + "layer.proj.input_scale": torch.tensor(3.0, dtype=torch.float32), | ||
| } | ||
|
|
||
| nvfp4_imp.load_hook(mock_state_dict, prefix, None, weight_name=weight_name) | ||
|
|
||
| assert mock_state_dict[prefix + "layer.proj.alpha"] == torch.tensor(6.0) | ||
| assert mock_state_dict[prefix + "layer.proj.input_scale"] == torch.tensor(1.0 / 3.0) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
if ! rg -n --glob '*.{yml,yaml}' \
'test_nvfp4_load_hook_maps_scale_only_shard_with_prefix|test_quantization_utils\.py' \
tests/integration/test_lists; then
echo "No matching test-list entry found."
fiRepository: NVIDIA/TensorRT-LLM
Length of output: 192
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- candidate test file ---'
sed -n '1,220p' tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py
echo '--- quantization implementation references ---'
rg -n -C 8 \
'class NVFP4LinearQuantizationFromConfig|block_scale_interleave|weight_scale|load_hook' \
tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py
echo '--- test-list files and matching references ---'
find tests/integration/test_lists -maxdepth 2 -type f \( -name '*.yml' -o -name '*.yaml' -o -name '*.sqlite' \) -print | sort
rg -n -i \
'quantization_utils|nvfp4|auto_deploy|singlegpu' \
tests/integration/test_lists || true
echo '--- coding guidance and test configuration ---'
find . -maxdepth 3 -iname 'CODING_GUIDELINES.md' -o -iname 'pyproject.toml' -o -iname 'setup.cfg' -o -iname 'tox.ini' | sortRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- candidate test file ---'
sed -n '1,220p' tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py
echo '--- quantization implementation references ---'
rg -n -C 8 \
'class NVFP4LinearQuantizationFromConfig|block_scale_interleave|weight_scale|load_hook' \
tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py
echo '--- test-list files and matching references ---'
find tests/integration/test_lists -maxdepth 2 -type f \( -name '*.yml' -o -name '*.yaml' -o -name '*.sqlite' \) -print | sort
rg -n -i \
'quantization_utils|nvfp4|auto_deploy|singlegpu' \
tests/integration/test_lists || true
echo '--- coding guidance and test configuration ---'
find . -maxdepth 3 \( -iname 'CODING_GUIDELINES.md' -o -iname 'pyproject.toml' -o -iname 'setup.cfg' -o -iname 'tox.ini' \) -print | sortRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py")
impl_path = Path("tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py")
test_source = test_path.read_text()
test_tree = ast.parse(test_source)
test_fn = next(
node for node in test_tree.body
if isinstance(node, ast.FunctionDef)
and node.name == "test_nvfp4_load_hook_maps_scale_only_shard_with_prefix"
)
lambda_nodes = [
node for node in ast.walk(test_fn)
if isinstance(node, ast.Lambda)
]
assert len(lambda_nodes) == 1
identity_lambda = (
len(lambda_nodes[0].args.args) == 1
and isinstance(lambda_nodes[0].body, ast.Name)
and lambda_nodes[0].body.id == lambda_nodes[0].args.args[0].arg
)
assert_texts = [
ast.get_source_segment(test_source, node)
for node in ast.walk(test_fn)
if isinstance(node, ast.Assert)
]
scale_asserts = [
text for text in assert_texts
if text and "_scale" in text
]
print(f"test_function_arguments={[arg.arg for arg in test_fn.args.args]}")
print(f"mock_is_identity={identity_lambda}")
print(f"assertions={assert_texts}")
print(f"assertions_observing_weight_scale={scale_asserts}")
impl_source = impl_path.read_text()
impl_tree = ast.parse(impl_source)
impl_fn = next(
node for node in ast.walk(impl_tree)
if isinstance(node, ast.FunctionDef)
and node.name == "load_hook"
and any(
isinstance(parent, ast.ClassDef)
and parent.name == "NVFP4LinearQuantizationFromConfig"
for parent in ast.walk(impl_tree)
if isinstance(parent, ast.ClassDef)
and node in parent.body
)
)
impl_text = ast.get_source_segment(impl_source, impl_fn)
assert impl_text is not None
print(f"implementation_calls_block_scale_interleave={'block_scale_interleave' in impl_text}")
print(f"implementation_stores_transformed_scale='weight_scale_swizzled' in state_dict={\"weight_scale_swizzled\" in impl_text}")
PY
echo '--- exact test-list membership ---'
rg -n -F \
'unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py' \
tests/integration/test_lists || true
rg -n \
'(^|[/ ])auto_deploy(/|$)|(^|[/ ])singlegpu(/|$)' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa \
--glob '*.yml' --glob '*.yaml' --glob '*.txt' || true
echo '--- repository guidance files ---'
find . -type f \( -name 'CODING_GUIDELINES.md' -o -name 'pyproject.toml' \) -print | sort
echo '--- changed-file summary ---'
git diff --numstat -- tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 514
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py")
impl_path = Path("tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py")
test_source = test_path.read_text()
test_tree = ast.parse(test_source)
test_fn = next(
node for node in test_tree.body
if isinstance(node, ast.FunctionDef)
and node.name == "test_nvfp4_load_hook_maps_scale_only_shard_with_prefix"
)
lambda_nodes = [node for node in ast.walk(test_fn) if isinstance(node, ast.Lambda)]
assert len(lambda_nodes) == 1
identity_lambda = (
len(lambda_nodes[0].args.args) == 1
and isinstance(lambda_nodes[0].body, ast.Name)
and lambda_nodes[0].body.id == lambda_nodes[0].args.args[0].arg
)
assert_texts = [
ast.get_source_segment(test_source, node)
for node in ast.walk(test_fn)
if isinstance(node, ast.Assert)
]
scale_asserts = [
text for text in assert_texts
if text and "_scale" in text
]
print(f"test_function_arguments={[arg.arg for arg in test_fn.args.args]}")
print(f"mock_is_identity={identity_lambda}")
print(f"assertions={assert_texts}")
print(f"assertions_observing_weight_scale={scale_asserts}")
impl_source = impl_path.read_text()
print(
"implementation_calls_block_scale_interleave="
f"{'block_scale_interleave' in impl_source}"
)
print(
"implementation_stores_transformed_scale="
f"{'weight_scale_swizzled' in impl_source}"
)
PY
echo '--- exact test-list membership ---'
rg -n -F \
'unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py' \
tests/integration/test_lists || true
rg -n \
'(^|[/ ])auto_deploy(/|$)|(^|[/ ])singlegpu(/|$)' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa \
--glob '*.yml' --glob '*.yaml' --glob '*.txt' || true
echo '--- repository guidance files ---'
find . -type f \( -name 'CODING_GUIDELINES.md' -o -name 'pyproject.toml' \) -print | sort
echo '--- changed-file summary ---'
git diff --numstat -- tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 6598
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- applicable Python and test guidance ---'
rg -n -C 3 \
'Annotate every function|test coverage|test list|test-db|qa/|directory|recursive' \
CODING_GUIDELINES.md
echo '--- CI entries covering the changed test directory ---'
sed -n '72,94p' tests/integration/test_lists/test-db/l0_a30.yml
sed -n '486,502p' tests/integration/test_lists/test-db/l0_h100.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 201
Strengthen the NVFP4 scale assertion.
test_nvfp4_load_hook_maps_scale_only_shard_with_prefix is covered by the utils directory entries in l0_a30.yml and l0_h100.yml. The identity mock and missing weight_scale assertion leave the conversion untested. Return distinguishable uint8 data from block_scale_interleave, then assert the stored scale shape and bytes. Add pytest.MonkeyPatch and None annotations to the test function.
Coverage verdict: insufficient.
🤖 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 `@tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.py` around
lines 135 - 161, The test_nvfp4_load_hook_maps_scale_only_shard_with_prefix test
only verifies alpha and input_scale, so it does not validate scale conversion.
Annotate the test’s monkeypatch parameter with pytest.MonkeyPatch and its unused
third parameter with None, make block_scale_interleave return distinguishable
uint8 data, and assert the resulting stored weight scale’s shape and byte values
after nvfp4_imp.load_hook.
Source: Path instructions
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
alpha, invertedinput_scale, and swizzledweight_scale.Fixes #11541
Test Plan
python -m ruff check tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.pypython -m py_compile tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py tests/unittest/auto_deploy/singlegpu/utils/test_quantization_utils.pyimpi.dll/msmpi.dll) and compiled TensorRT-LLM bindings.Dev Engineer Review
alpha, invertsinput_scale, and swizzlesweight_scalebefore requiring the packed weight.QA Engineer Review
test_nvfp4_load_hook_maps_scale_only_shard_with_prefix.