Skip to content

Add Gear model support - #1648

Open
crinex wants to merge 1 commit into
ml-explore:mainfrom
crinex:add-gear-model
Open

Add Gear model support#1648
crinex wants to merge 1 commit into
ml-explore:mainfrom
crinex:add-gear-model

Conversation

@crinex

@crinex crinex commented Aug 3, 2026

Copy link
Copy Markdown

Summary

This PR adds initial mlx-lm support for the Gear causal language model
architecture.

Included scope:

  • new mlx_lm.models.gear architecture module for model_type="gear"
  • full and sliding attention layers with Gear Q/K normalization
  • local conv_mixer layers with depthwise key/value convolution state
  • Gear RMSNorm semantics using (1 + weight)
  • scaled token embeddings matching the reference implementation
  • HF Conv1d weight layout sanitization for key_conv and value_conv
  • model smoke coverage in tests/test_models.py

Excluded from PR scope:

  • local validation scripts and model artifacts
  • bidirectional Gear attention mode
  • model-specific training examples

Validation

Local model artifact:

  • sft_final
  • HF architecture: GearForCausalLM
  • model_type: gear

Weight key check against the local safetensors header:

actual_count=362 expected_count=362
missing=0 extra=0

Commands run:

.venv-contrib/bin/python -m py_compile mlx_lm/models/gear.py tests/test_models.py
git diff --check
.venv-contrib/bin/pre-commit run --files mlx_lm/models/gear.py tests/test_models.py
.venv-contrib/bin/python -m unittest tests.test_models.TestModels.test_gear tests.test_evaluate tests.test_tokenizers.TestTokenizers.test_tokenizers
.venv-contrib/bin/python -m unittest discover tests/
.venv-contrib/bin/python - <<'PY'
from pathlib import Path
from mlx_lm.utils import load_model

model, config = load_model(Path('../contribution-prep/models/sft_final'), lazy=True)
print('model_type', model.model_type)
print('layers', len(model.layers))
print('vocab_size', config['vocab_size'])
print('loaded_lazy', True)
PY
.venv-contrib/bin/mlx_lm.generate \
  --model ../contribution-prep/models/sft_final \
  --prompt "What is 2+2?" \
  --max-tokens 16 \
  --trust-remote-code
.venv-contrib/bin/mlx_lm.convert \
  --hf-path ../contribution-prep/models/sft_final \
  --mlx-path /private/tmp/gear_mlx_q4_smoke_20260714_1635 \
  --quantize \
  --q-bits 4 \
  --q-group-size 64 \
  --trust-remote-code
.venv-contrib/bin/mlx_lm.generate \
  --model /private/tmp/gear_mlx_q4_smoke_20260714_1635 \
  --prompt "What is 2+2?" \
  --max-tokens 16 \
  --trust-remote-code

Result:

black....................................................................Passed
isort....................................................................Passed
Ran 3 tests in 8.347s
OK
Ran 210 tests in 41.611s
OK (skipped=1)
model_type gear
layers 26
vocab_size 125184
loaded_lazy True
Prompt: 18 tokens
Generation: 16 tokens
Peak memory: 1.720 GB
Quantized model with 4.502 bits per weight.
Quantized generation peak memory: 0.508 GB

Environment note:

  • Validation used a dedicated .venv-contrib environment with mlx==0.32.0,
    transformers==5.13.1, lm-eval==0.4.12, and datasets==5.0.0.
  • MLX runtime checks were run outside the sandbox because the sandbox cannot
    access the local Metal device.
  • The local sft_final artifact requires --trust-remote-code for
    tokenizer/config loading because it contains auto_map custom code.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant