Skip to content

[Spec] Build DFlash draft layers under their checkpoint names and refuse tensors the draft would drop #40884

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
twu3202 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Spec] Build DFlash draft layers under their checkpoint names and refuse tensors the draft would drop #40884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
123 changes: 92 additions & 31 deletions python/sglang/srt/models/dflash.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
Expand All @@ -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"
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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 "
Expand All @@ -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",
Expand Down Expand Up @@ -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. "
Expand All @@ -1029,15 +1090,15 @@ 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)
for i, norm in enumerate(self.aux_hidden_norms):
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)

Expand Down
Loading
Loading