diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index c8ce6af88a48..99f0cde583a7 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -626,6 +626,12 @@ class DFlashDraftModel(nn.Module): supports_fused_context_kv = True # Layer prefixes and a quant_config-aware `fc`, for quantized checkpoints. supports_quantization = False + # Checkpoints store q/k/v and gate/up as separate tensors; a quantization + # config that names those shards is matched against the fused modules. + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } def __init__(self, config, quant_config=None, prefix: str = "") -> None: super().__init__() @@ -642,7 +648,13 @@ def __init__(self, config, quant_config=None, prefix: str = "") -> None: self.candidate_selector: Optional[nn.Module] = None self.lilicorr: Optional[nn.Module] = None self.is_nemotron_35_draft = is_nemotron_35_draft_config(config) - self.quantizable = self.supports_quantization or self.is_nemotron_35_draft + # A draft whose own checkpoint is quantized also qualifies; a target + # --quantization the draft only inherits leaves it as before. + self.quantizable = ( + self.supports_quantization + or self.is_nemotron_35_draft + or getattr(config, "quantization_config", None) is not None + ) self.embed_tokens: Optional[VocabParallelEmbedding] = None if self.is_nemotron_35_draft: embed_prefix = f"{prefix}.embed_tokens" if prefix else "embed_tokens" @@ -832,30 +844,75 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ] params_dict = dict(self.named_parameters()) + modules_dict = dict(self.named_modules()) loaded_params = set() # Alias the native export's "encoder." names. _VENDOR_ENCODER_ALIASES = { - "encoder.fc.weight": "fc.weight", "encoder.output_norm_enc.weight": "hidden_norm.weight", } + # Checkpoint tensors that have no parameter on the draft by design. + _IGNORED_TENSOR_SUFFIXES = ( + ".bias", + ".inv_freq", + ".cos_cached", + ".sin_cached", + ".input_scale", + ".output_scale", + ".k_scale", + ".v_scale", + ".kv_scale", + ) - def resolve_param_name(name: str) -> Optional[str]: - if name in params_dict: + def resolve_name(name: str, known: dict) -> Optional[str]: + if name in known: return name if name.startswith("model."): stripped_name = name[len("model.") :] - if stripped_name in params_dict: + if stripped_name in known: return stripped_name else: prefixed_name = f"model.{name}" - if prefixed_name in params_dict: + if prefixed_name in known: return prefixed_name + return None + + def alias_vendor_name(name: str) -> str: aliased_name = _VENDOR_ENCODER_ALIASES.get(name) - if aliased_name is not None and aliased_name in params_dict: + if aliased_name is not None: + return aliased_name + if name.startswith("encoder.fc."): + return name[len("encoder.") :] + return name + + def resolve_param_name(name: str) -> Optional[str]: + resolved_name = resolve_name(name, params_dict) + if resolved_name is not None: + return resolved_name + aliased_name = alias_vendor_name(name) + if aliased_name != name and aliased_name in params_dict: return aliased_name return None + def reject_unmatched(*, name: str, param_name: str) -> None: + """Refuse a tensor for a module the draft has, under a parameter + name that module does not have (a dense tensor for a packed layer + or the reverse); tensors for absent modules stay ignored.""" + if name.endswith(_IGNORED_TENSOR_SUFFIXES): + return + module_stem = alias_vendor_name(param_name).rsplit(".", 1)[0] + module_name = resolve_name(module_stem, modules_dict) + if module_name is None: + return + module = modules_dict[module_name] + raise ValueError( + f"DFLASH checkpoint tensor {name!r} has no matching parameter on " + f"{module_name} ({type(module).__name__} with parameters " + f"{sorted(dict(module.named_parameters(recurse=False)))}). The " + "draft was built without it: its config or quantization config " + "does not describe this layer the way the checkpoint stores it." + ) + for name, loaded_weight in weights: unprefixed_name = name.removeprefix("model.") if self.projector_type != "domino" and unprefixed_name.startswith( @@ -871,6 +928,7 @@ def resolve_param_name(name: str) -> Optional[str]: mapped_name = name.replace(weight_name, param_name) resolved_name = resolve_param_name(mapped_name) if resolved_name is None: + reject_unmatched(name=name, param_name=mapped_name) continue param = params_dict[resolved_name] weight_loader = getattr(param, "weight_loader", default_weight_loader) @@ -880,28 +938,23 @@ def resolve_param_name(name: str) -> Optional[str]: else: resolved_name = resolve_param_name(name) if resolved_name is None: - # Ignore unexpected weights (e.g., HF rotary caches). + reject_unmatched(name=name, param_name=name) continue param = params_dict[resolved_name] if resolved_name.endswith("fc.weight"): - if self.is_nemotron_35_draft: - expected_shape = ( - int(self.config.hidden_size), - int(self.num_context_features * self.config.hidden_size), - ) - loaded_shape = _logical_linear_weight_shape( - param, - loaded_weight, - output_features=expected_shape[0], - ) - shape_matches = loaded_shape == expected_shape or ( - getattr(param, "pack_factor", None) is None - and tuple(loaded_weight.shape) == tuple(param.shape) - ) - else: - expected_shape = tuple(param.shape) - loaded_shape = tuple(loaded_weight.shape) - shape_matches = loaded_shape == expected_shape + expected_shape = ( + int(self.config.hidden_size), + int(self.num_context_features * self.config.hidden_size), + ) + loaded_shape = _logical_linear_weight_shape( + param, + loaded_weight, + output_features=expected_shape[0], + ) + shape_matches = loaded_shape == expected_shape or ( + getattr(param, "pack_factor", None) is None + and tuple(loaded_weight.shape) == tuple(param.shape) + ) if not shape_matches: raise ValueError( "DFLASH fc.weight shape mismatch. This usually means the draft checkpoint's " @@ -922,6 +975,16 @@ def resolve_param_name(name: str) -> Optional[str]: weight_loader(param, loaded_weight) loaded_params.add(resolved_name) + # A checkpoint that fills the decoder layers but never fc would leave + # fc at whatever memory its weight was allocated over. + if any(name.startswith("layers.") for name in loaded_params) and not any( + name.startswith("fc.") for name in loaded_params + ): + raise ValueError( + "DFLASH checkpoint has decoder-layer tensors but no fc tensor " + "(fc.weight, encoder.fc.weight, or their packed forms)." + ) + if self.projector_type == "domino": required = { "prefix_gru.weight_ih_l0", @@ -1015,9 +1078,7 @@ def prepare_context_hidden_for_kv( return layer.input_layernorm(ctx_hidden) def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: - expected = int( - self.fc.input_size if self.is_nemotron_35_draft else self.fc.in_features - ) + expected = int(self.fc.input_size if self.quantizable else self.fc.in_features) if target_hidden.ndim != 2 or int(target_hidden.shape[-1]) != expected: raise ValueError( "Laguna DFLASH target_hidden feature dim mismatch. " @@ -1029,7 +1090,7 @@ def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: num_slices = int(self.num_context_features) slice_size = int(target_hidden.shape[-1]) // num_slices slices = target_hidden.view(target_hidden.shape[0], num_slices, slice_size) - compute_dtype = self.fc.weight.dtype + compute_dtype = self.hidden_norm.weight.dtype if slices.dtype != compute_dtype: slices = slices.to(compute_dtype) normed = torch.empty_like(slices) @@ -1037,7 +1098,7 @@ def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: normed[:, i, :] = norm(slices[:, i, :]) fused = normed.reshape(target_hidden.shape[0], -1) projected = self.fc(fused) - if self.is_nemotron_35_draft: + if self.quantizable: projected = projected[0] return self.hidden_norm(projected) diff --git a/test/registered/unit/spec/test_dflash_quantized_draft_layout.py b/test/registered/unit/spec/test_dflash_quantized_draft_layout.py new file mode 100644 index 000000000000..42c0bc2c1803 --- /dev/null +++ b/test/registered/unit/spec/test_dflash_quantized_draft_layout.py @@ -0,0 +1,280 @@ +"""A quantized DFlash draft is built from its quantization config by module +name and loaded by tensor name, and both have to see the same names. With an +empty layer prefix every attention and MLP projection was matched as a bare +``qkv_proj`` / ``o_proj`` / ..., so name-based ``ignore`` entries never +applied and the checkpoint's tensors for those layers were dropped without a +word; ``fc`` was a plain ``nn.Linear`` that could not take packed tensors at +all. These cases pin the layer names, the quantizable ``fc``, and the refusal +to drop a tensor whose module the draft does have. Both apply only to a draft +whose own checkpoint is quantized; an unquantized draft that inherits the +target's ``--quantization`` keeps its dense ``fc``.""" + +import sys +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + CompressedTensorsLinearMethod, +) +from sglang.srt.layers.quantization.compressed_tensors.utils import ( + should_ignore_layer, +) +from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.models import dflash +from sglang.srt.models.dflash import DFlashDecoderLayer, DFlashDraftModel +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + +HIDDEN = 128 # the INT8 group size, so fc's packed weight needs no padding +TARGET_LAYER_IDS = [1, 3] +K = len(TARGET_LAYER_IDS) + + +def _draft_config(num_layers=2, quantized=False): + config = SimpleNamespace( + architectures=["DFlash2DraftModel"], + hidden_size=HIDDEN, + num_hidden_layers=num_layers, + rms_norm_eps=1e-6, + dflash_config={"block_size": 4, "target_layer_ids": TARGET_LAYER_IDS}, + ) + if quantized: + # Only a quantized checkpoint's config.json has this key. + config.quantization_config = {"quant_method": "compressed-tensors"} + return config + + +def _int8_config(ignore=()): + """compressed-tensors pack-quantized INT8 g128 on every Linear, as + llm-compressor writes it.""" + return CompressedTensorsConfig.from_config( + { + "format": "pack-quantized", + "quant_method": "compressed-tensors", + "ignore": list(ignore), + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 8, + "type": "int", + "group_size": 128, + "strategy": "group", + "symmetric": True, + }, + } + }, + } + ) + + +class _RecordingAttention(nn.Module): + """Stands in for DFlashAttention (whose rotary embedding needs a GPU + build): records the prefix it was built with and carries the module shapes + load_weights has to reason about -- a fused qkv_proj holding only a packed + weight, and a rotary module with a parameter next to its inv_freq buffer.""" + + prefixes = [] + + def __init__(self, config, layer_id, quant_config=None, prefix=""): + super().__init__() + type(self).prefixes.append(prefix) + self.qkv_proj = nn.Module() + self.qkv_proj.weight_packed = nn.Parameter( + torch.zeros(4, 4, dtype=torch.int32), requires_grad=False + ) + self.rotary_emb = nn.Module() + self.rotary_emb.cos_coef = nn.Parameter(torch.zeros(2)) + self.rotary_emb.register_buffer("inv_freq", torch.zeros(2)) + + +class _RecordingMLP(nn.Module): + prefixes = [] + + def __init__(self, config, quant_config=None, prefix=""): + super().__init__() + type(self).prefixes.append(prefix) + + +class _Layer(DFlashDecoderLayer): + attention_cls = _RecordingAttention + + +class _Draft(DFlashDraftModel): + decoder_layer_cls = _Layer + + +@pytest.fixture(autouse=True) +def _cpu_build(monkeypatch): + _RecordingAttention.prefixes.clear() + _RecordingMLP.prefixes.clear() + monkeypatch.setattr(dflash, "DFlashMLP", _RecordingMLP) + # Scheme resolution asks the GPU for its compute capability. + monkeypatch.setattr( + CompressedTensorsConfig, + "_check_scheme_supported", + lambda self, min_capability, error=True: True, + ) + + +def _fc_tensors(): + """Dense fc tensors for a draft whose fc is unquantized.""" + return [("fc.weight", torch.randn(HIDDEN, K * HIDDEN))] + + +@pytest.mark.parametrize("model_prefix", ["", "draft"]) +def test_every_layer_is_built_under_its_checkpoint_name(model_prefix): + """Names reach the quantization config through this chain; an empty + prefix at either link turns every projection into a bare name.""" + _Draft(_draft_config(quantized=True), quant_config=None, prefix=model_prefix) + root = f"{model_prefix}." if model_prefix else "" + assert _RecordingAttention.prefixes == [ + f"{root}layers.0.self_attn", + f"{root}layers.1.self_attn", + ] + assert _RecordingMLP.prefixes == [f"{root}layers.0.mlp", f"{root}layers.1.mlp"] + + +def test_shard_only_ignore_reaches_the_fused_projection(): + """An ignore list that names q_proj/k_proj/v_proj (how llm-compressor + writes it) applies to the fused qkv_proj only through the mapping the + loader takes from the model class.""" + ignore = [ + r"re:.*self_attn\.q_proj$", + r"re:.*self_attn\.k_proj$", + r"re:.*self_attn\.v_proj$", + ] + name = "layers.0.self_attn.qkv_proj" + assert should_ignore_layer( + name, ignore=ignore, fused_mapping=DFlashDraftModel.packed_modules_mapping + ) + assert not should_ignore_layer(name, ignore=ignore) + + +def test_fc_is_quantized_by_a_config_that_targets_linear(): + model = _Draft(_draft_config(quantized=True), quant_config=_int8_config()) + assert isinstance(model.fc, ReplicatedLinear) + assert isinstance(model.fc.quant_method, CompressedTensorsLinearMethod) + assert model.fc.weight_packed.shape == (HIDDEN, K * HIDDEN // 4) + assert not hasattr(model.fc, "weight") + + +def test_fc_stays_dense_when_ignored_or_unquantized(): + ignored = _Draft( + _draft_config(quantized=True), quant_config=_int8_config(ignore=["re:^fc$"]) + ) + assert isinstance(ignored.fc.quant_method, UnquantizedLinearMethod) + plain = _Draft(_draft_config(), quant_config=None) + assert plain.fc.weight.shape == (HIDDEN, K * HIDDEN) + + +def test_inherited_online_quantization_leaves_an_unquantized_draft_dense(): + """An unquantized draft inherits the target's --quantization (online FP8 + here). Its fc must stay a dense nn.Linear and its layers unnamed, as on + main; giving it a quantizable fc turned fc into FP8 too.""" + online_fp8 = Fp8Config(is_checkpoint_fp8_serialized=False) + model = _Draft(_draft_config(), quant_config=online_fp8, prefix="draft") + assert type(model.fc) is nn.Linear + assert model.fc.weight.shape == (HIDDEN, K * HIDDEN) + assert _RecordingAttention.prefixes == ["", ""] + assert _RecordingMLP.prefixes == ["", ""] + + +def test_packed_fc_tensors_load(): + model = _Draft(_draft_config(quantized=True), quant_config=_int8_config()) + packed = torch.randint( + -(2**31), 2**31 - 1, model.fc.weight_packed.shape, dtype=torch.int32 + ) + scale = torch.randn(HIDDEN, K, dtype=torch.bfloat16) + + model.load_weights( + [ + ("fc.weight_packed", packed), + ("fc.weight_scale", scale), + ("fc.weight_shape", torch.tensor([HIDDEN, K * HIDDEN])), + ] + ) + + torch.testing.assert_close(model.fc.weight_packed.data, packed) + torch.testing.assert_close( + model.fc.weight_scale.data, scale.to(model.fc.weight_scale.dtype) + ) + + +@pytest.mark.parametrize("name", ["fc.weight", "encoder.fc.weight"]) +def test_dense_fc_loads_under_its_own_and_the_vendor_name(name): + model = _Draft(_draft_config(), quant_config=None) + weight = torch.randn(HIDDEN, K * HIDDEN) + model.load_weights([(name, weight)]) + torch.testing.assert_close(model.fc.weight.data, weight) + + +def test_dense_fc_shape_is_still_checked(): + model = _Draft(_draft_config(), quant_config=None) + with pytest.raises(ValueError, match="fc.weight shape mismatch"): + model.load_weights([("fc.weight", torch.randn(HIDDEN, HIDDEN))]) + + +def test_loading_refuses_a_tensor_the_module_stores_differently(): + """A tensor for a module the draft has, under a parameter name that module + was not built with, must not be dropped: a BF16 q/k/v tensor against a + fused qkv_proj built packed left the packed weight uninitialized.""" + model = _Draft(_draft_config(), quant_config=None) + with pytest.raises(ValueError, match="layers.0.self_attn.qkv_proj"): + model.load_weights( + [("layers.0.self_attn.q_proj.weight", torch.randn(4, HIDDEN))] + ) + # The export's "model." prefix resolves to the same module. + with pytest.raises(ValueError, match="layers.0.self_attn.qkv_proj"): + model.load_weights( + [("model.layers.0.self_attn.k_proj.weight", torch.randn(4, HIDDEN))] + ) + # A packed fc tensor for a draft whose fc is dense, under either name. + for name in ("fc.weight_packed", "encoder.fc.weight_packed"): + with pytest.raises(ValueError, match=rf"'{name}'.*\bfc\b"): + model.load_weights([(name, torch.zeros(4, 4, dtype=torch.int32))]) + + +def test_loading_refuses_decoder_layers_without_an_fc(): + """A checkpoint that fills the layers but never fc would leave fc over + uninitialized memory; a projector-only load (Domino) stays allowed.""" + model = _Draft(_draft_config(), quant_config=None) + with pytest.raises(ValueError, match="no fc tensor"): + model.load_weights([("layers.0.input_layernorm.weight", torch.ones(HIDDEN))]) + model.load_weights( + _fc_tensors() + [("layers.0.input_layernorm.weight", torch.ones(HIDDEN))] + ) + + +def test_loading_still_ignores_tensors_the_draft_has_no_parameter_for(): + """Tensors that have no parameter on the draft by design: rotary caches, + KV-cache and activation scales, biases the layer was built without, and + modules the draft does not have at all.""" + model = _Draft(_draft_config(), quant_config=None) + model.load_weights( + _fc_tensors() + + [ + ("layers.0.self_attn.rotary_emb.inv_freq", torch.zeros(2)), + ("layers.0.self_attn.rotary_emb.cos_cached", torch.zeros(2, 2)), + ("layers.0.self_attn.rotary_emb.sin_cached", torch.zeros(2, 2)), + ("layers.0.self_attn.k_proj.output_scale", torch.ones(1)), + ("layers.0.self_attn.k_proj.k_scale", torch.ones(1)), + ("layers.0.self_attn.v_proj.v_scale", torch.ones(1)), + ("layers.0.self_attn.q_proj.input_scale", torch.ones(1)), + ("layers.0.self_attn.q_proj.bias", torch.zeros(4)), + ("lm_head.weight", torch.zeros(4, HIDDEN)), + ("embed_tokens.weight", torch.zeros(4, HIDDEN)), + ("layers.5.self_attn.q_proj.weight", torch.zeros(4, HIDDEN)), + ] + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))