Skip to content

[FP8] Add opt-in ParallelLMHead dispatch to Fp8Config #41000

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
webcodes-cz wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[FP8] Add opt-in ParallelLMHead dispatch to Fp8Config #41000

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
72238d4
[FP8] Add opt-in ParallelLMHead quantization to Fp8Config
webcodes-cz Apr 27, 2026
3520a71
Gap 3: weight loaders for FP8 ParallelLMHead
webcodes-cz Apr 27, 2026
5ca163e
Merge branch 'main' into fp8-parallel-lm-head-opt-in
webcodes-cz Apr 27, 2026
511412b
Merge branch 'main' into fp8-parallel-lm-head-opt-in
webcodes-cz Apr 28, 2026
49a8a6b
Merge branch 'main' into fp8-parallel-lm-head-opt-in
webcodes-cz Apr 28, 2026
de77662
Merge branch 'main' into fp8-parallel-lm-head-opt-in
webcodes-cz Apr 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion vllm/model_executor/layers/quantization/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@
LinearMethodBase,
UnquantizedLinearMethod,
)
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
UnquantizedEmbeddingMethod,
)
Comment on lines +41 to +44

@gemini-code-assist gemini-code-assist Bot Apr 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Importing ParallelLMHead and UnquantizedEmbeddingMethod at the top level of fp8.py from vllm.model_executor.layers.vocab_parallel_embedding may lead to circular import issues in the future, as quantization configs are often imported by the layers they configure. It is generally safer to perform these imports inside get_quant_method or use TYPE_CHECKING for type hints and importlib for runtime checks if necessary.

from vllm.model_executor.layers.quantization import QuantizationMethods
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
Expand Down Expand Up @@ -102,10 +106,12 @@ def __init__(
activation_scheme: str = "dynamic",
ignored_layers: list[str] | None = None,
weight_block_size: list[int] | None = None,
lm_head_quantized: bool = False,
) -> None:
super().__init__()

self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
self.lm_head_quantized = lm_head_quantized

if activation_scheme not in ACTIVATION_SCHEMES:
raise ValueError(f"Unsupported activation scheme {activation_scheme}")
Expand Down Expand Up @@ -162,22 +168,29 @@ def from_config(cls, config: dict[str, Any]) -> "Fp8Config":
ignored_layers = cls.get_from_keys_or(
config, ["modules_to_not_convert"], None
)
lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
return cls(
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
activation_scheme=activation_scheme,
ignored_layers=ignored_layers,
weight_block_size=weight_block_size,
lm_head_quantized=lm_head_quantized,
)

def get_quant_method(
self, layer: torch.nn.Module, prefix: str
) -> "QuantizeMethodBase | None":
if isinstance(layer, LinearBase):
is_parallel_lm_head = isinstance(layer, ParallelLMHead)
if isinstance(layer, LinearBase) or (
is_parallel_lm_head and self.lm_head_quantized
):
Comment on lines +183 to +186

@gemini-code-assist gemini-code-assist Bot Apr 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current implementation of get_quant_method returns Fp8LinearMethod (or Fp8OnlineLinearMethod) for ParallelLMHead. However, Fp8LinearMethod is designed for LinearBase modules and does not implement the embedding method required by VocabParallelEmbedding (the base class of ParallelLMHead). While ParallelLMHead overrides forward to raise a RuntimeError, any code path that might attempt to use it as a standard embedding layer (e.g., if weights are tied and accessed via the embedding interface) will fail with a NotImplementedError.

Furthermore, as noted in the PR description, VocabParallelEmbedding.weight_loader does not currently handle the companion parameters (like weight_scale) created by Fp8LinearMethod. Returning a linear method for an embedding-sharded layer without ensuring the loader and interface compatibility is a high-risk change.

if is_layer_skipped(
prefix=prefix,
ignored_layers=self.ignored_layers,
fused_mapping=self.packed_modules_mapping,
):
if is_parallel_lm_head:
return UnquantizedEmbeddingMethod()
return UnquantizedLinearMethod()
if not self.is_checkpoint_fp8_serialized:
online_method = Fp8OnlineLinearMethod(self)
Expand Down
1 change: 1 addition & 0 deletions vllm/model_executor/models/qwen3_5.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -501,6 +501,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.lm_head = ParallelLMHead(
config.vocab_size,
config.hidden_size,
quant_config=self.quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
else:
Expand Down
Loading