feat(example): Add Spark-X2.5 model example with hybrid attention support - #22865
dongjiang1989 wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22865
Note: Links to docs will display an error until the docs builds have been completed.
|
|
Hi @dongjiang1989! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
|
This PR needs a
|
Add examples/models/spark_x2_5/ with support for the Spark-X2.5-1.7B and Spark-X2.5-4B models from XHToken. These are hybrid-attention LLMs (3:1 sliding-window to full-attention layers) with native 1M-token context, per-layer-type RoPE configs, headwise sigmoid attention gates, and GELU activation. New module: - convert_weights.py: HF safetensors to Meta format conversion, including fused QKV projection splitting, sharded checkpoint support, and tied embeddings handling. - config/: JSON configs for 1.7B (28 layers) and 4B (36 layers) variants, plus XNNPack (fp32/q8da4w), CoreML, and MLX export configs. - test_spark_x2_5.py: config validation and model registration tests. - BUCK, README.md: build and usage documentation. Core changes to support Spark-X2.5 architecture: - model_args.py: Add headwise_attn_output_gate and rope_parameters fields. - feed_forward.py: FeedForward now accepts act_fn parameter (was hardcoded to SiLU). - attention.py: Add headwise attention output gate (dim->n_heads broadcast over head_dim), per-layer is_sliding detection, RingKVCache for sliding-window layers. - llama_transformer.py: Per-layer-type RoPE construction via _build_ropes(), freqs_by_type dispatch in _forward_layers(), act_fn plumbed to FeedForward. - export_llama_lib.py, llm_config.py: Register spark_x2_5_1_7b and spark_x2_5_4b model types. This commit was authored with AI assistance (Claude Code). Signed-off-by: dongjiang1989 <dongjiang1989@126.com>
c825ccf to
980ad58
Compare
Summary
Add
examples/models/spark_x2_5/with support for the Spark-X2.5-1.7B and Spark-X2.5-4B models from XHToken. These are compact, general-purpose language models with a hybrid attention architecture (3 sliding-window attention layers + 1 full-attention layer, repeating) that natively supports context windows up to 1M tokens.Architecture highlights
rope_theta=5M, partial_rotary_factor=0.25; sliding-attention layers userope_theta=10K, partial_rotary_factor=1.0use_attn_o_gate)RingKVCache, window=512 tokens)Files added (in
examples/models/spark_x2_5/)convert_weights.py— HF safetensors → Meta format conversion with fused QKV split, sharded checkpoint support, and tied embeddings handlingconfig/— JSON configs for 1.7B and 4B variants; XNNPack (fp32, q8da4w), CoreML, and MLX export configstest_spark_x2_5.py— config validation and model registration testsBUCK,README.md— build target and usage documentationCore changes to existing files
examples/models/llama/model_args.py— addheadwise_attn_output_gateandrope_parametersfieldsexamples/models/llama/feed_forward.py—FeedForwardnow acceptsact_fnparameter (was hardcoded to SiLU)examples/models/llama/attention.py— headwise attention output gate (dim→n_heads, broadcast over head_dim), per-layeris_slidingdetection,RingKVCachefor sliding-window layers, mutual-exclusion validation for gate flagsexamples/models/llama/llama_transformer.py— per-layer-type RoPE via_build_ropes()helper,freqs_by_typedispatch in_forward_layers(),act_fnplumbed to FeedForwardexamples/models/llama/export_llama_lib.py— registerspark_x2_5_1_7bandspark_x2_5_4binEXECUTORCH_DEFINED_MODELS,HUGGING_FACE_REPO_IDS, and weight-conversion dispatchextension/llm/export/config/llm_config.py— addspark_x2_5_1_7bandspark_x2_5_4btoModelTypeenumExample export
Test plan
Unit tests
Model construction + forward pass (both variants)
Sharded checkpoint conversion
Lint
This PR was authored with AI assistance (Claude Code).