File size: 3,174 Bytes
3599318 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | # Agnes 3.0 Flash configuration for sglang.
#
# The HF checkpoint says model_type "agnes", names its layer types
# agnes_delta_attention / agnes_global_attention and carries a parallel FFN
# branch per layer. The server runs it on its built-in hybrid
# (delta-rule + global attention) implementation, so this class maps those
# onto the fields that implementation reads; the checkpoint's tensor names
# are translated while loading (see the model file patched by apply_patch.py).
from sglang.srt.configs.qwen3_5 import Qwen3_5Config
AGNES_DELTA = "agnes_delta_attention"
AGNES_GLOBAL = "agnes_global_attention"
class AgnesConfig(Qwen3_5Config):
model_type = "agnes"
def __init__(self, text_config=None, vision_config=None, **kwargs):
kwargs.pop("auto_map", None) # the transformers remote code is not used in the server
if isinstance(text_config, dict):
text_config = dict(text_config)
text_config["model_type"] = "qwen3_5_text"
width = int(text_config.pop("parallel_ffn_intermediate_size", 0) or 0)
plan = text_config.pop("layer_types", None)
interval = text_config.pop("global_attention_interval", None)
if interval is None:
interval = text_config.pop("full_attention_interval", None)
if interval is None and plan:
interval = next(i + 1 for i, t in enumerate(plan) if t == AGNES_GLOBAL)
text_config["full_attention_interval"] = int(interval or 4)
main = int(text_config["intermediate_size"])
# the parallel branch is folded into the main MLP at load time
text_config["intermediate_size"] = main + width
text_config["agnes_main_intermediate_size"] = main
text_config["agnes_parallel_ffn_intermediate_size"] = width
if isinstance(vision_config, dict):
vision_config = dict(vision_config)
vision_config["model_type"] = "qwen3_5"
kwargs["architectures"] = ["Qwen3_5ForConditionalGeneration"]
super().__init__(text_config=text_config, vision_config=vision_config, **kwargs)
# every downstream check sees the built-in hybrid architecture
self.model_type = "qwen3_5"
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
# The weight loader needs the checkpoint directory to pick up the parallel
# branch tensors. The path is written into the config *dict* before the
# object is built: the config reaches the worker processes through a
# to_dict round trip, which keeps fields that came in through __init__
# and drops attributes set afterwards (from_pretrained's own kwargs only
# override known fields, so they cannot carry it either).
path = str(pretrained_model_name_or_path)
config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
config_dict["agnes_model_path"] = path
if isinstance(config_dict.get("text_config"), dict):
config_dict["text_config"]["agnes_model_path"] = path
return cls.from_dict(config_dict, **kwargs)
|