[feat] Add support to load fp8 Meta Llama4 weights #6268
Conversation
Signed-off-by: Aurelien Chartier <[email protected]>
WalkthroughThe changes introduce support for a new quantization mode, "compressed-tensors" with "float-quantized" format, across model configuration, weight loading, and quantization logic. Updates include new enum values, properties, and conditional branches to handle this mode, as well as adjustments to tensor shape handling during weight loading. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant ModelLoader
participant ModelConfig
participant QuantConfig
User->>ModelLoader: load(model_path, hf_quant_config)
ModelLoader->>ModelConfig: from_pretrained(..., hf_quant_config)
ModelConfig->>QuantConfig: process quant_method/format
alt quant_method == "compressed-tensors" and format == "float-quantized"
QuantConfig->>QuantConfig: set quant_algo = FP8_PER_CHANNEL_PER_TOKEN
QuantConfig->>QuantConfig: set use_meta_recipe = True
QuantConfig->>QuantConfig: process ignore list
else other quantization
QuantConfig->>QuantConfig: handle other quantization configs
end
ModelConfig-->>ModelLoader: return config
ModelLoader-->>User: return loaded model
Estimated code review effort3 (~45 minutes) Suggested reviewers
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![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
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_llama.py (1)
276-276: Fix line length violation and improve readability.The conditional assignment logic is correct, but the line exceeds the 120-character limit flagged by static analysis.
- weight_loading_mode = MoEWeightLoadingMode.META_FP8_RECIPE if model_config.quant_config.use_meta_recipe else MoEWeightLoadingMode.FUSED_GATE_UP_PROJ + weight_loading_mode = (MoEWeightLoadingMode.META_FP8_RECIPE + if model_config.quant_config.use_meta_recipe + else MoEWeightLoadingMode.FUSED_GATE_UP_PROJ)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
tensorrt_llm/_torch/model_config.py(1 hunks)tensorrt_llm/_torch/models/modeling_llama.py(2 hunks)tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py(3 hunks)tensorrt_llm/_torch/modules/fused_moe/interface.py(2 hunks)tensorrt_llm/_torch/modules/fused_moe/quantization.py(3 hunks)tensorrt_llm/_torch/modules/linear.py(2 hunks)tensorrt_llm/llmapi/llm_utils.py(1 hunks)
🧠 Learnings (1)
tensorrt_llm/_torch/modules/linear.py (1)
Learnt from: yechank-nvidia
PR: #6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using from_shared_tensor() is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call strip_for_generation() to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/models/modeling_llama.py
276-276: Line too long (156 > 120)
(E501)
🧰 Additional context used
🧠 Learnings (1)
tensorrt_llm/_torch/modules/linear.py (1)
Learnt from: yechank-nvidia
PR: #6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using from_shared_tensor() is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call strip_for_generation() to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/models/modeling_llama.py
276-276: Line too long (156 > 120)
(E501)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (11)
tensorrt_llm/_torch/modules/linear.py (2)
501-502: LGTM: Proper tensor shape normalization for FP8 row-wise quantization.The
.squeeze()call ensures that singleton dimensions are removed from the weight scale tensor after loading, which is necessary for proper FP8 row-wise quantization handling.
537-537: LGTM: Consistent tensor shape handling in fused operations.The
.squeeze()call on the concatenated scale tensor maintains consistency with the vanilla weight loading approach and ensures proper tensor dimensionality for FP8 operations.tensorrt_llm/_torch/modules/fused_moe/interface.py (2)
15-15: LGTM: New enum member for META FP8 recipe support.The addition of
META_FP8_RECIPE = 2extends the weight loading modes to support the new FP8 quantization recipe, following the existing enum pattern.
115-119: LGTM: Consistent property implementation for FP8 row-wise detection.The new
has_fp8_rowwiseproperty follows the same pattern as existing quantization detection properties and provides a clean interface for checking FP8 row-wise quantization mode.tensorrt_llm/_torch/models/modeling_llama.py (1)
285-285: LGTM: Proper use of dynamic weight loading mode.The use of the local
weight_loading_modevariable instead of a hardcoded value enables proper conditional selection based on the quantization configuration.tensorrt_llm/llmapi/llm_utils.py (1)
429-434: LGTM! Proper integration of Llama4 FP8 quantization support.The new conditional branch correctly handles the "compressed-tensors" with "float-quantized" format by setting the appropriate quantization algorithm. The implementation follows the existing pattern and maintains consistency with the DeepSeek V3 FP8 case above it.
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (3)
133-133: Good addition to quantization mode validation.The inclusion of
has_fp8_rowwise()in the supported quantization modes check is consistent with the existing validation pattern and ensures the new fp8_rowwise mode is properly recognized.
149-150: Appropriate method selection for fp8_rowwise quantization.The extension of the condition to include
has_fp8_rowwise()alongsidehas_fp8_qdq()makes sense, as both modes appear to use the sameFP8QDQFusedMoEMethod()for quantization handling.
231-231: Consistent input quantization handling for fp8_rowwise.The addition of
has_fp8_rowwiseto the condition for per-tensor static quantization aligns with treating this mode similarly tohas_fp8_qdq, maintaining consistency in the quantization flow.tensorrt_llm/_torch/model_config.py (1)
272-286: Comprehensive Llama4 FP8 quantization support implementation.The new conditional branch properly handles the "compressed-tensors" with "float-quantized" format. The exclude_modules processing correctly transforms HF module names to TRT-LLM equivalents by removing the "language_model." prefix and mapping "q_proj" to "qkv_proj". The assignment of
use_meta_recipe = Trueenables the META_FP8_RECIPE weight loading mode as intended.tensorrt_llm/_torch/modules/fused_moe/quantization.py (1)
103-106: Weight loading implementation follows established patterns.The new
META_FP8_RECIPEcase correctly implements the weight loading logic using Meta's naming conventions (gate_proj,up_proj,down_projvsw1,w3,w2).
| elif module.weight_loading_mode == MoEWeightLoadingMode.META_FP8_RECIPE: | ||
| # HACK: dynamic quantization so no input scale - create empty tensors for testing | ||
| w1_input_scale = torch.tensor(1.0, dtype=torch.float32) | ||
| w3_input_scale = torch.tensor(1.0, dtype=torch.float32) | ||
| w2_input_scale = torch.tensor(1.0, dtype=torch.float32) |
There was a problem hiding this comment.
Replace hardcoded scale values with proper dynamic quantization implementation.
The current implementation uses hardcoded 1.0 values as a workaround for dynamic quantization. This approach could lead to incorrect quantization behavior and reduced model accuracy.
Consider implementing proper dynamic quantization scale computation instead:
elif module.weight_loading_mode == MoEWeightLoadingMode.META_FP8_RECIPE:
- # HACK: dynamic quantization so no input scale - create empty tensors for testing
- w1_input_scale = torch.tensor(1.0, dtype=torch.float32)
- w3_input_scale = torch.tensor(1.0, dtype=torch.float32)
- w2_input_scale = torch.tensor(1.0, dtype=torch.float32)
+ # TODO: Implement proper dynamic quantization scale computation
+ # For now, compute scales dynamically or load from appropriate keys
+ w1_input_scale = self._compute_dynamic_input_scale(weights, expert_id, "gate_proj")
+ w3_input_scale = self._compute_dynamic_input_scale(weights, expert_id, "up_proj")
+ w2_input_scale = self._compute_dynamic_input_scale(weights, expert_id, "down_proj")Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/modules/fused_moe/quantization.py around lines 371 to
375, the code uses hardcoded scale values of 1.0 for dynamic quantization, which
is a placeholder and can cause incorrect quantization results. Replace these
hardcoded tensors with a proper dynamic quantization scale computation by
calculating the input scales based on the actual data or model parameters to
ensure accurate quantization behavior.
| elif module.weight_loading_mode == MoEWeightLoadingMode.META_FP8_RECIPE: | ||
| # HACK: weight scale is not the right shape in blockwise recipe - create empty tensors for testing | ||
| w1_weight_scale = torch.tensor(1.0, dtype=torch.float32) | ||
| w3_weight_scale = torch.tensor(1.0, dtype=torch.float32) | ||
| w2_weight_scale = torch.tensor(1.0, dtype=torch.float32) |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Address the underlying shape mismatch instead of using hardcoded values.
The hardcoded weight scales could lead to suboptimal quantization. The comment indicates a shape mismatch with the "blockwise recipe" that should be addressed properly.
Consider addressing the root cause:
elif module.weight_loading_mode == MoEWeightLoadingMode.META_FP8_RECIPE:
- # HACK: weight scale is not the right shape in blockwise recipe - create empty tensors for testing
- w1_weight_scale = torch.tensor(1.0, dtype=torch.float32)
- w3_weight_scale = torch.tensor(1.0, dtype=torch.float32)
- w2_weight_scale = torch.tensor(1.0, dtype=torch.float32)
+ # Load and reshape weight scales to match expected format
+ w1_weight_scale = self._load_and_reshape_weight_scale(weights, expert_id, "gate_proj")
+ w3_weight_scale = self._load_and_reshape_weight_scale(weights, expert_id, "up_proj")
+ w2_weight_scale = self._load_and_reshape_weight_scale(weights, expert_id, "down_proj")Let me verify if there are existing implementations that handle dynamic quantization properly:
🏁 Script executed:
#!/bin/bash
# Search for dynamic quantization implementations and scale computation methods
rg -A 10 -B 5 "dynamic.*quantization|compute.*scale" --type pyLength of output: 45171
Replace hardcoded FP8 weight‐scale hack with proper loading and reshaping
The current META_FP8_RECIPE branch (lines 409–413 in
tensorrt_llm/_torch/modules/fused_moe/quantization.py) uses placeholder scalars for all three weight scales. This bypasses actual per-expert, blockwise scale data and will degrade quantization accuracy. Please implement a loader that:
- Reads each expert’s stored weight scale (
weights["{expert_id}.w1.weight_scale"], etc.) - Reshapes it to match the blockwise partitioning used by the fused MoE kernels
- Registers or returns the correctly shaped tensor for
w1_weight_scale,w3_weight_scale, andw2_weight_scale
For example, instead of:
- # HACK: weight scale is not the right shape in blockwise recipe - create empty tensors for testing
- w1_weight_scale = torch.tensor(1.0, dtype=torch.float32)
- w3_weight_scale = torch.tensor(1.0, dtype=torch.float32)
- w2_weight_scale = torch.tensor(1.0, dtype=torch.float32)
+ # TODO: load and reshape the per-expert FP8 weight scales for blockwise META_FP8_RECIPE
+ raw_w1 = weights[f"{expert_id}.w1.weight_scale"]
+ raw_w3 = weights[f"{expert_id}.w3.weight_scale"]
+ raw_w2 = weights[f"{expert_id}.w2.weight_scale"]
+ w1_weight_scale = raw_w1.view(expected_block_shape).to(torch.float32)
+ w3_weight_scale = raw_w3.view(expected_block_shape).to(torch.float32)
+ w2_weight_scale = raw_w2.view(expected_block_shape).to(torch.float32)• File: tensorrt_llm/_torch/modules/fused_moe/quantization.py
• Lines: ~409–413
Implement and test this loader against your blockwise FP8 metadata to ensure correct shapes and quantization fidelity.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In tensorrt_llm/_torch/modules/fused_moe/quantization.py around lines 409 to
413, replace the hardcoded scalar tensors used as placeholders for
w1_weight_scale, w2_weight_scale, and w3_weight_scale with actual loading logic.
Implement code to read each expert's stored weight scale from weights using keys
like "{expert_id}.w1.weight_scale", reshape these tensors to match the blockwise
partitioning expected by the fused MoE kernels, and assign or return these
properly shaped tensors instead of the placeholders. Ensure the loader correctly
handles all experts and weight scales to maintain quantization accuracy.
Summary by CodeRabbit
New Features
Bug Fixes
Refactor
[feat] Add support to load fp8 Meta Llama4 weights
Description
TODO
Test Coverage
TODO
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.