diff --git a/examples/models/llama/BUCK b/examples/models/llama/BUCK index e4cdc2e3c12..0bfb74cd24f 100644 --- a/examples/models/llama/BUCK +++ b/examples/models/llama/BUCK @@ -32,6 +32,7 @@ fbcode_target(_kind = runtime.python_library, ":transformer_modules", "//caffe2:torch", "//executorch/examples/models/lfm2:lfm2", + "//executorch/examples/models/spark_x2_5:spark_x2_5", ], ) diff --git a/examples/models/llama/attention.py b/examples/models/llama/attention.py index 98a838bfa30..95321a03515 100644 --- a/examples/models/llama/attention.py +++ b/examples/models/llama/attention.py @@ -418,7 +418,17 @@ def __init__( self.enable_dynamic_shape = args.enable_dynamic_shape self.scale_query_by = args.scale_query_by self.use_attn_o_gate = args.use_attn_o_gate + self.headwise_attn_output_gate = args.headwise_attn_output_gate + if self.use_attn_o_gate and self.headwise_attn_output_gate: + raise ValueError( + "use_attn_o_gate and headwise_attn_output_gate are mutually exclusive" + ) self.use_attn_o_norm = args.use_attn_o_norm + self.is_sliding = ( + args.layer_types is not None + and layer_id < len(args.layer_types) + and args.layer_types[layer_id] == "sliding_attention" + ) q_out_dim = self.n_heads * self.head_dim * (2 if self.use_q_gate else 1) # YOCO: Determine if this is a KV shared layer (receives shared KV from donor). @@ -481,6 +491,8 @@ def _init_norms(self, args: ModelArgs) -> None: self.o_norm = ScalelessRMSNorm(self.head_dim, eps=args.norm_eps) if self.use_attn_o_gate: self.og = nn.Linear(args.dim, self.n_heads * self.head_dim, bias=False) + if self.headwise_attn_output_gate: + self.og = nn.Linear(args.dim, self.n_local_heads, bias=False) def _init_projections(self, args: ModelArgs, q_out_dim: int) -> None: """Initialize Q/K/V/O projection layers.""" @@ -509,13 +521,22 @@ def _init_projections(self, args: ModelArgs, q_out_dim: int) -> None: def _init_kv_cache(self, args: ModelArgs) -> None: """Initialize KV cache (only for non-shared layers).""" if self.has_kv_weights: - self.kv_cache = KVCache( - args.max_batch_size, - args.max_context_len, - self.n_kv_heads, - self.head_dim, - args.enable_dynamic_shape, - ) + if self.is_sliding and args.sliding_window: + self.kv_cache = RingKVCache( + args.max_batch_size, + args.sliding_window, + self.n_kv_heads, + self.head_dim, + args.enable_dynamic_shape, + ) + else: + self.kv_cache = KVCache( + args.max_batch_size, + args.max_context_len, + self.n_kv_heads, + self.head_dim, + args.enable_dynamic_shape, + ) else: self.kv_cache = None @@ -679,6 +700,11 @@ def _apply_output_transforms( og = self.og(x).view(bsz, seqlen, self.n_local_heads, self.head_dim) output_4d = torch.sigmoid(og) * output_4d output = output_4d.reshape(bsz, seqlen, -1) + if self.headwise_attn_output_gate: + output_4d = output.view(bsz, seqlen, self.n_local_heads, self.head_dim) + og = self.og(x).unsqueeze(-1).to(output_4d.dtype) + output_4d = torch.sigmoid(og) * output_4d + output = output_4d.reshape(bsz, seqlen, -1) if gate is not None: output = output * torch.sigmoid(gate) return output diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 07240b11d8c..ac9d07bcf50 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -121,6 +121,8 @@ "lfm2_1_2b", # hybrid "lfm2_5_350m", # hybrid "lfm2_5_1_2b", # hybrid + "spark_x2_5_1_7b", # hybrid + "spark_x2_5_4b", # hybrid ] TORCHTUNE_DEFINED_MODELS = ["llama3_2_vision"] HUGGING_FACE_REPO_IDS = { @@ -141,6 +143,8 @@ "lfm2_1_2b": "LiquidAI/LFM2-1.2B", "lfm2_5_350m": "LiquidAI/LFM2.5-350M", "lfm2_5_1_2b": "LiquidAI/LFM2.5-1.2B-Instruct", + "spark_x2_5_1_7b": "XHToken/Spark-X2.5-1.7B", + "spark_x2_5_4b": "XHToken/Spark-X2.5-4B", } @@ -718,6 +722,8 @@ def export_llama( # noqa: C901 from executorch.examples.models.smollm2 import convert_weights elif model_name.startswith("lfm2"): from executorch.examples.models.lfm2 import convert_weights + elif model_name.startswith("spark_x2_5"): + from executorch.examples.models.spark_x2_5 import convert_weights else: raise ValueError( f"Converting weights to meta format for {model_name} is not yet supported" diff --git a/examples/models/llama/feed_forward.py b/examples/models/llama/feed_forward.py index b23429b89f7..796894e73c3 100644 --- a/examples/models/llama/feed_forward.py +++ b/examples/models/llama/feed_forward.py @@ -6,14 +6,15 @@ class FeedForward(nn.Module): - def __init__(self, dim: int, hidden_dim: int): + def __init__(self, dim: int, hidden_dim: int, act_fn=F.silu): super().__init__() self.w1 = nn.Linear(dim, hidden_dim, bias=False) self.w2 = nn.Linear(hidden_dim, dim, bias=False) self.w3 = nn.Linear(dim, hidden_dim, bias=False) + self.act_fn = act_fn def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) + return self.w2(self.act_fn(self.w1(x)) * self.w3(x)) class LoRAFeedForward(nn.Module): diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py index 73117826708..c5b5b693a60 100644 --- a/examples/models/llama/llama_transformer.py +++ b/examples/models/llama/llama_transformer.py @@ -215,7 +215,11 @@ def __init__( ): self.feed_forward = LoRAFeedForward(args.dim, args.hidden_dim, args) else: - self.feed_forward = FeedForward(dim=args.dim, hidden_dim=args.hidden_dim) + self.feed_forward = FeedForward( + dim=args.dim, + hidden_dim=args.hidden_dim, + act_fn=args.act_fn.get_function(), + ) if isinstance(self.attention, AttentionSkip): self.attention_norm = nn.Identity() @@ -353,6 +357,7 @@ def __init__(self, params: ModelArgs, layers: nn.ModuleList, rope: Rope): self.output_prune_map = params.output_prune_map # YOCO (You Only Cache Once) KV sharing configuration. self.num_kv_shared_layers = params.num_kv_shared_layers + self.layer_types = params.layer_types def _forward_layers( self, @@ -361,6 +366,7 @@ def _forward_layers( freqs_sin: torch.Tensor, attn_options_: Dict, seqlen: int, + freqs_by_type: Optional[Dict[str, Tuple[torch.Tensor, torch.Tensor]]] = None, ) -> Tuple[torch.Tensor, Optional[Any]]: """Run transformer layers with YOCO KV sharing support.""" attn_options_update = None @@ -379,7 +385,14 @@ def _forward_layers( if donor_idx in shared_kv: attn_options_["shared_kv"] = shared_kv[donor_idx] - h, attn_options_update = layer(h, freqs_cos, freqs_sin, attn_options_) + # Per-layer-type RoPE: select freqs based on layer type when available. + l_cos, l_sin = freqs_cos, freqs_sin + if freqs_by_type is not None and self.layer_types is not None: + layer_type = self.layer_types[layer_idx] + if layer_type in freqs_by_type: + l_cos, l_sin = freqs_by_type[layer_type] + + h, attn_options_update = layer(h, l_cos, l_sin, attn_options_) if _is_kv_donor_layer(layer_idx, self.n_layers, self.num_kv_shared_layers): assert ( @@ -421,10 +434,23 @@ def forward( attn_options.get("input_pos"), seqlen ) + # Compute per-layer-type freqs when per-layer RoPE is configured. + freqs_by_type = None + if hasattr(self, "ropes"): + input_pos = attn_options.get("input_pos") + freqs_by_type = { + lt: r.get_freqs(input_pos, seqlen) for lt, r in self.ropes.items() + } + attn_options_ = attn_options.copy() if attn_options is not None else {} h, attn_options_update = self._forward_layers( - h, freqs_cos, freqs_sin, attn_options_, seqlen + h, + freqs_cos, + freqs_sin, + attn_options_, + seqlen, + freqs_by_type=freqs_by_type, ) if not self.generate_full_logits: @@ -463,11 +489,35 @@ def forward( return logits +def _build_ropes(model_args: ModelArgs) -> Tuple[Rope, Dict[str, Rope]]: + """Build Rope instances, creating per-layer-type ropes when rope_parameters is set. + + Returns (default_rope, ropes_by_type). ropes_by_type is empty when no + per-layer-type configuration is provided. + """ + import copy as _copy + + if not model_args.rope_parameters: + return Rope(model_args), {} + + ropes: Dict[str, Rope] = {} + for layer_type, rope_params in model_args.rope_parameters.items(): + rope_args = _copy.copy(model_args) + if "rope_theta" in rope_params: + rope_args.rope_theta = rope_params["rope_theta"] + rope_args.rope_freq_base = rope_params["rope_theta"] + if "partial_rotary_factor" in rope_params: + rope_args.partial_rotary_factor = rope_params["partial_rotary_factor"] + ropes[layer_type] = Rope(rope_args) + return next(iter(ropes.values())), ropes + + def construct_transformer(model_args: ModelArgs) -> Transformer: """ Construct a Transformer model from the given model arguments. """ - rope = Rope(model_args) + rope, ropes = _build_ropes(model_args) + if model_args.attention_type not in ATTENTION_REGISTRY: raise ValueError( f"Unknown attention type: {model_args.attention_type}. " @@ -517,12 +567,20 @@ def construct_transformer(model_args: ModelArgs) -> Transformer: ) layers.append(transformer_block) else: + # Select per-layer-type RoPE when available. + layer_rope = rope + if ropes and model_args.layer_types: + layer_type = model_args.layer_types[layer_id] + layer_rope = ropes.get(layer_type, rope) attention = cls( - model_args, layer_id, rope, **model_args.attention_kwargs + model_args, layer_id, layer_rope, **model_args.attention_kwargs ) # pyre-ignore[45] transformer_block = TransformerBlock( model_args, attention, layer_id=layer_id ) layers.append(transformer_block) - return Transformer(model_args, layers, rope) + transformer = Transformer(model_args, layers, rope) + if ropes: + transformer.ropes = torch.nn.ModuleDict(ropes) + return transformer diff --git a/examples/models/llama/model_args.py b/examples/models/llama/model_args.py index a71b9857dbf..24ca5480dee 100644 --- a/examples/models/llama/model_args.py +++ b/examples/models/llama/model_args.py @@ -126,6 +126,9 @@ class ModelArgs: local_rope_theta: Optional[float] = ( None # For sliding window attention. e.g., gemma3-1b ) + rope_parameters: Optional[Dict[str, Dict[str, Any]]] = ( + None # Per-layer-type RoPE configs. e.g., {"full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}} + ) rope_freq_base: float = 10000.0 # The base frequency for RoPE. Keep it for BC. use_scaled_rope: bool = False # Use scaled RoPE, introduced in llama3.1. rope_scale_factor: int = 8 @@ -184,6 +187,7 @@ class ModelArgs: normalize_tok_embeddings: bool = False scale_query_by: float = 1.0 use_attn_o_gate: bool = False + headwise_attn_output_gate: bool = False use_attn_o_norm: bool = False use_residual_gate: bool = False use_ffn_learnable_scales: bool = False diff --git a/examples/models/spark_x2_5/BUCK b/examples/models/spark_x2_5/BUCK new file mode 100644 index 00000000000..6eaa29212c0 --- /dev/null +++ b/examples/models/spark_x2_5/BUCK @@ -0,0 +1,22 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. This file can contain fbcode-only targets. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target(_kind = runtime.python_library, + name = "spark_x2_5", + srcs = [ + "__init__.py", + "convert_weights.py", + ], + base_module = "executorch.examples.models.spark_x2_5", + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/examples/models/llama:transformer_modules", + "fbsource//third-party/pypi/safetensors:safetensors", + ], +) diff --git a/examples/models/spark_x2_5/README.md b/examples/models/spark_x2_5/README.md new file mode 100644 index 00000000000..ebd673c68ff --- /dev/null +++ b/examples/models/spark_x2_5/README.md @@ -0,0 +1,100 @@ +## Summary +[Spark-X2.5](https://huggingface.co/collections/XHToken/spark-x25) is a compact, efficient general-purpose language model series developed by [XHToken](https://github.com/XHToken/Spark-X2.5), available in two variants — 1.7B and 4B. The models use a hybrid attention architecture that combines one full-attention layer with three sliding-window attention layers, natively supporting context windows of up to 1M tokens and 200+ languages. + +## Architecture highlights +- **Hybrid attention**: 3 sliding-window attention layers + 1 full-attention layer (repeating) +- **Per-layer-type RoPE**: full-attention layers use `rope_theta=5M, partial_rotary_factor=0.25`; sliding-attention layers use `rope_theta=10K, partial_rotary_factor=1.0` +- **Headwise attention output gate**: per-head sigmoid gate broadcast over head dim +- **Sliding window**: 512 tokens for sliding-attention layers +- **GELU activation** in the MLP +- **Tied word embeddings** + +## Instructions + +Spark-X2.5 uses the same export pipeline as the optimized Llama model. Please see the [Llama README](../llama/README.md) for general runner and mobile-app details. + +### Example export + +Export Spark-X2.5-1.7B to XNNPack, FP32: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_xnnpack_fp32.yaml \ + +base.model_class="spark_x2_5_1_7b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json" \ + +export.output_name="spark_x2_5_1_7b_fp32.pte" +``` + +Export Spark-X2.5-1.7B to XNNPack, quantized with 8da4w: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.yaml \ + +base.model_class="spark_x2_5_1_7b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json" \ + +export.output_name="spark_x2_5_1_7b_8da4w.pte" +``` + +Export Spark-X2.5-4B to XNNPack, quantized with 8da4w: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.yaml \ + +base.model_class="spark_x2_5_4b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_4b_config.json" \ + +export.output_name="spark_x2_5_4b_8da4w.pte" +``` + +Export Spark-X2.5-1.7B to MLX on Apple Silicon, 4-bit weights: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml \ + +base.model_class="spark_x2_5_1_7b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json" \ + +export.output_name="spark_x2_5_1_7b_mlx_4w.pte" +``` + +To export with extended context (e.g., 2048 tokens): +``` +python -m extension.llm.export.export_llm \ + --config examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.yaml \ + +base.model_class="spark_x2_5_1_7b" \ + +base.params="examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json" \ + +export.max_seq_length=2048 \ + +export.max_context_length=2048 \ + +export.output_name="spark_x2_5_1_7b_8da4w.pte" +``` + +### Example run + +With ExecuTorch pybindings: +``` +python -m examples.models.llama.runner.native \ + --model spark_x2_5_1_7b \ + --pte spark_x2_5_1_7b_8da4w.pte \ + --tokenizer ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer.json \ + --tokenizer_config ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer_config.json \ + --prompt="<|user|>\nWho are you?<|end|>\n<|assistant|>\n" \ + --params examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json \ + --max_len 128 \ + -kv \ + --temperature 0.3 +``` + +With ExecuTorch's sample C++ runner: +``` +cmake-out/examples/models/llama/llama_main \ + --model_path spark_x2_5_1_7b_8da4w.pte \ + --tokenizer_path ~/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots//tokenizer.json \ + --prompt="<|user|>\nWho are you?<|end|>\n<|assistant|>\n" \ + --temperature 0.3 +``` + +Find the Hugging Face cache snapshot directory with: +``` +python - <<'PY' +from pathlib import Path +root = Path.home() / ".cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/snapshots" +for path in root.glob("*/tokenizer.json"): + print(path.parent) +PY +``` + +To run the model on an example iOS or Android app, see the Llama README's [Step 5: Build Mobile apps](../llama/README.md#step-5-build-mobile-apps) section. diff --git a/examples/models/spark_x2_5/__init__.py b/examples/models/spark_x2_5/__init__.py new file mode 100644 index 00000000000..0387e1ca3ec --- /dev/null +++ b/examples/models/spark_x2_5/__init__.py @@ -0,0 +1,5 @@ +from executorch.examples.models.spark_x2_5.convert_weights import convert_weights + +__all__ = [ + "convert_weights", +] diff --git a/examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json b/examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json new file mode 100644 index 00000000000..811a77c1a5c --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_1_7b_config.json @@ -0,0 +1,39 @@ +{ + "dim": 2048, + "ffn_dim_multiplier": 1, + "hidden_dim": 6656, + "n_heads": 8, + "n_kv_heads": 2, + "head_dim": 256, + "n_layers": 28, + "norm_eps": 1e-06, + "vocab_size": 131072, + "rope_theta": 10000.0, + "use_scaled_rope": false, + "use_hf_rope": true, + "partial_rotary_factor": 1.0, + "max_position_embeddings": 1048576, + "sliding_window": 512, + "act_fn": "gelu", + "headwise_attn_output_gate": true, + "attention_qkv_bias": false, + "layer_types": [ + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention" + ], + "rope_parameters": { + "full_attention": { + "rope_theta": 5000000, + "partial_rotary_factor": 0.25 + }, + "sliding_attention": { + "rope_theta": 10000, + "partial_rotary_factor": 1.0 + } + } +} diff --git a/examples/models/spark_x2_5/config/spark_x2_5_4b_config.json b/examples/models/spark_x2_5/config/spark_x2_5_4b_config.json new file mode 100644 index 00000000000..7285153836c --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_4b_config.json @@ -0,0 +1,41 @@ +{ + "dim": 2560, + "ffn_dim_multiplier": 1, + "hidden_dim": 10240, + "n_heads": 16, + "n_kv_heads": 4, + "head_dim": 256, + "n_layers": 36, + "norm_eps": 1e-06, + "vocab_size": 131072, + "rope_theta": 10000.0, + "use_scaled_rope": false, + "use_hf_rope": true, + "partial_rotary_factor": 1.0, + "max_position_embeddings": 1048576, + "sliding_window": 512, + "act_fn": "gelu", + "headwise_attn_output_gate": true, + "attention_qkv_bias": false, + "layer_types": [ + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "full_attention" + ], + "rope_parameters": { + "full_attention": { + "rope_theta": 5000000, + "partial_rotary_factor": 0.25 + }, + "sliding_attention": { + "rope_theta": 10000, + "partial_rotary_factor": 1.0 + } + } +} diff --git a/examples/models/spark_x2_5/config/spark_x2_5_coreml_fp32.yaml b/examples/models/spark_x2_5/config/spark_x2_5_coreml_fp32.yaml new file mode 100644 index 00000000000..c56012b38ee --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_coreml_fp32.yaml @@ -0,0 +1,15 @@ +base: + metadata: '{"get_bos_id": 0, "get_eos_ids":[1]}' + +model: + use_kv_cache: True + enable_dynamic_shape: False + dtype_override: fp32 + +backend: + coreml: + enabled: True + ios: 18 + enable_state: True + preserve_sdpa: True + compute_units: cpu_and_ne diff --git a/examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml b/examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml new file mode 100644 index 00000000000..24122e9ad02 --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_mlx_4w.yaml @@ -0,0 +1,15 @@ +base: + metadata: '{"get_bos_id": 0, "get_eos_ids":[1]}' + +model: + use_kv_cache: True + use_sdpa_with_kv_cache: True + dtype_override: bf16 + +quantization: + qmode: 4w + group_size: 64 + +backend: + mlx: + enabled: True diff --git a/examples/models/spark_x2_5/config/spark_x2_5_xnnpack_fp32.yaml b/examples/models/spark_x2_5/config/spark_x2_5_xnnpack_fp32.yaml new file mode 100644 index 00000000000..72b8d3cc854 --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_xnnpack_fp32.yaml @@ -0,0 +1,12 @@ +base: + metadata: '{"get_bos_id": 0, "get_eos_ids":[1]}' + +model: + use_kv_cache: True + use_sdpa_with_kv_cache: True + dtype_override: fp32 + +backend: + xnnpack: + enabled: True + extended_ops: True diff --git a/examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.yaml b/examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.yaml new file mode 100644 index 00000000000..60532df51dd --- /dev/null +++ b/examples/models/spark_x2_5/config/spark_x2_5_xnnpack_q8da4w.yaml @@ -0,0 +1,15 @@ +base: + metadata: '{"get_bos_id": 0, "get_eos_ids":[1]}' + +model: + use_kv_cache: True + use_sdpa_with_kv_cache: True + dtype_override: fp32 + +quantization: + qmode: 8da4w + +backend: + xnnpack: + enabled: True + extended_ops: True diff --git a/examples/models/spark_x2_5/convert_weights.py b/examples/models/spark_x2_5/convert_weights.py new file mode 100644 index 00000000000..64037fa16f9 --- /dev/null +++ b/examples/models/spark_x2_5/convert_weights.py @@ -0,0 +1,136 @@ +import argparse +import json +import os +from typing import Dict + +import torch + +from executorch.examples.models.checkpoint import get_mapped_key +from safetensors.torch import load_file + +_SPARK_X2_5_TO_META = { + "model.embedding.weight": "tok_embeddings.weight", + "model.norm.weight": "norm.weight", + "model.layers.{}.self_attn.out_proj.weight": "layers.{}.attention.wo.weight", + "model.layers.{}.self_attn.g_proj.weight": "layers.{}.attention.og.weight", + "model.layers.{}.post_attention_layernorm.weight": "layers.{}.ffn_norm.weight", + "model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight", + "model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight", + "model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight", + "model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight", +} + + +def spark_x2_5_to_meta( + state_dict: Dict[str, torch.Tensor], + n_heads: int, + n_kv_heads: int, + head_dim: int, +) -> Dict[str, torch.Tensor]: + """Convert Spark-X2.5 HF state dict to Meta format. + + Splits the fused q_k_v_proj into separate wq/wk/wv projections and handles + tied word embeddings (no lm_head → output.weight = tok_embeddings.weight). + """ + converted: Dict[str, torch.Tensor] = {} + q_size = n_heads * head_dim + kv_size = n_kv_heads * head_dim + + for key, value in state_dict.items(): + # Split fused QKV projection. + if key.endswith(".self_attn.q_k_v_proj.weight"): + layer_idx = key.split(".")[2] + q, k, v = torch.split(value, [q_size, kv_size, kv_size], dim=0) + converted[f"layers.{layer_idx}.attention.wq.weight"] = q + converted[f"layers.{layer_idx}.attention.wk.weight"] = k + converted[f"layers.{layer_idx}.attention.wv.weight"] = v + continue + + try: + new_key = get_mapped_key(key, _SPARK_X2_5_TO_META) + except Exception: + new_key = key.removeprefix("model.") + + converted[new_key] = value + + # Tied embeddings: no lm_head in Spark-X2.5. + if "lm_head.weight" not in state_dict: + converted["output.weight"] = converted["tok_embeddings.weight"] + + return converted + + +def load_checkpoint(input_dir: str) -> Dict[str, torch.Tensor]: + """Load a safetensors checkpoint, supporting both single-file and sharded formats.""" + index_path = os.path.join(input_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + # Sharded checkpoint. + print("Loading checkpoint from sharded safetensors") + with open(index_path, "r") as f: + index = json.load(f) + weight_map = index["weight_map"] + checkpoint_shards = sorted(set(weight_map.values())) + + shard_to_keys: Dict[str, list] = {} + for weight_name, shard in weight_map.items(): + shard_to_keys.setdefault(shard, []).append(weight_name) + + merged: Dict[str, torch.Tensor] = {} + for shard in checkpoint_shards: + shard_data = load_file(os.path.join(input_dir, shard)) + for weight_name in shard_to_keys[shard]: + merged[weight_name] = shard_data[weight_name] + del shard_data + return merged + + # Single checkpoint. + model_path = os.path.join(input_dir, "model.safetensors") + if os.path.exists(model_path): + print("Loading checkpoint from safetensors directory") + return load_file(model_path) + + raise FileNotFoundError(f"Could not find safetensors checkpoint in {input_dir}") + + +def convert_weights(input_dir: str, output_file: str) -> None: + """Convert Spark-X2.5 HF weights to Meta format. + + Reads model hyperparameters from config.json in input_dir so the function + signature stays compatible with download_and_convert_hf_checkpoint. + """ + config_path = os.path.join(input_dir, "config.json") + if os.path.exists(config_path): + with open(config_path) as f: + config = json.load(f) + n_heads = config.get("num_attention_heads", 8) + n_kv_heads = config.get("num_key_value_heads", 2) + head_dim = config.get("head_dim", 256) + else: + n_heads, n_kv_heads, head_dim = 8, 2, 256 + + print("Loading checkpoint...") + sd = load_checkpoint(input_dir) + print("Converting checkpoint...") + sd = spark_x2_5_to_meta(sd, n_heads, n_kv_heads, head_dim) + print("Saving checkpoint...") + torch.save(sd, output_file) + print("Done.") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert Spark-X2.5 weights to Meta format." + ) + parser.add_argument( + "input_dir", + type=str, + help="Path to directory containing safetensor checkpoint files.", + ) + parser.add_argument("output", type=str, help="Path to the output checkpoint") + + args = parser.parse_args() + convert_weights(args.input_dir, args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/models/spark_x2_5/test_spark_x2_5.py b/examples/models/spark_x2_5/test_spark_x2_5.py new file mode 100644 index 00000000000..fda5e991186 --- /dev/null +++ b/examples/models/spark_x2_5/test_spark_x2_5.py @@ -0,0 +1,126 @@ +import ast +import json +from pathlib import Path + +from omegaconf import OmegaConf + + +REPO_ROOT = Path(__file__).resolve().parents[3] +CONFIG_DIR = REPO_ROOT / "examples" / "models" / "spark_x2_5" / "config" +EXPORT_LLAMA_LIB = REPO_ROOT / "examples" / "models" / "llama" / "export_llama_lib.py" +LLM_CONFIG = REPO_ROOT / "extension" / "llm" / "export" / "config" / "llm_config.py" + + +def _load_json_config(name: str) -> dict: + with open(CONFIG_DIR / name, "r") as f: + return json.load(f) + + +def _module_ast(path: Path) -> ast.Module: + return ast.parse(path.read_text()) + + +def _literal_assignment(module: ast.Module, name: str): + for node in module.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == name: + return ast.literal_eval(node.value) + raise AssertionError(f"{name} not found") + + +def _class_string_assignments(module: ast.Module, class_name: str) -> dict[str, str]: + for node in module.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + values = {} + for stmt in node.body: + if ( + isinstance(stmt, ast.Assign) + and len(stmt.targets) == 1 + and isinstance(stmt.targets[0], ast.Name) + ): + values[stmt.targets[0].id] = ast.literal_eval(stmt.value) + return values + raise AssertionError(f"{class_name} not found") + + +def test_spark_x2_5_models_are_registered() -> None: + export_module = _module_ast(EXPORT_LLAMA_LIB) + model_types = _class_string_assignments(_module_ast(LLM_CONFIG), "ModelType") + executor_defined_models = _literal_assignment( + export_module, "EXECUTORCH_DEFINED_MODELS" + ) + hf_repo_ids = _literal_assignment(export_module, "HUGGING_FACE_REPO_IDS") + + assert "spark_x2_5_1_7b" in executor_defined_models + assert "spark_x2_5_4b" in executor_defined_models + assert model_types["spark_x2_5_1_7b"] == "spark_x2_5_1_7b" + assert model_types["spark_x2_5_4b"] == "spark_x2_5_4b" + assert hf_repo_ids["spark_x2_5_1_7b"] == "XHToken/Spark-X2.5-1.7B" + assert hf_repo_ids["spark_x2_5_4b"] == "XHToken/Spark-X2.5-4B" + + +def test_spark_x2_5_architecture_configs_match_expected_shapes() -> None: + expected = { + "spark_x2_5_1_7b_config.json": { + "dim": 2048, + "hidden_dim": 6656, + "n_heads": 8, + "n_kv_heads": 2, + "head_dim": 256, + }, + "spark_x2_5_4b_config.json": { + "dim": 2560, + "hidden_dim": 10240, + "n_heads": 16, + "n_kv_heads": 4, + "head_dim": 256, + }, + } + + for filename, expected_fields in expected.items(): + cfg = _load_json_config(filename) + for key, value in expected_fields.items(): + assert ( + cfg[key] == value + ), f"{filename}: {key} expected {value}, got {cfg[key]}" + + # 1.7B: 28 layers = 7 groups of 4 + cfg_1_7b = _load_json_config("spark_x2_5_1_7b_config.json") + assert cfg_1_7b["n_layers"] == 28 + assert len(cfg_1_7b["layer_types"]) == 28 + assert cfg_1_7b["layer_types"].count("full_attention") == 7 + assert cfg_1_7b["layer_types"].count("sliding_attention") == 21 + + # 4B: 36 layers = 9 groups of 4 + cfg_4b = _load_json_config("spark_x2_5_4b_config.json") + assert cfg_4b["n_layers"] == 36 + assert len(cfg_4b["layer_types"]) == 36 + assert cfg_4b["layer_types"].count("full_attention") == 9 + assert cfg_4b["layer_types"].count("sliding_attention") == 27 + + # Shared architecture features. + for filename in expected: + cfg = _load_json_config(filename) + assert cfg["vocab_size"] == 131072 + assert cfg["sliding_window"] == 512 + assert cfg["use_hf_rope"] is True + assert cfg["headwise_attn_output_gate"] is True + assert cfg["act_fn"] == "gelu" + assert cfg["rope_parameters"]["full_attention"]["rope_theta"] == 5000000 + assert cfg["rope_parameters"]["full_attention"]["partial_rotary_factor"] == 0.25 + assert cfg["rope_parameters"]["sliding_attention"]["rope_theta"] == 10000 + assert ( + cfg["rope_parameters"]["sliding_attention"]["partial_rotary_factor"] == 1.0 + ) + + +def test_spark_x2_5_mlx_config_enables_mlx_backend() -> None: + cfg = OmegaConf.load(CONFIG_DIR / "spark_x2_5_mlx_4w.yaml") + assert cfg.base.metadata == '{"get_bos_id": 0, "get_eos_ids":[1]}' + assert cfg.model.use_kv_cache is True + assert cfg.model.use_sdpa_with_kv_cache is True + assert cfg.model.dtype_override == "bf16" + assert cfg.quantization.qmode == "4w" + assert cfg.quantization.group_size == 64 + assert cfg.backend.mlx.enabled is True diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index 3eb8b8a18f6..352ebf5e603 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -56,6 +56,8 @@ class ModelType(str, Enum): lfm2_1_2b = "lfm2_1_2b" lfm2_5_350m = "lfm2_5_350m" lfm2_5_1_2b = "lfm2_5_1_2b" + spark_x2_5_1_7b = "spark_x2_5_1_7b" + spark_x2_5_4b = "spark_x2_5_4b" class PreqMode(str, Enum):