[config] Recognize custom hybrid SWA models via hf_text_config.is_hybrid_swa (#23988)
This commit is contained in:
@@ -504,7 +504,7 @@ class ModelConfig:
|
|||||||
def _derive_hybrid_model(self):
|
def _derive_hybrid_model(self):
|
||||||
# Use self.context_len after it has been initialized to prevent using context_len which may be None.
|
# Use self.context_len after it has been initialized to prevent using context_len which may be None.
|
||||||
self.is_hybrid_swa = (
|
self.is_hybrid_swa = (
|
||||||
is_hybrid_swa_model(self.hf_config.architectures)
|
is_hybrid_swa_model(self.hf_config.architectures, self.hf_text_config)
|
||||||
and not self.disable_hybrid_swa_memory
|
and not self.disable_hybrid_swa_memory
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1671,7 +1671,10 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float
|
|||||||
return base_scaling * mscale * mscale
|
return base_scaling * mscale * mscale
|
||||||
|
|
||||||
|
|
||||||
def is_hybrid_swa_model(model_architectures: List[str]):
|
def is_hybrid_swa_model(
|
||||||
|
model_architectures: List[str],
|
||||||
|
hf_text_config: Optional[PretrainedConfig] = None,
|
||||||
|
):
|
||||||
|
|
||||||
hybrid_swa_archs = {
|
hybrid_swa_archs = {
|
||||||
"Llama4ForConditionalGeneration",
|
"Llama4ForConditionalGeneration",
|
||||||
@@ -1687,7 +1690,13 @@ def is_hybrid_swa_model(model_architectures: List[str]):
|
|||||||
"Gemma4ForConditionalGeneration",
|
"Gemma4ForConditionalGeneration",
|
||||||
"LagunaForCausalLM",
|
"LagunaForCausalLM",
|
||||||
}
|
}
|
||||||
return any(arch in hybrid_swa_archs for arch in model_architectures)
|
if any(arch in hybrid_swa_archs for arch in model_architectures):
|
||||||
|
return True
|
||||||
|
# Also recognize models that explicitly opt-in via their HF text config,
|
||||||
|
# so custom hybrid-SWA architectures don't need to be added to the allowlist.
|
||||||
|
if hf_text_config is not None and getattr(hf_text_config, "is_hybrid_swa", False):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def get_hybrid_layer_ids(
|
def get_hybrid_layer_ids(
|
||||||
@@ -1758,6 +1767,17 @@ def get_hybrid_layer_ids(
|
|||||||
full_attention_layer_ids = [
|
full_attention_layer_ids = [
|
||||||
i for i, x in enumerate(layer_types) if x == "full_attention"
|
i for i, x in enumerate(layer_types) if x == "full_attention"
|
||||||
]
|
]
|
||||||
|
elif getattr(hf_text_config, "hybrid_layer_pattern", None) is not None:
|
||||||
|
# Generic fallback for custom hybrid SWA models that opt in via
|
||||||
|
# hf_text_config.is_hybrid_swa and expose a hybrid_layer_pattern
|
||||||
|
# (1 = SWA, 0 = full) without needing to be added to the allowlist.
|
||||||
|
hybrid_layer_pattern = hf_text_config.hybrid_layer_pattern
|
||||||
|
swa_attention_layer_ids = [
|
||||||
|
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1
|
||||||
|
]
|
||||||
|
full_attention_layer_ids = [
|
||||||
|
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0
|
||||||
|
]
|
||||||
else:
|
else:
|
||||||
swa_attention_layer_ids = None
|
swa_attention_layer_ids = None
|
||||||
full_attention_layer_ids = None
|
full_attention_layer_ids = None
|
||||||
|
|||||||
Reference in New Issue
Block a user