Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Awni Hannun
|
f1b6fd63ec | ||
|
Awni Hannun
|
96cb7c3957 |
@@ -1,66 +0,0 @@
|
||||
version: 2.1
|
||||
|
||||
orbs:
|
||||
apple: ml-explore/pr-approval@0.1.0
|
||||
|
||||
jobs:
|
||||
linux_build_and_test:
|
||||
docker:
|
||||
- image: cimg/python:3.9
|
||||
|
||||
steps:
|
||||
- checkout
|
||||
- run:
|
||||
name: Run style checks
|
||||
command: |
|
||||
pip install pre-commit
|
||||
pre-commit run --all
|
||||
if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
|
||||
|
||||
mlx_lm_build_and_test:
|
||||
macos:
|
||||
xcode: "15.2.0"
|
||||
resource_class: macos.m1.large.gen1
|
||||
steps:
|
||||
- checkout
|
||||
- run:
|
||||
name: Install dependencies
|
||||
command: |
|
||||
brew install python@3.9
|
||||
python3.9 -m venv env
|
||||
source env/bin/activate
|
||||
pip install --upgrade pip
|
||||
pip install unittest-xml-reporting
|
||||
pip install -e ".[test]"
|
||||
- run:
|
||||
name: Run Python tests
|
||||
command: |
|
||||
source env/bin/activate
|
||||
python -m xmlrunner discover -v tests -o test-results/
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
workflows:
|
||||
build_and_test:
|
||||
when:
|
||||
matches:
|
||||
pattern: "^(?!pull/)[-\\w]+$"
|
||||
value: << pipeline.git.branch >>
|
||||
jobs:
|
||||
- mlx_lm_build_and_test
|
||||
- linux_build_and_test
|
||||
|
||||
prb:
|
||||
when:
|
||||
matches:
|
||||
pattern: "^pull/\\d+(/head)?$"
|
||||
value: << pipeline.git.branch >>
|
||||
jobs:
|
||||
- hold:
|
||||
type: approval
|
||||
- apple/authenticate:
|
||||
context: pr-approval
|
||||
- mlx_lm_build_and_test:
|
||||
requires: [ hold ]
|
||||
- linux_build_and_test:
|
||||
requires: [ hold ]
|
||||
-139
@@ -1,139 +0,0 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Vim
|
||||
*.swp
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
pip-wheel-metadata/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# IDE files
|
||||
.idea/
|
||||
.vscode/
|
||||
|
||||
# .DS_Store files
|
||||
.DS_Store
|
||||
@@ -1,11 +0,0 @@
|
||||
repos:
|
||||
- repo: https://github.com/psf/black-pre-commit-mirror
|
||||
rev: 25.1.0
|
||||
hooks:
|
||||
- id: black
|
||||
- repo: https://github.com/pycqa/isort
|
||||
rev: 6.0.0
|
||||
hooks:
|
||||
- id: isort
|
||||
args:
|
||||
- --profile=black
|
||||
+7
-3
@@ -5,8 +5,12 @@ with a short description of your contribution(s) below. For example:
|
||||
|
||||
- Jane Smith: Added the `foo` example.
|
||||
|
||||
MLX LM was developed with contributions from the following individuals:
|
||||
MLX Examples was developed with contributions from the following individuals:
|
||||
|
||||
- Juarez Bochi: Added support for T5 models.
|
||||
- Sarthak Yadav: Added the `cifar` and `speechcommands` examples.
|
||||
- Shunta Saito: Added support for PLaMo models.
|
||||
- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, Rednote `dots.llm1`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`, and `reporting training metrics to WandB (Weights & Biases)`.
|
||||
- Prince Canuma: Helped add support for the following model architectures: HuggingFace's `Starcoder2`, Cohere's `Cohere (1 and 2)`, Alibaba Qwen's `Qwen (2, 3 and MoE)`, Microsoft's `Phi (3 and 3.5 MoE)`, `BitNet1.58`, Meta's `Llama (3 and 4)`, Google DeepMind's `Gemma 3`, and InterLM's `InternLM 2.5`.
|
||||
- Gabrijel Boduljak: Implemented `CLIP`.
|
||||
- Markus Enzweiler: Added the `cvae` examples.
|
||||
- Prince Canuma: Helped add support for `Starcoder2` models.
|
||||
- Shiyu Li: Added the `Segment Anything Model`.
|
||||
|
||||
+8
-51
@@ -1,54 +1,11 @@
|
||||
# Contributing to MLX LM
|
||||
|
||||
We want to make contributing to this project as easy and transparent as
|
||||
possible.
|
||||
|
||||
## Pull Requests
|
||||
|
||||
1. Fork and submit pull requests to the repo.
|
||||
2. If you've added code that should be tested, add tests.
|
||||
3. Every PR should have passing tests and at least one review.
|
||||
4. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`.
|
||||
This should install hooks for running `black` and `clang-format` to ensure
|
||||
consistent style for C++ and python code.
|
||||
|
||||
You can also run the formatters manually as follows on individual files:
|
||||
|
||||
```bash
|
||||
clang-format -i file.cpp
|
||||
```
|
||||
|
||||
```bash
|
||||
black file.py
|
||||
```
|
||||
|
||||
or,
|
||||
|
||||
```bash
|
||||
# single file
|
||||
pre-commit run --files file1.py
|
||||
|
||||
# specific files
|
||||
pre-commit run --files file1.py file2.py
|
||||
```
|
||||
|
||||
or run `pre-commit run --all-files` to check all files in the repo.
|
||||
|
||||
## Issues
|
||||
|
||||
We use GitHub issues to track public bugs. Please ensure your description is
|
||||
clear and has sufficient instructions to be able to reproduce the issue.
|
||||
|
||||
## License
|
||||
|
||||
By contributing to mlx-lm, you agree that your contributions will be licensed
|
||||
under the LICENSE file in the root directory of this source tree.
|
||||
|
||||
## Adding New Models
|
||||
|
||||
Below are some tips to port LLMs available on Hugging Face to MLX.
|
||||
|
||||
From this directory, do an editable install:
|
||||
Before starting checkout the [general contribution
|
||||
guidelines](https://github.com/ml-explore/mlx-examples/blob/main/CONTRIBUTING.md).
|
||||
|
||||
Next, from this directory, do an editable install:
|
||||
|
||||
```shell
|
||||
pip install -e .
|
||||
@@ -60,7 +17,7 @@ Then check if the model has weights in the
|
||||
convert it.
|
||||
|
||||
After that, add the model file to the
|
||||
[`mlx_lm/models`](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/models)
|
||||
[`mlx_lm/models`](https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/models)
|
||||
directory. You can see other examples there. We recommend starting from a model
|
||||
that is similar to the model you are porting.
|
||||
|
||||
@@ -78,12 +35,12 @@ To determine the model layer names, we suggest either:
|
||||
in the Hugging Face repo.
|
||||
|
||||
To add LoRA support edit
|
||||
[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/utils.py#L27-L60)
|
||||
[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/tuner/utils.py#L27-L60)
|
||||
|
||||
Finally, add a test for the new modle type to the [model
|
||||
tests](https://github.com/ml-explore/mlx-lm/blob/main/tests/test_models.py).
|
||||
tests](https://github.com/ml-explore/mlx-examples/blob/main/llms/tests/test_models.py).
|
||||
|
||||
You can run the tests with:
|
||||
From the `llms/` directory, you can run the tests with:
|
||||
|
||||
```shell
|
||||
python -m unittest discover tests/
|
||||
|
||||
+1
-1
@@ -1,2 +1,2 @@
|
||||
include requirements.txt
|
||||
include mlx_lm/requirements.txt
|
||||
recursive-include mlx_lm/ *.py
|
||||
|
||||
@@ -1,17 +1,4 @@
|
||||
## MLX LM
|
||||
|
||||
MLX LM is a Python package for generating text and fine-tuning large language
|
||||
models on Apple silicon with MLX.
|
||||
|
||||
Some key features include:
|
||||
|
||||
* Integration with the Hugging Face Hub to easily use thousands of LLMs with a
|
||||
single command.
|
||||
* Support for quantizing and uploading models to the Hugging Face Hub.
|
||||
* [Low-rank and full model
|
||||
fine-tuning](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md)
|
||||
with support for quantized models.
|
||||
* Distributed inference and fine-tuning with `mx.distributed`
|
||||
## Generate Text with LLMs and MLX
|
||||
|
||||
The easiest way to get started is to install the `mlx-lm` package:
|
||||
|
||||
@@ -27,30 +14,11 @@ pip install mlx-lm
|
||||
conda install -c conda-forge mlx-lm
|
||||
```
|
||||
|
||||
### Quick Start
|
||||
The `mlx-lm` package also has:
|
||||
|
||||
To generate text with an LLM use:
|
||||
|
||||
```bash
|
||||
mlx_lm.generate --prompt "How tall is Mt Everest?"
|
||||
```
|
||||
|
||||
To chat with an LLM use:
|
||||
|
||||
```bash
|
||||
mlx_lm.chat
|
||||
```
|
||||
|
||||
This will give you a chat REPL that you can use to interact with the LLM. The
|
||||
chat context is preserved during the lifetime of the REPL.
|
||||
|
||||
Commands in `mlx-lm` typically take command line options which let you specify
|
||||
the model, sampling parameters, and more. Use `-h` to see a list of available
|
||||
options for a command, e.g.:
|
||||
|
||||
```bash
|
||||
mlx_lm.generate -h
|
||||
```
|
||||
- [LoRA and QLoRA fine-tuning](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md)
|
||||
- [Merging models](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/MERGE.md)
|
||||
- [HTTP model serving](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/SERVER.md)
|
||||
|
||||
### Python API
|
||||
|
||||
@@ -61,14 +29,7 @@ from mlx_lm import load, generate
|
||||
|
||||
model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
|
||||
|
||||
prompt = "Write a story about Einstein"
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=True
|
||||
)
|
||||
|
||||
text = generate(model, tokenizer, prompt=prompt, verbose=True)
|
||||
response = generate(model, tokenizer, prompt="hello", verbose=True)
|
||||
```
|
||||
|
||||
To see a description of all the arguments you can do:
|
||||
@@ -77,14 +38,10 @@ To see a description of all the arguments you can do:
|
||||
>>> help(generate)
|
||||
```
|
||||
|
||||
Check out the [generation
|
||||
example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/generate_response.py)
|
||||
to see how to use the API in more detail.
|
||||
|
||||
The `mlx-lm` package also comes with functionality to quantize and optionally
|
||||
upload models to the Hugging Face Hub.
|
||||
|
||||
You can convert models using the Python API:
|
||||
You can convert models in the Python API with:
|
||||
|
||||
```python
|
||||
from mlx_lm import convert
|
||||
@@ -107,10 +64,8 @@ To see a description of all the arguments you can do:
|
||||
|
||||
#### Streaming
|
||||
|
||||
For streaming generation, use the `stream_generate` function. This yields
|
||||
a generation response object.
|
||||
|
||||
For example,
|
||||
For streaming generation, use the `stream_generate` function. This returns a
|
||||
generator object which streams the output text. For example,
|
||||
|
||||
```python
|
||||
from mlx_lm import load, stream_generate
|
||||
@@ -120,28 +75,11 @@ model, tokenizer = load(repo)
|
||||
|
||||
prompt = "Write a story about Einstein"
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=True
|
||||
)
|
||||
|
||||
for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
|
||||
print(response.text, end="", flush=True)
|
||||
for t in stream_generate(model, tokenizer, prompt, max_tokens=512):
|
||||
print(t, end="", flush=True)
|
||||
print()
|
||||
```
|
||||
|
||||
#### Sampling
|
||||
|
||||
The `generate` and `stream_generate` functions accept `sampler` and
|
||||
`logits_processors` keyword arguments. A sampler is any callable which accepts
|
||||
a possibly batched logits array and returns an array of sampled tokens. The
|
||||
`logits_processors` must be a list of callables which take the token history
|
||||
and current logits as input and return the processed logits. The logits
|
||||
processors are applied in order.
|
||||
|
||||
Some standard sampling functions and logits processors are provided in
|
||||
`mlx_lm.sample_utils`.
|
||||
|
||||
### Command Line
|
||||
|
||||
You can also use `mlx-lm` from the command line with:
|
||||
@@ -182,55 +120,11 @@ mlx_lm.convert \
|
||||
--upload-repo mlx-community/my-4bit-mistral
|
||||
```
|
||||
|
||||
Models can also be converted and quantized directly in the
|
||||
[mlx-my-repo](https://huggingface.co/spaces/mlx-community/mlx-my-repo) Hugging
|
||||
Face Space.
|
||||
|
||||
### Long Prompts and Generations
|
||||
|
||||
`mlx-lm` has some tools to scale efficiently to long prompts and generations:
|
||||
|
||||
- A rotating fixed-size key-value cache.
|
||||
- Prompt caching
|
||||
|
||||
To use the rotating key-value cache pass the argument `--max-kv-size n` where
|
||||
`n` can be any integer. Smaller values like `512` will use very little RAM but
|
||||
result in worse quality. Larger values like `4096` or higher will use more RAM
|
||||
but have better quality.
|
||||
|
||||
Caching prompts can substantially speedup reusing the same long context with
|
||||
different queries. To cache a prompt use `mlx_lm.cache_prompt`. For example:
|
||||
|
||||
```bash
|
||||
cat prompt.txt | mlx_lm.cache_prompt \
|
||||
--model mistralai/Mistral-7B-Instruct-v0.3 \
|
||||
--prompt - \
|
||||
--prompt-cache-file mistral_prompt.safetensors
|
||||
```
|
||||
|
||||
Then use the cached prompt with `mlx_lm.generate`:
|
||||
|
||||
```
|
||||
mlx_lm.generate \
|
||||
--prompt-cache-file mistral_prompt.safetensors \
|
||||
--prompt "\nSummarize the above text."
|
||||
```
|
||||
|
||||
The cached prompt is treated as a prefix to the supplied prompt. Also notice
|
||||
when using a cached prompt, the model to use is read from the cache and need
|
||||
not be supplied explicitly.
|
||||
|
||||
Prompt caching can also be used in the Python API in order to avoid
|
||||
recomputing the prompt. This is useful in multi-turn dialogues or across
|
||||
requests that use the same context. See the
|
||||
[example](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/examples/chat.py)
|
||||
for more usage details.
|
||||
|
||||
### Supported Models
|
||||
|
||||
`mlx-lm` supports thousands of Hugging Face format LLMs. If the model you want to
|
||||
run is not supported, file an
|
||||
[issue](https://github.com/ml-explore/mlx-lm/issues/new) or better yet,
|
||||
The example supports Hugging Face format Mistral, Llama, and Phi-2 style
|
||||
models. If the model you want to run is not supported, file an
|
||||
[issue](https://github.com/ml-explore/mlx-examples/issues/new) or better yet,
|
||||
submit a pull request.
|
||||
|
||||
Here are a few examples of Hugging Face models that work with this example:
|
||||
@@ -246,7 +140,6 @@ Here are a few examples of Hugging Face models that work with this example:
|
||||
- [pfnet/plamo-13b-instruct](https://huggingface.co/pfnet/plamo-13b-instruct)
|
||||
- [stabilityai/stablelm-2-zephyr-1_6b](https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b)
|
||||
- [internlm/internlm2-7b](https://huggingface.co/internlm/internlm2-7b)
|
||||
- [tiiuae/falcon-mamba-7b-instruct](https://huggingface.co/tiiuae/falcon-mamba-7b-instruct)
|
||||
|
||||
Most
|
||||
[Mistral](https://huggingface.co/models?library=transformers,safetensors&other=mistral&sort=trending),
|
||||
@@ -274,28 +167,3 @@ model, tokenizer = load(
|
||||
tokenizer_config={"eos_token": "<|endoftext|>", "trust_remote_code": True},
|
||||
)
|
||||
```
|
||||
|
||||
### Large Models
|
||||
|
||||
> [!NOTE]
|
||||
This requires macOS 15.0 or higher to work.
|
||||
|
||||
Models which are large relative to the total RAM available on the machine can
|
||||
be slow. `mlx-lm` will attempt to make them faster by wiring the memory
|
||||
occupied by the model and cache. This requires macOS 15 or higher to
|
||||
work.
|
||||
|
||||
If you see the following warning message:
|
||||
|
||||
> [WARNING] Generating with a model that requires ...
|
||||
|
||||
then the model will likely be slow on the given machine. If the model fits in
|
||||
RAM then it can often be sped up by increasing the system wired memory limit.
|
||||
To increase the limit, set the following `sysctl`:
|
||||
|
||||
```bash
|
||||
sudo sysctl iogpu.wired_limit_mb=N
|
||||
```
|
||||
|
||||
The value `N` should be larger than the size of the model in megabytes but
|
||||
smaller than the memory size of the machine.
|
||||
|
||||
@@ -1,149 +0,0 @@
|
||||
# Learned Quantization
|
||||
|
||||
To reduce the quality loss from quantization MLX LM has several options:
|
||||
|
||||
- Distilled Weight Quantization (DWQ)
|
||||
- Activation-aware Weight Quantization (AWQ)[^1]
|
||||
- Dynamic quantization
|
||||
|
||||
All methods use calibration data to tune parameters or hyper-parameters of the
|
||||
model. DWQ fine-tunes non-quantized parameters (including quantization scales
|
||||
and biases) using the non-quantized model as a teacher. AWQ scales and clips
|
||||
the weights prior to quantization. Dynamic quantization estimates the
|
||||
sensitivity of a model's outputs to each layer and uses a higher precision for
|
||||
layers which have higher sensitivity.
|
||||
|
||||
Dynamic quantization is the fastest to run. DWQ takes longer but typically
|
||||
yields better results. You can also cascade methods. For example a dynamically
|
||||
quantized model can be further refined with DWQ.
|
||||
|
||||
To get started, first install the requirements:
|
||||
|
||||
```
|
||||
pip install mlx-lm[quant]
|
||||
```
|
||||
|
||||
### DWQ
|
||||
|
||||
Use `mlx_lm.dwq` to run DWQ on a given model. For example:
|
||||
|
||||
```bash
|
||||
mlx_lm.dwq --model mistralai/Mistral-7B-Instruct-v0.3
|
||||
```
|
||||
|
||||
Some important options, along with their default values are:
|
||||
|
||||
- `--mlx-path mlx_model`: The location to save the DWQ model.
|
||||
- `--bits 4`: Precision of the quantization.
|
||||
- `--num-samples 1024`: Number of samples to use. Using more samples can lead to
|
||||
better results but takes longer.
|
||||
- `--batch-size 8`: Use a smaller batch size to reduce the memory footprint.
|
||||
|
||||
For a full list of options run:
|
||||
|
||||
```bash
|
||||
mlx_lm.dwq --help
|
||||
```
|
||||
|
||||
#### Tips
|
||||
|
||||
- DWQ works best distilling to lower precision, anywhere from 2-bit to 4-bit
|
||||
models.
|
||||
- Distilling 16-bit precision to 8-bit and even 6-bit often doesn't work well.
|
||||
The loss starts out so low that it's difficult to reduce further.
|
||||
- Decreasing the quantization group size (e.g. `--group-size 32`) doubles the
|
||||
number of tunable parameters and can work much better.
|
||||
- If the loss is oscillating and not going down consistently, try reducing the
|
||||
learning rate. If it is decreasing but slowly, try increasing the learning
|
||||
rate.
|
||||
- As a rule of thumb, lower precision can benefit from a higher learning rate
|
||||
since the loss starts out higher. Conversely, higher precision needs a lower
|
||||
learning rate.
|
||||
|
||||
|
||||
#### Memory Use
|
||||
|
||||
A few options to reduce memory use for DWQ:
|
||||
|
||||
- Distill from an 8-bit model instead of a 16-bit model. The 8-bit
|
||||
models are usually as good as 16-bit precision models.
|
||||
- Use a shorter maximum sequence length. The default is 2048. Using
|
||||
`--max-seq-length 512` reduces the memory and still gets good results.
|
||||
- Use a smaller batch size, e.g. `--batch-size 1`
|
||||
|
||||
### Dynamic Quantization
|
||||
|
||||
Use `mlx_lm.dynamic_quant` to generate a dynamic quantization of given model.
|
||||
For example:
|
||||
|
||||
```bash
|
||||
mlx_lm.dynamic_quant --model mistralai/Mistral-7B-Instruct-v0.3
|
||||
```
|
||||
|
||||
The script will estimate the sensitivity for each quantizable layer in the
|
||||
model. It will then quantize the model using higher precision (default 5 bits)
|
||||
for the more sensitive layers and lower precision (default 4 bits) for the
|
||||
rest. The script also saves a JSON file with each layer's sensitivities which
|
||||
saves needing to compute it multiple times to make different precision quants
|
||||
of the same model.
|
||||
|
||||
Some important options are:
|
||||
|
||||
- `--target-bpw`: The target bits-per-weight. For a given set of quantization
|
||||
parameters only certain ranges are possible. For example, with the default
|
||||
parameters a BPW in the range `[4.5, 5.5]` is achievable.
|
||||
- `--sensitivities`: A path to a precomputed sensitivities file.
|
||||
- `--low-bits`: The number of bits to use for the less sensitive layers.
|
||||
- `--high-bits`: The number of bits to use for the more sensitive layers.
|
||||
|
||||
### AWQ
|
||||
|
||||
Use `mlx_lm.awq` to run AWQ on a given model. For example:
|
||||
|
||||
```bash
|
||||
mlx_lm.awq --model mistralai/Mistral-7B-Instruct-v0.3
|
||||
```
|
||||
|
||||
The script can take anywhere form a few minutes to several hours to run
|
||||
depending on the model size and the number of samples.
|
||||
|
||||
Some important options, along with their default values, are:
|
||||
|
||||
- `--mlx-path mlx_model`: The location to save the AWQ model.
|
||||
- `--bits 4`: Precision of the quantization.
|
||||
- `--num-samples 32`: Number of samples to use. Using more samples can lead to
|
||||
better results but takes longer.
|
||||
- `--n-grid 10`: The granularity of the AWQ search. A larger grid can lead to
|
||||
better results but takes longer.
|
||||
|
||||
For a full list of options run:
|
||||
|
||||
```bash
|
||||
mlx_lm.awq --help
|
||||
```
|
||||
|
||||
### Evaluate
|
||||
|
||||
Once the training script finishes, you can evaluate the quality of the model
|
||||
on downstream tasks using `mlx_lm.evaluate`. For example:
|
||||
|
||||
```bash
|
||||
mlx_lm.evaluate \
|
||||
--model mlx_model \
|
||||
--tasks winogrande boolq arc_challenge arc_easy hellaswag openbookqa piqa social_iqa
|
||||
```
|
||||
|
||||
### Upload to Hugging Face
|
||||
|
||||
Use `mlx_lm.upload` to upload the quantized model to the Hugging Face Hub. For
|
||||
example:
|
||||
|
||||
```bash
|
||||
mlx_lm.upload \
|
||||
--path mlx_model \
|
||||
--upload-repo mlx-community/Mistral-7B-Instruct-v0.3-3bit-DWQ
|
||||
```
|
||||
|
||||
[^1]: Refer to the [paper](https://arxiv.org/abs/2306.00978)
|
||||
and [github repository](https://github.com/mit-han-lab/llm-awq) for more
|
||||
details.
|
||||
+38
-169
@@ -57,9 +57,6 @@ mlx_lm.lora \
|
||||
--iters 600
|
||||
```
|
||||
|
||||
To fine-tune the full model weights, add the `--fine-tune-type full` flag.
|
||||
Currently supported fine-tuning types are `lora` (default), `dora`, and `full`.
|
||||
|
||||
The `--data` argument must specify a path to a `train.jsonl`, `valid.jsonl`
|
||||
when using `--train` and a path to a `test.jsonl` when using `--test`. For more
|
||||
details on the data format see the section on [Data](#Data).
|
||||
@@ -70,25 +67,12 @@ mistralai/Mistral-7B-v0.1`.
|
||||
If `--model` points to a quantized model, then the training will use QLoRA,
|
||||
otherwise it will use regular LoRA.
|
||||
|
||||
By default, the adapter config and learned weights are saved in `adapters/`.
|
||||
You can specify the output location with `--adapter-path`.
|
||||
By default, the adapter config and weights are saved in `adapters/`. You can
|
||||
specify the output location with `--adapter-path`.
|
||||
|
||||
You can resume fine-tuning with an existing adapter with
|
||||
`--resume-adapter-file <path_to_adapters.safetensors>`.
|
||||
|
||||
#### Logging
|
||||
|
||||
You can log training metrics to Weights & Biases by passing a project name with
|
||||
the `--wandb` flag. Make sure to install wandb with `pip install wandb`.
|
||||
|
||||
#### Prompt Masking
|
||||
|
||||
The default training computes a loss for every token in the sample. You can
|
||||
ignore the prompt and compute loss for just the completion by passing
|
||||
`--mask-prompt`. Note this is only supported for `chat` and `completion`
|
||||
datasets. For `chat` datasets the final message in the message list is
|
||||
considered the completion. See the [dataset section](#Data) for more details.
|
||||
|
||||
### Evaluate
|
||||
|
||||
To compute test set perplexity use:
|
||||
@@ -134,7 +118,7 @@ mlx_lm.fuse --model <path_to_model>
|
||||
```
|
||||
|
||||
This will by default load the adapters from `adapters/`, and save the fused
|
||||
model in the path `fused_model/`. All of these are configurable.
|
||||
model in the path `lora_fused_model/`. All of these are configurable.
|
||||
|
||||
To upload a fused model, supply the `--upload-repo` and `--hf-path` arguments
|
||||
to `mlx_lm.fuse`. The latter is the repo name of the original model, which is
|
||||
@@ -157,7 +141,7 @@ mlx_lm.fuse \
|
||||
--export-gguf
|
||||
```
|
||||
|
||||
This will save the GGUF model in `fused_model/ggml-model-f16.gguf`. You
|
||||
This will save the GGUF model in `lora_fused_model/ggml-model-f16.gguf`. You
|
||||
can specify the file name with `--gguf-path`.
|
||||
|
||||
## Data
|
||||
@@ -167,173 +151,59 @@ Examples GitHub repo has an [example of the WikiSQL
|
||||
data](https://github.com/ml-explore/mlx-examples/tree/main/lora/data) in the
|
||||
correct format.
|
||||
|
||||
Datasets can be specified in `*.jsonl` files locally or loaded from Hugging
|
||||
Face.
|
||||
|
||||
### Local Datasets
|
||||
|
||||
For fine-tuning (`--train`), the data loader expects a `train.jsonl` and a
|
||||
`valid.jsonl` to be in the data directory. For evaluation (`--test`), the data
|
||||
loader expects a `test.jsonl` in the data directory.
|
||||
loader expects a `test.jsonl` in the data directory.
|
||||
|
||||
Currently, `*.jsonl` files support `chat`, `tools`, `completions`, and `text`
|
||||
data formats. Here are examples of these formats:
|
||||
Currently, `*.jsonl` files support three data formats: `chat`,
|
||||
`completions`, and `text`. Here are three examples of these formats:
|
||||
|
||||
`chat`:
|
||||
|
||||
```jsonl
|
||||
{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello."}, {"role": "assistant", "content": "How can I assistant you today."}]}
|
||||
```
|
||||
|
||||
`tools`:
|
||||
|
||||
```jsonl
|
||||
{"messages":[{"role":"user","content":"What is the weather in San Francisco?"},{"role":"assistant","tool_calls":[{"id":"call_id","type":"function","function":{"name":"get_current_weather","arguments":"{\"location\": \"San Francisco, USA\", \"format\": \"celsius\"}"}}]}],"tools":[{"type":"function","function":{"name":"get_current_weather","description":"Get the current weather","parameters":{"type":"object","properties":{"location":{"type":"string","description":"The city and country, eg. San Francisco, USA"},"format":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location","format"]}}}]}
|
||||
```
|
||||
|
||||
<details>
|
||||
<summary>View the expanded single data tool format</summary>
|
||||
|
||||
```jsonl
|
||||
{
|
||||
"messages": [
|
||||
{ "role": "user", "content": "What is the weather in San Francisco?" },
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_id",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"arguments": "{\"location\": \"San Francisco, USA\", \"format\": \"celsius\"}"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and country, eg. San Francisco, USA"
|
||||
},
|
||||
"format": { "type": "string", "enum": ["celsius", "fahrenheit"] }
|
||||
},
|
||||
"required": ["location", "format"]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant."
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hello."
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "How can I assistant you today."
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
The format for the `arguments` field in a function varies for different models.
|
||||
Common formats include JSON strings and dictionaries. The example provided
|
||||
follows the format used by
|
||||
[OpenAI](https://platform.openai.com/docs/guides/fine-tuning/fine-tuning-examples)
|
||||
and [Mistral
|
||||
AI](https://github.com/mistralai/mistral-finetune?tab=readme-ov-file#instruct).
|
||||
A dictionary format is used in Hugging Face's [chat
|
||||
templates](https://huggingface.co/docs/transformers/main/en/chat_templating#a-complete-tool-use-example).
|
||||
Refer to the documentation for the model you are fine-tuning for more details.
|
||||
|
||||
</details>
|
||||
|
||||
`completions`:
|
||||
|
||||
```jsonl
|
||||
{"prompt": "What is the capital of France?", "completion": "Paris."}
|
||||
{
|
||||
"prompt": "What is the capital of France?",
|
||||
"completion": "Paris."
|
||||
}
|
||||
```
|
||||
|
||||
For the `completions` data format, a different key can be used for the prompt
|
||||
and completion by specifying the following in the YAML config:
|
||||
|
||||
```yaml
|
||||
prompt_feature: "input"
|
||||
completion_feature: "output"
|
||||
```
|
||||
|
||||
Here, `"input"` is the expected key instead of the default `"prompt"`, and
|
||||
`"output"` is the expected key instead of `"completion"`.
|
||||
|
||||
`text`:
|
||||
|
||||
```jsonl
|
||||
{"text": "This is an example for the model."}
|
||||
{
|
||||
"text": "This is an example for the model."
|
||||
}
|
||||
```
|
||||
|
||||
Note, the format is automatically determined by the dataset. Note also, keys
|
||||
in each line not expected by the loader will be ignored.
|
||||
Note, the format is automatically determined by the dataset. Note also, keys in
|
||||
each line not expected by the loader will be ignored.
|
||||
|
||||
> [!NOTE]
|
||||
> Each example in the datasets must be on a single line. Do not put more than
|
||||
> one example per line and do not split an example across multiple lines.
|
||||
|
||||
### Hugging Face Datasets
|
||||
|
||||
To use Hugging Face datasets, first install the `datasets` package:
|
||||
|
||||
```
|
||||
pip install datasets
|
||||
```
|
||||
|
||||
If the Hugging Face dataset is already in a supported format, you can specify
|
||||
it on the command line. For example, pass `--data mlx-community/wikisql` to
|
||||
train on the pre-formatted WikiwSQL data.
|
||||
|
||||
Otherwise, provide a mapping of keys in the dataset to the features MLX LM
|
||||
expects. Use a YAML config to specify the Hugging Face dataset arguments. For
|
||||
example:
|
||||
|
||||
```yaml
|
||||
hf_dataset:
|
||||
path: "billsum"
|
||||
prompt_feature: "text"
|
||||
completion_feature: "summary"
|
||||
```
|
||||
|
||||
- Use `prompt_feature` and `completion_feature` to specify keys for a
|
||||
`completions` dataset. Use `text_feature` to specify the key for a `text`
|
||||
dataset. Use `chat_feature` to specify the key for a chat dataset.
|
||||
|
||||
- To specify the train, valid, or test splits, set the corresponding
|
||||
`{train,valid,test}_split` argument.
|
||||
|
||||
You can specify a list of Hugging Face datasets with a list of records each
|
||||
with the same structure as above. For example:
|
||||
|
||||
```yaml
|
||||
hf_dataset:
|
||||
- path: "Open-Orca/OpenOrca"
|
||||
train_split: "train[:90%]"
|
||||
valid_split: "train[-10%:]"
|
||||
prompt_feature: "question"
|
||||
completion_feature: "response"
|
||||
- path: "trl-lib/ultrafeedback_binarized"
|
||||
train_split: "train[:90%]"
|
||||
valid_split: "train[-10%:]"
|
||||
chat_feature: "chosen"
|
||||
```
|
||||
|
||||
- Arguments specified in `config` will be passed as keyword arguments to
|
||||
[`datasets.load_dataset`](https://huggingface.co/docs/datasets/v2.20.0/en/package_reference/loading_methods#datasets.load_dataset).
|
||||
|
||||
In general, for the `chat`, `tools` and `completions` formats, Hugging Face
|
||||
[chat
|
||||
templates](https://huggingface.co/docs/transformers/main/en/chat_templating)
|
||||
are used. This applies the model's chat template by default. If the model does
|
||||
not have a chat template, then Hugging Face will use a default. For example,
|
||||
the final text in the `chat` example above with Hugging Face's default template
|
||||
becomes:
|
||||
For the `chat` and `completions` formats, Hugging Face [chat
|
||||
templates](https://huggingface.co/blog/chat-templates) are used. This applies
|
||||
the model's chat template by default. If the model does not have a chat
|
||||
template, then Hugging Face will use a default. For example, the final text in
|
||||
the `chat` example above with Hugging Face's default template becomes:
|
||||
|
||||
```text
|
||||
<|im_start|>system
|
||||
@@ -361,7 +231,7 @@ of memory. Here are some tips to reduce memory use should you need to do so:
|
||||
setting this to `2` or `1` will reduce memory consumption. This may slow
|
||||
things down a little, but will also reduce the memory use.
|
||||
|
||||
3. Reduce the number of layers to fine-tune with `--num-layers`. The default
|
||||
3. Reduce the number of layers to fine-tune with `--lora-layers`. The default
|
||||
is `16`, so you can try `8` or `4`. This reduces the amount of memory
|
||||
needed for back propagation. It may also reduce the quality of the
|
||||
fine-tuned model if you are fine-tuning with a lot of data.
|
||||
@@ -383,8 +253,8 @@ mlx_lm.lora \
|
||||
--model mistralai/Mistral-7B-v0.1 \
|
||||
--train \
|
||||
--batch-size 1 \
|
||||
--num-layers 4 \
|
||||
--data mlx-community/wikisql
|
||||
--lora-layers 4 \
|
||||
--data wikisql
|
||||
```
|
||||
|
||||
The above command on an M1 Max with 32 GB runs at about 250
|
||||
@@ -393,5 +263,4 @@ tokens-per-second, using the MLX Example
|
||||
data set.
|
||||
|
||||
[^lora]: Refer to the [arXiv paper](https://arxiv.org/abs/2106.09685) for more details on LoRA.
|
||||
|
||||
[^qlora]: Refer to the paper [QLoRA: Efficient Finetuning of Quantized LLMs](https://arxiv.org/abs/2305.14314)
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
# Model Merging
|
||||
|
||||
You can use `mlx-lm` to merge models and upload them to the Hugging
|
||||
Face hub or save them locally for LoRA fine tuning.
|
||||
|
||||
The main command is `mlx_lm.merge`:
|
||||
|
||||
```shell
|
||||
mlx_lm.merge --config config.yaml
|
||||
```
|
||||
|
||||
The merged model will be saved by default in `mlx_merged_model`. To see a
|
||||
full list of options run:
|
||||
|
||||
```shell
|
||||
mlx_lm.merge --help
|
||||
```
|
||||
|
||||
Here is an example `config.yaml`:
|
||||
|
||||
```yaml
|
||||
models:
|
||||
- OpenPipe/mistral-ft-optimized-1218
|
||||
- mlabonne/NeuralHermes-2.5-Mistral-7B
|
||||
method: slerp
|
||||
parameters:
|
||||
t:
|
||||
- filter: self_attn
|
||||
value: [0, 0.5, 0.3, 0.7, 1]
|
||||
- filter: mlp
|
||||
value: [1, 0.5, 0.7, 0.3, 0]
|
||||
- value: 0.5
|
||||
```
|
||||
|
||||
The `models` field is a list of Hugging Face repo ids. The first model in the
|
||||
list is treated as the base model into which the remaining models are merged.
|
||||
|
||||
The `method` field is the merging method. Right now `slerp` is the only
|
||||
supported method.
|
||||
|
||||
The `parameters` are the corresponding parameters for the given `method`.
|
||||
Each parameter is a list with `filter` determining which layer the parameter
|
||||
applies to and `value` determining the actual value used. The last item in
|
||||
the list without a `filter` field is the default.
|
||||
|
||||
If `value` is a list, it specifies the start and end values for the
|
||||
corresponding segment of blocks. In the example above, the models have 32
|
||||
blocks. For blocks 1-8, the layers with `self_attn` in the name will use the
|
||||
values `np.linspace(0, 0.5, 8)`, the same layers in the next 8 blocks (9-16)
|
||||
will use `np.linspace(0.5, 0.3, 8)`, and so on.
|
||||
+5
-72
@@ -17,7 +17,7 @@ mlx_lm.server --model <path_to_model_or_hf_repo>
|
||||
For example:
|
||||
|
||||
```shell
|
||||
mlx_lm.server --model mlx-community/Mistral-7B-Instruct-v0.3-4bit
|
||||
mlx_lm.server --model mistralai/Mistral-7B-Instruct-v0.1
|
||||
```
|
||||
|
||||
This will start a text generation server on port `8080` of the `localhost`
|
||||
@@ -50,28 +50,22 @@ curl localhost:8080/v1/chat/completions \
|
||||
- `role_mapping`: (Optional) A dictionary to customize the role prefixes in
|
||||
the generated prompt. If not provided, the default mappings are used.
|
||||
|
||||
- `stop`: (Optional) An array of strings or a single string. These are
|
||||
- `stop`: (Optional) An array of strings or a single string. Thesse are
|
||||
sequences of tokens on which the generation should stop.
|
||||
|
||||
- `max_tokens`: (Optional) An integer specifying the maximum number of tokens
|
||||
to generate. Defaults to `512`.
|
||||
to generate. Defaults to `100`.
|
||||
|
||||
- `stream`: (Optional) A boolean indicating if the response should be
|
||||
streamed. If true, responses are sent as they are generated. Defaults to
|
||||
false.
|
||||
|
||||
- `temperature`: (Optional) A float specifying the sampling temperature.
|
||||
Defaults to `0.0`.
|
||||
Defaults to `1.0`.
|
||||
|
||||
- `top_p`: (Optional) A float specifying the nucleus sampling parameter.
|
||||
Defaults to `1.0`.
|
||||
|
||||
- `top_k`: (Optional) An integer specifying the top-k sampling parameter.
|
||||
Defaults to `0` (disabled).
|
||||
|
||||
- `min_p`: (Optional) A float specifying the min-p sampling parameter.
|
||||
Defaults to `0.0` (disabled).
|
||||
|
||||
- `repetition_penalty`: (Optional) Applies a penalty to repeated tokens.
|
||||
Defaults to `1.0`.
|
||||
|
||||
@@ -79,65 +73,4 @@ curl localhost:8080/v1/chat/completions \
|
||||
applying repetition penalty. Defaults to `20`.
|
||||
|
||||
- `logit_bias`: (Optional) A dictionary mapping token IDs to their bias
|
||||
values. Defaults to `None`.
|
||||
|
||||
- `logprobs`: (Optional) An integer specifying the number of top tokens and
|
||||
corresponding log probabilities to return for each output in the generated
|
||||
sequence. If set, this can be any value between 1 and 10, inclusive.
|
||||
|
||||
- `model`: (Optional) A string path to a local model or Hugging Face repo id.
|
||||
If the path is local is must be relative to the directory the server was
|
||||
started in.
|
||||
|
||||
- `adapters`: (Optional) A string path to low-rank adapters. The path must be
|
||||
relative to the directory the server was started in.
|
||||
|
||||
- `draft_model`: (Optional) Specifies a smaller model to use for speculative
|
||||
decoding. Set to `null` to unload.
|
||||
|
||||
- `num_draft_tokens`: (Optional) The number of draft tokens the draft model
|
||||
should predict at once. Defaults to `3`.
|
||||
|
||||
### Response Fields
|
||||
|
||||
- `id`: A unique identifier for the chat.
|
||||
|
||||
- `system_fingerprint`: A unique identifier for the system.
|
||||
|
||||
- `object`: Any of "chat.completion", "chat.completion.chunk" (for
|
||||
streaming), or "text.completion".
|
||||
|
||||
- `model`: The model repo or path (e.g. `"mlx-community/Llama-3.2-3B-Instruct-4bit"`).
|
||||
|
||||
- `created`: A time-stamp for when the request was processed.
|
||||
|
||||
- `choices`: A list of outputs. Each output is a dictionary containing the fields:
|
||||
- `index`: The index in the list.
|
||||
- `logprobs`: A dictionary containing the fields:
|
||||
- `token_logprobs`: A list of the log probabilities for the generated
|
||||
tokens.
|
||||
- `tokens`: A list of the generated token ids.
|
||||
- `top_logprobs`: A list of lists. Each list contains the `logprobs`
|
||||
top tokens (if requested) with their corresponding probabilities.
|
||||
- `finish_reason`: The reason the completion ended. This can be either of
|
||||
`"stop"` or `"length"`.
|
||||
- `message`: The text response from the model.
|
||||
|
||||
- `usage`: A dictionary containing the fields:
|
||||
- `prompt_tokens`: The number of prompt tokens processed.
|
||||
- `completion_tokens`: The number of tokens generated.
|
||||
- `total_tokens`: The total number of tokens, i.e. the sum of the above two fields.
|
||||
|
||||
### List Models
|
||||
|
||||
Use the `v1/models` endpoint to list available models:
|
||||
|
||||
```shell
|
||||
curl localhost:8080/v1/models -H "Content-Type: application/json"
|
||||
```
|
||||
|
||||
This will return a list of locally available models where each model in the
|
||||
list contains the following fields:
|
||||
|
||||
- `id`: The Hugging Face repo id.
|
||||
- `created`: A time-stamp representing the model creation time.
|
||||
values. Defaults to `None`.
|
||||
+2
-9
@@ -1,11 +1,4 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import os
|
||||
|
||||
from ._version import __version__
|
||||
|
||||
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
|
||||
|
||||
from .convert import convert
|
||||
from .generate import generate, stream_generate
|
||||
from .utils import load
|
||||
from .utils import convert, generate, load, stream_generate
|
||||
from .version import __version__
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
|
||||
if __name__ == "__main__":
|
||||
subcommands = {
|
||||
"quant.awq",
|
||||
"quant.dwq",
|
||||
"quant.dynamic_quant",
|
||||
"cache_prompt",
|
||||
"chat",
|
||||
"convert",
|
||||
"evaluate",
|
||||
"fuse",
|
||||
"generate",
|
||||
"lora",
|
||||
"server",
|
||||
"manage",
|
||||
"upload",
|
||||
}
|
||||
if len(sys.argv) < 2:
|
||||
raise ValueError(f"CLI requires a subcommand in {subcommands}")
|
||||
subcommand = sys.argv.pop(1)
|
||||
if subcommand not in subcommands:
|
||||
raise ValueError(f"CLI requires a subcommand in {subcommands}")
|
||||
submodule = importlib.import_module(f"mlx_lm.{subcommand}")
|
||||
submodule.main()
|
||||
@@ -1,166 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
from .generate import generate_step
|
||||
from .models.cache import make_prompt_cache, save_prompt_cache
|
||||
from .utils import load
|
||||
|
||||
DEFAULT_QUANTIZED_KV_START = 5000
|
||||
|
||||
|
||||
def setup_arg_parser():
|
||||
"""Set up and return the argument parser."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Cache the state of a prompt to be reused with mlx_lm.generate"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="mlx_model",
|
||||
help="The path to the local model directory or Hugging Face repo.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--adapter-path",
|
||||
type=str,
|
||||
help="Optional path for the trained adapter weights and config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--trust-remote-code",
|
||||
action="store_true",
|
||||
help="Enable trusting remote code for tokenizer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eos-token",
|
||||
type=str,
|
||||
default=None,
|
||||
help="End of sequence token for tokenizer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ignore-chat-template",
|
||||
action="store_true",
|
||||
help="Use the raw prompt without the tokenizer's chat template.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-default-chat-template",
|
||||
action="store_true",
|
||||
help="Use the default chat template",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-kv-size",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Set the maximum key-value cache size",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-cache-file",
|
||||
help="The file to save the prompt cache in",
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
required=True,
|
||||
help="Message to be processed by the model ('-' reads from stdin)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-bits",
|
||||
type=int,
|
||||
help="Number of bits for KV cache quantization. "
|
||||
"Defaults to no quantization.",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-group-size",
|
||||
type=int,
|
||||
help="Group size for KV cache quantization.",
|
||||
default=64,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantized-kv-start",
|
||||
help="When --kv-bits is set, start quantizing the KV cache "
|
||||
"from this step onwards.",
|
||||
type=int,
|
||||
default=DEFAULT_QUANTIZED_KV_START,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main():
|
||||
parser = setup_arg_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
# Building tokenizer_config
|
||||
tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
|
||||
if args.eos_token is not None:
|
||||
tokenizer_config["eos_token"] = args.eos_token
|
||||
|
||||
model, tokenizer = load(
|
||||
args.model,
|
||||
adapter_path=args.adapter_path,
|
||||
tokenizer_config=tokenizer_config,
|
||||
)
|
||||
|
||||
args.prompt = sys.stdin.read() if args.prompt == "-" else args.prompt
|
||||
|
||||
if args.use_default_chat_template:
|
||||
if tokenizer.chat_template is None:
|
||||
tokenizer.chat_template = tokenizer.default_chat_template
|
||||
|
||||
if not args.ignore_chat_template and tokenizer.chat_template is not None:
|
||||
messages = [{"role": "user", "content": args.prompt}]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=False, continue_final_message=True
|
||||
)
|
||||
|
||||
else:
|
||||
prompt = tokenizer.encode(args.prompt)
|
||||
|
||||
cache = make_prompt_cache(model, args.max_kv_size)
|
||||
y = mx.array(prompt)
|
||||
|
||||
# Process the prompt
|
||||
start = time.time()
|
||||
max_msg_len = 0
|
||||
|
||||
def callback(processed, total_tokens):
|
||||
current = time.time()
|
||||
speed = processed / (current - start)
|
||||
msg = f"\rProcessed {processed:6d} tokens ({speed:6.2f} tok/s)"
|
||||
nonlocal max_msg_len
|
||||
max_msg_len = max(max_msg_len, len(msg))
|
||||
print(msg + " " * (max_msg_len - len(msg)), end="", flush=True)
|
||||
|
||||
for _ in generate_step(
|
||||
y,
|
||||
model,
|
||||
max_tokens=0,
|
||||
prompt_cache=cache,
|
||||
kv_bits=args.kv_bits,
|
||||
kv_group_size=args.kv_group_size,
|
||||
quantized_kv_start=args.quantized_kv_start,
|
||||
prompt_progress_callback=callback,
|
||||
):
|
||||
pass
|
||||
|
||||
print()
|
||||
print(f"Peak memory: {mx.get_peak_memory() / 1e9:.3f} GB")
|
||||
|
||||
print("Saving...")
|
||||
metadata = {}
|
||||
metadata["model"] = args.model
|
||||
metadata["chat_template"] = json.dumps(tokenizer.chat_template)
|
||||
metadata["tokenizer_config"] = json.dumps(tokenizer_config)
|
||||
save_prompt_cache(args.prompt_cache_file, cache, metadata)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.cache_prompt...` directly is deprecated."
|
||||
" Use `mlx_lm.cache_prompt...` or `python -m mlx_lm cache_prompt ...` instead."
|
||||
)
|
||||
main()
|
||||
-134
@@ -1,134 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import argparse
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
from .generate import stream_generate
|
||||
from .models.cache import make_prompt_cache
|
||||
from .sample_utils import make_sampler
|
||||
from .utils import load
|
||||
|
||||
DEFAULT_TEMP = 0.0
|
||||
DEFAULT_TOP_P = 1.0
|
||||
DEFAULT_XTC_PROBABILITY = 0.0
|
||||
DEFAULT_XTC_THRESHOLD = 0.0
|
||||
DEFAULT_SEED = None
|
||||
DEFAULT_MAX_TOKENS = 256
|
||||
DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
|
||||
|
||||
|
||||
def setup_arg_parser():
|
||||
"""Set up and return the argument parser."""
|
||||
parser = argparse.ArgumentParser(description="Chat with an LLM")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
help="The path to the local model directory or Hugging Face repo.",
|
||||
default=DEFAULT_MODEL,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--adapter-path",
|
||||
type=str,
|
||||
help="Optional path for the trained adapter weights and config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temp", type=float, default=DEFAULT_TEMP, help="Sampling temperature"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--xtc-probability",
|
||||
type=float,
|
||||
default=DEFAULT_XTC_PROBABILITY,
|
||||
help="Probability of XTC sampling to happen each next token",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--xtc-threshold",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Thresold the probs of each next token candidate to be sampled by XTC",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=DEFAULT_SEED,
|
||||
help="PRNG seed",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-kv-size",
|
||||
type=int,
|
||||
help="Set the maximum key-value cache size",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
"-m",
|
||||
type=int,
|
||||
default=DEFAULT_MAX_TOKENS,
|
||||
help="Maximum number of tokens to generate",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main():
|
||||
parser = setup_arg_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.seed is not None:
|
||||
mx.random.seed(args.seed)
|
||||
|
||||
model, tokenizer = load(
|
||||
args.model,
|
||||
adapter_path=args.adapter_path,
|
||||
tokenizer_config={"trust_remote_code": True},
|
||||
)
|
||||
|
||||
def print_help():
|
||||
print("The command list:")
|
||||
print("- 'q' to exit")
|
||||
print("- 'r' to reset the chat")
|
||||
print("- 'h' to display these commands")
|
||||
|
||||
print(f"[INFO] Starting chat session with {args.model}.")
|
||||
print_help()
|
||||
prompt_cache = make_prompt_cache(model, args.max_kv_size)
|
||||
while True:
|
||||
query = input(">> ")
|
||||
if query == "q":
|
||||
break
|
||||
if query == "r":
|
||||
prompt_cache = make_prompt_cache(model, args.max_kv_size)
|
||||
continue
|
||||
if query == "h":
|
||||
print_help()
|
||||
continue
|
||||
messages = [{"role": "user", "content": query}]
|
||||
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
|
||||
for response in stream_generate(
|
||||
model,
|
||||
tokenizer,
|
||||
prompt,
|
||||
max_tokens=args.max_tokens,
|
||||
sampler=make_sampler(
|
||||
args.temp,
|
||||
args.top_p,
|
||||
xtc_threshold=args.xtc_threshold,
|
||||
xtc_probability=args.xtc_probability,
|
||||
xtc_special_tokens=(
|
||||
tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
|
||||
),
|
||||
),
|
||||
prompt_cache=prompt_cache,
|
||||
):
|
||||
print(response.text, flush=True, end="")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.chat...` directly is deprecated."
|
||||
" Use `mlx_lm.chat...` or `python -m mlx_lm chat ...` instead."
|
||||
)
|
||||
main()
|
||||
+4
-181
@@ -1,174 +1,8 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_map_with_path
|
||||
|
||||
from .utils import (
|
||||
dequantize_model,
|
||||
fetch_from_hub,
|
||||
get_model_path,
|
||||
quantize_model,
|
||||
save,
|
||||
upload_to_hub,
|
||||
)
|
||||
|
||||
|
||||
def mixed_quant_predicate_builder(
|
||||
recipe: str, model: nn.Module
|
||||
) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
|
||||
|
||||
high_bits = 6
|
||||
group_size = 64
|
||||
|
||||
if recipe == "mixed_2_6":
|
||||
low_bits = 2
|
||||
elif recipe == "mixed_3_4":
|
||||
low_bits = 3
|
||||
high_bits = 4
|
||||
elif recipe == "mixed_3_6":
|
||||
low_bits = 3
|
||||
elif recipe == "mixed_4_6":
|
||||
low_bits = 4
|
||||
else:
|
||||
raise ValueError("Invalid quant recipe {recipe}")
|
||||
|
||||
down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
|
||||
if len(down_keys) == 0:
|
||||
raise ValueError("Model does not have expected keys for mixed quant.")
|
||||
|
||||
# Look for the layer index location in the path:
|
||||
for layer_location, k in enumerate(down_keys[0].split(".")):
|
||||
if k.isdigit():
|
||||
break
|
||||
num_layers = len(model.layers)
|
||||
|
||||
def mixed_quant_predicate(
|
||||
path: str,
|
||||
module: nn.Module,
|
||||
config: dict,
|
||||
) -> Union[bool, dict]:
|
||||
"""Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
|
||||
Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
|
||||
By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
|
||||
"""
|
||||
|
||||
if not hasattr(module, "to_quantized"):
|
||||
return False
|
||||
if module.weight.shape[1] % group_size != 0:
|
||||
return False
|
||||
|
||||
index = (
|
||||
int(path.split(".")[layer_location])
|
||||
if len(path.split(".")) > layer_location
|
||||
else 0
|
||||
)
|
||||
use_more_bits = (
|
||||
index < num_layers // 8
|
||||
or index >= 7 * num_layers // 8
|
||||
or (index - num_layers // 8) % 3 == 2
|
||||
)
|
||||
if "v_proj" in path and use_more_bits:
|
||||
return {"group_size": group_size, "bits": high_bits}
|
||||
if "down_proj" in path and use_more_bits:
|
||||
return {"group_size": group_size, "bits": high_bits}
|
||||
if "lm_head" in path:
|
||||
return {"group_size": group_size, "bits": high_bits}
|
||||
|
||||
return {"group_size": group_size, "bits": low_bits}
|
||||
|
||||
return mixed_quant_predicate
|
||||
|
||||
|
||||
QUANT_RECIPES = ["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"]
|
||||
|
||||
MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]
|
||||
|
||||
|
||||
def convert(
|
||||
hf_path: str,
|
||||
mlx_path: str = "mlx_model",
|
||||
quantize: bool = False,
|
||||
q_group_size: int = 64,
|
||||
q_bits: int = 4,
|
||||
dtype: Optional[str] = None,
|
||||
upload_repo: str = None,
|
||||
revision: Optional[str] = None,
|
||||
dequantize: bool = False,
|
||||
quant_predicate: Optional[
|
||||
Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
|
||||
] = None,
|
||||
):
|
||||
# Check the save path is empty
|
||||
if isinstance(mlx_path, str):
|
||||
mlx_path = Path(mlx_path)
|
||||
|
||||
if mlx_path.exists():
|
||||
raise ValueError(
|
||||
f"Cannot save to the path {mlx_path} as it already exists."
|
||||
" Please delete the file/directory or specify a new path to save to."
|
||||
)
|
||||
|
||||
print("[INFO] Loading")
|
||||
model_path, hf_path = get_model_path(hf_path, revision=revision)
|
||||
model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
|
||||
|
||||
def base_quant_predicate(path, module, config):
|
||||
if not hasattr(module, "to_quantized"):
|
||||
return False
|
||||
if module.weight.shape[1] % q_group_size != 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
if isinstance(quant_predicate, str):
|
||||
quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
|
||||
quant_predicate = quant_predicate or base_quant_predicate
|
||||
|
||||
if dtype is None:
|
||||
dtype = config.get("torch_dtype", None)
|
||||
if dtype in MODEL_CONVERSION_DTYPES:
|
||||
print("[INFO] Using dtype:", dtype)
|
||||
dtype = getattr(mx, dtype)
|
||||
cast_predicate = getattr(model, "cast_predicate", lambda _: True)
|
||||
|
||||
def set_dtype(k, v):
|
||||
if cast_predicate(k) and mx.issubdtype(v.dtype, mx.floating):
|
||||
return v.astype(dtype)
|
||||
else:
|
||||
return v
|
||||
|
||||
model.update(tree_map_with_path(set_dtype, model.parameters()))
|
||||
|
||||
if quantize and dequantize:
|
||||
raise ValueError("Choose either quantize or dequantize, not both.")
|
||||
|
||||
if quantize:
|
||||
print("[INFO] Quantizing")
|
||||
model, config = quantize_model(
|
||||
model, config, q_group_size, q_bits, quant_predicate=quant_predicate
|
||||
)
|
||||
|
||||
if dequantize:
|
||||
print("[INFO] Dequantizing")
|
||||
config.pop("quantization", None)
|
||||
config.pop("quantization_config", None)
|
||||
model = dequantize_model(model)
|
||||
|
||||
save(
|
||||
mlx_path,
|
||||
model_path,
|
||||
model,
|
||||
tokenizer,
|
||||
config,
|
||||
hf_repo=hf_path,
|
||||
)
|
||||
|
||||
if upload_repo is not None:
|
||||
upload_to_hub(mlx_path, upload_repo)
|
||||
from .utils import convert
|
||||
|
||||
|
||||
def configure_parser() -> argparse.ArgumentParser:
|
||||
@@ -195,19 +29,12 @@ def configure_parser() -> argparse.ArgumentParser:
|
||||
parser.add_argument(
|
||||
"--q-bits", help="Bits per weight for quantization.", type=int, default=4
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quant-predicate",
|
||||
help=f"Mixed-bit quantization recipe.",
|
||||
choices=QUANT_RECIPES,
|
||||
type=str,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
help="Type to save the non-quantized parameters. Defaults to config.json's `torch_dtype` or the current model weights dtype.",
|
||||
help="Type to save the parameters, ignored if -q is given.",
|
||||
type=str,
|
||||
choices=MODEL_CONVERSION_DTYPES,
|
||||
default=None,
|
||||
choices=["float16", "bfloat16", "float32"],
|
||||
default="float16",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-repo",
|
||||
@@ -232,8 +59,4 @@ def main():
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.convert ...` directly is deprecated."
|
||||
" Use `mlx_lm.convert ...` or `python -m mlx_lm convert ...` instead."
|
||||
)
|
||||
main()
|
||||
|
||||
@@ -1,415 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
"""
|
||||
Adapted from a PyTorch implementation by David Grangier
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import collections
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from importlib.metadata import version
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
import lm_eval
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
from lm_eval.api.model import LM
|
||||
from lm_eval.api.registry import register_model
|
||||
from tqdm import tqdm
|
||||
|
||||
from .generate import stream_generate
|
||||
from .models.base import create_causal_mask
|
||||
from .models.cache import make_prompt_cache
|
||||
from .utils import common_prefix_len, load
|
||||
|
||||
|
||||
def _rstrip_until(s, untils):
|
||||
"""Limit a string <s> to the first occurrence of any substring in untils."""
|
||||
l = len(s)
|
||||
f = [s.find(u) for u in untils]
|
||||
f = [l if x < 0 else x for x in f]
|
||||
return s[: min(f)]
|
||||
|
||||
|
||||
def _pad_inputs(inputs):
|
||||
lengths = np.array([len(x) for x in inputs])
|
||||
maxlen = lengths.max()
|
||||
padded = np.stack(
|
||||
[np.pad(x, (0, maxlen - len(x))) for x in inputs],
|
||||
axis=0,
|
||||
)
|
||||
return mx.array(padded), mx.array(lengths)
|
||||
|
||||
|
||||
def chat_template_fn(**extra_kwargs):
|
||||
def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
|
||||
return self.tokenizer.apply_chat_template(
|
||||
chat_history,
|
||||
tokenize=False,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
continue_final_message=not add_generation_prompt,
|
||||
**extra_kwargs,
|
||||
)
|
||||
|
||||
return apply_chat_template
|
||||
|
||||
|
||||
@register_model("mlxlm")
|
||||
class MLXLM(LM):
|
||||
|
||||
tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
|
||||
apply_chat_template = chat_template_fn()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path_or_hf_repo: str,
|
||||
max_tokens: Optional[int] = None,
|
||||
use_chat_template: Optional[bool] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._model, self.tokenizer = load(path_or_hf_repo)
|
||||
self._max_tokens = max_tokens or self.tokenizer.model_max_length
|
||||
self._batch_size = 8
|
||||
self.use_chat_template = use_chat_template
|
||||
if use_chat_template is None:
|
||||
self.use_chat_template = self.tokenizer.chat_template is not None
|
||||
|
||||
def _process_prompt(self, prompt, step_size: int = 2048):
|
||||
prompt = mx.array(prompt)[None]
|
||||
cache = make_prompt_cache(self._model)
|
||||
for i in range(0, prompt.shape[1], step_size):
|
||||
logits = self._model(prompt[:, i : i + step_size], cache=cache)
|
||||
mx.eval([c.state for c in cache])
|
||||
mx.clear_cache()
|
||||
logprobs = nn.log_softmax(logits[:, -1, :].astype(mx.float32))
|
||||
return logprobs, cache
|
||||
|
||||
def _score_fn(self, inputs, cache: Optional[Any] = None, step_size: int = 2048):
|
||||
inputs, lengths = _pad_inputs(inputs)
|
||||
inputs, targets = inputs[..., :-1], inputs[..., 1:]
|
||||
|
||||
cache = cache or make_prompt_cache(self._model)
|
||||
lengths += cache[0].offset
|
||||
|
||||
scores, is_greedy = [], []
|
||||
for i in range(0, inputs.shape[1], step_size):
|
||||
inp = inputs[:, i : i + step_size]
|
||||
T = inp.shape[1]
|
||||
|
||||
offset = cache[0].offset
|
||||
mask = create_causal_mask(T, offset, lengths=lengths)
|
||||
|
||||
logits = self._model(inp, cache=cache, mask=mask)
|
||||
log_probs = nn.log_softmax(logits.astype(mx.float32))
|
||||
|
||||
score = mx.take_along_axis(
|
||||
log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
|
||||
)[..., 0]
|
||||
ig = targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
|
||||
ig = mx.where(mx.arange(T) + offset < lengths[:, None], ig, False)
|
||||
|
||||
mx.eval(score, ig)
|
||||
mx.clear_cache()
|
||||
|
||||
is_greedy.append(ig)
|
||||
scores.append(score)
|
||||
|
||||
scores = mx.concatenate(scores, axis=1)
|
||||
is_greedy = mx.concatenate(is_greedy, axis=1)
|
||||
|
||||
return scores, lengths, is_greedy
|
||||
|
||||
def _tokenize(self, texts):
|
||||
return [
|
||||
tuple(
|
||||
self.tokenizer.encode(t, add_special_tokens=not self.use_chat_template)
|
||||
)
|
||||
for t in texts
|
||||
]
|
||||
|
||||
def loglikelihood(self, requests) -> list[tuple[float, bool]]:
|
||||
"""Compute log-likelihood of generating a continuation from a context.
|
||||
Downstream tasks should attempt to use loglikelihood instead of other
|
||||
LM calls whenever possible.
|
||||
:param requests: list[Instance]
|
||||
A list of Instance objects, with property `args` which returns a tuple (context, continuation).
|
||||
`context: str`
|
||||
Context string. Implementations of LM must be able to handle an
|
||||
empty context string.
|
||||
`continuation: str`
|
||||
The continuation over which log likelihood will be calculated. If
|
||||
there is a word boundary, the space should be in the continuation.
|
||||
For example, context="hello" continuation=" world" is correct.
|
||||
:return: list[tuple[float, bool]]
|
||||
A list of pairs (logprob, isgreedy)
|
||||
`logprob: float`
|
||||
The log probability of `continuation`.
|
||||
`isgreedy`:
|
||||
Whether `continuation` would be generated by greedy sampling from `context`.
|
||||
"""
|
||||
logging.info("Estimating loglikelihood for %d pairs." % len(requests))
|
||||
|
||||
group = mx.distributed.init()
|
||||
|
||||
# Group by common prefix
|
||||
group_reqs = collections.defaultdict(list)
|
||||
for idx, req in enumerate(requests):
|
||||
group_reqs[req.args[0]].append((idx, req.args[1]))
|
||||
questions = list(group_reqs.keys())
|
||||
responses = []
|
||||
indices = []
|
||||
for v in group_reqs.values():
|
||||
idx, resp = zip(*v)
|
||||
indices.extend(idx)
|
||||
responses.append(resp)
|
||||
|
||||
# split data accross ranks
|
||||
questions = questions[group.rank() :: group.size()]
|
||||
responses = responses[group.rank() :: group.size()]
|
||||
|
||||
long_completions = 0
|
||||
scores, is_greedy = [], []
|
||||
for q, rs in tqdm(zip(questions, responses), total=len(questions)):
|
||||
prefix = self._tokenize([q])[0]
|
||||
full_sequences = self._tokenize([q + r for r in rs])
|
||||
max_completed_l = max(len(s) for s in full_sequences)
|
||||
|
||||
# compute truncation length
|
||||
truncation = max(0, max_completed_l - self._max_tokens - 1)
|
||||
orig_prefix_l = len(prefix)
|
||||
prefix_l = max(len(prefix) - truncation, 0)
|
||||
prefix = prefix[len(prefix) - prefix_l :]
|
||||
|
||||
# If the entire prompt got truncated ignore the question
|
||||
if prefix_l == 0:
|
||||
long_completions += 1
|
||||
all_scores.extend([-float("inf")] * len(rs))
|
||||
all_is_greedy.extend([False] * len(rs))
|
||||
continue
|
||||
|
||||
# model scoring, returns num_requests x (logp, is_greedy, length).
|
||||
logprobs, cache = self._process_prompt(prefix)
|
||||
max_idx = mx.argmax(logprobs).item()
|
||||
|
||||
for s in full_sequences:
|
||||
inputs = s[len(prefix) :]
|
||||
# The logprobs from the last token of the prompt are
|
||||
# for the first input token
|
||||
scores.append(logprobs[0, inputs[0]].item())
|
||||
is_greedy.append((inputs[0] == max_idx))
|
||||
|
||||
if len(inputs) == 1:
|
||||
continue
|
||||
score, _, ig = self._score_fn(
|
||||
mx.array(inputs)[None, :], cache=copy.deepcopy(cache)
|
||||
)
|
||||
scores[-1] += mx.sum(score).item()
|
||||
is_greedy[-1] &= mx.all(ig).item()
|
||||
|
||||
scores = mx.array(scores)
|
||||
is_greedy = mx.array(is_greedy)
|
||||
|
||||
if long_completions > 0:
|
||||
logging.info(
|
||||
f"Prefix eliminated for {long_completions} requests with "
|
||||
+ "completion longer than context."
|
||||
)
|
||||
|
||||
num_results = len(requests)
|
||||
|
||||
# all gather the results across groups
|
||||
if group.size() > 1:
|
||||
per_group = int(np.ceil(num_results / group.size()))
|
||||
scores = mx.pad(scores, ((0, per_group - len(scores)),))
|
||||
is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
|
||||
scores = mx.distributed.all_gather(scores[mx.newaxis], stream=mx.cpu)
|
||||
is_greedy = mx.distributed.all_gather(is_greedy[mx.newaxis], stream=mx.cpu)
|
||||
mx.eval(scores, is_greedy)
|
||||
scores = scores.T.reshape(-1)
|
||||
is_greedy = is_greedy.T.reshape(-1)
|
||||
|
||||
inv_sort = mx.argsort(mx.array(indices))
|
||||
scores = scores[:num_results][inv_sort]
|
||||
is_greedy = is_greedy[:num_results][inv_sort]
|
||||
return list(zip(scores.tolist(), is_greedy.tolist()))
|
||||
|
||||
def loglikelihood_rolling(self, requests) -> list[float]:
|
||||
"""Compute full log-likelihood of a string, with no truncation, for perplexity computation
|
||||
- We will use the full max context length of the model.
|
||||
- For inputs that exceed the max context length, we divide the tokenized string into chunks of up to
|
||||
the max context length.
|
||||
- IMPORTANT: Each document's loglikelihood/perplexity is computed *separately*, unlike other implementations
|
||||
which may simply concatenate multiple documents together.
|
||||
- IMPORTANT: We maximize the amount of context for each prediction. Specifically, for inputs that we break into
|
||||
multiple chunks, the last input will still a full-sized context.
|
||||
Example:
|
||||
Input tokens: [ 0 1 2 3 4 5 6 7 8 9 ]
|
||||
Prefix: EOT
|
||||
Max context length: 4
|
||||
Resulting input/prediction pairs:
|
||||
INPUT: EOT 0 1 2
|
||||
PRED: 0 1 2 3
|
||||
INPUT: 3 4 5 6
|
||||
PRED: 4 5 6 7
|
||||
INPUT: 5 6 7 8
|
||||
PRED: 8 9
|
||||
Observe that:
|
||||
1. Each token is predicted exactly once
|
||||
2. For the last pair, we provide the full context, but only score the last two tokens
|
||||
:param requests: list[Instance]
|
||||
A list of Instance objects with property `args` which returns a tuple (context,).
|
||||
string: str
|
||||
String for which we are computing overall loglikelihood
|
||||
:return: list[tuple[float]]
|
||||
A list of tuples (logprob,)
|
||||
logprob: float
|
||||
The log probability of `context` conditioned on the EOT token.
|
||||
"""
|
||||
logging.info(
|
||||
"Estimating loglikelihood rolling for %d sequences." % len(requests)
|
||||
)
|
||||
inputs = self._tokenize([req.args[0] for req in requests])
|
||||
all_scores = []
|
||||
for i in tqdm(range(0, len(texts), self._batch_size)):
|
||||
batch = texts[i : i + self._batch_size]
|
||||
scores, lengths, _ = self._score_fn(batch)
|
||||
mask = mx.arange(scores.shape[-1]) < lengths[:, None]
|
||||
all_scores.extend((mask * scores).sum(axis=-1).tolist())
|
||||
|
||||
return all_scores
|
||||
|
||||
def generate_until(self, requests) -> list[str]:
|
||||
"""Generate greedily until a stopping sequence
|
||||
:param requests: list[Instance]
|
||||
A list of Instance objects with property `args` which returns a tuple (context, until).
|
||||
context: str
|
||||
Context string
|
||||
until: [str]
|
||||
The string sequences to generate until. These string sequences
|
||||
may each span across multiple tokens, or may be part of one token.
|
||||
:return: list[str]
|
||||
A list of strings continuation
|
||||
continuation: str
|
||||
The generated continuation.
|
||||
"""
|
||||
logging.info("Generating continuation for %d sequences." % len(requests))
|
||||
contexts, options = zip(*[req.args for req in requests])
|
||||
# contrary to the doc the second element of the tuple contains
|
||||
# {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
|
||||
completions = []
|
||||
|
||||
for context, opt in tqdm(zip(contexts, options), total=len(contexts)):
|
||||
until = opt["until"]
|
||||
context = self.tokenizer.encode(
|
||||
context, add_special_tokens=not self.use_chat_template
|
||||
)
|
||||
max_tokens = min(
|
||||
opt.get("max_gen_tokens", self._max_tokens),
|
||||
self.tokenizer.model_max_length - len(context),
|
||||
)
|
||||
text = ""
|
||||
for response in stream_generate(
|
||||
self._model, self.tokenizer, prompt=context, max_tokens=max_tokens
|
||||
):
|
||||
text += response.text
|
||||
if any(u in text for u in until):
|
||||
text = _rstrip_until(text, until)
|
||||
completions.append(text)
|
||||
break
|
||||
else:
|
||||
completions.append(text)
|
||||
return completions
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
"Evaluate an MLX model using lm-evaluation-harness."
|
||||
)
|
||||
parser.add_argument("--model", help="Model to evaluate", required=True)
|
||||
parser.add_argument("--tasks", nargs="+", required=True)
|
||||
parser.add_argument(
|
||||
"--output-dir", default=".", help="Output directory for result files."
|
||||
)
|
||||
parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
|
||||
parser.add_argument("--num-shots", type=int, default=None, help="Number of shots")
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
help="Maximum nunber of tokens to generate. Defaults to the model's max context length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--limit",
|
||||
default=None,
|
||||
help="Limit the number of examples per task.",
|
||||
type=int,
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=123, help="Random seed.")
|
||||
parser.add_argument(
|
||||
"--fewshot-as-multiturn",
|
||||
action="store_true",
|
||||
help="Whether to provide the fewshot examples as a multiturn "
|
||||
"conversation or a single user turn.",
|
||||
default=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--apply-chat-template",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
help="Specifies whether to apply a chat template to the prompt. If "
|
||||
"the model has a chat template, this defaults to `True`, "
|
||||
"otherwise `False`.",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chat-template-args",
|
||||
type=json.loads,
|
||||
help="""A JSON formatted string of arguments for the tokenizer's "
|
||||
"apply_chat_template, e.g. '{"enable_thinking":false}'""",
|
||||
default="{}",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Silence tokenizer warnings
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
mx.random.seed(args.seed)
|
||||
|
||||
lm = MLXLM(
|
||||
args.model,
|
||||
max_tokens=args.max_tokens,
|
||||
use_chat_template=args.apply_chat_template,
|
||||
)
|
||||
MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
|
||||
|
||||
results = lm_eval.simple_evaluate(
|
||||
model=lm,
|
||||
tasks=args.tasks,
|
||||
fewshot_as_multiturn=args.fewshot_as_multiturn,
|
||||
apply_chat_template=lm.use_chat_template,
|
||||
num_fewshot=args.num_shots,
|
||||
limit=args.limit,
|
||||
random_seed=args.seed,
|
||||
numpy_random_seed=args.seed,
|
||||
torch_random_seed=args.seed,
|
||||
fewshot_random_seed=args.seed,
|
||||
)
|
||||
|
||||
file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
|
||||
if args.num_shots is not None:
|
||||
file_keys += [f"{args.num_shots:02d}"]
|
||||
file_keys += args.tasks
|
||||
filename = "_".join(file_keys)
|
||||
if mx.distributed.init().rank() == 0:
|
||||
output_path = output_dir / filename
|
||||
output_path.write_text(json.dumps(results["results"], indent=4))
|
||||
print("Results:")
|
||||
for result in results["results"].values():
|
||||
print(json.dumps(result, indent=4))
|
||||
@@ -1,47 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
"""
|
||||
An example of a multi-turn chat with prompt caching.
|
||||
"""
|
||||
|
||||
from mlx_lm import generate, load
|
||||
from mlx_lm.models.cache import load_prompt_cache, make_prompt_cache, save_prompt_cache
|
||||
|
||||
model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
|
||||
|
||||
# Make the initial prompt cache for the model
|
||||
prompt_cache = make_prompt_cache(model)
|
||||
|
||||
# User turn
|
||||
prompt = "Hi my name is <Name>."
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
|
||||
|
||||
# Assistant response
|
||||
response = generate(
|
||||
model,
|
||||
tokenizer,
|
||||
prompt=prompt,
|
||||
verbose=True,
|
||||
prompt_cache=prompt_cache,
|
||||
)
|
||||
|
||||
# User turn
|
||||
prompt = "What's my name?"
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
|
||||
|
||||
# Assistant response
|
||||
response = generate(
|
||||
model,
|
||||
tokenizer,
|
||||
prompt=prompt,
|
||||
verbose=True,
|
||||
prompt_cache=prompt_cache,
|
||||
)
|
||||
|
||||
# Save the prompt cache to disk to reuse it at a later time
|
||||
save_prompt_cache("mistral_prompt.safetensors", prompt_cache)
|
||||
|
||||
# Load the prompt cache from disk
|
||||
prompt_cache = load_prompt_cache("mistral_prompt.safetensors")
|
||||
@@ -1,33 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
from mlx_lm import generate, load
|
||||
|
||||
# Specify the checkpoint
|
||||
checkpoint = "mistralai/Mistral-7B-Instruct-v0.3"
|
||||
|
||||
# Load the corresponding model and tokenizer
|
||||
model, tokenizer = load(path_or_hf_repo=checkpoint)
|
||||
|
||||
# Specify the prompt and conversation history
|
||||
prompt = "Why is the sky blue?"
|
||||
conversation = [{"role": "user", "content": prompt}]
|
||||
|
||||
# Transform the prompt into the chat template
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
conversation=conversation, add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Specify the maximum number of tokens
|
||||
max_tokens = 1_000
|
||||
|
||||
# Specify if tokens and timing information will be printed
|
||||
verbose = True
|
||||
|
||||
# Generate a response with the specified settings
|
||||
response = generate(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
prompt=prompt,
|
||||
max_tokens=max_tokens,
|
||||
verbose=verbose,
|
||||
)
|
||||
@@ -1,29 +1,16 @@
|
||||
# The path to the local model directory or Hugging Face repo.
|
||||
model: "mlx-community/Llama-3.2-1B-Instruct"
|
||||
|
||||
model: "mlx_model"
|
||||
# Whether or not to train (boolean)
|
||||
train: true
|
||||
|
||||
# The fine-tuning method: "lora", "dora", or "full".
|
||||
fine_tune_type: lora
|
||||
|
||||
# The Optimizer with its possible inputs
|
||||
optimizer: adamw
|
||||
# optimizer_config:
|
||||
# adamw:
|
||||
# betas: [0.9, 0.98]
|
||||
# eps: 1e-6
|
||||
# weight_decay: 0.05
|
||||
# bias_correction: true
|
||||
|
||||
# Directory with {train, valid, test}.jsonl files
|
||||
data: "mlx-community/WikiSQL"
|
||||
data: "/path/to/training/data"
|
||||
|
||||
# The PRNG seed
|
||||
seed: 0
|
||||
|
||||
# Number of layers to fine-tune
|
||||
num_layers: 16
|
||||
lora_layers: 16
|
||||
|
||||
# Minibatch size.
|
||||
batch_size: 4
|
||||
@@ -37,9 +24,6 @@ val_batches: 25
|
||||
# Adam learning rate.
|
||||
learning_rate: 1e-5
|
||||
|
||||
# Whether to report the logs to WandB
|
||||
# wand: "wandb-project"
|
||||
|
||||
# Number of training steps between loss reporting.
|
||||
steps_per_report: 10
|
||||
|
||||
@@ -67,6 +51,9 @@ max_seq_length: 2048
|
||||
# Use gradient checkpointing to reduce memory use.
|
||||
grad_checkpoint: false
|
||||
|
||||
# Use DoRA instead of LoRA.
|
||||
use_dora: false
|
||||
|
||||
# LoRA parameters can only be specified in a config file
|
||||
lora_parameters:
|
||||
# The layer keys to apply LoRA to.
|
||||
@@ -82,11 +69,3 @@ lora_parameters:
|
||||
# warmup: 100 # 0 for no warmup
|
||||
# warmup_init: 1e-7 # 0 if not specified
|
||||
# arguments: [1e-5, 1000, 1e-7] # passed to scheduler
|
||||
|
||||
#hf_dataset:
|
||||
# path: "billsum"
|
||||
# train_split: "train[:1000]"
|
||||
# valid_split: "train[-100:]"
|
||||
# prompt_feature: "text"
|
||||
# completion_feature: "summary"
|
||||
|
||||
|
||||
@@ -1,65 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
"""
|
||||
This is an example of tool use with mlx_lm and the OpenAI client.
|
||||
|
||||
To run, first start the server:
|
||||
|
||||
>>> mlx_lm.server
|
||||
|
||||
Then run this script.
|
||||
"""
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8080/v1", api_key="not-needed")
|
||||
|
||||
model = "mlx-community/qwen3-4b-4bit-DWQ"
|
||||
messages = [{"role": "user", "content": "What's the weather in Boston?"}]
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
},
|
||||
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def get_current_weather(**kwargs):
|
||||
return "51 Farenheit, clear skies"
|
||||
|
||||
|
||||
functions = {"get_current_weather": get_current_weather}
|
||||
|
||||
# The first query generates a tool call
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
# Call the function
|
||||
function = response.choices[0].message.tool_calls[0].function
|
||||
tool_result = functions[function.name](**function.arguments)
|
||||
|
||||
# Put the result of the function in the messages and generate the final
|
||||
# response:
|
||||
messages.append({"role": "tool", "name": function.name, "content": tool_result})
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
)
|
||||
print(response.choices[0].message.content)
|
||||
@@ -1,131 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
"""
|
||||
Run with:
|
||||
|
||||
```
|
||||
mlx.launch \
|
||||
--hostfile /path/to/hosts.json \
|
||||
/path/to/pipeline_generate.py \
|
||||
--prompt "hello world"
|
||||
```
|
||||
|
||||
Make sure you can run MLX over MPI on two hosts. For more information see the
|
||||
documentation:
|
||||
|
||||
https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import resource
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
from huggingface_hub import snapshot_download
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
from mlx_lm import load, stream_generate
|
||||
from mlx_lm.utils import load_model, load_tokenizer
|
||||
|
||||
# Needed for 8 bit model
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
|
||||
|
||||
|
||||
def download(repo: str, allow_patterns: list[str]) -> Path:
|
||||
return Path(
|
||||
snapshot_download(
|
||||
repo,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def shard_and_load(repo):
|
||||
# Get model path with everything but weight safetensors
|
||||
model_path = download(
|
||||
args.model,
|
||||
allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
|
||||
)
|
||||
|
||||
# Lazy load and shard model to figure out
|
||||
# which weights we need
|
||||
model, _ = load_model(model_path, lazy=True, strict=False)
|
||||
|
||||
group = mx.distributed.init()
|
||||
rank = group.rank()
|
||||
model.model.pipeline(group)
|
||||
|
||||
# Figure out which files we need for the local shard
|
||||
with open(model_path / "model.safetensors.index.json", "r") as fid:
|
||||
weight_index = json.load(fid)["weight_map"]
|
||||
|
||||
local_files = set()
|
||||
for k, _ in tree_flatten(model.parameters()):
|
||||
local_files.add(weight_index[k])
|
||||
|
||||
# Download weights for local shard
|
||||
download(args.model, allow_patterns=local_files)
|
||||
|
||||
# Load and shard the model, and load the weights
|
||||
tokenizer = load_tokenizer(model_path)
|
||||
model, _ = load_model(model_path, lazy=True, strict=False)
|
||||
model.model.pipeline(group)
|
||||
mx.eval(model.parameters())
|
||||
|
||||
# Synchronize processes before generation to avoid timeout if downloading
|
||||
# model for the first time.
|
||||
mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="LLM pipelined inference example")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="mlx-community/DeepSeek-R1-3bit",
|
||||
help="HF repo or path to local model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
default="Write a quicksort in C++.",
|
||||
help="Message to be processed by the model ('-' reads from stdin)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
"-m",
|
||||
type=int,
|
||||
default=256,
|
||||
help="Maximum number of tokens to generate",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
group = mx.distributed.init()
|
||||
rank = group.rank()
|
||||
|
||||
def rprint(*args, **kwargs):
|
||||
if rank == 0:
|
||||
print(*args, **kwargs)
|
||||
|
||||
model, tokenizer = shard_and_load(args.model)
|
||||
|
||||
messages = [{"role": "user", "content": args.prompt}]
|
||||
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
|
||||
|
||||
for response in stream_generate(
|
||||
model, tokenizer, prompt, max_tokens=args.max_tokens
|
||||
):
|
||||
rprint(response.text, end="", flush=True)
|
||||
|
||||
rprint()
|
||||
rprint("=" * 10)
|
||||
rprint(
|
||||
f"Prompt: {response.prompt_tokens} tokens, "
|
||||
f"{response.prompt_tps:.3f} tokens-per-sec"
|
||||
)
|
||||
rprint(
|
||||
f"Generation: {response.generation_tokens} tokens, "
|
||||
f"{response.generation_tps:.3f} tokens-per-sec"
|
||||
)
|
||||
rprint(f"Peak memory: {response.peak_memory:.3f} GB")
|
||||
@@ -1,73 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import json
|
||||
|
||||
from mlx_lm import generate, load
|
||||
from mlx_lm.models.cache import make_prompt_cache
|
||||
|
||||
# Specify the checkpoint
|
||||
checkpoint = "mlx-community/Qwen2.5-32B-Instruct-4bit"
|
||||
|
||||
# Load the corresponding model and tokenizer
|
||||
model, tokenizer = load(path_or_hf_repo=checkpoint)
|
||||
|
||||
|
||||
# An example tool, make sure to include a docstring and type hints
|
||||
def multiply(a: float, b: float):
|
||||
"""
|
||||
A function that multiplies two numbers
|
||||
|
||||
Args:
|
||||
a: The first number to multiply
|
||||
b: The second number to multiply
|
||||
"""
|
||||
return a * b
|
||||
|
||||
|
||||
tools = {"multiply": multiply}
|
||||
|
||||
# Specify the prompt and conversation history
|
||||
prompt = "Multiply 12234585 and 48838483920."
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=True, tools=list(tools.values())
|
||||
)
|
||||
|
||||
prompt_cache = make_prompt_cache(model)
|
||||
|
||||
# Generate the initial tool call:
|
||||
response = generate(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
prompt=prompt,
|
||||
max_tokens=2048,
|
||||
verbose=True,
|
||||
prompt_cache=prompt_cache,
|
||||
)
|
||||
|
||||
# Parse the tool call:
|
||||
# (Note, the tool call format is model specific)
|
||||
tool_open = "<tool_call>"
|
||||
tool_close = "</tool_call>"
|
||||
start_tool = response.find(tool_open) + len(tool_open)
|
||||
end_tool = response.find(tool_close)
|
||||
tool_call = json.loads(response[start_tool:end_tool].strip())
|
||||
tool_result = tools[tool_call["name"]](**tool_call["arguments"])
|
||||
|
||||
# Put the tool result in the prompt
|
||||
messages = [{"role": "tool", "name": tool_call["name"], "content": tool_result}]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
|
||||
# Generate the final response:
|
||||
response = generate(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
prompt=prompt,
|
||||
max_tokens=2048,
|
||||
verbose=True,
|
||||
prompt_cache=prompt_cache,
|
||||
)
|
||||
+38
-25
@@ -1,14 +1,19 @@
|
||||
import argparse
|
||||
import glob
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from .gguf import convert_to_gguf
|
||||
from .tuner.utils import dequantize, load_adapters
|
||||
from .tuner.dora import DoRALinear
|
||||
from .tuner.lora import LoRALinear, LoRASwitchLinear
|
||||
from .tuner.utils import apply_lora_layers, dequantize
|
||||
from .utils import (
|
||||
fetch_from_hub,
|
||||
get_model_path,
|
||||
save,
|
||||
save_config,
|
||||
save_weights,
|
||||
upload_to_hub,
|
||||
)
|
||||
|
||||
@@ -24,7 +29,7 @@ def parse_arguments() -> argparse.Namespace:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-path",
|
||||
default="fused_model",
|
||||
default="lora_fused_model",
|
||||
help="The path to save the fused model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -33,6 +38,12 @@ def parse_arguments() -> argparse.Namespace:
|
||||
default="adapters",
|
||||
help="Path to the trained adapter weights and config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to the original Hugging Face model. Required for upload if --model is a local directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-repo",
|
||||
help="The Hugging Face repo to upload the model to.",
|
||||
@@ -62,36 +73,40 @@ def main() -> None:
|
||||
print("Loading pretrained model")
|
||||
args = parse_arguments()
|
||||
|
||||
model_path, hf_path = get_model_path(args.model)
|
||||
model_path = get_model_path(args.model)
|
||||
model, config, tokenizer = fetch_from_hub(model_path)
|
||||
|
||||
model.freeze()
|
||||
model = load_adapters(model, args.adapter_path)
|
||||
model = apply_lora_layers(model, args.adapter_path)
|
||||
|
||||
fused_linears = [
|
||||
(n, m.fuse(de_quantize=args.de_quantize))
|
||||
(n, m.to_linear())
|
||||
for n, m in model.named_modules()
|
||||
if hasattr(m, "fuse")
|
||||
if isinstance(m, (LoRASwitchLinear, LoRALinear, DoRALinear))
|
||||
]
|
||||
|
||||
if fused_linears:
|
||||
model.update_modules(tree_unflatten(fused_linears))
|
||||
model.update_modules(tree_unflatten(fused_linears))
|
||||
|
||||
if args.de_quantize:
|
||||
print("De-quantizing model")
|
||||
model = dequantize(model)
|
||||
config.pop("quantization", None)
|
||||
|
||||
weights = dict(tree_flatten(model.parameters()))
|
||||
|
||||
save_path = Path(args.save_path)
|
||||
save(
|
||||
save_path,
|
||||
model_path,
|
||||
model,
|
||||
tokenizer,
|
||||
config,
|
||||
hf_repo=hf_path,
|
||||
donate_model=False,
|
||||
)
|
||||
|
||||
save_weights(save_path, weights)
|
||||
|
||||
py_files = glob.glob(str(model_path / "*.py"))
|
||||
for file in py_files:
|
||||
shutil.copy(file, save_path)
|
||||
|
||||
tokenizer.save_pretrained(save_path)
|
||||
|
||||
if args.de_quantize:
|
||||
config.pop("quantization", None)
|
||||
|
||||
save_config(config, config_path=save_path / "config.json")
|
||||
|
||||
if args.export_gguf:
|
||||
model_type = config["model_type"]
|
||||
@@ -99,20 +114,18 @@ def main() -> None:
|
||||
raise ValueError(
|
||||
f"Model type {model_type} not supported for GGUF conversion."
|
||||
)
|
||||
weights = dict(tree_flatten(model.parameters()))
|
||||
convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
|
||||
|
||||
if args.upload_repo is not None:
|
||||
hf_path = args.hf_path or (
|
||||
args.model if not Path(args.model).exists() else None
|
||||
)
|
||||
if hf_path is None:
|
||||
raise ValueError(
|
||||
"Must provide original Hugging Face repo to upload local model."
|
||||
)
|
||||
upload_to_hub(args.save_path, args.upload_repo)
|
||||
upload_to_hub(args.save_path, args.upload_repo, hf_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.fuse...` directly is deprecated."
|
||||
" Use `mlx_lm.fuse...` or `python -m mlx_lm fuse ...` instead."
|
||||
)
|
||||
main()
|
||||
|
||||
+65
-784
@@ -1,52 +1,17 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import functools
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Generator,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_reduce
|
||||
from transformers import PreTrainedTokenizer
|
||||
|
||||
from .models import cache
|
||||
from .models.cache import (
|
||||
QuantizedKVCache,
|
||||
load_prompt_cache,
|
||||
)
|
||||
from .sample_utils import make_sampler
|
||||
from .tokenizer_utils import TokenizerWrapper
|
||||
from .utils import does_model_support_input_embeddings, load
|
||||
from .utils import generate, load
|
||||
|
||||
DEFAULT_MODEL_PATH = "mlx_model"
|
||||
DEFAULT_PROMPT = "hello"
|
||||
DEFAULT_MAX_TOKENS = 100
|
||||
DEFAULT_TEMP = 0.0
|
||||
DEFAULT_TEMP = 0.6
|
||||
DEFAULT_TOP_P = 1.0
|
||||
DEFAULT_MIN_P = 0.0
|
||||
DEFAULT_TOP_K = 0
|
||||
DEFAULT_XTC_PROBABILITY = 0.0
|
||||
DEFAULT_XTC_THRESHOLD = 0.0
|
||||
DEFAULT_MIN_TOKENS_TO_KEEP = 1
|
||||
DEFAULT_SEED = None
|
||||
DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
|
||||
DEFAULT_QUANTIZED_KV_START = 5000
|
||||
|
||||
|
||||
def str2bool(string):
|
||||
return string.lower() not in ["false", "f"]
|
||||
DEFAULT_SEED = 0
|
||||
|
||||
|
||||
def setup_arg_parser():
|
||||
@@ -55,11 +20,8 @@ def setup_arg_parser():
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
help=(
|
||||
"The path to the local model directory or Hugging Face repo. "
|
||||
f"If no model is specified, then {DEFAULT_MODEL} is used."
|
||||
),
|
||||
default=None,
|
||||
default="mlx_model",
|
||||
help="The path to the local model directory or Hugging Face repo.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--adapter-path",
|
||||
@@ -67,27 +29,18 @@ def setup_arg_parser():
|
||||
help="Optional path for the trained adapter weights and config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-eos-token",
|
||||
"--trust-remote-code",
|
||||
action="store_true",
|
||||
help="Enable trusting remote code for tokenizer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eos-token",
|
||||
type=str,
|
||||
default=(),
|
||||
nargs="+",
|
||||
help="Add tokens in the list of eos tokens that stop generation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--system-prompt",
|
||||
default=None,
|
||||
help="System prompt to be used for the chat template",
|
||||
help="End of sequence token for tokenizer",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
default=DEFAULT_PROMPT,
|
||||
help="Message to be processed by the model ('-' reads from stdin)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefill-response",
|
||||
default=None,
|
||||
help="Prefill response to be used for the chat template",
|
||||
"--prompt", default=DEFAULT_PROMPT, help="Message to be processed by the model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
@@ -102,36 +55,7 @@ def setup_arg_parser():
|
||||
parser.add_argument(
|
||||
"--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-p", type=float, default=DEFAULT_MIN_P, help="Sampling min-p"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-k", type=int, default=DEFAULT_TOP_K, help="Sampling top-k"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--xtc-probability",
|
||||
type=float,
|
||||
default=DEFAULT_XTC_PROBABILITY,
|
||||
help="Probability of XTC sampling to happen each next token",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--xtc-threshold",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Thresold the probs of each next token candidate to be sampled by XTC",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-tokens-to-keep",
|
||||
type=int,
|
||||
default=DEFAULT_MIN_TOKENS_TO_KEEP,
|
||||
help="Minimum tokens to keep for min-p sampling.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=DEFAULT_SEED,
|
||||
help="PRNG seed",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
|
||||
parser.add_argument(
|
||||
"--ignore-chat-template",
|
||||
action="store_true",
|
||||
@@ -143,738 +67,95 @@ def setup_arg_parser():
|
||||
help="Use the default chat template",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chat-template-config",
|
||||
help="Additional config for `apply_chat_template`. Should be a dictionary of"
|
||||
" string keys to values represented as a JSON decodable string.",
|
||||
default=None,
|
||||
"--colorize",
|
||||
action="store_true",
|
||||
help="Colorize output based on T[0] probability",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
type=str2bool,
|
||||
default=True,
|
||||
help="Log verbose output when 'True' or 'T' or only print the response when 'False' or 'F'",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-kv-size",
|
||||
"--cache-limit-gb",
|
||||
type=int,
|
||||
help="Set the maximum key-value cache size",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-cache-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="A file containing saved KV caches to avoid recomputing them",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-bits",
|
||||
type=int,
|
||||
help="Number of bits for KV cache quantization. "
|
||||
"Defaults to no quantization.",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-group-size",
|
||||
type=int,
|
||||
help="Group size for KV cache quantization.",
|
||||
default=64,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantized-kv-start",
|
||||
help="When --kv-bits is set, start quantizing the KV cache "
|
||||
"from this step onwards.",
|
||||
type=int,
|
||||
default=DEFAULT_QUANTIZED_KV_START,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--draft-model",
|
||||
type=str,
|
||||
help="A model to be used for speculative decoding.",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-draft-tokens",
|
||||
type=int,
|
||||
help="Number of tokens to draft when using speculative decoding.",
|
||||
default=3,
|
||||
help="Set the MLX cache limit in GB",
|
||||
required=False,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
# A stream on the default device just for generation
|
||||
generation_stream = mx.new_stream(mx.default_device())
|
||||
def colorprint(color, s):
|
||||
color_codes = {
|
||||
"black": 30,
|
||||
"red": 31,
|
||||
"green": 32,
|
||||
"yellow": 33,
|
||||
"blue": 34,
|
||||
"magenta": 35,
|
||||
"cyan": 36,
|
||||
"white": 39,
|
||||
}
|
||||
ccode = color_codes.get(color, 30)
|
||||
print(f"\033[1m\033[{ccode}m{s}\033[0m", end="", flush=True)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
|
||||
"""
|
||||
A context manager to temporarily change the wired limit.
|
||||
|
||||
Note, the wired limit should not be changed during an async eval. If an
|
||||
async eval could be running pass in the streams to synchronize with prior
|
||||
to exiting the context manager.
|
||||
"""
|
||||
if not mx.metal.is_available():
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
return
|
||||
|
||||
model_bytes = tree_reduce(
|
||||
lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
|
||||
)
|
||||
max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
|
||||
if model_bytes > 0.9 * max_rec_size:
|
||||
model_mb = model_bytes // 2**20
|
||||
max_rec_mb = max_rec_size // 2**20
|
||||
print(
|
||||
f"[WARNING] Generating with a model that requires {model_mb} MB "
|
||||
f"which is close to the maximum recommended size of {max_rec_mb} "
|
||||
"MB. This can be slow. See the documentation for possible work-arounds: "
|
||||
"https://github.com/ml-explore/mlx-lm/tree/main#large-models"
|
||||
)
|
||||
old_limit = mx.set_wired_limit(max_rec_size)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if streams is not None:
|
||||
for s in streams:
|
||||
mx.synchronize(s)
|
||||
else:
|
||||
mx.synchronize()
|
||||
mx.set_wired_limit(old_limit)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResponse:
|
||||
"""
|
||||
The output of :func:`stream_generate`.
|
||||
|
||||
Args:
|
||||
text (str): The next segment of decoded text. This can be an empty string.
|
||||
token (int): The next token.
|
||||
from_draft (bool): Whether the token was generated by the draft model.
|
||||
logprobs (mx.array): A vector of log probabilities.
|
||||
prompt_tokens (int): The number of tokens in the prompt.
|
||||
prompt_tps (float): The prompt processing tokens-per-second.
|
||||
generation_tokens (int): The number of generated tokens.
|
||||
generation_tps (float): The tokens-per-second for generation.
|
||||
peak_memory (float): The peak memory used so far in GB.
|
||||
finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
|
||||
"""
|
||||
|
||||
text: str
|
||||
token: int
|
||||
logprobs: mx.array
|
||||
from_draft: bool
|
||||
prompt_tokens: int
|
||||
prompt_tps: float
|
||||
generation_tokens: int
|
||||
generation_tps: float
|
||||
peak_memory: float
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
|
||||
if (
|
||||
kv_bits is not None
|
||||
and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
|
||||
and prompt_cache[0].offset > quantized_kv_start
|
||||
):
|
||||
for i in range(len(prompt_cache)):
|
||||
if isinstance(prompt_cache[i], cache.KVCache):
|
||||
prompt_cache[i] = prompt_cache[i].to_quantized(
|
||||
group_size=kv_group_size, bits=kv_bits
|
||||
)
|
||||
|
||||
|
||||
def generate_step(
|
||||
prompt: mx.array,
|
||||
model: nn.Module,
|
||||
*,
|
||||
max_tokens: int = 256,
|
||||
sampler: Optional[Callable[mx.array, mx.array]] = None,
|
||||
logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
|
||||
max_kv_size: Optional[int] = None,
|
||||
prompt_cache: Optional[Any] = None,
|
||||
prefill_step_size: int = 2048,
|
||||
kv_bits: Optional[int] = None,
|
||||
kv_group_size: int = 64,
|
||||
quantized_kv_start: int = 0,
|
||||
prompt_progress_callback: Optional[Callable[int, int]] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> Generator[Tuple[mx.array, mx.array], None, None]:
|
||||
"""
|
||||
A generator producing token ids based on the given prompt from the model.
|
||||
|
||||
Args:
|
||||
prompt (mx.array): The input prompt.
|
||||
model (nn.Module): The model to use for generation.
|
||||
max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
|
||||
generator. Default: ``256``.
|
||||
sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
|
||||
token from a vector of log probabilities. Default: ``None``.
|
||||
logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
|
||||
A list of functions that take tokens and logits and return the processed
|
||||
logits. Default: ``None``.
|
||||
max_kv_size (int, optional): Maximum size of the key-value cache. Old
|
||||
entries (except the first 4 tokens) will be overwritten.
|
||||
prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
|
||||
provided, the cache will be updated in place.
|
||||
prefill_step_size (int): Step size for processing the prompt.
|
||||
kv_bits (int, optional): Number of bits to use for KV cache quantization.
|
||||
None implies no cache quantization. Default: ``None``.
|
||||
kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
|
||||
quantized_kv_start (int): Step to begin using a quantized KV cache.
|
||||
when ``kv_bits`` is non-None. Default: ``0``.
|
||||
prompt_progress_callback (Callable[int, int]): A call-back which takes the
|
||||
prompt tokens processed so far and the total number of prompt tokens.
|
||||
input_embeddings (mx.array, optional): Input embeddings to use in place of
|
||||
prompt tokens. Default: ``None``.
|
||||
|
||||
Yields:
|
||||
Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
|
||||
"""
|
||||
if input_embeddings is not None:
|
||||
if not does_model_support_input_embeddings(model):
|
||||
raise ValueError("Model does not support input embeddings.")
|
||||
if len(prompt) != 0:
|
||||
raise ValueError(
|
||||
"If using input embeddings, prompt tokens must be an empty array."
|
||||
)
|
||||
|
||||
tokens = None
|
||||
|
||||
# Create the KV cache for generation
|
||||
if prompt_cache is None:
|
||||
prompt_cache = cache.make_prompt_cache(
|
||||
model,
|
||||
max_kv_size=max_kv_size,
|
||||
)
|
||||
|
||||
prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
|
||||
|
||||
quantize_cache_fn = functools.partial(
|
||||
maybe_quantize_kv_cache,
|
||||
quantized_kv_start=quantized_kv_start,
|
||||
kv_group_size=kv_group_size,
|
||||
kv_bits=kv_bits,
|
||||
)
|
||||
|
||||
sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
|
||||
|
||||
def _model_call(y):
|
||||
if y.ndim == 3:
|
||||
return model(None, cache=prompt_cache, input_embeddings=y)
|
||||
else:
|
||||
return model(y, cache=prompt_cache)
|
||||
|
||||
def _step(y):
|
||||
nonlocal tokens
|
||||
|
||||
with mx.stream(generation_stream):
|
||||
logits = _model_call(y[None])
|
||||
|
||||
logits = logits[:, -1, :]
|
||||
|
||||
if logits_processors and input_embeddings is None:
|
||||
tokens = mx.concat([tokens, y]) if tokens is not None else y
|
||||
for processor in logits_processors:
|
||||
logits = processor(tokens, logits)
|
||||
|
||||
quantize_cache_fn(prompt_cache)
|
||||
|
||||
logprobs = logits - mx.logsumexp(logits, keepdims=True)
|
||||
y = sampler(logprobs)
|
||||
return y, logprobs.squeeze(0)
|
||||
|
||||
using_embeddings = input_embeddings is not None
|
||||
|
||||
y = input_embeddings if using_embeddings else prompt
|
||||
with mx.stream(generation_stream):
|
||||
total_prompt_tokens = y.shape[0]
|
||||
prompt_processed_tokens = 0
|
||||
while y.shape[0] > prefill_step_size:
|
||||
_model_call(y[:prefill_step_size][None])
|
||||
quantize_cache_fn(prompt_cache)
|
||||
mx.eval([c.state for c in prompt_cache])
|
||||
prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
|
||||
prompt_processed_tokens += prefill_step_size
|
||||
y = y[prefill_step_size:]
|
||||
mx.clear_cache()
|
||||
|
||||
y, logprobs = _step(y)
|
||||
|
||||
mx.async_eval(y, logprobs)
|
||||
n = 0
|
||||
while True:
|
||||
if n != max_tokens:
|
||||
next_y, next_logprobs = _step(y)
|
||||
mx.async_eval(next_y, next_logprobs)
|
||||
if n == 0:
|
||||
mx.eval(y)
|
||||
prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
|
||||
if n == max_tokens:
|
||||
break
|
||||
yield y.item(), logprobs
|
||||
if n % 256 == 0:
|
||||
mx.clear_cache()
|
||||
y, logprobs = next_y, next_logprobs
|
||||
n += 1
|
||||
|
||||
|
||||
def speculative_generate_step(
|
||||
prompt: mx.array,
|
||||
model: nn.Module,
|
||||
draft_model: nn.Module,
|
||||
*,
|
||||
num_draft_tokens=2,
|
||||
max_tokens: int = 256,
|
||||
sampler: Optional[Callable[mx.array, mx.array]] = None,
|
||||
logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
|
||||
prompt_cache: Optional[Any] = None,
|
||||
prefill_step_size: int = 512,
|
||||
kv_bits: Optional[int] = None,
|
||||
kv_group_size: int = 64,
|
||||
quantized_kv_start: int = 0,
|
||||
) -> Generator[Tuple[mx.array, mx.array, bool], None, None]:
|
||||
"""
|
||||
A generator producing token ids based on the given prompt from the model.
|
||||
|
||||
Args:
|
||||
prompt (mx.array): The input prompt.
|
||||
model (nn.Module): The model to use for generation.
|
||||
draft_model (nn.Module): The draft model for speculative decoding.
|
||||
num_draft_tokens (int, optional): The number of draft tokens for
|
||||
speculative decoding. Default: ``2``.
|
||||
max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
|
||||
generator. Default: ``256``.
|
||||
sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
|
||||
token from a vector of log probabilities. Default: ``None``.
|
||||
logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
|
||||
A list of functions that take tokens and logits and return the processed
|
||||
logits. Default: ``None``.
|
||||
prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
|
||||
provided, the cache will be updated in place. The cache must be trimmable.
|
||||
prefill_step_size (int): Step size for processing the prompt.
|
||||
kv_bits (int, optional): Number of bits to use for KV cache quantization.
|
||||
None implies no cache quantization. Default: ``None``.
|
||||
kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
|
||||
quantized_kv_start (int): Step to begin using a quantized KV cache.
|
||||
when ``kv_bits`` is non-None. Default: ``0``.
|
||||
|
||||
Yields:
|
||||
Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
|
||||
and a bool indicating if the token was generated by the draft model
|
||||
"""
|
||||
|
||||
y = prompt.astype(mx.uint32)
|
||||
prev_tokens = None
|
||||
|
||||
# Create the KV cache for generation
|
||||
if prompt_cache is None:
|
||||
model_cache = cache.make_prompt_cache(model)
|
||||
draft_cache = cache.make_prompt_cache(draft_model)
|
||||
def colorprint_by_t0(s, t0):
|
||||
if t0 > 0.95:
|
||||
color = "white"
|
||||
elif t0 > 0.70:
|
||||
color = "green"
|
||||
elif t0 > 0.30:
|
||||
color = "yellow"
|
||||
else:
|
||||
model_cache = prompt_cache[: len(model.layers)]
|
||||
draft_cache = prompt_cache[len(model.layers) :]
|
||||
|
||||
sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
|
||||
|
||||
quantize_cache_fn = functools.partial(
|
||||
maybe_quantize_kv_cache,
|
||||
quantized_kv_start=quantized_kv_start,
|
||||
kv_group_size=kv_group_size,
|
||||
kv_bits=kv_bits,
|
||||
)
|
||||
|
||||
def _process_and_sample(tokens, logits):
|
||||
if logits_processors:
|
||||
for processor in logits_processors:
|
||||
logits = processor(tokens, logits)
|
||||
|
||||
logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
|
||||
y = sampler(logprobs)
|
||||
return y, logprobs
|
||||
|
||||
def _step(model, cache, y, n_predict=1):
|
||||
with mx.stream(generation_stream):
|
||||
logits = model(y[None], cache=cache)
|
||||
logits = logits[:, -n_predict:, :]
|
||||
|
||||
quantize_cache_fn(cache)
|
||||
if logits_processors:
|
||||
nonlocal prev_tokens
|
||||
out_y, out_logprobs = [], []
|
||||
if n_predict > 1:
|
||||
y = y[: -(n_predict - 1)]
|
||||
for i in range(n_predict):
|
||||
prev_tokens = (
|
||||
mx.concat([prev_tokens, y]) if prev_tokens is not None else y
|
||||
)
|
||||
y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
|
||||
out_y.append(y)
|
||||
out_logprobs.append(logprobs)
|
||||
return mx.concatenate(out_y, axis=0), mx.concatenate(
|
||||
out_logprobs, axis=0
|
||||
)
|
||||
else:
|
||||
return _process_and_sample(None, logits.squeeze(0))
|
||||
|
||||
def _prefill(model, cache, y):
|
||||
while y.size > prefill_step_size:
|
||||
model(y[:prefill_step_size][None], cache=cache)
|
||||
quantize_cache_fn(cache)
|
||||
mx.eval([c.state for c in cache])
|
||||
y = y[prefill_step_size:]
|
||||
mx.clear_cache()
|
||||
return y
|
||||
|
||||
def _rewind_cache(num_draft, num_accept):
|
||||
cache.trim_prompt_cache(model_cache, num_draft - num_accept)
|
||||
cache.trim_prompt_cache(draft_cache, max(num_draft - num_accept - 1, 0))
|
||||
|
||||
def _draft_generate(y, num_draft):
|
||||
if num_draft == 0:
|
||||
return mx.array([], mx.uint32)
|
||||
ys = []
|
||||
for _ in range(num_draft):
|
||||
y, _ = _step(draft_model, draft_cache, y)
|
||||
mx.async_eval(y)
|
||||
ys.append(y)
|
||||
return mx.concatenate(ys)
|
||||
|
||||
with mx.stream(generation_stream):
|
||||
draft_y = _prefill(draft_model, draft_cache, y)
|
||||
y = _prefill(model, model_cache, y)
|
||||
|
||||
ntoks = 0
|
||||
# Set these so the finally block doesn't raise
|
||||
num_draft = 0
|
||||
n = 0
|
||||
try:
|
||||
while True:
|
||||
num_draft = min(max_tokens - ntoks, num_draft_tokens)
|
||||
draft_tokens = _draft_generate(draft_y, num_draft)
|
||||
if prev_tokens is not None:
|
||||
prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1]
|
||||
y = mx.concatenate([y, draft_tokens])
|
||||
tokens, logprobs = _step(model, model_cache, y, num_draft + 1)
|
||||
mx.eval(tokens, draft_tokens)
|
||||
draft_tokens = draft_tokens.tolist()
|
||||
tokens = tokens.tolist()
|
||||
n = 0
|
||||
while n < num_draft:
|
||||
tn, dtn, lpn = tokens[n], draft_tokens[n], logprobs[n]
|
||||
if tn != dtn:
|
||||
break
|
||||
n += 1
|
||||
ntoks += 1
|
||||
yield tn, lpn, True
|
||||
if ntoks == max_tokens:
|
||||
break
|
||||
if ntoks < max_tokens:
|
||||
ntoks += 1
|
||||
yield tokens[n], logprobs[n], False
|
||||
|
||||
if ntoks == max_tokens:
|
||||
break
|
||||
|
||||
y = mx.array([tokens[n]], mx.uint32)
|
||||
draft_y = y
|
||||
|
||||
# If we accepted all the draft tokens, include the last
|
||||
# draft token in the next draft step since it hasn't been
|
||||
# processed yet by the draft model
|
||||
if n == num_draft:
|
||||
draft_y = mx.concatenate(
|
||||
[mx.array(draft_tokens[-1:], mx.uint32), draft_y]
|
||||
)
|
||||
|
||||
if prev_tokens is not None:
|
||||
prev_tokens = prev_tokens[: -max(num_draft - n, 1)]
|
||||
_rewind_cache(num_draft, n)
|
||||
finally:
|
||||
_rewind_cache(num_draft, n)
|
||||
|
||||
|
||||
def stream_generate(
|
||||
model: nn.Module,
|
||||
tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
|
||||
prompt: Union[str, mx.array, List[int]],
|
||||
draft_model: Optional[nn.Module] = None,
|
||||
**kwargs,
|
||||
) -> Generator[GenerationResponse, None, None]:
|
||||
"""
|
||||
A generator producing text based on the given prompt from the model.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The model to use for generation.
|
||||
tokenizer (PreTrainedTokenizer): The tokenizer.
|
||||
prompt (Union[str, mx.array, List[int]]): The input prompt string or
|
||||
integer tokens.
|
||||
draft_model (Optional[nn.Module]): An optional draft model. If provided
|
||||
then speculative decoding is used. The draft model must use the same
|
||||
tokenizer as the main model. Default: ``None``.
|
||||
kwargs: The remaining options get passed to :func:`generate_step`.
|
||||
See :func:`generate_step` for more details.
|
||||
|
||||
Yields:
|
||||
GenerationResponse: An instance containing the generated text segment and
|
||||
associated metadata. See :class:`GenerationResponse` for details.
|
||||
"""
|
||||
if not isinstance(tokenizer, TokenizerWrapper):
|
||||
tokenizer = TokenizerWrapper(tokenizer)
|
||||
|
||||
if not isinstance(prompt, mx.array):
|
||||
if isinstance(prompt, str):
|
||||
# Try to infer if special tokens are needed
|
||||
add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
|
||||
tokenizer.bos_token
|
||||
)
|
||||
prompt = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
|
||||
prompt = mx.array(prompt)
|
||||
|
||||
detokenizer = tokenizer.detokenizer
|
||||
|
||||
if draft_model is None:
|
||||
kwargs.pop("num_draft_tokens", None)
|
||||
token_generator = generate_step(prompt, model, **kwargs)
|
||||
# from_draft always false for non-speculative generation
|
||||
token_generator = (
|
||||
(token, logprobs, False) for token, logprobs in token_generator
|
||||
)
|
||||
else:
|
||||
kwargs.pop("max_kv_size", None)
|
||||
token_generator = speculative_generate_step(
|
||||
prompt, model, draft_model, **kwargs
|
||||
)
|
||||
with wired_limit(model, [generation_stream]):
|
||||
detokenizer.reset()
|
||||
tic = time.perf_counter()
|
||||
for n, (token, logprobs, from_draft) in enumerate(token_generator):
|
||||
if n == 0:
|
||||
prompt_time = time.perf_counter() - tic
|
||||
prompt_tps = prompt.size / prompt_time
|
||||
tic = time.perf_counter()
|
||||
if token in tokenizer.eos_token_ids:
|
||||
break
|
||||
|
||||
detokenizer.add_token(token)
|
||||
|
||||
yield GenerationResponse(
|
||||
text=detokenizer.last_segment,
|
||||
token=token,
|
||||
logprobs=logprobs,
|
||||
from_draft=from_draft,
|
||||
prompt_tokens=prompt.size,
|
||||
prompt_tps=prompt_tps,
|
||||
generation_tokens=n + 1,
|
||||
generation_tps=(n + 1) / (time.perf_counter() - tic),
|
||||
peak_memory=mx.get_peak_memory() / 1e9,
|
||||
finish_reason=None,
|
||||
)
|
||||
|
||||
detokenizer.finalize()
|
||||
yield GenerationResponse(
|
||||
text=detokenizer.last_segment,
|
||||
token=token,
|
||||
logprobs=logprobs,
|
||||
from_draft=from_draft,
|
||||
prompt_tokens=prompt.size,
|
||||
prompt_tps=prompt_tps,
|
||||
generation_tokens=n + 1,
|
||||
generation_tps=(n + 1) / (time.perf_counter() - tic),
|
||||
peak_memory=mx.get_peak_memory() / 1e9,
|
||||
finish_reason="stop" if token in tokenizer.eos_token_ids else "length",
|
||||
)
|
||||
|
||||
|
||||
def generate(
|
||||
model: nn.Module,
|
||||
tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
|
||||
prompt: Union[str, List[int]],
|
||||
verbose: bool = False,
|
||||
formatter: Optional[Callable] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""
|
||||
Generate a complete response from the model.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The language model.
|
||||
tokenizer (PreTrainedTokenizer): The tokenizer.
|
||||
prompt (Union[str, List[int]]): The input prompt string or integer tokens.
|
||||
verbose (bool): If ``True``, print tokens and timing information.
|
||||
Default: ``False``.
|
||||
kwargs: The remaining options get passed to :func:`stream_generate`.
|
||||
See :func:`stream_generate` for more details.
|
||||
"""
|
||||
if formatter is not None:
|
||||
print(
|
||||
"[Warning] Text formatting is deprecated and no longer used. "
|
||||
"The argument will be removed in a future version."
|
||||
)
|
||||
if verbose:
|
||||
print("=" * 10)
|
||||
|
||||
text = ""
|
||||
for response in stream_generate(model, tokenizer, prompt, **kwargs):
|
||||
if verbose:
|
||||
print(response.text, end="", flush=True)
|
||||
text += response.text
|
||||
|
||||
if verbose:
|
||||
print()
|
||||
print("=" * 10)
|
||||
if len(text) == 0:
|
||||
print("No text generated for this prompt")
|
||||
return
|
||||
print(
|
||||
f"Prompt: {response.prompt_tokens} tokens, "
|
||||
f"{response.prompt_tps:.3f} tokens-per-sec"
|
||||
)
|
||||
print(
|
||||
f"Generation: {response.generation_tokens} tokens, "
|
||||
f"{response.generation_tps:.3f} tokens-per-sec"
|
||||
)
|
||||
print(f"Peak memory: {response.peak_memory:.3f} GB")
|
||||
return text
|
||||
color = "red"
|
||||
colorprint(color, s)
|
||||
|
||||
|
||||
def main():
|
||||
parser = setup_arg_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.seed is not None:
|
||||
mx.random.seed(args.seed)
|
||||
mx.random.seed(args.seed)
|
||||
|
||||
# Load the prompt cache and metadata if a cache file is provided
|
||||
using_cache = args.prompt_cache_file is not None
|
||||
if using_cache:
|
||||
prompt_cache, metadata = load_prompt_cache(
|
||||
args.prompt_cache_file,
|
||||
return_metadata=True,
|
||||
)
|
||||
if isinstance(prompt_cache[0], QuantizedKVCache):
|
||||
if args.kv_bits is not None and args.kv_bits != prompt_cache[0].bits:
|
||||
raise ValueError(
|
||||
"--kv-bits does not match the kv cache loaded from --prompt-cache-file."
|
||||
)
|
||||
if args.kv_group_size != prompt_cache[0].group_size:
|
||||
raise ValueError(
|
||||
"--kv-group-size does not match the kv cache loaded from --prompt-cache-file."
|
||||
)
|
||||
if args.cache_limit_gb is not None:
|
||||
mx.metal.set_cache_limit(args.cache_limit_gb * 1024 * 1024 * 1024)
|
||||
|
||||
# Building tokenizer_config
|
||||
tokenizer_config = (
|
||||
{} if not using_cache else json.loads(metadata["tokenizer_config"])
|
||||
)
|
||||
tokenizer_config["trust_remote_code"] = True
|
||||
|
||||
model_path = args.model
|
||||
if using_cache:
|
||||
if model_path is None:
|
||||
model_path = metadata["model"]
|
||||
elif model_path != metadata["model"]:
|
||||
raise ValueError(
|
||||
f"Providing a different model ({model_path}) than that "
|
||||
f"used to create the prompt cache ({metadata['model']}) "
|
||||
"is an error."
|
||||
)
|
||||
model_path = model_path or DEFAULT_MODEL
|
||||
tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
|
||||
if args.eos_token is not None:
|
||||
tokenizer_config["eos_token"] = args.eos_token
|
||||
|
||||
model, tokenizer = load(
|
||||
model_path,
|
||||
args.model,
|
||||
adapter_path=args.adapter_path,
|
||||
tokenizer_config=tokenizer_config,
|
||||
)
|
||||
for eos_token in args.extra_eos_token:
|
||||
tokenizer.add_eos_token(eos_token)
|
||||
|
||||
template_kwargs = {}
|
||||
if args.chat_template_config is not None:
|
||||
template_kwargs = json.loads(args.chat_template_config)
|
||||
|
||||
if args.use_default_chat_template:
|
||||
if tokenizer.chat_template is None:
|
||||
tokenizer.chat_template = tokenizer.default_chat_template
|
||||
elif using_cache:
|
||||
tokenizer.chat_template = json.loads(metadata["chat_template"])
|
||||
|
||||
prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
|
||||
prompt = sys.stdin.read() if prompt == "-" else prompt
|
||||
if not args.ignore_chat_template and tokenizer.chat_template is not None:
|
||||
if args.system_prompt is not None:
|
||||
messages = [{"role": "system", "content": args.system_prompt}]
|
||||
else:
|
||||
messages = []
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
has_prefill = args.prefill_response is not None
|
||||
if has_prefill:
|
||||
messages.append({"role": "assistant", "content": args.prefill_response})
|
||||
if not args.ignore_chat_template and (
|
||||
hasattr(tokenizer, "apply_chat_template")
|
||||
and tokenizer.chat_template is not None
|
||||
):
|
||||
messages = [{"role": "user", "content": args.prompt}]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
continue_final_message=has_prefill,
|
||||
add_generation_prompt=not has_prefill,
|
||||
**template_kwargs,
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Treat the prompt as a suffix assuming that the prefix is in the
|
||||
# stored kv cache.
|
||||
if using_cache:
|
||||
messages[-1]["content"] = "<query>"
|
||||
test_prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
continue_final_message=has_prefill,
|
||||
add_generation_prompt=not has_prefill,
|
||||
)
|
||||
prompt = prompt[test_prompt.index("<query>") :]
|
||||
prompt = tokenizer.encode(prompt, add_special_tokens=False)
|
||||
else:
|
||||
prompt = tokenizer.encode(prompt)
|
||||
prompt = args.prompt
|
||||
|
||||
if args.draft_model is not None:
|
||||
draft_model, draft_tokenizer = load(args.draft_model)
|
||||
if draft_tokenizer.vocab_size != tokenizer.vocab_size:
|
||||
raise ValueError("Draft model tokenizer does not match model tokenizer.")
|
||||
else:
|
||||
draft_model = None
|
||||
sampler = make_sampler(
|
||||
args.temp,
|
||||
args.top_p,
|
||||
args.min_p,
|
||||
args.min_tokens_to_keep,
|
||||
top_k=args.top_k,
|
||||
xtc_probability=args.xtc_probability,
|
||||
xtc_threshold=args.xtc_threshold,
|
||||
xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids),
|
||||
)
|
||||
response = generate(
|
||||
formatter = colorprint_by_t0 if args.colorize else None
|
||||
|
||||
generate(
|
||||
model,
|
||||
tokenizer,
|
||||
prompt,
|
||||
max_tokens=args.max_tokens,
|
||||
verbose=args.verbose,
|
||||
sampler=sampler,
|
||||
max_kv_size=args.max_kv_size,
|
||||
prompt_cache=prompt_cache if using_cache else None,
|
||||
kv_bits=args.kv_bits,
|
||||
kv_group_size=args.kv_group_size,
|
||||
quantized_kv_start=args.quantized_kv_start,
|
||||
draft_model=draft_model,
|
||||
num_draft_tokens=args.num_draft_tokens,
|
||||
args.max_tokens,
|
||||
verbose=True,
|
||||
formatter=formatter,
|
||||
temp=args.temp,
|
||||
top_p=args.top_p,
|
||||
)
|
||||
if not args.verbose:
|
||||
print(response)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.generate...` directly is deprecated."
|
||||
" Use `mlx_lm.generate...` or `python -m mlx_lm generate ...` instead."
|
||||
)
|
||||
main()
|
||||
|
||||
+15
-16
@@ -59,7 +59,7 @@ class HfVocab:
|
||||
for token_id in range(self.vocab_size_base):
|
||||
if token_id in self.added_tokens_ids:
|
||||
continue
|
||||
token_text = reverse_vocab[token_id]
|
||||
token_text = reverse_vocab[token_id].encode("utf-8")
|
||||
yield token_text, self.get_token_score(token_id), self.get_token_type(
|
||||
token_id, token_text, self.special_ids
|
||||
)
|
||||
@@ -67,7 +67,7 @@ class HfVocab:
|
||||
def get_token_type(
|
||||
self, token_id: int, token_text: bytes, special_ids: Set[int]
|
||||
) -> TokenType:
|
||||
if re.fullmatch(r"<0x[0-9A-Fa-f]{2}>", token_text):
|
||||
if re.fullmatch(rb"<0x[0-9A-Fa-f]{2}>", token_text):
|
||||
return TokenType.BYTE
|
||||
return TokenType.CONTROL if token_id in special_ids else TokenType.NORMAL
|
||||
|
||||
@@ -77,12 +77,14 @@ class HfVocab:
|
||||
def added_tokens(self) -> Iterable[Tuple[bytes, float, TokenType]]:
|
||||
for text in self.added_tokens_list:
|
||||
if text in self.specials:
|
||||
toktype = self.get_token_type(self.specials[text], "", self.special_ids)
|
||||
toktype = self.get_token_type(
|
||||
self.specials[text], b"", self.special_ids
|
||||
)
|
||||
score = self.get_token_score(self.specials[text])
|
||||
else:
|
||||
toktype = TokenType.USER_DEFINED
|
||||
score = -1000.0
|
||||
yield text, score, toktype
|
||||
yield text.encode("utf-8"), score, toktype
|
||||
|
||||
def has_newline_token(self):
|
||||
return "<0x0A>" in self.tokenizer.vocab or "\n" in self.tokenizer.vocab
|
||||
@@ -241,18 +243,15 @@ def prepare_metadata(config, vocab):
|
||||
metadata["tokenizer.ggml.tokens"] = tokens
|
||||
metadata["tokenizer.ggml.scores"] = mx.array(scores, dtype=mx.float32)
|
||||
metadata["tokenizer.ggml.token_type"] = mx.array(toktypes, dtype=mx.uint32)
|
||||
if vocab.tokenizer.bos_token_id is not None:
|
||||
metadata["tokenizer.ggml.bos_token_id"] = mx.array(
|
||||
vocab.tokenizer.bos_token_id, dtype=mx.uint32
|
||||
)
|
||||
if vocab.tokenizer.eos_token_id is not None:
|
||||
metadata["tokenizer.ggml.eos_token_id"] = mx.array(
|
||||
vocab.tokenizer.eos_token_id, dtype=mx.uint32
|
||||
)
|
||||
if vocab.tokenizer.unk_token_id is not None:
|
||||
metadata["tokenizer.ggml.unknown_token_id"] = mx.array(
|
||||
vocab.tokenizer.unk_token_id, dtype=mx.uint32
|
||||
)
|
||||
metadata["tokenizer.ggml.bos_token_id"] = mx.array(
|
||||
vocab.tokenizer.bos_token_id, dtype=mx.uint32
|
||||
)
|
||||
metadata["tokenizer.ggml.eos_token_id"] = mx.array(
|
||||
vocab.tokenizer.eos_token_id, dtype=mx.uint32
|
||||
)
|
||||
metadata["tokenizer.ggml.unknown_token_id"] = mx.array(
|
||||
vocab.tokenizer.unk_token_id, dtype=mx.uint32
|
||||
)
|
||||
|
||||
metadata = {k: v for k, v in metadata.items() if v is not None}
|
||||
return metadata
|
||||
|
||||
+38
-109
@@ -1,23 +1,23 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import mlx.optimizers as optim
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
from .tuner.callbacks import WandBCallback
|
||||
from .tuner.datasets import CacheDataset, load_dataset
|
||||
from .tokenizer_utils import TokenizerWrapper
|
||||
from .tuner.datasets import load_dataset
|
||||
from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
|
||||
from .tuner.utils import (
|
||||
apply_lora_layers,
|
||||
build_schedule,
|
||||
linear_to_lora_layers,
|
||||
load_adapters,
|
||||
print_trainable_parameters,
|
||||
)
|
||||
from .utils import load, save_config
|
||||
@@ -41,15 +41,9 @@ yaml_loader.add_implicit_resolver(
|
||||
CONFIG_DEFAULTS = {
|
||||
"model": "mlx_model",
|
||||
"train": False,
|
||||
"fine_tune_type": "lora",
|
||||
"optimizer": "adam",
|
||||
"optimizer_config": {
|
||||
"adam": {},
|
||||
"adamw": {},
|
||||
},
|
||||
"data": "data/",
|
||||
"seed": 0,
|
||||
"num_layers": 16,
|
||||
"lora_layers": 16,
|
||||
"batch_size": 4,
|
||||
"iters": 1000,
|
||||
"val_batches": 25,
|
||||
@@ -62,12 +56,9 @@ CONFIG_DEFAULTS = {
|
||||
"test": False,
|
||||
"test_batches": 500,
|
||||
"max_seq_length": 2048,
|
||||
"config": None,
|
||||
"grad_checkpoint": False,
|
||||
"lr_schedule": None,
|
||||
"lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
|
||||
"mask_prompt": False,
|
||||
"wandb": None,
|
||||
"lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
|
||||
"use_dora": False,
|
||||
}
|
||||
|
||||
|
||||
@@ -75,7 +66,6 @@ def build_parser():
|
||||
parser = argparse.ArgumentParser(description="LoRA or QLoRA finetuning.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
help="The path to the local model directory or Hugging Face repo.",
|
||||
)
|
||||
|
||||
@@ -89,32 +79,10 @@ def build_parser():
|
||||
parser.add_argument(
|
||||
"--data",
|
||||
type=str,
|
||||
help=(
|
||||
"Directory with {train, valid, test}.jsonl files or the name "
|
||||
"of a Hugging Face dataset (e.g., 'mlx-community/wikisql')"
|
||||
),
|
||||
help="Directory with {train, valid, test}.jsonl files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fine-tune-type",
|
||||
type=str,
|
||||
choices=["lora", "dora", "full"],
|
||||
help="Type of fine-tuning to perform: lora, dora, or full.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--optimizer",
|
||||
type=str,
|
||||
choices=["adam", "adamw"],
|
||||
default=None,
|
||||
help="Optimizer to use for training: adam or adamw",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mask-prompt",
|
||||
action="store_true",
|
||||
help="Mask the prompt in the loss when training",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-layers",
|
||||
"--lora-layers",
|
||||
type=int,
|
||||
help="Number of layers to fine-tune. Default is 16, use -1 for all.",
|
||||
)
|
||||
@@ -139,12 +107,12 @@ def build_parser():
|
||||
parser.add_argument(
|
||||
"--resume-adapter-file",
|
||||
type=str,
|
||||
help="Load path to resume training from the given fine-tuned weights.",
|
||||
help="Load path to resume training with the given adapters.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--adapter-path",
|
||||
type=str,
|
||||
help="Save/load path for the fine-tuned weights.",
|
||||
help="Save/load path for the adapters.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-every",
|
||||
@@ -170,7 +138,7 @@ def build_parser():
|
||||
parser.add_argument(
|
||||
"-c",
|
||||
"--config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="A YAML configuration file with the training options",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -179,57 +147,36 @@ def build_parser():
|
||||
help="Use gradient checkpointing to reduce memory use.",
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="The PRNG seed")
|
||||
parser.add_argument(
|
||||
"--wandb",
|
||||
type=str,
|
||||
default=None,
|
||||
help="WandB project name to report training metrics. Disabled if None.",
|
||||
"--use-dora", action="store_true", default=None, help="Use DoRA to finetune."
|
||||
)
|
||||
parser.add_argument("--seed", type=int, help="The PRNG seed")
|
||||
return parser
|
||||
|
||||
|
||||
def train_model(
|
||||
args,
|
||||
model: nn.Module,
|
||||
tokenizer: TokenizerWrapper,
|
||||
train_set,
|
||||
valid_set,
|
||||
training_callback: TrainingCallback = None,
|
||||
):
|
||||
mx.random.seed(args.seed)
|
||||
# Freeze all layers
|
||||
model.freeze()
|
||||
if args.num_layers > len(model.layers):
|
||||
raise ValueError(
|
||||
f"Requested to train {args.num_layers} layers "
|
||||
f"but the model only has {len(model.layers)} layers."
|
||||
)
|
||||
|
||||
if args.fine_tune_type == "full":
|
||||
for l in model.layers[-max(args.num_layers, 0) :]:
|
||||
l.unfreeze()
|
||||
# Convert linear layers to lora layers and unfreeze in the process
|
||||
linear_to_lora_layers(model, args.lora_layers, args.lora_parameters)
|
||||
|
||||
args.lora_parameters = None
|
||||
elif args.fine_tune_type in ["lora", "dora"]:
|
||||
# Convert linear layers to lora/dora layers and unfreeze in the process
|
||||
linear_to_lora_layers(
|
||||
model,
|
||||
args.num_layers,
|
||||
args.lora_parameters,
|
||||
use_dora=(args.fine_tune_type == "dora"),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Received unknown fine-tune-type {args.fine_tune_type}")
|
||||
|
||||
# Resume from weights if provided
|
||||
# Resume training the given adapters.
|
||||
if args.resume_adapter_file is not None:
|
||||
print(f"Loading fine-tuned weights from {args.resume_adapter_file}")
|
||||
print(f"Loading pretrained adapters from {args.resume_adapter_file}")
|
||||
model.load_weights(args.resume_adapter_file, strict=False)
|
||||
|
||||
print_trainable_parameters(model)
|
||||
|
||||
adapter_path = Path(args.adapter_path)
|
||||
adapter_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
adapter_file = adapter_path / "adapters.safetensors"
|
||||
save_config(vars(args), adapter_path / "adapter_config.json")
|
||||
|
||||
@@ -246,36 +193,31 @@ def train_model(
|
||||
grad_checkpoint=args.grad_checkpoint,
|
||||
)
|
||||
|
||||
# Initialize the selected optimizer
|
||||
lr = build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
|
||||
|
||||
optimizer_name = args.optimizer.lower()
|
||||
optimizer_config = args.optimizer_config.get(optimizer_name, {})
|
||||
|
||||
if optimizer_name == "adam":
|
||||
opt_class = optim.Adam
|
||||
elif optimizer_name == "adamw":
|
||||
opt_class = optim.AdamW
|
||||
else:
|
||||
raise ValueError(f"Unsupported optimizer: {optimizer_name}")
|
||||
|
||||
opt = opt_class(learning_rate=lr, **optimizer_config)
|
||||
|
||||
model.train()
|
||||
opt = optim.Adam(
|
||||
learning_rate=(
|
||||
build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
|
||||
)
|
||||
)
|
||||
# Train model
|
||||
train(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
args=training_args,
|
||||
optimizer=opt,
|
||||
train_dataset=CacheDataset(train_set),
|
||||
val_dataset=CacheDataset(valid_set),
|
||||
train_dataset=train_set,
|
||||
val_dataset=valid_set,
|
||||
training_callback=training_callback,
|
||||
)
|
||||
|
||||
|
||||
def evaluate_model(args, model: nn.Module, test_set):
|
||||
def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
|
||||
model.eval()
|
||||
|
||||
test_loss = evaluate(
|
||||
model=model,
|
||||
dataset=CacheDataset(test_set),
|
||||
dataset=test_set,
|
||||
tokenizer=tokenizer,
|
||||
batch_size=args.batch_size,
|
||||
num_batches=args.test_batches,
|
||||
max_seq_length=args.max_seq_length,
|
||||
@@ -289,14 +231,6 @@ def evaluate_model(args, model: nn.Module, test_set):
|
||||
def run(args, training_callback: TrainingCallback = None):
|
||||
np.random.seed(args.seed)
|
||||
|
||||
if args.wandb is not None:
|
||||
training_callback = WandBCallback(
|
||||
project_name=args.wandb,
|
||||
log_dir=args.adapter_path,
|
||||
config=vars(args),
|
||||
wrapped_callback=training_callback,
|
||||
)
|
||||
|
||||
print("Loading pretrained model")
|
||||
model, tokenizer = load(args.model)
|
||||
|
||||
@@ -306,21 +240,20 @@ def run(args, training_callback: TrainingCallback = None):
|
||||
if args.test and not args.train:
|
||||
# Allow testing without LoRA layers by providing empty path
|
||||
if args.adapter_path != "":
|
||||
load_adapters(model, args.adapter_path)
|
||||
apply_lora_layers(model, args.adapter_path)
|
||||
|
||||
elif args.train:
|
||||
print("Training")
|
||||
train_model(args, model, train_set, valid_set, training_callback)
|
||||
train_model(args, model, tokenizer, train_set, valid_set, training_callback)
|
||||
else:
|
||||
raise ValueError("Must provide at least one of --train or --test")
|
||||
|
||||
if args.test:
|
||||
print("Testing")
|
||||
evaluate_model(args, model, test_set)
|
||||
evaluate_model(args, model, tokenizer, test_set)
|
||||
|
||||
|
||||
def main():
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
config = args.config
|
||||
@@ -342,8 +275,4 @@ def main():
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.lora...` directly is deprecated."
|
||||
" Use `mlx_lm.lora...` or `python -m mlx_lm lora ...` instead."
|
||||
)
|
||||
main()
|
||||
|
||||
+14
-36
@@ -2,37 +2,23 @@ import argparse
|
||||
from typing import List, Union
|
||||
|
||||
from huggingface_hub import scan_cache_dir
|
||||
|
||||
|
||||
def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
|
||||
"""
|
||||
Inspired by:
|
||||
- stackoverflow.com/a/8356620/593036
|
||||
- stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
|
||||
"""
|
||||
col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
|
||||
row_format = ("{{:{}}} " * len(headers)).format(*col_widths)
|
||||
lines = []
|
||||
lines.append(row_format.format(*headers))
|
||||
lines.append(row_format.format(*["-" * w for w in col_widths]))
|
||||
for row in rows:
|
||||
lines.append(row_format.format(*row))
|
||||
return "\n".join(lines)
|
||||
from transformers.commands.user import tabulate
|
||||
|
||||
|
||||
def ask_for_confirmation(message: str) -> bool:
|
||||
"""Ask user for confirmation with Y/N prompt.
|
||||
Returns True for Y/yes, False for N/no/empty."""
|
||||
y = ("y", "yes", "1")
|
||||
n = ("n", "no", "0", "")
|
||||
full_message = f"{message} (y/n) "
|
||||
n = ("n", "no", "0")
|
||||
all_values = y + n + ("",)
|
||||
full_message = f"{message} (Y/n) "
|
||||
while True:
|
||||
answer = input(full_message).lower()
|
||||
if answer == "":
|
||||
return False
|
||||
if answer in y:
|
||||
return True
|
||||
if answer in n:
|
||||
return False
|
||||
print(f"Invalid input. Must be one of: yes/no/y/n or empty for no")
|
||||
print(f"Invalid input. Must be one of {all_values}")
|
||||
|
||||
|
||||
def main():
|
||||
@@ -57,7 +43,9 @@ def main():
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.scan:
|
||||
print(f'Scanning Hugging Face cache for models with pattern "{args.pattern}".')
|
||||
print(
|
||||
"Scanning Hugging Face cache for models with" f'pattern "{args.pattern}".'
|
||||
)
|
||||
hf_cache_info = scan_cache_dir()
|
||||
print(
|
||||
tabulate(
|
||||
@@ -98,46 +86,36 @@ def main():
|
||||
if args.pattern in repo.repo_id
|
||||
]
|
||||
if repos:
|
||||
print("\nFound the following models:")
|
||||
print(
|
||||
tabulate(
|
||||
rows=[
|
||||
[
|
||||
repo.repo_id,
|
||||
repo.size_on_disk_str, # Added size information
|
||||
str(repo.repo_path),
|
||||
]
|
||||
for repo in repos
|
||||
],
|
||||
headers=[
|
||||
"REPO ID",
|
||||
"SIZE", # Added size header
|
||||
"LOCAL PATH",
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
confirmed = ask_for_confirmation(
|
||||
"\nAre you sure you want to delete these models?"
|
||||
)
|
||||
confirmed = ask_for_confirmation(f"Confirm deletion ?")
|
||||
if confirmed:
|
||||
for model_info in repos:
|
||||
print(f"\nDeleting {model_info.repo_id}...")
|
||||
for revision in sorted(
|
||||
model_info.revisions, key=lambda revision: revision.commit_hash
|
||||
):
|
||||
strategy = hf_cache_info.delete_revisions(revision.commit_hash)
|
||||
strategy.execute()
|
||||
print("\nModel(s) deleted successfully.")
|
||||
print("Model(s) deleted.")
|
||||
else:
|
||||
print("\nDeletion cancelled - no changes made.")
|
||||
print("Deletion is cancelled. Do nothing.")
|
||||
else:
|
||||
print(f'No models found matching pattern "{args.pattern}"')
|
||||
print(f"No models found.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(
|
||||
"Calling `python -m mlx_lm.manage...` directly is deprecated."
|
||||
" Use `mlx_lm.manage...` or `python -m mlx_lm manage ...` instead."
|
||||
)
|
||||
main()
|
||||
|
||||
+172
@@ -0,0 +1,172 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
import yaml
|
||||
from mlx.utils import tree_flatten, tree_map
|
||||
|
||||
from .utils import (
|
||||
fetch_from_hub,
|
||||
get_model_path,
|
||||
save_config,
|
||||
save_weights,
|
||||
upload_to_hub,
|
||||
)
|
||||
|
||||
|
||||
def configure_parser() -> argparse.ArgumentParser:
|
||||
"""
|
||||
Configures and returns the argument parser for the script.
|
||||
|
||||
Returns:
|
||||
argparse.ArgumentParser: Configured argument parser.
|
||||
"""
|
||||
parser = argparse.ArgumentParser(description="Merge multiple models.")
|
||||
|
||||
parser.add_argument("--config", type=str, help="Path to the YAML config.")
|
||||
parser.add_argument(
|
||||
"--mlx-path",
|
||||
type=str,
|
||||
default="mlx_merged_model",
|
||||
help="Path to save the MLX model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload-repo",
|
||||
help="The Hugging Face repo to upload the model to.",
|
||||
type=str,
|
||||
default=None,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def slerp(t, w1, w2, eps=1e-5):
|
||||
"""
|
||||
Spherical linear interpolation
|
||||
|
||||
Args:
|
||||
t (float): Interpolation weight in [0.0, 1.0]
|
||||
w1 (mx.array): First input
|
||||
w2 (mx.array): Second input
|
||||
eps (float): Constant for numerical stability
|
||||
Returns:
|
||||
mx.array: Interpolated result
|
||||
"""
|
||||
t = float(t)
|
||||
if t == 0:
|
||||
return w1
|
||||
elif t == 1:
|
||||
return w2
|
||||
# Normalize
|
||||
v1 = w1 / mx.linalg.norm(w1)
|
||||
v2 = w2 / mx.linalg.norm(w2)
|
||||
# Angle
|
||||
dot = mx.clip((v1 * v2).sum(), 0.0, 1.0)
|
||||
theta = mx.arccos(dot)
|
||||
sin_theta = mx.sin(theta + eps)
|
||||
s1 = mx.sin(theta * (1 - t)) / sin_theta
|
||||
s2 = mx.sin(theta * t) / sin_theta
|
||||
return s1 * w1 + s2 * w2
|
||||
|
||||
|
||||
def merge_models(base_model: nn.Module, model: nn.Module, config: dict):
|
||||
method = config.get("method", None)
|
||||
if method != "slerp":
|
||||
raise ValueError(f"Merge method {method} not supported")
|
||||
|
||||
num_layers = len(model.layers)
|
||||
|
||||
def unpack_values(vals):
|
||||
if isinstance(vals, (int, float)):
|
||||
return np.full(num_layers, vals)
|
||||
bins = len(vals) - 1
|
||||
sizes = [num_layers // bins] * bins
|
||||
sizes[-1] = num_layers - sum(sizes[:-1])
|
||||
return np.concatenate(
|
||||
[np.linspace(v1, v2, s) for v1, v2, s in zip(vals[:-1], vals[1:], sizes)]
|
||||
)
|
||||
|
||||
param_list = config["parameters"]["t"]
|
||||
params = {}
|
||||
filter_keys = set()
|
||||
for pl in param_list[:-1]:
|
||||
params[pl["filter"]] = unpack_values(pl["value"])
|
||||
filter_keys.add(pl["filter"])
|
||||
default = unpack_values(param_list[-1]["value"])
|
||||
|
||||
for e in range(num_layers):
|
||||
bl = base_model.layers[e]
|
||||
l = model.layers[e]
|
||||
base_weights = bl.parameters()
|
||||
weights = l.parameters()
|
||||
for k, w1 in base_weights.items():
|
||||
w2 = weights[k]
|
||||
t = params.get(k, default)[e]
|
||||
base_weights[k] = tree_map(lambda x, y: slerp(t, x, y), w1, w2)
|
||||
base_model.update(base_weights)
|
||||
|
||||
|
||||
def merge(
|
||||
config: str,
|
||||
mlx_path: str = "mlx_model",
|
||||
upload_repo: Optional[str] = None,
|
||||
):
|
||||
with open(config, "r") as fid:
|
||||
merge_conf = yaml.safe_load(fid)
|
||||
print("[INFO] Loading")
|
||||
|
||||
model_paths = merge_conf.get("models", [])
|
||||
if len(model_paths) < 2:
|
||||
raise ValueError(f"Expected at least 2 models, got {len(model_paths)}.")
|
||||
|
||||
# Load all models
|
||||
base_hf_path = model_paths[0]
|
||||
base_path = get_model_path(base_hf_path)
|
||||
base_model, base_config, tokenizer = fetch_from_hub(base_path, lazy=True)
|
||||
models = []
|
||||
for mp in model_paths[1:]:
|
||||
model, model_config, _ = fetch_from_hub(get_model_path(mp), lazy=True)
|
||||
base_type = base_config["model_type"]
|
||||
model_type = model_config["model_type"]
|
||||
if base_type != model_type:
|
||||
raise ValueError(
|
||||
f"Can only merge models of the same type,"
|
||||
f" but got {base_type} and {model_type}."
|
||||
)
|
||||
models.append(model)
|
||||
|
||||
# Merge models into base model
|
||||
for m in models:
|
||||
merge_models(base_model, m, merge_conf)
|
||||
|
||||
# Save base model
|
||||
mlx_path = Path(mlx_path)
|
||||
weights = dict(tree_flatten(base_model.parameters()))
|
||||
del models, base_model
|
||||
save_weights(mlx_path, weights, donate_weights=True)
|
||||
py_files = glob.glob(str(base_path / "*.py"))
|
||||
for file in py_files:
|
||||
shutil.copy(file, mlx_path)
|
||||
|
||||
tokenizer.save_pretrained(mlx_path)
|
||||
|
||||
save_config(config, config_path=mlx_path / "config.json")
|
||||
|
||||
if upload_repo is not None:
|
||||
upload_to_hub(mlx_path, upload_repo, base_hf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = configure_parser()
|
||||
args = parser.parse_args()
|
||||
merge(**vars(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,397 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from itertools import accumulate
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import ConcatenateKVCache, KVCache
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
vocab_size: int
|
||||
hidden_dim: int
|
||||
num_layers: int
|
||||
num_kv_reuse_layers: int
|
||||
num_heads: int
|
||||
num_kv_heads: int
|
||||
hidden_dim_scale_factor: float = 3.25
|
||||
rope_theta: float = 50000
|
||||
rms_norm_eps: float = 1e-5
|
||||
|
||||
|
||||
class FusedLoRALinear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dims: int,
|
||||
output_dims: list[int],
|
||||
r: int = 8,
|
||||
dropout: float = 0.0,
|
||||
scale: float = 20.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.linear = FusedLinear(input_dims, output_dims)
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
self.scale = scale
|
||||
|
||||
scale = 1 / math.sqrt(input_dims)
|
||||
self.lora_a = [
|
||||
mx.random.uniform(low=-scale, high=scale, shape=(input_dims, r))
|
||||
for _ in output_dims
|
||||
]
|
||||
self.lora_b = [mx.zeros((r, od)) for od in output_dims]
|
||||
|
||||
def fuse(self, de_quantize: bool = False):
|
||||
linear = self.linear
|
||||
weight = linear.weight
|
||||
is_quantized = isinstance(linear, FusedQuantizedLinear)
|
||||
|
||||
# Use the same type as the linear weight if not quantized
|
||||
dtype = weight.dtype
|
||||
|
||||
if is_quantized:
|
||||
dtype = linear.scales.dtype
|
||||
weight = mx.dequantize(
|
||||
weight,
|
||||
linear.scales,
|
||||
linear.biases,
|
||||
linear.group_size,
|
||||
linear.bits,
|
||||
)
|
||||
|
||||
input_dims = weight.shape[-1]
|
||||
output_dims = linear.output_dims
|
||||
fused_linear = FusedLinear(input_dims, output_dims)
|
||||
fused_linear.weight = weight
|
||||
deltas = [
|
||||
((self.scale * b.T) @ a.T).astype(dtype)
|
||||
for a, b in zip(self.lora_a, self.lora_b)
|
||||
]
|
||||
delta = mx.concatenate(deltas, axis=0)
|
||||
fused_linear.weight = weight + delta
|
||||
|
||||
if is_quantized and not de_quantize:
|
||||
fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
|
||||
|
||||
return fused_linear
|
||||
|
||||
def __call__(self, x):
|
||||
dt = x.dtype
|
||||
y = self.linear(x)
|
||||
x = self.dropout(x)
|
||||
z = [(x @ a) @ b for a, b in zip(self.lora_a, self.lora_b)]
|
||||
return tuple(yi + (self.scale * zi).astype(dt) for yi, zi in zip(y, z))
|
||||
|
||||
|
||||
class FusedQuantizedLinear(nn.QuantizedLinear):
|
||||
def __init__(self, input_dims, output_dims, group_size: int = 64, bits: int = 4):
|
||||
*indices, output_dims = accumulate(output_dims)
|
||||
self.indices = indices
|
||||
super().__init__(
|
||||
input_dims, output_dims, bias=False, group_size=group_size, bits=bits
|
||||
)
|
||||
|
||||
@property
|
||||
def input_dims(self):
|
||||
return self.scales.shape[-1] * self.group_size
|
||||
|
||||
@property
|
||||
def output_dims(self):
|
||||
indices = [0] + self.indices + [self.weight.shape[0]]
|
||||
return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
|
||||
|
||||
def __call__(self, x):
|
||||
x = super().__call__(x)
|
||||
return x.split(self.indices, axis=-1)
|
||||
|
||||
def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
|
||||
lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
|
||||
lora_lin.linear = self
|
||||
return lora_lin
|
||||
|
||||
|
||||
class FusedLinear(nn.Linear):
|
||||
def __init__(self, input_dims, output_dims):
|
||||
*indices, output_dims = accumulate(output_dims)
|
||||
self.indices = indices
|
||||
super().__init__(input_dims, output_dims, bias=False)
|
||||
|
||||
@property
|
||||
def input_dims(self):
|
||||
return self.weight.shape[-1]
|
||||
|
||||
@property
|
||||
def output_dims(self):
|
||||
indices = [0] + self.indices + [self.weight.shape[0]]
|
||||
return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
|
||||
|
||||
def __call__(self, x):
|
||||
x = super().__call__(x)
|
||||
return x.split(self.indices, axis=-1)
|
||||
|
||||
def to_quantized(self, group_size: int = 64, bits: int = 4):
|
||||
input_dims = self.input_dims
|
||||
output_dims = self.output_dims
|
||||
ql = FusedQuantizedLinear(input_dims, output_dims, group_size, bits)
|
||||
ql.weight, ql.scales, ql.biases = mx.quantize(self.weight, group_size, bits)
|
||||
|
||||
return ql
|
||||
|
||||
def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
|
||||
lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
|
||||
lora_lin.linear = self
|
||||
return lora_lin
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def fake_8bit_quant(x, scale):
|
||||
dt = x.dtype
|
||||
x = x.astype(mx.float32)
|
||||
x = (x / scale).round()
|
||||
x = mx.clip(x, -128, 127)
|
||||
return (x * scale).astype(dt)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_dim
|
||||
self.n_heads = n_heads = args.num_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_kv_heads
|
||||
self.head_dim = head_dim = args.hidden_dim // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
qkv_dim = (n_heads + 2 * n_kv_heads) * head_dim
|
||||
self.qkv_proj = FusedLinear(
|
||||
dim, [n_heads * head_dim] + 2 * [n_kv_heads * head_dim]
|
||||
)
|
||||
self.out_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
True,
|
||||
)
|
||||
self.q_norm = nn.RMSNorm(head_dim)
|
||||
self.k_norm = nn.RMSNorm(head_dim)
|
||||
self.quant_key_scale = mx.array(1.0)
|
||||
self.quant_value_scale = mx.array(1.0)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
# Get the queries, keys and values
|
||||
queries, keys, values = self.qkv_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.q_norm(self.rope(queries, offset=cache.offset))
|
||||
keys = self.k_norm(self.rope(keys, offset=cache.offset))
|
||||
keys = fake_8bit_quant(keys, self.quant_key_scale)
|
||||
values = fake_8bit_quant(values, self.quant_value_scale)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.q_norm(self.rope(queries))
|
||||
keys = self.k_norm(self.rope(keys))
|
||||
keys = fake_8bit_quant(keys, self.quant_key_scale)
|
||||
values = fake_8bit_quant(values, self.quant_value_scale)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.out_proj(output)
|
||||
|
||||
|
||||
class KVReuseAttention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_dim
|
||||
self.n_heads = n_heads = args.num_heads
|
||||
self.head_dim = head_dim = args.hidden_dim // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.out_proj = nn.Linear(dim, dim, bias=False)
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
True,
|
||||
)
|
||||
self.q_norm = nn.RMSNorm(head_dim)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
keys: mx.array,
|
||||
values: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
_, _, S, _ = keys.shape
|
||||
|
||||
queries = self.q_proj(x)
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
queries = self.q_norm(self.rope(queries, offset=S - L))
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=None, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.out_proj(output)
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def _swiglu(g, x):
|
||||
return nn.silu(g) * x
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_dim
|
||||
hidden_dim = int(dim * args.hidden_dim_scale_factor)
|
||||
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
g = self.gate_proj(x)
|
||||
x = self.up_proj(x)
|
||||
return self.down_proj(_swiglu(g, x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_dim, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class KVReuseTransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.self_attn = KVReuseAttention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_dim, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
keys: mx.array,
|
||||
values: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), keys, values, mask)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class AFMModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
|
||||
self.embedding = nn.Embedding(args.vocab_size, args.hidden_dim)
|
||||
self.layers = [
|
||||
TransformerBlock(args)
|
||||
for _ in range(args.num_layers - args.num_kv_reuse_layers)
|
||||
]
|
||||
self.kv_reuse_layers = [
|
||||
KVReuseTransformerBlock(args) for _ in range(args.num_kv_reuse_layers)
|
||||
]
|
||||
self.output_norm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embedding(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
cache[-1] = ConcatenateKVCache()
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
keys, values = cache[-1].state
|
||||
for layer in self.kv_reuse_layers:
|
||||
h = layer(h, keys, values, mask)
|
||||
|
||||
return self.output_norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = AFMModel(args)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model.embedding.as_linear(out)
|
||||
return out
|
||||
|
||||
def make_cache(self):
|
||||
return [KVCache() for _ in range(len(self.model.layers))]
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers + self.model.kv_reuse_layers
|
||||
@@ -1,226 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import CacheList, KVCache, MambaCache, RotatingKVCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
vocab_size: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
rope_theta: float
|
||||
sliding_window: int
|
||||
sliding_window_layers: List[int]
|
||||
conv_window: int
|
||||
rms_norm_eps: float
|
||||
model_type: str = "baichuan_m1"
|
||||
num_swa_attention_heads: Optional[int] = None
|
||||
num_swa_key_value_heads: Optional[int] = None
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, config: ModelArgs, layer_idx: Optional[int] = None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layer_idx = layer_idx
|
||||
if layer_idx is None:
|
||||
raise ValueError("Layer index must be provided to Attention module.")
|
||||
|
||||
self.is_swa = layer_idx in config.sliding_window_layers
|
||||
self.num_heads = (
|
||||
config.num_swa_attention_heads
|
||||
if self.is_swa and config.num_swa_attention_heads
|
||||
else config.num_attention_heads
|
||||
)
|
||||
self.num_kv_heads = (
|
||||
config.num_swa_key_value_heads
|
||||
if self.is_swa and config.num_swa_key_value_heads
|
||||
else config.num_key_value_heads
|
||||
)
|
||||
|
||||
self.hidden_size = config.hidden_size
|
||||
self.head_dim = self.hidden_size // self.num_heads
|
||||
assert self.head_dim * self.num_heads == self.hidden_size
|
||||
|
||||
self.scale = self.head_dim**-0.5
|
||||
|
||||
self.W_pack = nn.Linear(
|
||||
config.hidden_size,
|
||||
self.hidden_size + 2 * self.num_kv_heads * self.head_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.o_proj = nn.Linear(
|
||||
self.num_heads * self.head_dim, config.hidden_size, bias=False
|
||||
)
|
||||
|
||||
self.rope = nn.RoPE(self.head_dim, traditional=False, base=config.rope_theta)
|
||||
|
||||
self.conv_window = config.conv_window
|
||||
assert self.conv_window == 2
|
||||
self.conv_k = mx.zeros((1, 1, self.num_kv_heads, 1, self.conv_window))
|
||||
self.conv_v = mx.zeros((1, 1, self.num_kv_heads, 1, self.conv_window))
|
||||
|
||||
def _custom_convolution(self, u, weights, state=None):
|
||||
B, H, L, D = u.shape
|
||||
weights = weights.reshape((1, H, self.conv_window, 1, 1))
|
||||
w0 = weights[:, :, 0]
|
||||
w1 = weights[:, :, 1]
|
||||
if state is None:
|
||||
state = mx.zeros((B, H, 1, D), u.dtype)
|
||||
if L > 1:
|
||||
u_prev = mx.concatenate([state, u[:, :, :-1]], axis=2)
|
||||
else:
|
||||
u_prev = state
|
||||
return u_prev * w0 + u * w1
|
||||
|
||||
def __call__(
|
||||
self, x: mx.array, mask: mx.array = None, cache: Any = None
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
proj = self.W_pack(x)
|
||||
q, k, v = mx.split(proj, (D, D + self.num_kv_heads * self.head_dim), axis=-1)
|
||||
|
||||
q = q.reshape(B, L, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
|
||||
k = k.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
|
||||
v = v.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
offset = cache[1].offset
|
||||
last_k, last_v = cache[0][0], cache[0][1]
|
||||
else:
|
||||
offset = 0
|
||||
last_k, last_v = None, None
|
||||
|
||||
k_init = k
|
||||
v_init = v
|
||||
k = self._custom_convolution(k, self.conv_k, state=last_k)
|
||||
v = self._custom_convolution(v, self.conv_v, state=last_v)
|
||||
q = self.rope(q, offset=offset)
|
||||
k = self.rope(k, offset=offset)
|
||||
|
||||
if cache is not None:
|
||||
k, v = cache[1].update_and_fetch(k, v)
|
||||
if L > 0:
|
||||
cache[0][0] = k_init[:, :, -1:, :]
|
||||
cache[0][1] = v_init[:, :, -1:, :]
|
||||
|
||||
out = scaled_dot_product_attention(
|
||||
q, k, v, cache=cache[1], scale=self.scale, mask=mask
|
||||
)
|
||||
out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(out)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(
|
||||
config.hidden_size, config.intermediate_size, bias=False
|
||||
)
|
||||
self.up_proj = nn.Linear(
|
||||
config.hidden_size, config.intermediate_size, bias=False
|
||||
)
|
||||
self.down_proj = nn.Linear(
|
||||
config.intermediate_size, config.hidden_size, bias=False
|
||||
)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, config: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(config, layer_idx)
|
||||
self.mlp = MLP(config)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self, x: mx.array, mask: mx.array = None, cache: Any = None
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
x = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(x))
|
||||
return x + r
|
||||
|
||||
|
||||
class BaichuanModel(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = config
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
|
||||
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self, inputs: mx.array, mask: mx.array = None, cache: Any = None
|
||||
) -> mx.array:
|
||||
x = self.embed_tokens(inputs)
|
||||
if mask is None:
|
||||
if cache is not None:
|
||||
c = [cache[0][1]]
|
||||
mask = create_attention_mask(x, c)
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
for layer, c in zip(self.layers, cache):
|
||||
x = layer(x, mask, c)
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.model_type = config.model_type
|
||||
self.model = BaichuanModel(config)
|
||||
self.tie_word_embeddings = config.tie_word_embeddings
|
||||
if not config.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
def make_cache(self) -> List[Any]:
|
||||
caches = []
|
||||
for i, layer in enumerate(self.model.layers):
|
||||
is_swa = i in self.config.sliding_window_layers
|
||||
conv_cache = MambaCache()
|
||||
if is_swa:
|
||||
kv_cache = RotatingKVCache(max_size=self.config.sliding_window)
|
||||
else:
|
||||
kv_cache = KVCache()
|
||||
caches.append(CacheList(conv_cache, kv_cache))
|
||||
return caches
|
||||
|
||||
def sanitize(self, weights: dict) -> dict:
|
||||
is_quantized = "lm_head.scales" in weights
|
||||
if not is_quantized and "lm_head.weight" in weights:
|
||||
w = weights["lm_head.weight"]
|
||||
dtype = w.dtype
|
||||
w = w.astype(mx.float32)
|
||||
norm = mx.linalg.norm(w, axis=-1, keepdims=True)
|
||||
w = (w / (norm + 1e-7)).astype(dtype)
|
||||
weights["lm_head.weight"] = w
|
||||
return weights
|
||||
|
||||
def __call__(
|
||||
self, inputs: mx.array, mask: mx.array = None, cache: Any = None
|
||||
) -> mx.array:
|
||||
outputs = self.model(inputs, mask, cache)
|
||||
return self.lm_head(outputs)
|
||||
|
||||
@property
|
||||
def layers(self) -> List[nn.Module]:
|
||||
return self.model.layers
|
||||
+38
-115
@@ -1,13 +1,46 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
from mlx.utils import tree_map
|
||||
|
||||
from .cache import QuantizedKVCache
|
||||
|
||||
def create_additive_causal_mask(N: int, offset: int = 0):
|
||||
rinds = mx.arange(offset + N)
|
||||
linds = mx.arange(offset, offset + N) if offset else rinds
|
||||
mask = linds[:, None] < rinds[None]
|
||||
return mask * -1e9
|
||||
|
||||
|
||||
class KVCache:
|
||||
|
||||
def __init__(self, head_dim, n_kv_heads):
|
||||
self.n_kv_heads = n_kv_heads
|
||||
self.head_dim = head_dim
|
||||
self.keys = None
|
||||
self.values = None
|
||||
self.offset = 0
|
||||
self.step = 256
|
||||
|
||||
def update_and_fetch(self, keys, values):
|
||||
prev = self.offset
|
||||
if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
|
||||
n_steps = (self.step + keys.shape[2] - 1) // self.step
|
||||
shape = (1, self.n_kv_heads, n_steps * self.step, self.head_dim)
|
||||
new_k = mx.zeros(shape, keys.dtype)
|
||||
new_v = mx.zeros(shape, values.dtype)
|
||||
if self.keys is not None:
|
||||
if prev % self.step != 0:
|
||||
self.keys = self.keys[..., :prev, :]
|
||||
self.values = self.values[..., :prev, :]
|
||||
self.keys = mx.concatenate([self.keys, new_k], axis=2)
|
||||
self.values = mx.concatenate([self.values, new_v], axis=2)
|
||||
else:
|
||||
self.keys, self.values = new_k, new_v
|
||||
|
||||
self.offset += keys.shape[2]
|
||||
self.keys[..., prev : self.offset, :] = keys
|
||||
self.values[..., prev : self.offset, :] = values
|
||||
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -21,113 +54,3 @@ class BaseModelArgs:
|
||||
if k in inspect.signature(cls).parameters
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def create_causal_mask(
|
||||
N: int,
|
||||
offset: int = 0,
|
||||
window_size: Optional[int] = None,
|
||||
lengths: Optional[mx.array] = None,
|
||||
):
|
||||
rinds = mx.arange(offset + N)
|
||||
linds = mx.arange(offset, offset + N) if offset else rinds
|
||||
linds = linds[:, None]
|
||||
rinds = rinds[None]
|
||||
mask = linds >= rinds
|
||||
if window_size is not None:
|
||||
mask = mask & (linds <= rinds + window_size)
|
||||
if lengths is not None:
|
||||
lengths = lengths[:, None, None, None]
|
||||
mask = mask & (rinds < lengths)
|
||||
return mask
|
||||
|
||||
|
||||
def create_attention_mask(
|
||||
h: mx.array, cache: Optional[Any] = None, return_array: bool = False
|
||||
):
|
||||
T = h.shape[1]
|
||||
if T > 1:
|
||||
offset = 0
|
||||
window_size = None
|
||||
if cache is not None and cache[0] is not None:
|
||||
c = cache[0]
|
||||
offset = c.offset
|
||||
if hasattr(c, "max_size"):
|
||||
window_size = c.max_size
|
||||
offset = min(window_size, offset)
|
||||
return_array = return_array or offset + T > window_size
|
||||
if return_array:
|
||||
return create_causal_mask(T, offset, window_size=window_size)
|
||||
else:
|
||||
return "causal"
|
||||
else:
|
||||
mask = None
|
||||
return mask
|
||||
|
||||
|
||||
def quantized_scaled_dot_product_attention(
|
||||
queries: mx.array,
|
||||
q_keys: tuple[mx.array, mx.array, mx.array],
|
||||
q_values: tuple[mx.array, mx.array, mx.array],
|
||||
scale: float,
|
||||
mask: Optional[mx.array],
|
||||
group_size: int = 64,
|
||||
bits: int = 8,
|
||||
) -> mx.array:
|
||||
B, n_q_heads, L, D = queries.shape
|
||||
n_kv_heads = q_keys[0].shape[-3]
|
||||
n_repeats = n_q_heads // n_kv_heads
|
||||
|
||||
queries *= scale
|
||||
|
||||
if n_repeats > 1:
|
||||
queries = mx.reshape(queries, (B, n_kv_heads, n_repeats, L, D))
|
||||
q_keys = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_keys)
|
||||
q_values = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_values)
|
||||
|
||||
scores = mx.quantized_matmul(
|
||||
queries, *q_keys, transpose=True, group_size=group_size, bits=bits
|
||||
)
|
||||
if mask is not None:
|
||||
if isinstance(mask, str):
|
||||
qL, kL = scores.shape[-2:]
|
||||
q_indices = mx.arange(kL - qL, kL)
|
||||
k_indices = mx.arange(kL)
|
||||
mask = q_indices[:, None] >= k_indices[None]
|
||||
if mask.dtype == mx.bool_:
|
||||
scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
|
||||
else:
|
||||
scores += mask
|
||||
scores = mx.softmax(scores, axis=-1, precise=True)
|
||||
out = mx.quantized_matmul(
|
||||
scores, *q_values, transpose=False, group_size=group_size, bits=bits
|
||||
)
|
||||
|
||||
if n_repeats > 1:
|
||||
out = mx.reshape(out, (B, n_q_heads, L, D))
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def scaled_dot_product_attention(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
cache,
|
||||
scale: float,
|
||||
mask: Optional[mx.array],
|
||||
) -> mx.array:
|
||||
if isinstance(cache, QuantizedKVCache):
|
||||
return quantized_scaled_dot_product_attention(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
scale=scale,
|
||||
mask=mask,
|
||||
group_size=cache.group_size,
|
||||
bits=cache.bits,
|
||||
)
|
||||
else:
|
||||
return mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=scale, mask=mask
|
||||
)
|
||||
|
||||
@@ -1,131 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.nn.layers.quantized import QuantizedLinear
|
||||
|
||||
|
||||
def make_bitlinear_kernel():
|
||||
"""
|
||||
Custom Metal kernel that performs matrix multiplication directly on
|
||||
packed weights and scales the output. This eliminates the need to
|
||||
store unpacked weights in memory.
|
||||
"""
|
||||
source = """
|
||||
constexpr int M = 4;
|
||||
constexpr int BLOCK = 32;
|
||||
|
||||
uint tid = thread_position_in_grid.y;
|
||||
uint in_offset = thread_position_in_grid.x;
|
||||
|
||||
uint batch_idx = tid / (out_features / 4);
|
||||
uint row_idx = tid % (out_features / 4);
|
||||
|
||||
float sum[4] = {0.0};
|
||||
|
||||
for (uint i = in_offset * M; i < in_features; i += BLOCK * M) {
|
||||
float v[M];
|
||||
for (int j=0; j<M; j++) {
|
||||
v[j] = x[batch_idx * in_features + i + j];
|
||||
}
|
||||
|
||||
for (int j=0; j<M; j++) {
|
||||
uint8_t w = packed_weights[row_idx * in_features + i + j];
|
||||
sum[0] += v[j] * ((w & 3) - 1);
|
||||
sum[1] += v[j] * (((w >> 2) & 3) - 1);
|
||||
sum[2] += v[j] * (((w >> 4) & 3) - 1);
|
||||
sum[3] += v[j] * (((w >> 6) & 3) - 1);
|
||||
}
|
||||
}
|
||||
|
||||
for (int j=0; j<4; j++) {
|
||||
sum[j] = simd_sum(sum[j]);
|
||||
}
|
||||
|
||||
// Apply weight scaling by diving them or multiplying them
|
||||
if (in_offset == 0) {
|
||||
float scale = invert_weight_scales ? 1 / weight_scale[0] : weight_scale[0];
|
||||
for (int i=0; i<4; i++) {
|
||||
out[batch_idx * out_features + row_idx + i * (out_features/4)] = static_cast<T>(sum[i] * scale);
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
return mx.fast.metal_kernel(
|
||||
name="bitlinear_matmul",
|
||||
input_names=["x", "packed_weights", "weight_scale"],
|
||||
output_names=["out"],
|
||||
source=source,
|
||||
)
|
||||
|
||||
|
||||
_bitlinear_kernel = make_bitlinear_kernel()
|
||||
|
||||
|
||||
class BitLinear(nn.Module):
|
||||
"""
|
||||
BitLinear module with memory-efficient weight handling.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
bias=True,
|
||||
invert_weight_scales=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
# Calculate packed dimensions - the first dimension gets packed 4:1
|
||||
# The weights are ternary so can be represented with 2 bits, and they
|
||||
# are packed in uint8 tensors, hence the number of values per item is 4
|
||||
packed_out_features = (out_features + 3) // 4
|
||||
self.weight = mx.zeros((packed_out_features, in_features), dtype=mx.uint8)
|
||||
|
||||
self.invert_weight_scales = invert_weight_scales
|
||||
self.weight_scale = mx.array([1.0])
|
||||
|
||||
if bias:
|
||||
self.bias = mx.zeros((out_features,))
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
def execute_matmul_kernel(self, x, packed_weights):
|
||||
original_shape = x.shape
|
||||
if len(original_shape) > 2:
|
||||
x = x.reshape(-1, original_shape[-1])
|
||||
total_batch_elements, in_features = x.shape
|
||||
|
||||
out_features = self.out_features
|
||||
|
||||
dtype = self.weight_scale.dtype
|
||||
assert x.dtype == dtype, "Wrong type for input."
|
||||
out = _bitlinear_kernel(
|
||||
inputs=[
|
||||
x,
|
||||
packed_weights,
|
||||
self.weight_scale,
|
||||
],
|
||||
template=[
|
||||
("T", dtype),
|
||||
("invert_weight_scales", self.invert_weight_scales),
|
||||
("in_features", in_features),
|
||||
("out_features", out_features),
|
||||
],
|
||||
grid=(32, total_batch_elements * out_features // 4, 1),
|
||||
threadgroup=(32, 1, 1), # SIMD width is 32 threads
|
||||
output_shapes=[(total_batch_elements, out_features)],
|
||||
output_dtypes=[dtype],
|
||||
)[0]
|
||||
|
||||
if len(original_shape) > 2:
|
||||
out = out.reshape(*original_shape[:-1], out_features)
|
||||
return out
|
||||
|
||||
def __call__(self, x):
|
||||
y = self.execute_matmul_kernel(x, self.weight)
|
||||
|
||||
if self.bias is not None:
|
||||
y = mx.add(y, self.bias)
|
||||
return y
|
||||
@@ -1,215 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .bitlinear_layers import BitLinear
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
head_dim: Optional[int] = None
|
||||
max_position_embeddings: Optional[int] = None
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = True
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
|
||||
self.scale = head_dim**-0.5
|
||||
attention_bias = args.attention_bias
|
||||
|
||||
self.q_proj = BitLinear(dim, n_heads * head_dim, bias=attention_bias)
|
||||
self.k_proj = BitLinear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.v_proj = BitLinear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.o_proj = BitLinear(n_heads * head_dim, dim, bias=attention_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
args.rope_traditional,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
)
|
||||
self.attn_sub_norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
output = self.attn_sub_norm(output)
|
||||
output = self.o_proj(output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def relu2(x):
|
||||
return mx.square(nn.relu(x))
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
hidden_dim = args.intermediate_size
|
||||
if hasattr(args, "mlp_bias"):
|
||||
mlp_bias = args.mlp_bias
|
||||
else:
|
||||
mlp_bias = False
|
||||
|
||||
self.gate_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
|
||||
self.down_proj = BitLinear(hidden_dim, dim, bias=mlp_bias)
|
||||
self.up_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
|
||||
self.ffn_sub_norm = nn.RMSNorm(args.intermediate_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
x = relu2(self.gate_proj(x)) * self.up_proj(x)
|
||||
x = self.ffn_sub_norm(x)
|
||||
x = self.down_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class LlamaModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = LlamaModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
weights = {
|
||||
k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
|
||||
}
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,545 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_map, tree_unflatten
|
||||
|
||||
|
||||
def make_prompt_cache(
|
||||
model: nn.Module,
|
||||
max_kv_size: Optional[int] = None,
|
||||
) -> List[Any]:
|
||||
"""
|
||||
Construct the model's cache for use in generation.
|
||||
|
||||
This function will defer the cache construction to the model if it has a
|
||||
``make_cache`` method, otherwise it will make a default KV cache.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The language model.
|
||||
max_kv_size (Optional[int]): If provided and the model does not have a
|
||||
``make_cache`` method, a ``RotatingKVCache`` is used with a maximum
|
||||
size of ``max_kv_size``
|
||||
"""
|
||||
if hasattr(model, "make_cache"):
|
||||
return model.make_cache()
|
||||
|
||||
num_layers = len(model.layers)
|
||||
if max_kv_size is not None:
|
||||
return [
|
||||
RotatingKVCache(max_size=max_kv_size, keep=4) for _ in range(num_layers)
|
||||
]
|
||||
else:
|
||||
return [KVCache() for _ in range(num_layers)]
|
||||
|
||||
|
||||
def save_prompt_cache(file_name: str, cache: List[Any], metadata: Dict[str, str] = {}):
|
||||
"""
|
||||
Save a pre-computed prompt cache to a file.
|
||||
|
||||
Args:
|
||||
file_name (str): The ``.safetensors`` file name.
|
||||
cache (List[Any]): The model state.
|
||||
metadata (Dict[str, str]): Optional metadata to save along with model
|
||||
state.
|
||||
"""
|
||||
cache_data = [c.state for c in cache]
|
||||
cache_info = [c.meta_state for c in cache]
|
||||
cache_data = dict(tree_flatten(cache_data))
|
||||
cache_classes = [type(c).__name__ for c in cache]
|
||||
cache_metadata = [cache_info, metadata, cache_classes]
|
||||
cache_metadata = dict(tree_flatten(cache_metadata))
|
||||
mx.save_safetensors(file_name, cache_data, cache_metadata)
|
||||
|
||||
|
||||
def load_prompt_cache(file_name, return_metadata=False):
|
||||
"""
|
||||
Load a prompt cache from a file.
|
||||
|
||||
Args:
|
||||
file_name (str): The ``.safetensors`` file name.
|
||||
return_metadata (bool): Whether or not to return metadata.
|
||||
Default: ``False``.
|
||||
|
||||
Returns:
|
||||
List[Any] or Tuple[List[Any], Dict[str, str]]: The prompt cache and
|
||||
the metadata if requested.
|
||||
"""
|
||||
arrays, cache_metadata = mx.load(file_name, return_metadata=True)
|
||||
arrays = tree_unflatten(list(arrays.items()))
|
||||
cache_metadata = tree_unflatten(list(cache_metadata.items()))
|
||||
info, metadata, classes = cache_metadata
|
||||
cache = [globals()[c]() for c in classes]
|
||||
for c, state, meta_state in zip(cache, arrays, info):
|
||||
c.state = state
|
||||
c.meta_state = meta_state
|
||||
if return_metadata:
|
||||
return cache, metadata
|
||||
return cache
|
||||
|
||||
|
||||
def can_trim_prompt_cache(cache: List[Any]) -> bool:
|
||||
"""
|
||||
Check if model's cache can be trimmed.
|
||||
"""
|
||||
return all(c.is_trimmable() for c in cache)
|
||||
|
||||
|
||||
def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
|
||||
"""
|
||||
Trim the model's cache by the given number of tokens.
|
||||
|
||||
This function will trim the cache if possible (in-place) and return the
|
||||
number of tokens that were trimmed.
|
||||
|
||||
Args:
|
||||
cache (List[Any]): The model's cache.
|
||||
num_tokens (int): The number of tokens to trim.
|
||||
|
||||
Returns:
|
||||
(int): The number of tokens that were trimmed.
|
||||
"""
|
||||
if not can_trim_prompt_cache(cache) or len(cache) == 0:
|
||||
return 0
|
||||
return [c.trim(num_tokens) for c in cache][0]
|
||||
|
||||
|
||||
class _BaseCache:
|
||||
@property
|
||||
def state(self):
|
||||
return []
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
if v is not None and v:
|
||||
raise ValueError("This cache has no state but a state was set.")
|
||||
|
||||
@property
|
||||
def meta_state(self):
|
||||
return ""
|
||||
|
||||
@meta_state.setter
|
||||
def meta_state(self, v):
|
||||
if v is not None and v:
|
||||
raise ValueError("This cache has no meta_state but a meta_state was set.")
|
||||
|
||||
def is_trimmable(self):
|
||||
return False
|
||||
|
||||
|
||||
class ConcatenateKVCache(_BaseCache):
|
||||
"""ConcatenateKVCache the simplest KV cache implementation.
|
||||
|
||||
Can be used as a mock KV cache or when large blocks are being processed at
|
||||
a time in which case KVCache isn't necessarily faster. Consider using the
|
||||
KVCache with a larger step size before using this cache.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.keys = None
|
||||
self.values = None
|
||||
self.offset = 0
|
||||
|
||||
def update_and_fetch(self, keys, values):
|
||||
if self.keys is None:
|
||||
self.keys = keys
|
||||
self.values = values
|
||||
else:
|
||||
self.keys = mx.concatenate([self.keys, keys], axis=-2)
|
||||
self.values = mx.concatenate([self.values, values], axis=-2)
|
||||
self.offset = self.keys.shape[-2]
|
||||
|
||||
return self.keys, self.values
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return self.keys, self.values
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
self.keys, self.values = v
|
||||
self.offset = self.keys.shape[-2]
|
||||
|
||||
|
||||
class QuantizedKVCache(_BaseCache):
|
||||
def __init__(self, group_size: int = 64, bits: int = 8):
|
||||
self.keys = None
|
||||
self.values = None
|
||||
self.offset = 0
|
||||
self.step = 256
|
||||
self.group_size = group_size
|
||||
self.bits = bits
|
||||
|
||||
def update_and_fetch(self, keys, values):
|
||||
B, n_kv_heads, num_steps, k_head_dim = keys.shape
|
||||
v_head_dim = values.shape[-1]
|
||||
prev = self.offset
|
||||
|
||||
if self.keys is None or (prev + num_steps) > self.keys[0].shape[-2]:
|
||||
el_per_int = 8 * mx.uint32.size // self.bits
|
||||
new_steps = (self.step + num_steps - 1) // self.step * self.step
|
||||
shape = (B, n_kv_heads, new_steps)
|
||||
|
||||
def init_quant(dim):
|
||||
return (
|
||||
mx.zeros((*shape, dim // el_per_int), dtype=mx.uint32),
|
||||
mx.zeros((*shape, dim // self.group_size), dtype=keys.dtype),
|
||||
mx.zeros((*shape, dim // self.group_size), dtype=keys.dtype),
|
||||
)
|
||||
|
||||
def expand_quant(x):
|
||||
new_x = mx.zeros((*shape, x.shape[-1]), dtype=x.dtype)
|
||||
return mx.concatenate([x, new_x], axis=-2)
|
||||
|
||||
if self.keys is not None:
|
||||
if prev % self.step != 0:
|
||||
self.keys, self.values = tree_map(
|
||||
lambda x: x[..., :prev, :], (self.keys, self.values)
|
||||
)
|
||||
|
||||
self.keys, self.values = tree_map(
|
||||
expand_quant, (self.keys, self.values)
|
||||
)
|
||||
else:
|
||||
self.keys, self.values = init_quant(k_head_dim), init_quant(v_head_dim)
|
||||
|
||||
self.offset += num_steps
|
||||
|
||||
keys = mx.quantize(keys, group_size=self.group_size, bits=self.bits)
|
||||
values = mx.quantize(values, group_size=self.group_size, bits=self.bits)
|
||||
for i in range(len(self.keys)):
|
||||
self.keys[i][..., prev : self.offset, :] = keys[i]
|
||||
self.values[i][..., prev : self.offset, :] = values[i]
|
||||
|
||||
return tree_map(lambda x: x[..., : self.offset, :], (self.keys, self.values))
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
if self.offset == self.keys[0].shape[2]:
|
||||
return self.keys, self.values
|
||||
else:
|
||||
return tree_map(
|
||||
lambda x: x[..., : self.offset, :], (self.keys, self.values)
|
||||
)
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
self.keys, self.values = v
|
||||
|
||||
@property
|
||||
def meta_state(self):
|
||||
return tuple(map(str, (self.step, self.offset, self.group_size, self.bits)))
|
||||
|
||||
@meta_state.setter
|
||||
def meta_state(self, v):
|
||||
self.step, self.offset, self.group_size, self.bits = map(int, v)
|
||||
|
||||
def is_trimmable(self):
|
||||
return True
|
||||
|
||||
def trim(self, n):
|
||||
n = min(self.offset, n)
|
||||
self.offset -= n
|
||||
return n
|
||||
|
||||
|
||||
class KVCache(_BaseCache):
|
||||
def __init__(self):
|
||||
self.keys = None
|
||||
self.values = None
|
||||
self.offset = 0
|
||||
self.step = 256
|
||||
|
||||
def update_and_fetch(self, keys, values):
|
||||
prev = self.offset
|
||||
if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
|
||||
B, n_kv_heads, _, k_head_dim = keys.shape
|
||||
v_head_dim = values.shape[3]
|
||||
n_steps = (self.step + keys.shape[2] - 1) // self.step
|
||||
k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
|
||||
v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
|
||||
new_k = mx.zeros(k_shape, keys.dtype)
|
||||
new_v = mx.zeros(v_shape, values.dtype)
|
||||
if self.keys is not None:
|
||||
if prev % self.step != 0:
|
||||
self.keys = self.keys[..., :prev, :]
|
||||
self.values = self.values[..., :prev, :]
|
||||
self.keys = mx.concatenate([self.keys, new_k], axis=2)
|
||||
self.values = mx.concatenate([self.values, new_v], axis=2)
|
||||
else:
|
||||
self.keys, self.values = new_k, new_v
|
||||
|
||||
self.offset += keys.shape[2]
|
||||
self.keys[..., prev : self.offset, :] = keys
|
||||
self.values[..., prev : self.offset, :] = values
|
||||
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
if self.offset == self.keys.shape[2]:
|
||||
return self.keys, self.values
|
||||
else:
|
||||
return (
|
||||
self.keys[..., : self.offset, :],
|
||||
self.values[..., : self.offset, :],
|
||||
)
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
self.keys, self.values = v
|
||||
self.offset = self.keys.shape[2]
|
||||
|
||||
def is_trimmable(self):
|
||||
return True
|
||||
|
||||
def trim(self, n):
|
||||
n = min(self.offset, n)
|
||||
self.offset -= n
|
||||
return n
|
||||
|
||||
def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
|
||||
quant_cache = QuantizedKVCache(group_size=group_size, bits=bits)
|
||||
quant_cache.offset = self.offset
|
||||
if self.keys is not None:
|
||||
quant_cache.keys = mx.quantize(self.keys, group_size=group_size, bits=bits)
|
||||
quant_cache.values = mx.quantize(
|
||||
self.values, group_size=group_size, bits=bits
|
||||
)
|
||||
return quant_cache
|
||||
|
||||
|
||||
class RotatingKVCache(_BaseCache):
|
||||
|
||||
def __init__(self, max_size=None, keep=0, step=256):
|
||||
self.keep = keep
|
||||
self.keys = None
|
||||
self.values = None
|
||||
self.offset = 0
|
||||
self.max_size = max_size
|
||||
self.step = step
|
||||
self._idx = 0
|
||||
|
||||
def _trim(self, trim_size, v, append=None):
|
||||
to_cat = []
|
||||
if trim_size > 0:
|
||||
to_cat = [v[..., : self.keep, :], v[..., trim_size + self.keep :, :]]
|
||||
else:
|
||||
to_cat = [v]
|
||||
if append is not None:
|
||||
to_cat.append(append)
|
||||
return mx.concatenate(to_cat, axis=2)
|
||||
|
||||
def _temporal_order(self, v):
|
||||
"""
|
||||
Rearrange the cache into temporal order, slicing off the end if unused.
|
||||
"""
|
||||
if self._idx == v.shape[2]:
|
||||
return v
|
||||
elif self._idx < self.offset:
|
||||
return mx.concatenate(
|
||||
[
|
||||
v[..., : self.keep, :],
|
||||
v[..., self._idx :, :],
|
||||
v[..., self.keep : self._idx, :],
|
||||
],
|
||||
axis=2,
|
||||
)
|
||||
else:
|
||||
return v[..., : self._idx, :]
|
||||
|
||||
def _update_concat(self, keys, values):
|
||||
if self.keys is None:
|
||||
self.keys = keys
|
||||
self.values = values
|
||||
else:
|
||||
# Put the keys/values in temporal order to
|
||||
# preserve context
|
||||
self.keys = self._temporal_order(self.keys)
|
||||
self.values = self._temporal_order(self.values)
|
||||
|
||||
# The largest size is self.max_size + S to ensure
|
||||
# every token gets at least self.max_size context
|
||||
trim_size = self._idx - self.max_size
|
||||
self.keys = self._trim(trim_size, self.keys, keys)
|
||||
self.values = self._trim(trim_size, self.values, values)
|
||||
self.offset += keys.shape[2]
|
||||
self._idx = self.keys.shape[2]
|
||||
return self.keys, self.values
|
||||
|
||||
def _update_in_place(self, keys, values):
|
||||
# May not have hit the max size yet, so potentially
|
||||
# keep growing the cache
|
||||
B, n_kv_heads, S, k_head_dim = keys.shape
|
||||
prev = self.offset
|
||||
if self.keys is None or (
|
||||
prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size
|
||||
):
|
||||
v_head_dim = values.shape[3]
|
||||
new_size = min(self.step, self.max_size - prev)
|
||||
k_shape = (B, n_kv_heads, new_size, k_head_dim)
|
||||
v_shape = (B, n_kv_heads, new_size, v_head_dim)
|
||||
new_k = mx.zeros(k_shape, keys.dtype)
|
||||
new_v = mx.zeros(v_shape, values.dtype)
|
||||
if self.keys is not None:
|
||||
self.keys = mx.concatenate([self.keys, new_k], axis=2)
|
||||
self.values = mx.concatenate([self.values, new_v], axis=2)
|
||||
else:
|
||||
self.keys, self.values = new_k, new_v
|
||||
self._idx = prev
|
||||
|
||||
# Trim if needed
|
||||
trim_size = self.keys.shape[2] - self.max_size
|
||||
if trim_size > 0:
|
||||
self.keys = self._trim(trim_size, self.keys)
|
||||
self.values = self._trim(trim_size, self.values)
|
||||
self._idx = self.max_size
|
||||
|
||||
# Rotate
|
||||
if self._idx == self.max_size:
|
||||
self._idx = self.keep
|
||||
|
||||
# Assign
|
||||
self.keys[..., self._idx : self._idx + S, :] = keys
|
||||
self.values[..., self._idx : self._idx + S, :] = values
|
||||
self.offset += S
|
||||
self._idx += S
|
||||
|
||||
# If the buffer is not full, slice off the end
|
||||
if self.offset < self.max_size:
|
||||
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
|
||||
return self.keys, self.values
|
||||
|
||||
def update_and_fetch(self, keys, values):
|
||||
if keys.shape[2] == 1:
|
||||
return self._update_in_place(keys, values)
|
||||
return self._update_concat(keys, values)
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
if self.offset < self.keys.shape[2]:
|
||||
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
|
||||
else:
|
||||
return self.keys, self.values
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
self.keys, self.values = v
|
||||
|
||||
@property
|
||||
def meta_state(self):
|
||||
return tuple(
|
||||
map(str, (self.keep, self.max_size, self.step, self.offset, self._idx))
|
||||
)
|
||||
|
||||
@meta_state.setter
|
||||
def meta_state(self, v):
|
||||
self.keep, self.max_size, self.step, self.offset, self._idx = map(
|
||||
int,
|
||||
v,
|
||||
)
|
||||
|
||||
def is_trimmable(self):
|
||||
return self.offset < self.max_size
|
||||
|
||||
def trim(self, n):
|
||||
n = min(self.offset, n)
|
||||
self.offset -= n
|
||||
self._idx -= n
|
||||
return n
|
||||
|
||||
def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
|
||||
raise NotImplementedError("RotatingKVCache Quantization NYI")
|
||||
|
||||
|
||||
class MambaCache(_BaseCache):
|
||||
def __init__(self):
|
||||
self.cache = [None, None]
|
||||
|
||||
def __setitem__(self, idx, value):
|
||||
self.cache[idx] = value
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.cache[idx]
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return self.cache
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
self.cache = v
|
||||
|
||||
|
||||
class ChunkedKVCache(KVCache):
|
||||
def __init__(self, chunk_size=None):
|
||||
super().__init__()
|
||||
self.chunk_size = chunk_size
|
||||
self.start_position = 0
|
||||
|
||||
def maybe_trim_front(self):
|
||||
# Maintain the cache below the chunk size
|
||||
if self.keys is not None and self.keys.shape[2] >= self.chunk_size:
|
||||
self.start_position += self.keys.shape[2] - self.chunk_size
|
||||
self.keys = self.keys[..., -self.chunk_size :, :]
|
||||
self.values = self.values[..., -self.chunk_size :, :]
|
||||
|
||||
def update_and_fetch(self, keys, values):
|
||||
prev = self.offset - self.start_position
|
||||
if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
|
||||
B, n_kv_heads, _, k_head_dim = keys.shape
|
||||
v_head_dim = values.shape[3]
|
||||
n_steps = (self.step + keys.shape[2] - 1) // self.step
|
||||
k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
|
||||
v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
|
||||
new_k = mx.zeros(k_shape, keys.dtype)
|
||||
new_v = mx.zeros(v_shape, values.dtype)
|
||||
if self.keys is not None:
|
||||
if prev % self.step != 0:
|
||||
self.keys = self.keys[..., :prev, :]
|
||||
self.values = self.values[..., :prev, :]
|
||||
self.keys = mx.concatenate([self.keys, new_k], axis=2)
|
||||
self.values = mx.concatenate([self.values, new_v], axis=2)
|
||||
else:
|
||||
self.keys, self.values = new_k, new_v
|
||||
|
||||
self.offset += keys.shape[2]
|
||||
end = self.offset - self.start_position
|
||||
self.keys[..., prev:end, :] = keys
|
||||
self.values[..., prev:end, :] = values
|
||||
return self.keys[..., :end, :], self.values[..., :end, :]
|
||||
|
||||
def trim(self, n):
|
||||
n = min(self.offset - self.start_position, n)
|
||||
self.offset -= n
|
||||
return n
|
||||
|
||||
@property
|
||||
def meta_state(self):
|
||||
return tuple(map(str, (self.chunk_size, self.start_position)))
|
||||
|
||||
@meta_state.setter
|
||||
def meta_state(self, v):
|
||||
self.chunk_size, self.start_position = map(int, v)
|
||||
|
||||
|
||||
class CacheList(KVCache):
|
||||
def __init__(self, *caches):
|
||||
self.caches = caches
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.caches[idx]
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return [s for c in self.caches for s in c.state]
|
||||
|
||||
@state.setter
|
||||
def state(self, v):
|
||||
state_lens = [len(c.state) for c in self.caches]
|
||||
start = 0
|
||||
for c in self.caches:
|
||||
l = len(c.state)
|
||||
c.state = v[start : start + l]
|
||||
start += l
|
||||
+19
-13
@@ -1,12 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -69,7 +67,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -93,8 +91,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
@@ -129,7 +127,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
h = self.input_layernorm(x)
|
||||
attn_h = self.self_attn(h, mask, cache)
|
||||
@@ -155,13 +153,14 @@ class CohereModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -182,10 +181,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
out = out * self.model.args.logit_scale
|
||||
return out
|
||||
@@ -193,3 +191,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,227 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import KVCache, RotatingKVCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int = 4096
|
||||
head_dim: int = 128
|
||||
num_hidden_layers: int = 32
|
||||
intermediate_size: int = 14336
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 8
|
||||
rope_theta: float = 50000.0
|
||||
vocab_size: int = 256000
|
||||
layer_norm_eps: float = 1e-05
|
||||
logit_scale: float = 0.0625
|
||||
attention_bias: bool = False
|
||||
layer_norm_bias: bool = False
|
||||
sliding_window: int = 4096
|
||||
sliding_window_pattern: int = 4
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
self.head_dim = head_dim = args.head_dim
|
||||
if (head_dim * n_heads) != dim:
|
||||
raise ValueError(
|
||||
f"hidden_size must be divisible by num_heads (got `hidden_size`: {dim}"
|
||||
f" and `num_heads`: {n_heads})."
|
||||
)
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
attetion_bias = args.attention_bias
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attetion_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attetion_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attetion_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attetion_bias)
|
||||
|
||||
self.rope = nn.RoPE(head_dim, traditional=True, base=args.rope_theta)
|
||||
|
||||
self.use_sliding_window = (layer_idx + 1) % args.sliding_window_pattern != 0
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
# Apply RoPE only if sliding window is enabled
|
||||
if self.use_sliding_window:
|
||||
if cache is None:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
else:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
if self.use_sliding_window and isinstance(mask, mx.array):
|
||||
key_len = keys.shape[-2]
|
||||
if mask.shape[-1] != key_len:
|
||||
mask = mask[..., -key_len:]
|
||||
|
||||
# TODO: maybe remove cast once fused mask is supported since attention
|
||||
# may be in higher precision
|
||||
sdpa_type = mx.float32 if queries.dtype == mx.float16 else queries.dtype
|
||||
output = scaled_dot_product_attention(
|
||||
queries.astype(sdpa_type),
|
||||
keys,
|
||||
values,
|
||||
cache=cache,
|
||||
scale=self.scale,
|
||||
mask=mask,
|
||||
).astype(queries.dtype)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
self.n_heads = args.num_attention_heads
|
||||
|
||||
self.self_attn = Attention(args, layer_idx)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
self.input_layernorm = nn.LayerNorm(
|
||||
args.hidden_size, eps=args.layer_norm_eps, bias=args.layer_norm_bias
|
||||
)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
|
||||
h = self.input_layernorm(x)
|
||||
attn_h = self.self_attn(h, mask, cache)
|
||||
ff_h = self.mlp(h)
|
||||
|
||||
return attn_h + ff_h + x
|
||||
|
||||
|
||||
class CohereModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args, layer_idx=i)
|
||||
for i in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.LayerNorm(
|
||||
args.hidden_size, eps=args.layer_norm_eps, bias=args.layer_norm_bias
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
if mask is None:
|
||||
j = self.args.sliding_window_pattern
|
||||
full_mask = create_attention_mask(h, cache[j - 1 : j])
|
||||
sliding_window_mask = create_attention_mask(h, cache)
|
||||
|
||||
for i, (layer, c) in enumerate(zip(self.layers, cache)):
|
||||
is_global = (
|
||||
i % self.args.sliding_window_pattern
|
||||
== self.args.sliding_window_pattern - 1
|
||||
)
|
||||
|
||||
local_mask = mask
|
||||
if mask is None and is_global:
|
||||
local_mask = full_mask
|
||||
elif mask is None:
|
||||
local_mask = sliding_window_mask
|
||||
|
||||
h = layer(h, local_mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.model_type = args.model_type
|
||||
self.model = CohereModel(args)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
out = out * self.model.args.logit_scale
|
||||
return out
|
||||
|
||||
def make_cache(self):
|
||||
caches = []
|
||||
for i in range(self.args.num_hidden_layers):
|
||||
if (
|
||||
i % self.args.sliding_window_pattern
|
||||
== self.args.sliding_window_pattern - 1
|
||||
):
|
||||
caches.append(KVCache())
|
||||
else:
|
||||
caches.append(
|
||||
RotatingKVCache(max_size=self.args.sliding_window, keep=0)
|
||||
)
|
||||
return caches
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+23
-15
@@ -1,13 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -49,7 +47,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
|
||||
qkv = self.Wqkv(x)
|
||||
@@ -74,8 +72,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.out_proj(output)
|
||||
@@ -92,7 +90,7 @@ class NormAttnNorm(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
h = self.attn(self.norm_1(x), mask=mask, cache=cache)
|
||||
x = h + x
|
||||
@@ -105,9 +103,10 @@ class MLP(nn.Module):
|
||||
self.v1 = nn.Linear(d_model, ffn_dim, bias=False)
|
||||
self.w1 = nn.Linear(d_model, ffn_dim, bias=False)
|
||||
self.w2 = nn.Linear(ffn_dim, d_model, bias=False)
|
||||
self.act_fn = nn.silu
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
current_hidden_states = nn.silu(self.w1(x)) * self.v1(x)
|
||||
current_hidden_states = self.act_fn(self.w1(x)) * self.v1(x)
|
||||
current_hidden_states = self.w2(current_hidden_states)
|
||||
return current_hidden_states
|
||||
|
||||
@@ -178,7 +177,7 @@ class DecoderLayer(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r, h = self.norm_attn_norm(x, mask, cache)
|
||||
out = self.ffn(h) + r
|
||||
@@ -196,13 +195,15 @@ class DBRX(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.wte(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
T = h.shape[1]
|
||||
if T > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(T)
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.blocks)
|
||||
@@ -224,10 +225,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.transformer(inputs, mask, cache)
|
||||
out = self.transformer(inputs, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
@property
|
||||
@@ -251,3 +251,11 @@ class Model(nn.Module):
|
||||
experts = [(s, sv.T) for s, sv in experts]
|
||||
new_weights.update(experts)
|
||||
return new_weights
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.d_model // self.args.n_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.attn_config["kv_n_heads"]
|
||||
|
||||
@@ -1,260 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str = "deepseek"
|
||||
vocab_size: int = 102400
|
||||
hidden_size: int = 4096
|
||||
intermediate_size: int = 11008
|
||||
moe_intermediate_size: int = 1407
|
||||
num_hidden_layers: int = 30
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 32
|
||||
n_shared_experts: Optional[int] = None
|
||||
n_routed_experts: Optional[int] = None
|
||||
num_experts_per_tok: Optional[int] = None
|
||||
moe_layer_freq: int = 1
|
||||
first_k_dense_replace: int = 0
|
||||
max_position_embeddings: int = 2048
|
||||
rms_norm_eps: float = 1e-6
|
||||
rope_theta: float = 10000.0
|
||||
rope_scaling: Optional[Dict] = None
|
||||
attention_bias: bool = False
|
||||
|
||||
|
||||
class DeepseekAttention(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.num_kv_heads = config.num_key_value_heads
|
||||
self.head_dim = config.hidden_size // config.num_attention_heads
|
||||
self.scale = self.head_dim**-0.5
|
||||
|
||||
attention_bias = getattr(config, "attention_bias", False)
|
||||
|
||||
self.q_proj = nn.Linear(
|
||||
self.hidden_size,
|
||||
config.num_attention_heads * self.head_dim,
|
||||
bias=attention_bias,
|
||||
)
|
||||
self.k_proj = nn.Linear(
|
||||
self.hidden_size,
|
||||
config.num_key_value_heads * self.head_dim,
|
||||
bias=attention_bias,
|
||||
)
|
||||
self.v_proj = nn.Linear(
|
||||
self.hidden_size,
|
||||
config.num_key_value_heads * self.head_dim,
|
||||
bias=attention_bias,
|
||||
)
|
||||
self.o_proj = nn.Linear(
|
||||
self.hidden_size,
|
||||
config.num_attention_heads * self.head_dim,
|
||||
bias=attention_bias,
|
||||
)
|
||||
|
||||
rope_scale = 1.0
|
||||
if config.rope_scaling and config.rope_scaling["type"] == "linear":
|
||||
assert isinstance(config.rope_scaling["factor"], float)
|
||||
rope_scale = 1 / config.rope_scaling["factor"]
|
||||
self.rope = nn.RoPE(
|
||||
self.head_dim,
|
||||
base=config.rope_theta,
|
||||
scale=rope_scale,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, _ = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
keys = keys.reshape(B, L, self.num_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.num_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class DeepseekMLP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: ModelArgs,
|
||||
hidden_size: Optional[int] = None,
|
||||
intermediate_size: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = hidden_size or config.hidden_size
|
||||
self.intermediate_size = intermediate_size or config.intermediate_size
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
|
||||
|
||||
def __call__(self, x):
|
||||
gates = x @ self.weight.T
|
||||
scores = mx.softmax(gates, axis=-1, precise=True)
|
||||
k = self.top_k
|
||||
inds = mx.stop_gradient(mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k])
|
||||
scores = mx.take_along_axis(scores, inds, axis=-1)
|
||||
return inds, scores
|
||||
|
||||
|
||||
class DeepseekMoE(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.switch_mlp = SwitchGLU(
|
||||
config.hidden_size, config.moe_intermediate_size, config.n_routed_experts
|
||||
)
|
||||
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts is not None:
|
||||
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
||||
self.shared_experts = DeepseekMLP(
|
||||
config=config, intermediate_size=intermediate_size
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
inds, scores = self.gate(x)
|
||||
y = self.switch_mlp(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2)
|
||||
if self.config.n_shared_experts is not None:
|
||||
y = y + self.shared_experts(x)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class DeepseekDecoderLayer(nn.Module):
|
||||
def __init__(self, config: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.self_attn = DeepseekAttention(config)
|
||||
self.mlp = (
|
||||
DeepseekMoE(config)
|
||||
if (
|
||||
config.n_routed_experts is not None
|
||||
and layer_idx >= config.first_k_dense_replace
|
||||
and layer_idx % config.moe_layer_freq == 0
|
||||
)
|
||||
else DeepseekMLP(config)
|
||||
)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class DeepseekModel(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = [
|
||||
DeepseekDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
h = self.embed_tokens(x)
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = config
|
||||
self.model_type = config.model_type
|
||||
self.model = DeepseekModel(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, cache, mask)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for m in ["gate_proj", "down_proj", "up_proj"]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
|
||||
for e in range(self.args.n_routed_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,458 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str = "deepseek_v2"
|
||||
vocab_size: int = 102400
|
||||
hidden_size: int = 4096
|
||||
intermediate_size: int = 11008
|
||||
moe_intermediate_size: int = 1407
|
||||
num_hidden_layers: int = 30
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 32
|
||||
n_shared_experts: Optional[int] = None
|
||||
n_routed_experts: Optional[int] = None
|
||||
routed_scaling_factor: float = 1.0
|
||||
kv_lora_rank: int = 512
|
||||
q_lora_rank: int = 1536
|
||||
qk_rope_head_dim: int = 64
|
||||
v_head_dim: int = 128
|
||||
qk_nope_head_dim: int = 128
|
||||
topk_method: str = "gready"
|
||||
n_group: Optional[int] = None
|
||||
topk_group: Optional[int] = None
|
||||
num_experts_per_tok: Optional[int] = None
|
||||
moe_layer_freq: int = 1
|
||||
first_k_dense_replace: int = 0
|
||||
max_position_embeddings: int = 2048
|
||||
rms_norm_eps: float = 1e-6
|
||||
rope_theta: float = 10000.0
|
||||
rope_scaling: Dict = None
|
||||
attention_bias: bool = False
|
||||
|
||||
|
||||
def yarn_find_correction_dim(
|
||||
num_rotations, dim, base=10000, max_position_embeddings=2048
|
||||
):
|
||||
return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
|
||||
2 * math.log(base)
|
||||
)
|
||||
|
||||
|
||||
def yarn_find_correction_range(
|
||||
low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
|
||||
):
|
||||
low = math.floor(
|
||||
yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
|
||||
)
|
||||
high = math.ceil(
|
||||
yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
|
||||
)
|
||||
return max(low, 0), min(high, dim - 1)
|
||||
|
||||
|
||||
def yarn_get_mscale(scale=1, mscale=1):
|
||||
if scale <= 1:
|
||||
return 1.0
|
||||
return 0.1 * mscale * math.log(scale) + 1.0
|
||||
|
||||
|
||||
def yarn_linear_ramp_mask(min_val, max_val, dim):
|
||||
if min_val == max_val:
|
||||
max_val += 0.001 # Prevent singularity
|
||||
|
||||
linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
|
||||
return mx.clip(linear_func, 0, 1)
|
||||
|
||||
|
||||
class DeepseekV2YarnRotaryEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
max_position_embeddings=2048,
|
||||
base=10000,
|
||||
scaling_factor=1.0,
|
||||
original_max_position_embeddings=4096,
|
||||
beta_fast=32,
|
||||
beta_slow=1,
|
||||
mscale=1,
|
||||
mscale_all_dim=0,
|
||||
):
|
||||
super().__init__()
|
||||
self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
|
||||
scaling_factor, mscale_all_dim
|
||||
)
|
||||
freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
|
||||
freq_inter = scaling_factor * base ** (
|
||||
mx.arange(0, dim, 2, dtype=mx.float32) / dim
|
||||
)
|
||||
low, high = yarn_find_correction_range(
|
||||
beta_fast,
|
||||
beta_slow,
|
||||
dim,
|
||||
base,
|
||||
original_max_position_embeddings,
|
||||
)
|
||||
freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
|
||||
self._freqs = (freq_inter * freq_extra) / (
|
||||
freq_inter * freq_mask + freq_extra * (1 - freq_mask)
|
||||
)
|
||||
|
||||
def __call__(self, x, offset=0):
|
||||
if self.mscale != 1.0:
|
||||
x = self.mscale * x
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
x.shape[-1],
|
||||
traditional=True,
|
||||
base=None,
|
||||
scale=1.0,
|
||||
offset=offset,
|
||||
freqs=self._freqs,
|
||||
)
|
||||
|
||||
|
||||
class DeepseekV2Attention(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
self.rope_theta = config.rope_theta
|
||||
self.q_lora_rank = config.q_lora_rank
|
||||
self.qk_rope_head_dim = config.qk_rope_head_dim
|
||||
self.kv_lora_rank = config.kv_lora_rank
|
||||
self.v_head_dim = config.v_head_dim
|
||||
self.qk_nope_head_dim = config.qk_nope_head_dim
|
||||
self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
|
||||
|
||||
self.scale = self.q_head_dim**-0.5
|
||||
|
||||
if self.q_lora_rank is None:
|
||||
self.q_proj = nn.Linear(
|
||||
self.hidden_size, self.num_heads * self.q_head_dim, bias=False
|
||||
)
|
||||
else:
|
||||
self.q_a_proj = nn.Linear(
|
||||
self.hidden_size, self.q_lora_rank, bias=config.attention_bias
|
||||
)
|
||||
self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
|
||||
self.q_b_proj = nn.Linear(
|
||||
self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
|
||||
)
|
||||
|
||||
self.kv_a_proj_with_mqa = nn.Linear(
|
||||
self.hidden_size,
|
||||
self.kv_lora_rank + self.qk_rope_head_dim,
|
||||
bias=config.attention_bias,
|
||||
)
|
||||
self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
|
||||
self.kv_b_proj = nn.Linear(
|
||||
self.kv_lora_rank,
|
||||
self.num_heads
|
||||
* (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.o_proj = nn.Linear(
|
||||
self.num_heads * self.v_head_dim,
|
||||
self.hidden_size,
|
||||
bias=config.attention_bias,
|
||||
)
|
||||
|
||||
mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
|
||||
scaling_factor = self.config.rope_scaling["factor"]
|
||||
if mscale_all_dim:
|
||||
mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
|
||||
self.scale = self.scale * mscale * mscale
|
||||
|
||||
rope_kwargs = {
|
||||
key: self.config.rope_scaling[key]
|
||||
for key in [
|
||||
"original_max_position_embeddings",
|
||||
"beta_fast",
|
||||
"beta_slow",
|
||||
"mscale",
|
||||
"mscale_all_dim",
|
||||
]
|
||||
if key in self.config.rope_scaling
|
||||
}
|
||||
self.rope = DeepseekV2YarnRotaryEmbedding(
|
||||
dim=self.qk_rope_head_dim,
|
||||
max_position_embeddings=self.max_position_embeddings,
|
||||
scaling_factor=scaling_factor,
|
||||
base=self.rope_theta,
|
||||
**rope_kwargs,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
if self.q_lora_rank is None:
|
||||
q = self.q_proj(x)
|
||||
else:
|
||||
q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
|
||||
|
||||
q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
|
||||
q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
|
||||
compressed_kv = self.kv_a_proj_with_mqa(x)
|
||||
compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
|
||||
k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
|
||||
kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
|
||||
kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
|
||||
|
||||
if cache is not None:
|
||||
q_pe = self.rope(q_pe, cache.offset)
|
||||
k_pe = self.rope(k_pe, cache.offset)
|
||||
k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
|
||||
keys, values = cache.update_and_fetch(
|
||||
mx.concatenate([k_nope, k_pe], axis=-1), values
|
||||
)
|
||||
else:
|
||||
q_pe = self.rope(q_pe)
|
||||
k_pe = self.rope(k_pe)
|
||||
k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
|
||||
keys = mx.concatenate([k_nope, k_pe], axis=-1)
|
||||
|
||||
queries = mx.concatenate([q_nope, q_pe], axis=-1)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class DeepseekV2MLP(nn.Module):
|
||||
def __init__(
|
||||
self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
|
||||
):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
|
||||
self.intermediate_size = (
|
||||
config.intermediate_size if intermediate_size is None else intermediate_size
|
||||
)
|
||||
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
|
||||
def __call__(self, x):
|
||||
down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
return down_proj
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
self.routed_scaling_factor = config.routed_scaling_factor
|
||||
self.topk_method = config.topk_method
|
||||
self.n_group = config.n_group
|
||||
self.topk_group = config.topk_group
|
||||
self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
|
||||
|
||||
def __call__(self, x):
|
||||
gates = x @ self.weight.T
|
||||
|
||||
scores = mx.softmax(gates, axis=-1, precise=True)
|
||||
|
||||
if self.topk_method == "group_limited_greedy":
|
||||
bsz, seq_len = x.shape[:2]
|
||||
scores = scores.reshape(bsz, seq_len, self.n_group, -1)
|
||||
group_scores = scores.max(axis=-1, keepdims=True)
|
||||
k = self.n_group - self.topk_group
|
||||
group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
|
||||
scores = mx.put_along_axis(
|
||||
scores, group_idx, mx.array(0.0, scores.dtype), axis=-2
|
||||
)
|
||||
scores = scores.reshape(bsz, seq_len, -1)
|
||||
|
||||
k = self.top_k
|
||||
inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
|
||||
scores = mx.take_along_axis(scores, inds, axis=-1)
|
||||
scores = scores * self.routed_scaling_factor
|
||||
|
||||
return inds, scores
|
||||
|
||||
|
||||
class DeepseekV2MoE(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.num_experts_per_tok = config.num_experts_per_tok
|
||||
self.switch_mlp = SwitchGLU(
|
||||
config.hidden_size, config.moe_intermediate_size, config.n_routed_experts
|
||||
)
|
||||
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts is not None:
|
||||
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
||||
self.shared_experts = DeepseekV2MLP(
|
||||
config=config, intermediate_size=intermediate_size
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
inds, scores = self.gate(x)
|
||||
y = self.switch_mlp(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2)
|
||||
if self.config.n_shared_experts is not None:
|
||||
y = y + self.shared_experts(x)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class DeepseekV2DecoderLayer(nn.Module):
|
||||
def __init__(self, config: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.self_attn = DeepseekV2Attention(config)
|
||||
self.mlp = (
|
||||
DeepseekV2MoE(config)
|
||||
if (
|
||||
config.n_routed_experts is not None
|
||||
and layer_idx >= config.first_k_dense_replace
|
||||
and layer_idx % config.moe_layer_freq == 0
|
||||
)
|
||||
else DeepseekV2MLP(config)
|
||||
)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class DeepseekV2Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.vocab_size = config.vocab_size
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = [
|
||||
DeepseekV2DecoderLayer(config, idx)
|
||||
for idx in range(config.num_hidden_layers)
|
||||
]
|
||||
self.start_idx = 0
|
||||
self.end_idx = len(self.layers)
|
||||
self.num_layers = self.end_idx
|
||||
|
||||
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
self.pipeline_rank = 0
|
||||
self.pipeline_size = 1
|
||||
|
||||
def pipeline(self, group):
|
||||
# Split layers in reverse so rank=0 gets the last layers and
|
||||
# rank=pipeline_size-1 gets the first
|
||||
self.pipeline_rank = group.rank()
|
||||
self.pipeline_size = group.size()
|
||||
layers_per_rank = len(self.layers) // self.pipeline_size
|
||||
extra = len(self.layers) - layers_per_rank * self.pipeline_size
|
||||
if self.pipeline_rank < extra:
|
||||
layers_per_rank += 1
|
||||
|
||||
self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
|
||||
self.end_idx = self.start_idx + layers_per_rank
|
||||
self.num_layers = layers_per_rank
|
||||
self.layers = self.layers[: self.end_idx]
|
||||
self.layers[: self.start_idx] = [None] * self.start_idx
|
||||
self.num_layers = len(self.layers) - self.start_idx
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
h = self.embed_tokens(x)
|
||||
|
||||
pipeline_rank = self.pipeline_rank
|
||||
pipeline_size = self.pipeline_size
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * self.num_layers
|
||||
|
||||
# Receive from the previous process in the pipeline
|
||||
if pipeline_rank < pipeline_size - 1:
|
||||
h = mx.distributed.recv_like(h, (pipeline_rank + 1))
|
||||
|
||||
for i in range(self.num_layers):
|
||||
h = self.layers[self.start_idx + i](h, mask, cache[i])
|
||||
|
||||
# Send to the next process in the pipeline
|
||||
if pipeline_rank != 0:
|
||||
h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
|
||||
|
||||
# Broadcast h while keeping it in the graph
|
||||
h = mx.distributed.all_gather(h)[: h.shape[0]]
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = config
|
||||
self.model_type = config.model_type
|
||||
self.model = DeepseekV2Model(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, cache, mask)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
|
||||
for e in range(self.args.n_routed_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers[self.model.start_idx : self.model.end_idx]
|
||||
@@ -1,545 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str = "deepseek_v3"
|
||||
vocab_size: int = 102400
|
||||
hidden_size: int = 4096
|
||||
intermediate_size: int = 11008
|
||||
moe_intermediate_size: int = 1407
|
||||
num_hidden_layers: int = 30
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 32
|
||||
n_shared_experts: Optional[int] = None
|
||||
n_routed_experts: Optional[int] = None
|
||||
routed_scaling_factor: float = 1.0
|
||||
kv_lora_rank: int = 512
|
||||
q_lora_rank: int = 1536
|
||||
qk_rope_head_dim: int = 64
|
||||
v_head_dim: int = 128
|
||||
qk_nope_head_dim: int = 128
|
||||
topk_method: str = "noaux_tc"
|
||||
scoring_func: str = "sigmoid"
|
||||
norm_topk_prob: bool = True
|
||||
n_group: Optional[int] = None
|
||||
topk_group: Optional[int] = None
|
||||
num_experts_per_tok: Optional[int] = None
|
||||
moe_layer_freq: int = 1
|
||||
first_k_dense_replace: int = 0
|
||||
max_position_embeddings: int = 2048
|
||||
rms_norm_eps: float = 1e-6
|
||||
rope_theta: float = 10000.0
|
||||
rope_scaling: Dict = None
|
||||
attention_bias: bool = False
|
||||
|
||||
|
||||
def yarn_find_correction_dim(
|
||||
num_rotations, dim, base=10000, max_position_embeddings=2048
|
||||
):
|
||||
return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
|
||||
2 * math.log(base)
|
||||
)
|
||||
|
||||
|
||||
def yarn_find_correction_range(
|
||||
low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
|
||||
):
|
||||
low = math.floor(
|
||||
yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
|
||||
)
|
||||
high = math.ceil(
|
||||
yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
|
||||
)
|
||||
return max(low, 0), min(high, dim - 1)
|
||||
|
||||
|
||||
def yarn_get_mscale(scale=1, mscale=1):
|
||||
if scale <= 1:
|
||||
return 1.0
|
||||
return 0.1 * mscale * math.log(scale) + 1.0
|
||||
|
||||
|
||||
def yarn_linear_ramp_mask(min_val, max_val, dim):
|
||||
if min_val == max_val:
|
||||
max_val += 0.001 # Prevent singularity
|
||||
|
||||
linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
|
||||
return mx.clip(linear_func, 0, 1)
|
||||
|
||||
|
||||
class DeepseekV3YarnRotaryEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
max_position_embeddings=2048,
|
||||
base=10000,
|
||||
scaling_factor=1.0,
|
||||
original_max_position_embeddings=4096,
|
||||
beta_fast=32,
|
||||
beta_slow=1,
|
||||
mscale=1,
|
||||
mscale_all_dim=0,
|
||||
):
|
||||
super().__init__()
|
||||
self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
|
||||
scaling_factor, mscale_all_dim
|
||||
)
|
||||
freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
|
||||
freq_inter = scaling_factor * freq_extra
|
||||
low, high = yarn_find_correction_range(
|
||||
beta_fast,
|
||||
beta_slow,
|
||||
dim,
|
||||
base,
|
||||
original_max_position_embeddings,
|
||||
)
|
||||
freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
|
||||
self._freqs = (freq_inter * freq_extra) / (
|
||||
freq_inter * freq_mask + freq_extra * (1 - freq_mask)
|
||||
)
|
||||
|
||||
def __call__(self, x, offset=0):
|
||||
if self.mscale != 1.0:
|
||||
x = self.mscale * x
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
x.shape[-1],
|
||||
traditional=True,
|
||||
base=None,
|
||||
scale=1.0,
|
||||
offset=offset,
|
||||
freqs=self._freqs,
|
||||
)
|
||||
|
||||
|
||||
# A clipped silu to prevent fp16 from overflowing
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def clipped_silu(x):
|
||||
return mx.clip(x * mx.sigmoid(x), a_min=-100, a_max=100)
|
||||
|
||||
|
||||
class ClippedSilu(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def __call__(self, x):
|
||||
return clipped_silu(x)
|
||||
|
||||
|
||||
class DeepseekV3Attention(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
self.rope_theta = config.rope_theta
|
||||
self.q_lora_rank = config.q_lora_rank
|
||||
self.qk_rope_head_dim = config.qk_rope_head_dim
|
||||
self.kv_lora_rank = config.kv_lora_rank
|
||||
self.v_head_dim = config.v_head_dim
|
||||
self.qk_nope_head_dim = config.qk_nope_head_dim
|
||||
self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
|
||||
|
||||
self.scale = self.q_head_dim**-0.5
|
||||
|
||||
if self.q_lora_rank is None:
|
||||
self.q_proj = nn.Linear(
|
||||
self.hidden_size, self.num_heads * self.q_head_dim, bias=False
|
||||
)
|
||||
else:
|
||||
self.q_a_proj = nn.Linear(
|
||||
self.hidden_size, self.q_lora_rank, bias=config.attention_bias
|
||||
)
|
||||
self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
|
||||
self.q_b_proj = nn.Linear(
|
||||
self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
|
||||
)
|
||||
|
||||
self.kv_a_proj_with_mqa = nn.Linear(
|
||||
self.hidden_size,
|
||||
self.kv_lora_rank + self.qk_rope_head_dim,
|
||||
bias=config.attention_bias,
|
||||
)
|
||||
self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
|
||||
self.kv_b_proj = nn.Linear(
|
||||
self.kv_lora_rank,
|
||||
self.num_heads
|
||||
* (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.o_proj = nn.Linear(
|
||||
self.num_heads * self.v_head_dim,
|
||||
self.hidden_size,
|
||||
bias=config.attention_bias,
|
||||
)
|
||||
|
||||
if self.config.rope_scaling is not None:
|
||||
mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
|
||||
scaling_factor = self.config.rope_scaling["factor"]
|
||||
if mscale_all_dim:
|
||||
mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
|
||||
self.scale = self.scale * mscale * mscale
|
||||
|
||||
rope_kwargs = {
|
||||
key: self.config.rope_scaling[key]
|
||||
for key in [
|
||||
"original_max_position_embeddings",
|
||||
"beta_fast",
|
||||
"beta_slow",
|
||||
"mscale",
|
||||
"mscale_all_dim",
|
||||
]
|
||||
if key in self.config.rope_scaling
|
||||
}
|
||||
self.rope = DeepseekV3YarnRotaryEmbedding(
|
||||
dim=self.qk_rope_head_dim,
|
||||
max_position_embeddings=self.max_position_embeddings,
|
||||
scaling_factor=scaling_factor,
|
||||
base=self.rope_theta,
|
||||
**rope_kwargs,
|
||||
)
|
||||
else:
|
||||
self.rope = nn.RoPE(
|
||||
dims=self.qk_rope_head_dim,
|
||||
base=self.rope_theta,
|
||||
traditional=True,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
if self.q_lora_rank is None:
|
||||
q = self.q_proj(x)
|
||||
else:
|
||||
q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
|
||||
|
||||
q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
|
||||
q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
|
||||
compressed_kv = self.kv_a_proj_with_mqa(x)
|
||||
compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
|
||||
k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
|
||||
kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
|
||||
kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
|
||||
|
||||
if cache is not None:
|
||||
q_pe = self.rope(q_pe, cache.offset)
|
||||
k_pe = self.rope(k_pe, cache.offset)
|
||||
k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
|
||||
keys, values = cache.update_and_fetch(
|
||||
mx.concatenate([k_nope, k_pe], axis=-1), values
|
||||
)
|
||||
else:
|
||||
q_pe = self.rope(q_pe)
|
||||
k_pe = self.rope(k_pe)
|
||||
k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
|
||||
keys = mx.concatenate([k_nope, k_pe], axis=-1)
|
||||
|
||||
queries = mx.concatenate([q_nope, q_pe], axis=-1)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class DeepseekV3MLP(nn.Module):
|
||||
def __init__(
|
||||
self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
|
||||
):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
|
||||
self.intermediate_size = (
|
||||
config.intermediate_size if intermediate_size is None else intermediate_size
|
||||
)
|
||||
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
|
||||
def __call__(self, x):
|
||||
down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
return down_proj
|
||||
|
||||
|
||||
@mx.compile
|
||||
def group_expert_select(
|
||||
gates,
|
||||
e_score_correction_bias,
|
||||
top_k,
|
||||
n_group,
|
||||
topk_group,
|
||||
routed_scaling_factor,
|
||||
norm_topk_prob,
|
||||
):
|
||||
|
||||
k = top_k
|
||||
scores = mx.sigmoid(gates.astype(mx.float32))
|
||||
orig_scores = scores
|
||||
scores = scores + e_score_correction_bias
|
||||
scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
|
||||
group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
|
||||
k = n_group - topk_group
|
||||
group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
|
||||
scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
|
||||
scores = mx.flatten(scores, -2, -1)
|
||||
|
||||
k = top_k
|
||||
inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
|
||||
scores = mx.take_along_axis(orig_scores, inds, axis=-1)
|
||||
if top_k > 1 and norm_topk_prob:
|
||||
denominator = scores.sum(axis=-1, keepdims=True)
|
||||
scores = scores / denominator
|
||||
scores = scores * routed_scaling_factor
|
||||
|
||||
return inds, scores
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.norm_topk_prob = config.norm_topk_prob
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
self.routed_scaling_factor = config.routed_scaling_factor
|
||||
self.n_group = config.n_group
|
||||
self.topk_group = config.topk_group
|
||||
self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
|
||||
self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
|
||||
assert config.topk_method == "noaux_tc", "Unsupported topk method."
|
||||
|
||||
def __call__(self, x):
|
||||
return group_expert_select(
|
||||
x @ self.weight.T,
|
||||
self.e_score_correction_bias,
|
||||
self.top_k,
|
||||
self.n_group,
|
||||
self.topk_group,
|
||||
self.routed_scaling_factor,
|
||||
self.norm_topk_prob,
|
||||
)
|
||||
|
||||
|
||||
class DeepseekV3MoE(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.num_experts_per_tok = config.num_experts_per_tok
|
||||
self.switch_mlp = SwitchGLU(
|
||||
config.hidden_size,
|
||||
config.moe_intermediate_size,
|
||||
config.n_routed_experts,
|
||||
activation=ClippedSilu(),
|
||||
)
|
||||
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts is not None:
|
||||
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
||||
self.shared_experts = DeepseekV3MLP(
|
||||
config=config, intermediate_size=intermediate_size
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
inds, scores = self.gate(x)
|
||||
y = self.switch_mlp(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
|
||||
if self.config.n_shared_experts is not None:
|
||||
y = y + self.shared_experts(x)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class DeepseekV3DecoderLayer(nn.Module):
|
||||
def __init__(self, config: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.self_attn = DeepseekV3Attention(config)
|
||||
self.mlp = (
|
||||
DeepseekV3MoE(config)
|
||||
if (
|
||||
config.n_routed_experts is not None
|
||||
and layer_idx >= config.first_k_dense_replace
|
||||
and layer_idx % config.moe_layer_freq == 0
|
||||
)
|
||||
else DeepseekV3MLP(config)
|
||||
)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
return h + r
|
||||
|
||||
|
||||
class DeepseekV3Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.vocab_size = config.vocab_size
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = [
|
||||
DeepseekV3DecoderLayer(config, idx)
|
||||
for idx in range(config.num_hidden_layers)
|
||||
]
|
||||
self.start_idx = 0
|
||||
self.end_idx = len(self.layers)
|
||||
self.num_layers = self.end_idx
|
||||
|
||||
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.pipeline_rank = 0
|
||||
self.pipeline_size = 1
|
||||
|
||||
def pipeline(self, group):
|
||||
# Split layers in reverse so rank=0 gets the last layers and
|
||||
# rank=pipeline_size-1 gets the first
|
||||
self.pipeline_rank = group.rank()
|
||||
self.pipeline_size = group.size()
|
||||
layers_per_rank = len(self.layers) // self.pipeline_size
|
||||
extra = len(self.layers) - layers_per_rank * self.pipeline_size
|
||||
if self.pipeline_rank < extra:
|
||||
layers_per_rank += 1
|
||||
self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
|
||||
self.end_idx = self.start_idx + layers_per_rank
|
||||
self.layers = self.layers[: self.end_idx]
|
||||
self.layers[: self.start_idx] = [None] * self.start_idx
|
||||
self.num_layers = len(self.layers) - self.start_idx
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
h = self.embed_tokens(x)
|
||||
|
||||
pipeline_rank = self.pipeline_rank
|
||||
pipeline_size = self.pipeline_size
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * self.num_layers
|
||||
|
||||
# Receive from the previous process in the pipeline
|
||||
|
||||
if pipeline_rank < pipeline_size - 1:
|
||||
h = mx.distributed.recv_like(h, (pipeline_rank + 1))
|
||||
|
||||
for i in range(self.num_layers):
|
||||
h = self.layers[self.start_idx + i](h, mask, cache[i])
|
||||
|
||||
# Send to the next process in the pipeline
|
||||
if pipeline_rank != 0:
|
||||
h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
|
||||
|
||||
# Broadcast h while keeping it in the graph
|
||||
h = mx.distributed.all_gather(h)[: h.shape[0]]
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = config
|
||||
self.model_type = config.model_type
|
||||
self.model = DeepseekV3Model(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, cache, mask)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
def dequant(weight, scale_inv):
|
||||
dtype = weight.dtype
|
||||
bs = 128 # block size
|
||||
m, n = weight.shape
|
||||
pad_bottom = (-m) % bs
|
||||
pad_side = (-n) % bs
|
||||
weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
|
||||
weight = weight.reshape(
|
||||
((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
|
||||
)
|
||||
weight = (weight * scale_inv[:, None, :, None]).reshape(
|
||||
m + pad_bottom, n + pad_side
|
||||
)
|
||||
return weight[:m, :n].astype(dtype)
|
||||
|
||||
# Dequantize
|
||||
new_weights = {}
|
||||
for k, v in weights.items():
|
||||
if "weight_scale_inv" in k:
|
||||
scale_inv = v
|
||||
wk = k.replace("_scale_inv", "")
|
||||
weight = weights[wk]
|
||||
weight = dequant(weight, scale_inv)
|
||||
new_weights[wk] = weight
|
||||
elif k not in new_weights:
|
||||
new_weights[k] = v
|
||||
weights = new_weights
|
||||
|
||||
# Stack experts
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
|
||||
for e in range(self.args.n_routed_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
|
||||
|
||||
# Remove multi-token prediction layer and any unused precomputed rotary freqs
|
||||
return {
|
||||
k: v
|
||||
for k, v in weights.items()
|
||||
if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
|
||||
}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers[self.model.start_idx : self.model.end_idx]
|
||||
|
||||
@property
|
||||
def cast_predicate(self):
|
||||
def predicate(k):
|
||||
return "e_score_correction_bias" not in k
|
||||
|
||||
return predicate
|
||||
@@ -1,320 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
max_position_embeddings: Optional[int]
|
||||
num_key_value_heads: Optional[int]
|
||||
first_k_dense_replace: int
|
||||
moe_intermediate_size: int
|
||||
moe_layer_freq: int
|
||||
n_routed_experts: int
|
||||
n_shared_experts: int
|
||||
norm_topk_prob: bool
|
||||
num_experts_per_tok: int
|
||||
rope_theta: float
|
||||
routed_scaling_factor: float
|
||||
head_dim: Optional[int] = None
|
||||
scoring_func: str = ("noaux_tc",)
|
||||
n_group: Optional[int] = 1
|
||||
topk_group: Optional[int] = 1
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
|
||||
class Dots1Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
|
||||
self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
|
||||
self.rope = initialize_rope(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
traditional=False,
|
||||
scaling_config=args.rope_scaling,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
@mx.compile
|
||||
def group_expert_select(
|
||||
gates,
|
||||
e_score_correction_bias,
|
||||
top_k,
|
||||
n_group,
|
||||
topk_group,
|
||||
routed_scaling_factor,
|
||||
norm_topk_prob,
|
||||
):
|
||||
|
||||
k = top_k
|
||||
scores = mx.sigmoid(gates.astype(mx.float32))
|
||||
orig_scores = scores
|
||||
scores = scores + e_score_correction_bias
|
||||
k = n_group - topk_group
|
||||
if k != 0:
|
||||
scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
|
||||
group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
|
||||
group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
|
||||
scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
|
||||
scores = mx.flatten(scores, -2, -1)
|
||||
|
||||
k = top_k
|
||||
inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
|
||||
scores = mx.take_along_axis(orig_scores, inds, axis=-1)
|
||||
if top_k > 1 and norm_topk_prob:
|
||||
denominator = scores.sum(axis=-1, keepdims=True)
|
||||
scores = scores / denominator
|
||||
scores = scores * routed_scaling_factor
|
||||
|
||||
return inds, scores
|
||||
|
||||
|
||||
class Dots1TopkRouter(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.top_k = args.num_experts_per_tok
|
||||
self.norm_topk_prob = args.norm_topk_prob
|
||||
self.n_routed_experts = args.n_routed_experts
|
||||
self.routed_scaling_factor = args.routed_scaling_factor
|
||||
self.n_group = args.n_group
|
||||
self.topk_group = args.topk_group
|
||||
self.weight = mx.zeros((self.n_routed_experts, args.hidden_size))
|
||||
self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
|
||||
|
||||
def __call__(self, x):
|
||||
return group_expert_select(
|
||||
x @ self.weight.T,
|
||||
self.e_score_correction_bias,
|
||||
self.top_k,
|
||||
self.n_group,
|
||||
self.topk_group,
|
||||
self.routed_scaling_factor,
|
||||
self.norm_topk_prob,
|
||||
)
|
||||
|
||||
|
||||
class Dots1MLP(nn.Module):
|
||||
def __init__(
|
||||
self, args: ModelArgs, hidden_size: int = None, intermediate_size: int = None
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = args.hidden_size if hidden_size is None else hidden_size
|
||||
self.intermediate_size = (
|
||||
args.intermediate_size if intermediate_size is None else intermediate_size
|
||||
)
|
||||
|
||||
self.gate_proj = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size, bias=args.mlp_bias
|
||||
)
|
||||
self.up_proj = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size, bias=args.mlp_bias
|
||||
)
|
||||
self.down_proj = nn.Linear(
|
||||
self.intermediate_size, self.hidden_size, bias=args.mlp_bias
|
||||
)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class Dots1MoE(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_experts_per_tok = args.num_experts_per_tok
|
||||
self.n_shared_experts = args.n_shared_experts
|
||||
|
||||
self.experts = SwitchGLU(
|
||||
args.hidden_size,
|
||||
args.moe_intermediate_size,
|
||||
args.n_routed_experts,
|
||||
)
|
||||
|
||||
self.gate = Dots1TopkRouter(args)
|
||||
|
||||
self.shared_experts = Dots1MLP(
|
||||
args=args,
|
||||
intermediate_size=args.moe_intermediate_size * args.n_shared_experts,
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
inds, scores = self.gate(x)
|
||||
y = self.experts(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
|
||||
if self.n_shared_experts is not None:
|
||||
y = y + self.shared_experts(x)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class Dots1DecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.self_attn = Dots1Attention(args)
|
||||
|
||||
if layer_idx >= args.first_k_dense_replace:
|
||||
self.mlp = Dots1MoE(args)
|
||||
else:
|
||||
self.mlp = Dots1MLP(args)
|
||||
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
return h + r
|
||||
|
||||
|
||||
class Dots1Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
Dots1DecoderLayer(args, layer_idx)
|
||||
for layer_idx in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = Dots1Model(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
if l >= self.args.first_k_dense_replace:
|
||||
for n, m in [
|
||||
("w1", "gate_proj"),
|
||||
("w2", "down_proj"),
|
||||
("w3", "up_proj"),
|
||||
]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
|
||||
for e in range(self.args.n_routed_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.experts.{m}.{k}"] = mx.stack(to_join)
|
||||
|
||||
return {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,167 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
model_type: str
|
||||
max_position_embeddings: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
head_dim: Optional[int]
|
||||
num_hidden_layers: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
rope_theta: float
|
||||
use_bias: bool
|
||||
tie_word_embeddings: bool
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or dim // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.use_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.use_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
traditional=True,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim, use_bias=False):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=use_bias)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size, args.use_bias)
|
||||
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
return h + r
|
||||
|
||||
|
||||
class Ernie45Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [DecoderLayer(args) for _ in range(args.num_hidden_layers)]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = Ernie45Model(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,166 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
vocab_size: int
|
||||
rope_theta: float
|
||||
layer_norm_epsilon: float
|
||||
num_key_value_heads: int
|
||||
head_dim: Optional[int] = None
|
||||
max_position_embeddings: Optional[int] = None
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = True
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
|
||||
|
||||
class AttentionModule(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
self.head_dim = head_dim = args.head_dim or (dim // n_heads)
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
args.rope_traditional,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
q = self.q_proj(x).reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
k = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
v = self.v_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
q = self.rope(q, offset=cache.offset)
|
||||
k = self.rope(k, offset=cache.offset)
|
||||
k, v = cache.update_and_fetch(k, v)
|
||||
else:
|
||||
q = self.rope(q)
|
||||
k = self.rope(k)
|
||||
|
||||
out = scaled_dot_product_attention(
|
||||
q, k, v, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
out = out.transpose(0, 2, 1, 3).reshape(B, L, D)
|
||||
return self.out_proj(out)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.attention = AttentionModule(args)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
dim = args.hidden_size
|
||||
hidden_dim = args.intermediate_size
|
||||
self.c_fc_0 = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
|
||||
self.c_fc_1 = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
|
||||
self.c_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
return self.c_proj(nn.silu(self.c_fc_0(x)) * self.c_fc_1(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.ln_1 = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
|
||||
self.attn = Attention(args)
|
||||
self.ln_2 = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
|
||||
self.mlp = MLP(args)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
h = x + self.attn.attention(self.ln_1(x), mask, cache)
|
||||
out = h + self.mlp(self.ln_2(h))
|
||||
return out
|
||||
|
||||
|
||||
class ExaoneModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.wte = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.h = [TransformerBlock(args) for _ in range(args.num_layers)]
|
||||
self.ln_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.wte(inputs)
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.h)
|
||||
|
||||
for layer, c in zip(self.h, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.ln_f(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.transformer = ExaoneModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.transformer(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.transformer.wte.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.transformer.h
|
||||
+19
-13
@@ -1,12 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -60,7 +58,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -79,8 +77,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
@@ -113,7 +111,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -138,14 +136,15 @@ class GemmaModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
h = h * (self.args.hidden_size**0.5)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -166,13 +165,20 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.head_dim
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,208 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
head_dim: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
attn_logit_softcapping: float = 50.0
|
||||
final_logit_softcapping: float = 30.0
|
||||
query_pre_attn_scalar: float = 144.0
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dims: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.weight = mx.ones((dims,))
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, x):
|
||||
return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
self.repeats = n_heads // n_kv_heads
|
||||
self.head_dim = head_dim = args.head_dim
|
||||
|
||||
self.scale = 1.0 / (args.query_pre_attn_scalar**0.5)
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
self.attn_logit_softcapping = args.attn_logit_softcapping
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
traditional=args.rope_traditional,
|
||||
base=args.rope_theta,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
queries = queries * self.scale
|
||||
|
||||
if self.repeats > 1:
|
||||
queries = queries.reshape(
|
||||
B, self.n_kv_heads, self.repeats, L, self.head_dim
|
||||
)
|
||||
keys = mx.expand_dims(keys, 2)
|
||||
values = mx.expand_dims(values, 2)
|
||||
|
||||
scores = queries @ keys.swapaxes(-1, -2)
|
||||
scores = mx.tanh(scores / self.attn_logit_softcapping)
|
||||
scores *= self.attn_logit_softcapping
|
||||
|
||||
if mask is not None:
|
||||
if mask.dtype == mx.bool_:
|
||||
scores = mx.where(
|
||||
mask, scores, mx.array(mx.finfo(scores.dtype).min, scores.dtype)
|
||||
)
|
||||
else:
|
||||
scores = scores + mask
|
||||
scores = mx.softmax(scores, precise=True, axis=-1)
|
||||
output = scores @ values
|
||||
if self.repeats > 1:
|
||||
output = output.reshape(B, self.n_heads, L, self.head_dim)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.pre_feedforward_layernorm = RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.post_feedforward_layernorm = RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + self.post_attention_layernorm(r)
|
||||
r = self.mlp(self.pre_feedforward_layernorm(h))
|
||||
out = h + self.post_feedforward_layernorm(r)
|
||||
return out
|
||||
|
||||
|
||||
class GemmaModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
h = h * (self.args.hidden_size**0.5)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache, return_array=True)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.model_type = args.model_type
|
||||
self.final_logit_softcapping = args.final_logit_softcapping
|
||||
self.model = GemmaModel(args)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
out = mx.tanh(out / self.final_logit_softcapping)
|
||||
out = out * self.final_logit_softcapping
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,64 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from . import gemma3_text
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
text_config: dict
|
||||
vocab_size: int = 262208
|
||||
|
||||
def __post_init__(self):
|
||||
self.text_config["vocab_size"] = self.vocab_size
|
||||
self.text_config["num_attention_heads"] = self.text_config.get(
|
||||
"num_attention_heads", 8
|
||||
)
|
||||
self.text_config["num_key_value_heads"] = self.text_config.get(
|
||||
"num_key_value_heads", 4
|
||||
)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.language_model = gemma3_text.Model(
|
||||
gemma3_text.ModelArgs.from_dict(args.text_config)
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask: Optional[mx.array] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
return self.language_model(
|
||||
inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
|
||||
)
|
||||
|
||||
def sanitize(self, weights):
|
||||
weights = tree_unflatten(list(weights.items()))
|
||||
weights.pop("vision_tower", None)
|
||||
weights.pop("multi_modal_projector", None)
|
||||
lm_weights = dict(tree_flatten(weights["language_model"]))
|
||||
lm_weights = self.language_model.sanitize(lm_weights)
|
||||
weights["language_model"] = tree_unflatten(list(lm_weights.items()))
|
||||
return dict(tree_flatten(weights))
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.language_model.layers
|
||||
|
||||
def make_cache(self):
|
||||
return self.language_model.make_cache()
|
||||
@@ -1,252 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import KVCache, RotatingKVCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int = 1152
|
||||
num_hidden_layers: int = 26
|
||||
intermediate_size: int = 6912
|
||||
num_attention_heads: int = 4
|
||||
head_dim: int = 256
|
||||
rms_norm_eps: float = 1.0e-6
|
||||
vocab_size: int = 262144
|
||||
num_key_value_heads: int = 1
|
||||
rope_global_base_freq: float = 1_000_000.0
|
||||
rope_local_base_freq: float = 10_000.0
|
||||
rope_traditional: bool = False
|
||||
query_pre_attn_scalar: float = 256
|
||||
sliding_window: int = 512
|
||||
sliding_window_pattern: int = 6
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
self.repeats = n_heads // n_kv_heads
|
||||
self.head_dim = head_dim = args.head_dim
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
self.scale = args.query_pre_attn_scalar**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.q_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
|
||||
self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
|
||||
self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
|
||||
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
traditional=args.rope_traditional,
|
||||
base=(
|
||||
args.rope_local_base_freq
|
||||
if self.is_sliding
|
||||
else args.rope_global_base_freq
|
||||
),
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, _ = x.shape
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
queries = self.q_norm(queries)
|
||||
keys = self.k_norm(keys)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
# Sliding window
|
||||
if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
|
||||
mask = mask[..., -keys.shape[-2] :]
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dims: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.weight = mx.ones((dims,))
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, x):
|
||||
return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def clip_residual(x, y):
|
||||
if x.dtype != mx.float16:
|
||||
return x + y
|
||||
bound = mx.finfo(mx.float16).max
|
||||
return mx.clip(x.astype(mx.float32) + y.astype(mx.float32), -bound, bound).astype(
|
||||
mx.float16
|
||||
)
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args, layer_idx)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.pre_feedforward_layernorm = RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.post_feedforward_layernorm = RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = clip_residual(x, self.post_attention_layernorm(r))
|
||||
r = self.mlp(self.pre_feedforward_layernorm(h))
|
||||
out = clip_residual(h, self.post_feedforward_layernorm(r))
|
||||
return out
|
||||
|
||||
|
||||
class Gemma3Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args, layer_idx=layer_idx)
|
||||
for layer_idx in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
if input_embeddings is not None:
|
||||
h = input_embeddings
|
||||
else:
|
||||
h = self.embed_tokens(inputs)
|
||||
h *= mx.array(self.args.hidden_size**0.5, mx.bfloat16).astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
if mask is None:
|
||||
j = self.args.sliding_window_pattern
|
||||
full_mask = create_attention_mask(h, cache[j - 1 : j])
|
||||
sliding_window_mask = create_attention_mask(h, cache)
|
||||
|
||||
for i, (layer, c) in enumerate(zip(self.layers, cache)):
|
||||
is_global = (
|
||||
i % self.args.sliding_window_pattern
|
||||
== self.args.sliding_window_pattern - 1
|
||||
)
|
||||
|
||||
local_mask = mask
|
||||
if mask is None and is_global:
|
||||
local_mask = full_mask
|
||||
elif mask is None:
|
||||
local_mask = sliding_window_mask
|
||||
|
||||
h = layer(h, local_mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = Gemma3Model(args)
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask: Optional[mx.array] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache, input_embeddings)
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
weights = dict(weights)
|
||||
if "lm_head.weight" not in weights:
|
||||
weights["lm_head.weight"] = weights["model.embed_tokens.weight"]
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
def make_cache(self):
|
||||
caches = []
|
||||
for i in range(self.args.num_hidden_layers):
|
||||
if (
|
||||
i % self.args.sliding_window_pattern
|
||||
== self.args.sliding_window_pattern - 1
|
||||
):
|
||||
caches.append(KVCache())
|
||||
else:
|
||||
caches.append(
|
||||
RotatingKVCache(max_size=self.args.sliding_window, keep=0)
|
||||
)
|
||||
return caches
|
||||
@@ -1,621 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import KVCache, RotatingKVCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextConfig(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
head_dim: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
num_kv_shared_layers: int
|
||||
query_pre_attn_scalar: float
|
||||
vocab_size_per_layer_input: int
|
||||
sliding_window: int
|
||||
max_position_embeddings: int
|
||||
rope_local_base_freq: float
|
||||
rope_theta: float
|
||||
final_logit_softcapping: float
|
||||
layer_types: List[str]
|
||||
activation_sparsity_pattern: List[float]
|
||||
hidden_size_per_layer_input: int
|
||||
altup_num_inputs: int
|
||||
altup_coef_clip: float
|
||||
altup_correct_scale: bool
|
||||
altup_active_idx: int
|
||||
laurel_rank: int
|
||||
rope_scaling: Optional[Dict] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
text_config: dict
|
||||
|
||||
|
||||
class RMSNoScale(nn.Module):
|
||||
def __init__(self, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, x):
|
||||
return mx.fast.rms_norm(x, None, self.eps)
|
||||
|
||||
|
||||
class Gemma3nLaurelBlock(nn.Module):
|
||||
"""Learned Augmented Residual Layer"""
|
||||
|
||||
def __init__(self, config: TextConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
self.linear_left = nn.Linear(
|
||||
self.config.hidden_size, self.config.laurel_rank, bias=False
|
||||
)
|
||||
self.linear_right = nn.Linear(
|
||||
self.config.laurel_rank, self.config.hidden_size, bias=False
|
||||
)
|
||||
self.post_laurel_norm = nn.RMSNorm(
|
||||
dims=self.config.hidden_size,
|
||||
eps=self.config.rms_norm_eps,
|
||||
)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
laurel_x = self.linear_left(x)
|
||||
laurel_x = self.linear_right(laurel_x)
|
||||
normed_laurel_x = self.post_laurel_norm(laurel_x)
|
||||
return x + normed_laurel_x
|
||||
|
||||
|
||||
class Gemma3nAttention(nn.Module):
|
||||
def __init__(self, config: TextConfig, layer_idx: int, is_kv_shared_layer: bool):
|
||||
super().__init__()
|
||||
self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
|
||||
|
||||
dim = config.hidden_size
|
||||
self.n_heads = n_heads = config.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = config.num_key_value_heads
|
||||
self.repeats = n_heads // n_kv_heads
|
||||
self.head_dim = head_dim = config.head_dim
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
self.scale = 1.0
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.q_norm = nn.RMSNorm(dims=config.head_dim, eps=config.rms_norm_eps)
|
||||
self.k_norm = nn.RMSNorm(dims=config.head_dim, eps=config.rms_norm_eps)
|
||||
self.v_norm = RMSNoScale(eps=config.rms_norm_eps)
|
||||
|
||||
self.is_kv_shared_layer = is_kv_shared_layer
|
||||
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
traditional=False,
|
||||
base=(
|
||||
config.rope_local_base_freq if self.is_sliding else config.rope_theta
|
||||
),
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, _ = x.shape
|
||||
|
||||
queries = self.q_proj(x)
|
||||
queries = queries.reshape(B, L, -1, self.head_dim)
|
||||
queries = self.q_norm(queries)
|
||||
|
||||
offset = 0
|
||||
if self.is_kv_shared_layer and cache is not None:
|
||||
# For shared layers, retrieve KV from the designated cache layer
|
||||
keys, values = cache.state
|
||||
offset = cache.offset
|
||||
|
||||
else:
|
||||
if cache is not None:
|
||||
offset = cache.offset
|
||||
keys = self.k_proj(x).reshape(B, L, -1, self.head_dim)
|
||||
keys = self.k_norm(keys)
|
||||
keys = keys.transpose(0, 2, 1, 3)
|
||||
keys = self.rope(keys, offset=offset)
|
||||
|
||||
values = self.v_proj(x).reshape(B, L, -1, self.head_dim)
|
||||
values = self.v_norm(values)
|
||||
values = values.transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
queries = queries.transpose(0, 2, 1, 3)
|
||||
queries = self.rope(queries, offset=offset)
|
||||
|
||||
if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
|
||||
mask = mask[:, : keys.shape[-2]]
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def gelu_topk(inputs, std_multiplier):
|
||||
inputs_mean = mx.mean(inputs, axis=-1, keepdims=True)
|
||||
inputs_std = mx.std(inputs, axis=-1, keepdims=True)
|
||||
cutoff_x = inputs_mean + inputs_std * std_multiplier.astype(inputs_std.dtype)
|
||||
return nn.gelu_approx(mx.maximum(0, inputs - cutoff_x))
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: TextConfig, layer_idx: int = 0):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
if config.activation_sparsity_pattern is not None:
|
||||
self.activation_sparsity = config.activation_sparsity_pattern[layer_idx]
|
||||
else:
|
||||
self.activation_sparsity = 0.0
|
||||
if self.activation_sparsity > 0:
|
||||
self._std_multiplier = math.sqrt(2.0) * mx.erfinv(
|
||||
2 * self.activation_sparsity - 1
|
||||
)
|
||||
|
||||
def __call__(self, x: mx.array):
|
||||
gate_proj = self.gate_proj(x)
|
||||
if self.activation_sparsity > 0.0:
|
||||
activations = gelu_topk(gate_proj, self._std_multiplier)
|
||||
else:
|
||||
activations = nn.gelu_approx(gate_proj)
|
||||
up_proj = self.up_proj(x)
|
||||
down_proj = self.down_proj(activations * up_proj)
|
||||
return down_proj
|
||||
|
||||
|
||||
class Gemma3nAltUp(nn.Module):
|
||||
"""Alternating Updates (AltUp)"""
|
||||
|
||||
def __init__(self, config: TextConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
self.correct_output_scale = mx.zeros((self.config.hidden_size,))
|
||||
self.correction_coefs = nn.Linear(
|
||||
self.config.altup_num_inputs, self.config.altup_num_inputs, bias=False
|
||||
)
|
||||
self.prediction_coefs = nn.Linear(
|
||||
self.config.altup_num_inputs, self.config.altup_num_inputs**2, bias=False
|
||||
)
|
||||
self.modality_router = nn.Linear(
|
||||
self.config.hidden_size, self.config.altup_num_inputs, bias=False
|
||||
)
|
||||
self.router_norm = nn.RMSNorm(
|
||||
dims=self.config.hidden_size,
|
||||
eps=self.config.rms_norm_eps,
|
||||
)
|
||||
|
||||
def compute_router_modalities(self, x: mx.array) -> mx.array:
|
||||
router_inputs = self.router_norm(x) * (self.config.hidden_size**-1.0)
|
||||
routed = self.modality_router(router_inputs).astype(mx.float32)
|
||||
return mx.tanh(routed)
|
||||
|
||||
def predict(self, x: mx.array) -> mx.array:
|
||||
modalities = self.compute_router_modalities(x[self.config.altup_active_idx])
|
||||
|
||||
self.prediction_coefs.weight = self.prediction_coefs.weight.astype(mx.float32)
|
||||
|
||||
if self.config.altup_coef_clip is not None:
|
||||
self.prediction_coefs.weight = mx.clip(
|
||||
self.prediction_coefs.weight,
|
||||
-self.config.altup_coef_clip,
|
||||
self.config.altup_coef_clip,
|
||||
)
|
||||
|
||||
all_coefs = (
|
||||
self.prediction_coefs(modalities)
|
||||
.reshape(
|
||||
*modalities.shape[:-1],
|
||||
self.config.altup_num_inputs,
|
||||
self.config.altup_num_inputs,
|
||||
)
|
||||
.transpose(0, 1, 3, 2)
|
||||
)
|
||||
|
||||
x_up = x.astype(mx.float32)
|
||||
x_permuted = x_up.transpose(1, 2, 3, 0)
|
||||
predictions = mx.matmul(x_permuted, all_coefs)
|
||||
predictions = predictions.transpose(3, 0, 1, 2)
|
||||
predictions += x_up
|
||||
return predictions.astype(x.dtype)
|
||||
|
||||
def correct(self, predictions: mx.array, activated: mx.array):
|
||||
modalities = self.compute_router_modalities(activated)
|
||||
|
||||
self.correction_coefs.weight = self.correction_coefs.weight.astype(mx.float32)
|
||||
|
||||
if self.config.altup_coef_clip is not None:
|
||||
self.correction_coefs.weight = mx.clip(
|
||||
self.correction_coefs.weight,
|
||||
-self.config.altup_coef_clip,
|
||||
self.config.altup_coef_clip,
|
||||
)
|
||||
|
||||
all_coefs = self.correction_coefs(modalities) + 1.0
|
||||
|
||||
active_x = predictions[self.config.altup_active_idx]
|
||||
innovation = activated - active_x
|
||||
|
||||
all_coefs = all_coefs.transpose(2, 1, 0)
|
||||
corrected = innovation[None] * all_coefs[:, None]
|
||||
corrected += predictions
|
||||
|
||||
return corrected.astype(activated.dtype)
|
||||
|
||||
|
||||
class Gemma3nDecoderLayer(nn.Module):
|
||||
def __init__(self, config: TextConfig, layer_idx: int, is_kv_shared_layer: bool):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.layer_idx = layer_idx
|
||||
self.self_attn = Gemma3nAttention(config, layer_idx, is_kv_shared_layer)
|
||||
self.mlp = MLP(config, layer_idx=layer_idx)
|
||||
self.input_layernorm = nn.RMSNorm(
|
||||
self.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
self.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
self.pre_feedforward_layernorm = nn.RMSNorm(
|
||||
self.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
self.post_feedforward_layernorm = nn.RMSNorm(
|
||||
self.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
self.is_sliding = self.self_attn.is_sliding
|
||||
self.sliding_window = config.sliding_window
|
||||
|
||||
self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
|
||||
|
||||
self.altup = Gemma3nAltUp(config)
|
||||
self.laurel = Gemma3nLaurelBlock(config)
|
||||
self.per_layer_input_gate = nn.Linear(
|
||||
self.hidden_size, self.hidden_size_per_layer_input, bias=False
|
||||
)
|
||||
self.per_layer_projection = nn.Linear(
|
||||
self.hidden_size_per_layer_input, self.hidden_size, bias=False
|
||||
)
|
||||
self.post_per_layer_input_norm = nn.RMSNorm(
|
||||
self.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
per_layer_input: Optional[mx.array] = None,
|
||||
):
|
||||
predictions = self.altup.predict(x)
|
||||
active_prediction = predictions[self.config.altup_active_idx]
|
||||
|
||||
active_prediction_normed = self.input_layernorm(active_prediction)
|
||||
laurel_output = self.laurel(active_prediction_normed)
|
||||
|
||||
attn = self.self_attn(
|
||||
active_prediction_normed,
|
||||
mask,
|
||||
cache,
|
||||
)
|
||||
|
||||
attn = self.post_attention_layernorm(attn)
|
||||
|
||||
attn_gated = active_prediction + attn
|
||||
attn_laurel = (attn_gated + laurel_output) * (2.0**-0.5)
|
||||
|
||||
attn_norm = self.pre_feedforward_layernorm(attn_laurel)
|
||||
attn_ffw = self.mlp(attn_norm)
|
||||
attn_ffw_norm = self.post_feedforward_layernorm(attn_ffw)
|
||||
attn_ffw_laurel_gated = attn_laurel + attn_ffw_norm
|
||||
|
||||
corrected_predictions = self.altup.correct(predictions, attn_ffw_laurel_gated)
|
||||
|
||||
first_prediction = corrected_predictions[self.config.altup_active_idx]
|
||||
if self.config.altup_correct_scale:
|
||||
first_prediction = first_prediction * self.altup.correct_output_scale
|
||||
|
||||
first_prediction = self.per_layer_input_gate(first_prediction)
|
||||
first_prediction = nn.gelu_approx(first_prediction)
|
||||
|
||||
first_prediction = mx.multiply(first_prediction, per_layer_input)
|
||||
|
||||
first_prediction = self.per_layer_projection(first_prediction)
|
||||
first_prediction = self.post_per_layer_input_norm(first_prediction)
|
||||
|
||||
corrected_predictions[1:] = corrected_predictions[1:] + first_prediction
|
||||
|
||||
return corrected_predictions
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def logit_softcap(softcap, x):
|
||||
out = mx.tanh(x / softcap)
|
||||
out = out * softcap
|
||||
return out
|
||||
|
||||
|
||||
class LanguageModel(nn.Module):
|
||||
def __init__(self, config: TextConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
|
||||
self.vocab_size = config.vocab_size
|
||||
self.vocab_size_per_layer_input = config.vocab_size_per_layer_input
|
||||
self.num_hidden_layers = config.num_hidden_layers
|
||||
self.final_logit_softcapping = config.final_logit_softcapping
|
||||
self.first_kv_shared_layer_idx = (
|
||||
config.num_hidden_layers - config.num_kv_shared_layers
|
||||
)
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = [
|
||||
Gemma3nDecoderLayer(
|
||||
config=config,
|
||||
layer_idx=layer_idx,
|
||||
is_kv_shared_layer=layer_idx >= self.first_kv_shared_layer_idx,
|
||||
)
|
||||
for layer_idx in range(config.num_hidden_layers)
|
||||
]
|
||||
|
||||
self.embed_tokens_per_layer = nn.Embedding(
|
||||
config.vocab_size_per_layer_input,
|
||||
config.num_hidden_layers * config.hidden_size_per_layer_input,
|
||||
)
|
||||
|
||||
self.per_layer_model_projection = nn.Linear(
|
||||
config.hidden_size,
|
||||
config.num_hidden_layers * config.hidden_size_per_layer_input,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.per_layer_projection_norm = nn.RMSNorm(
|
||||
dims=config.hidden_size_per_layer_input,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
|
||||
self.altup_projections = [
|
||||
nn.Linear(config.hidden_size, config.hidden_size, bias=False)
|
||||
for _ in range(1, self.config.altup_num_inputs)
|
||||
]
|
||||
|
||||
self.altup_unembed_projections = [
|
||||
nn.Linear(config.hidden_size, config.hidden_size, bias=False)
|
||||
for _ in range(1, self.config.altup_num_inputs)
|
||||
]
|
||||
|
||||
self.norm = nn.RMSNorm(
|
||||
config.hidden_size,
|
||||
eps=config.rms_norm_eps,
|
||||
)
|
||||
|
||||
self.first_sliding_idx = self.config.layer_types.index("sliding_attention")
|
||||
self.first_full_idx = self.config.layer_types.index("full_attention")
|
||||
|
||||
concrete_layers = self.config.layer_types[: self.first_kv_shared_layer_idx]
|
||||
shared_full_idx = (
|
||||
len(concrete_layers) - 1 - concrete_layers[::-1].index("full_attention")
|
||||
)
|
||||
shared_sliding_idx = (
|
||||
len(concrete_layers) - 1 - concrete_layers[::-1].index("sliding_attention")
|
||||
)
|
||||
|
||||
self.layer_idx_to_cache_idx = []
|
||||
for i, layer_type in enumerate(self.config.layer_types):
|
||||
if i < self.first_kv_shared_layer_idx:
|
||||
self.layer_idx_to_cache_idx.append(i)
|
||||
else:
|
||||
if layer_type == "full_attention":
|
||||
self.layer_idx_to_cache_idx.append(shared_full_idx)
|
||||
elif layer_type == "sliding_attention":
|
||||
self.layer_idx_to_cache_idx.append(shared_sliding_idx)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown layer type: {layer_type}")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array = None,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
input_embeddings: mx.array = None,
|
||||
):
|
||||
if input_embeddings is None:
|
||||
h = self.embed_tokens(inputs) * (self.hidden_size**0.5)
|
||||
else:
|
||||
h = input_embeddings
|
||||
|
||||
per_layer_inputs = self.get_per_layer_inputs(inputs)
|
||||
per_layer_inputs = self.project_per_layer_inputs(h, per_layer_inputs)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
if mask is None:
|
||||
full_mask = create_attention_mask(
|
||||
h,
|
||||
cache[self.first_full_idx :],
|
||||
)
|
||||
sliding_window_mask = create_attention_mask(
|
||||
h,
|
||||
cache[self.first_sliding_idx :],
|
||||
)
|
||||
h0 = h
|
||||
|
||||
# Expand hidden_states to support per-layer inputs
|
||||
target_magnitude = mx.mean(h0**2, axis=-1, keepdims=True) ** 0.5
|
||||
|
||||
h_list = [h0]
|
||||
h_list.extend([proj(h0) for proj in self.altup_projections])
|
||||
h = mx.stack(h_list, axis=0)
|
||||
mags = mx.mean(h[1:] ** 2, axis=-1, keepdims=True) ** 0.5
|
||||
h[1:] = h[1:] * (target_magnitude / mx.maximum(mags, mx.finfo(h0.dtype).min))
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
per_layer_input = per_layer_inputs[:, :, i, :]
|
||||
|
||||
is_global = self.config.layer_types[i] == "full_attention"
|
||||
|
||||
local_mask = mask
|
||||
if mask is None and is_global:
|
||||
local_mask = full_mask
|
||||
elif mask is None:
|
||||
local_mask = sliding_window_mask
|
||||
|
||||
h = layer(
|
||||
h,
|
||||
local_mask,
|
||||
cache[self.layer_idx_to_cache_idx[i]],
|
||||
per_layer_input,
|
||||
)
|
||||
|
||||
# Per-layer inputs to single output
|
||||
target_magnitude = mx.mean(h[0] ** 2, axis=-1, keepdims=True) ** 0.5
|
||||
for i, proj in enumerate(self.altup_unembed_projections):
|
||||
h[i + 1] = proj(h[i + 1])
|
||||
mags = mx.mean(h[1:] ** 2, axis=-1, keepdims=True) ** 0.5
|
||||
h[1:] = h[1:] * (target_magnitude / mx.maximum(mags, mx.finfo(h0.dtype).min))
|
||||
|
||||
h = mx.mean(h, axis=0)
|
||||
|
||||
out = self.norm(h)
|
||||
out = self.embed_tokens.as_linear(out)
|
||||
if self.final_logit_softcapping is not None:
|
||||
out = logit_softcap(self.final_logit_softcapping, out)
|
||||
return out
|
||||
|
||||
def get_per_layer_inputs(self, input_ids: mx.array) -> mx.array:
|
||||
per_layer_inputs_mask = input_ids < self.vocab_size_per_layer_input
|
||||
tokens = mx.where(per_layer_inputs_mask, input_ids, mx.zeros_like(input_ids))
|
||||
result = self.embed_tokens_per_layer(tokens) * (
|
||||
self.hidden_size_per_layer_input**0.5
|
||||
)
|
||||
return result.reshape(
|
||||
*input_ids.shape,
|
||||
self.num_hidden_layers,
|
||||
self.hidden_size_per_layer_input,
|
||||
)
|
||||
|
||||
def project_per_layer_inputs(
|
||||
self,
|
||||
inputs_embeds: mx.array,
|
||||
per_layer_inputs: mx.array,
|
||||
) -> mx.array:
|
||||
per_layer_projection = self.per_layer_model_projection(inputs_embeds) * (
|
||||
self.hidden_size**-0.5
|
||||
)
|
||||
per_layer_projection = per_layer_projection.reshape(
|
||||
*inputs_embeds.shape[:-1],
|
||||
self.config.num_hidden_layers,
|
||||
self.config.hidden_size_per_layer_input,
|
||||
)
|
||||
per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
|
||||
return (per_layer_projection + per_layer_inputs) * (2.0**-0.5)
|
||||
|
||||
def make_cache(self):
|
||||
caches = []
|
||||
for layer_type in self.config.layer_types[: self.first_kv_shared_layer_idx]:
|
||||
if layer_type == "full_attention":
|
||||
caches.append(KVCache())
|
||||
elif layer_type == "sliding_attention":
|
||||
caches.append(
|
||||
RotatingKVCache(max_size=self.config.sliding_window, keep=0)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown layer type: {layer_type}")
|
||||
return caches
|
||||
|
||||
|
||||
class Gemma3n(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.language_model = LanguageModel(TextConfig.from_dict(args.text_config))
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask: Optional[mx.array] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
return self.language_model(
|
||||
inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
|
||||
)
|
||||
|
||||
def make_cache(self):
|
||||
return self.language_model.make_cache()
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model = Gemma3n(args)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask: Optional[mx.array] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
return self.model(
|
||||
inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
|
||||
)
|
||||
|
||||
def sanitize(self, weights):
|
||||
weights = tree_unflatten(list(weights.items()))
|
||||
for k in ["vision_tower", "audio_tower", "embed_audio", "embed_vision"]:
|
||||
weights["model"].pop(k, None)
|
||||
return dict(tree_flatten(weights))
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.language_model.layers
|
||||
|
||||
def make_cache(self):
|
||||
return self.model.make_cache()
|
||||
@@ -1,183 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
attention_bias: bool
|
||||
head_dim: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
partial_rotary_factor: float
|
||||
rope_theta: float
|
||||
rope_traditional: bool = True
|
||||
max_position_embeddings: int = 32768
|
||||
|
||||
|
||||
class Glm4MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.gate_up_proj = nn.Linear(
|
||||
args.hidden_size, 2 * args.intermediate_size, bias=False
|
||||
)
|
||||
self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
x = self.gate_up_proj(x)
|
||||
gate, up_states = mx.split(x, 2, axis=-1)
|
||||
return self.down_proj(nn.silu(gate) * up_states)
|
||||
|
||||
|
||||
class Glm4Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.head_dim = getattr(
|
||||
args, "head_dim", args.hidden_size // args.num_attention_heads
|
||||
)
|
||||
self.n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = args.num_key_value_heads
|
||||
self.scale = self.head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(
|
||||
args.hidden_size,
|
||||
args.num_attention_heads * self.head_dim,
|
||||
bias=args.attention_bias,
|
||||
)
|
||||
self.k_proj = nn.Linear(
|
||||
args.hidden_size,
|
||||
args.num_key_value_heads * self.head_dim,
|
||||
bias=args.attention_bias,
|
||||
)
|
||||
self.v_proj = nn.Linear(
|
||||
args.hidden_size,
|
||||
args.num_key_value_heads * self.head_dim,
|
||||
bias=args.attention_bias,
|
||||
)
|
||||
self.o_proj = nn.Linear(
|
||||
args.num_attention_heads * self.head_dim, args.hidden_size, bias=False
|
||||
)
|
||||
|
||||
self.rope = nn.RoPE(
|
||||
dims=int(self.head_dim * args.partial_rotary_factor),
|
||||
base=args.rope_theta,
|
||||
traditional=args.rope_traditional,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class Glm4DecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.self_attn = Glm4Attention(args=args)
|
||||
|
||||
self.mlp = Glm4MLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.post_self_attn_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.post_mlp_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None
|
||||
) -> mx.array:
|
||||
x = x + self.post_self_attn_layernorm(
|
||||
self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
)
|
||||
residual = x
|
||||
x = (
|
||||
self.post_mlp_layernorm(self.mlp(self.post_attention_layernorm(x)))
|
||||
+ residual
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class Glm4Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
Glm4DecoderLayer(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = Glm4Model(args)
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+24
-18
@@ -1,12 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs, create_additive_causal_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -45,7 +44,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -60,8 +59,8 @@ class Attention(nn.Module):
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
@@ -99,7 +98,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.attn(self.ln_1(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -125,22 +124,22 @@ class GPT2Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
_, L = inputs.shape
|
||||
|
||||
hidden_states = self.wte(inputs)
|
||||
|
||||
offset = 0
|
||||
if cache is not None and len(cache) > 0 and cache[0] is not None:
|
||||
offset = cache[0].offset
|
||||
mask = None
|
||||
if hidden_states.shape[1] > 1:
|
||||
|
||||
position_ids = mx.arange(offset, offset + L)
|
||||
hidden_states += self.wpe(position_ids)
|
||||
position_ids = mx.array(np.arange(L))
|
||||
hidden_states += self.wpe(position_ids)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(hidden_states, cache)
|
||||
mask = create_additive_causal_mask(
|
||||
hidden_states.shape[1], cache[0].offset if cache is not None else 0
|
||||
)
|
||||
mask = mask.astype(hidden_states.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.h)
|
||||
@@ -161,10 +160,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
out = self.model.wte.as_linear(out)
|
||||
return out
|
||||
|
||||
@@ -199,3 +197,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.h
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.n_embd // self.args.n_head
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,13 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs, create_additive_causal_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -57,7 +55,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -74,8 +72,8 @@ class Attention(nn.Module):
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.c_proj(output)
|
||||
@@ -114,7 +112,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.attn(self.ln_1(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -137,7 +135,6 @@ class GPTBigCodeModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
B, L = inputs.shape
|
||||
@@ -145,16 +142,18 @@ class GPTBigCodeModel(nn.Module):
|
||||
hidden_states = self.wte(inputs)
|
||||
|
||||
mask = None
|
||||
if mask is not None and hidden_states.shape[1] > 1:
|
||||
mask = create_attention_mask(hidden_states, cache)
|
||||
if hidden_states.shape[1] > 1:
|
||||
|
||||
position_ids = mx.array(np.arange(L))
|
||||
hidden_states += self.wpe(position_ids)
|
||||
|
||||
mask = create_additive_causal_mask(
|
||||
hidden_states.shape[1], cache[0].offset if cache is not None else 0
|
||||
)
|
||||
mask = mask.astype(hidden_states.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.h)
|
||||
position_ids = mx.array(np.arange(L))
|
||||
else:
|
||||
position_ids = mx.array(np.arange(cache[0].offset, cache[0].offset + L))
|
||||
|
||||
hidden_states += self.wpe(position_ids)
|
||||
|
||||
for layer, c in zip(self.h, cache):
|
||||
hidden_states = layer(hidden_states, mask, cache=c)
|
||||
@@ -174,10 +173,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.transformer(inputs, mask, cache)
|
||||
out = self.transformer(inputs, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.transformer.wte.as_linear(out)
|
||||
else:
|
||||
@@ -187,3 +185,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.transformer.h
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.n_embd // self.args.n_head
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,218 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
|
||||
# Based on the transformers implementation at:
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
max_position_embeddings: int
|
||||
hidden_size: int
|
||||
num_attention_heads: int
|
||||
num_hidden_layers: int
|
||||
layer_norm_eps: float
|
||||
vocab_size: int
|
||||
rotary_emb_base: int
|
||||
rotary_pct: float
|
||||
num_key_value_heads: int = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.num_key_value_heads is None:
|
||||
self.num_key_value_heads = self.num_attention_heads
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
assert (
|
||||
args.hidden_size % args.num_attention_heads == 0
|
||||
), "hidden_size must be divisible by num_attention_heads"
|
||||
|
||||
self.hidden_size = args.hidden_size
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.head_dim = self.hidden_size // self.num_attention_heads
|
||||
|
||||
self.rope = nn.RoPE(
|
||||
dims=int(self.head_dim * args.rotary_pct),
|
||||
traditional=False,
|
||||
base=args.rotary_emb_base,
|
||||
)
|
||||
|
||||
self.scale = self.head_dim**-0.5
|
||||
|
||||
self.query_key_value = nn.Linear(
|
||||
self.hidden_size, 3 * self.hidden_size, bias=True
|
||||
)
|
||||
self.dense = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
qkv = self.query_key_value(x)
|
||||
|
||||
new_qkv_shape = qkv.shape[:-1] + (self.num_attention_heads, 3 * self.head_dim)
|
||||
qkv = qkv.reshape(*new_qkv_shape)
|
||||
|
||||
queries, keys, values = [x.transpose(0, 2, 1, 3) for x in qkv.split(3, -1)]
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.dense(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = args.hidden_size
|
||||
self.dense_h_to_4h = nn.Linear(self.hidden_size, 4 * self.hidden_size)
|
||||
self.dense_4h_to_h = nn.Linear(4 * self.hidden_size, self.hidden_size)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
# gelu_approx corresponds to FastGELUActivation in transformers.
|
||||
return self.dense_4h_to_h(nn.gelu_approx(self.dense_h_to_4h(x)))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = args.hidden_size
|
||||
self.layer_norm_eps = args.layer_norm_eps
|
||||
self.attention = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = nn.LayerNorm(
|
||||
self.hidden_size,
|
||||
eps=self.layer_norm_eps,
|
||||
)
|
||||
self.post_attention_layernorm = nn.LayerNorm(
|
||||
self.hidden_size, eps=self.layer_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
residual = x
|
||||
# NeoX runs attention and feedforward network in parallel.
|
||||
attn = self.attention(self.input_layernorm(x), mask, cache)
|
||||
ffn = self.mlp(self.post_attention_layernorm(x))
|
||||
out = attn + ffn + residual
|
||||
return out
|
||||
|
||||
|
||||
class GPTNeoXModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
self.layer_norm_eps = args.layer_norm_eps
|
||||
assert self.vocab_size > 0
|
||||
self.embed_in = nn.Embedding(self.vocab_size, self.hidden_size)
|
||||
self.embed_out = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
self.h = [TransformerBlock(args=args) for _ in range(self.num_hidden_layers)]
|
||||
self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
_, L = inputs.shape
|
||||
|
||||
hidden_states = self.embed_in(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(hidden_states, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.h)
|
||||
|
||||
for layer, c in zip(self.h, cache):
|
||||
hidden_states = layer(hidden_states, mask, cache=c)
|
||||
|
||||
out = self.final_layer_norm(hidden_states)
|
||||
out = self.embed_out(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = GPTNeoXModel(args)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
new_weights = {}
|
||||
|
||||
for w_key, w_value in weights.items():
|
||||
# Created through register_buffer in Pytorch, not needed here.
|
||||
ignore_suffixes = [
|
||||
".attention.bias",
|
||||
".attention.masked_bias",
|
||||
".attention.rotary_emb.inv_freq",
|
||||
]
|
||||
|
||||
skip_weight = False
|
||||
for ignored_suffix in ignore_suffixes:
|
||||
if w_key.endswith(ignored_suffix):
|
||||
skip_weight = True
|
||||
break
|
||||
|
||||
if skip_weight:
|
||||
continue
|
||||
|
||||
if not w_key.startswith("model."):
|
||||
w_key = f"model.{w_key}"
|
||||
|
||||
w_key = w_key.replace(".gpt_neox.layers.", ".h.")
|
||||
w_key = w_key.replace(".gpt_neox.", ".")
|
||||
|
||||
new_weights[w_key] = w_value
|
||||
|
||||
return new_weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.h
|
||||
@@ -1,195 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
logits_scaling: float
|
||||
attention_multiplier: float
|
||||
embedding_multiplier: float
|
||||
residual_multiplier: float
|
||||
max_position_embeddings: int
|
||||
num_key_value_heads: int
|
||||
attention_bias: bool
|
||||
mlp_bias: bool
|
||||
rope_theta: float
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = True
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.hidden_size // n_heads
|
||||
|
||||
self.scale = args.attention_multiplier
|
||||
attention_bias = args.attention_bias
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
False,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
hidden_dim = args.intermediate_size
|
||||
if hasattr(args, "mlp_bias"):
|
||||
mlp_bias = args.mlp_bias
|
||||
else:
|
||||
mlp_bias = False
|
||||
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.residual_multiplier = args.residual_multiplier
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r * self.residual_multiplier
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r * self.residual_multiplier
|
||||
return out
|
||||
|
||||
|
||||
class GraniteModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.embedding_multiplier = args.embedding_multiplier
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs) * self.embedding_multiplier
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = GraniteModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
self.logits_scaling = args.logits_scaling
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out / self.logits_scaling
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,185 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
attention_bias: bool
|
||||
head_dim: int
|
||||
max_position_embeddings: int
|
||||
mlp_bias: bool
|
||||
model_type: str
|
||||
rope_theta: float
|
||||
tie_word_embeddings: bool
|
||||
|
||||
|
||||
class HeliumAttention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
self.rope = nn.RoPE(head_dim, traditional=True, base=args.rope_theta)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class HeliumMLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
self.intermediate_size = args.intermediate_size
|
||||
|
||||
self.gate_proj = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size, bias=args.mlp_bias
|
||||
)
|
||||
self.up_proj = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size, bias=args.mlp_bias
|
||||
)
|
||||
self.down_proj = nn.Linear(
|
||||
self.intermediate_size, self.hidden_size, bias=args.mlp_bias
|
||||
)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class HeliumDecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
|
||||
self.self_attn = HeliumAttention(args)
|
||||
self.mlp = HeliumMLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class HeliumModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
self.vocab_size = args.vocab_size
|
||||
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
|
||||
self.layers = [HeliumDecoderLayer(args) for _ in range(args.num_hidden_layers)]
|
||||
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
|
||||
self.model = HeliumModel(args)
|
||||
|
||||
self.vocab_size = args.vocab_size
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,320 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
vocab_size: int
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
attention_bias: bool
|
||||
moe_topk: int
|
||||
num_experts: int
|
||||
num_shared_expert: int
|
||||
use_mixed_mlp_moe: bool
|
||||
use_qk_norm: bool
|
||||
rms_norm_eps: float
|
||||
rope_theta: float
|
||||
use_cla: bool
|
||||
cla_share_factor: 2
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
|
||||
if self.rope_scaling:
|
||||
required_keys = {"factor", "type"}
|
||||
if not all(key in self.rope_scaling for key in required_keys):
|
||||
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
||||
|
||||
|
||||
class DynamicNTKAlphaRoPE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
base: float = 10000,
|
||||
scaling_alpha: float = 1.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dims = dims
|
||||
base = base * scaling_alpha ** (dims / (dims - 2))
|
||||
self._freqs = base ** (mx.arange(0, self.dims, 2) / self.dims)
|
||||
|
||||
def __call__(self, x, offset: int = 0):
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
self.dims,
|
||||
traditional=False,
|
||||
base=None,
|
||||
scale=1.0,
|
||||
offset=offset,
|
||||
freqs=self._freqs,
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, kv_proj: bool, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
|
||||
if kv_proj:
|
||||
self.k_proj = nn.Linear(
|
||||
dim, n_kv_heads * head_dim, bias=args.attention_bias
|
||||
)
|
||||
self.v_proj = nn.Linear(
|
||||
dim, n_kv_heads * head_dim, bias=args.attention_bias
|
||||
)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
|
||||
self.use_qk_norm = args.use_qk_norm
|
||||
if self.use_qk_norm:
|
||||
self.query_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
|
||||
self.key_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
|
||||
|
||||
self.rope = DynamicNTKAlphaRoPE(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
scaling_alpha=args.rope_scaling["alpha"],
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
kv_states=None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries = self.q_proj(x)
|
||||
if kv_states is None:
|
||||
keys, values = self.k_proj(x), self.v_proj(x)
|
||||
kv_states = keys, values
|
||||
else:
|
||||
keys, values = kv_states
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
offset = cache.offset if cache else 0
|
||||
queries = self.rope(queries, offset=offset)
|
||||
keys = self.rope(keys, offset=offset)
|
||||
if self.use_qk_norm:
|
||||
queries = self.query_layernorm(queries)
|
||||
keys = self.key_layernorm(keys)
|
||||
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output), kv_states
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class Gate(nn.Module):
|
||||
def __init__(self, dim, num_experts):
|
||||
super().__init__()
|
||||
self.wg = nn.Linear(dim, num_experts, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.wg(x)
|
||||
|
||||
|
||||
class MoeBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
dim = args.hidden_size
|
||||
intermediate_size = args.intermediate_size
|
||||
self.use_shared_mlp = args.use_mixed_mlp_moe
|
||||
|
||||
if args.use_mixed_mlp_moe:
|
||||
self.shared_mlp = MLP(dim, intermediate_size * args.num_shared_expert)
|
||||
|
||||
self.num_experts = num_experts = args.num_experts
|
||||
self.top_k = args.moe_topk
|
||||
|
||||
self.gate = Gate(dim, num_experts)
|
||||
self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
):
|
||||
gates = self.gate(x)
|
||||
gates = mx.softmax(gates, axis=-1, precise=True)
|
||||
|
||||
k = self.top_k
|
||||
inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
|
||||
scores = mx.take_along_axis(gates, inds, axis=-1)
|
||||
|
||||
y = self.switch_mlp(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2)
|
||||
|
||||
if self.use_shared_mlp:
|
||||
shared_expert_output = self.shared_mlp(x)
|
||||
y = y + shared_expert_output
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs, kv_proj: bool):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(kv_proj, args)
|
||||
if args.num_experts == 1:
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
else:
|
||||
self.mlp = MoeBlock(args)
|
||||
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
shared_kv_states: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
):
|
||||
r, shared_kv_states = self.self_attn(
|
||||
self.input_layernorm(x), mask, cache, shared_kv_states
|
||||
)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out, shared_kv_states
|
||||
|
||||
|
||||
class HunYuanModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
DecoderLayer(
|
||||
args=args,
|
||||
kv_proj=(not args.use_cla) or (i % args.cla_share_factor) == 0,
|
||||
)
|
||||
for i in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for i, (layer, c) in enumerate(zip(self.layers, cache)):
|
||||
if (not self.args.use_cla) or i % self.args.cla_share_factor == 0:
|
||||
shared_kv_states = None
|
||||
h, shared_kv_states = layer(h, mask, c, shared_kv_states)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = HunYuanModel(args)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
return self.model.embed_tokens.as_linear(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
|
||||
if "model.layers.0.mlp.gate_and_up_proj.weight" in weights:
|
||||
new_weights = {}
|
||||
D = self.args.hidden_size
|
||||
n_kv_heads = self.args.num_key_value_heads
|
||||
n_kv_groups = self.args.num_attention_heads // n_kv_heads
|
||||
head_dim = D // self.args.num_attention_heads
|
||||
for k, v in weights.items():
|
||||
if "qkv_proj" in k:
|
||||
v = v.reshape(n_kv_heads, n_kv_groups + 2, head_dim, -1)
|
||||
splits = v.split([n_kv_groups, n_kv_groups + 1], axis=1)
|
||||
for k_up, v_new in zip(["q_proj", "k_proj", "v_proj"], splits):
|
||||
k_new = k.replace("qkv_proj", k_up)
|
||||
new_weights[k_new] = mx.flatten(v_new, 0, 2)
|
||||
elif "gate_and_up_proj" in k:
|
||||
splits = v.split(2, axis=0)
|
||||
for k_up, v_new in zip(["up_proj", "gate_proj"], splits):
|
||||
k_new = k.replace("gate_and_up_proj", k_up)
|
||||
new_weights[k_new] = v_new
|
||||
else:
|
||||
new_weights[k] = v
|
||||
weights = new_weights
|
||||
|
||||
if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
|
||||
return weights
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for n in ["up_proj", "down_proj", "gate_proj"]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
|
||||
for e in range(self.args.num_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{n}.{k}"] = mx.stack(to_join)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+23
-66
@@ -1,12 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -19,7 +17,6 @@ class ModelArgs(BaseModelArgs):
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
bias: bool = True
|
||||
max_position_embeddings: int = 32768
|
||||
num_key_value_heads: int = None
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
@@ -35,50 +32,8 @@ class ModelArgs(BaseModelArgs):
|
||||
if not all(key in self.rope_scaling for key in required_keys):
|
||||
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
||||
|
||||
if self.rope_scaling["type"] not in ["linear", "dynamic"]:
|
||||
raise ValueError(
|
||||
"rope_scaling 'type' currently only supports 'linear' or 'dynamic"
|
||||
)
|
||||
|
||||
|
||||
class DynamicNTKScalingRoPE(nn.Module):
|
||||
"""Implements the rotary positional encoding with Dynamic NTK scaling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
max_position_embeddings: int = 2048,
|
||||
traditional: bool = False,
|
||||
base: float = 10000,
|
||||
scale: float = 1.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.original_base = base
|
||||
self.dims = dims
|
||||
self.traditional = traditional
|
||||
self.scale = scale
|
||||
|
||||
def extra_repr(self):
|
||||
return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, scaling_factor={self.scaling_factor}"
|
||||
|
||||
def __call__(self, x, offset: int = 0):
|
||||
seq_len = x.shape[1] + offset
|
||||
if seq_len > self.max_position_embeddings:
|
||||
base = self.original_base * (
|
||||
(self.scale * seq_len / self.max_position_embeddings) - (self.scale - 1)
|
||||
) ** (self.dims / (self.dims - 2))
|
||||
else:
|
||||
base = self.original_base
|
||||
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
self.dims,
|
||||
traditional=self.traditional,
|
||||
base=base,
|
||||
scale=self.scale,
|
||||
offset=offset,
|
||||
)
|
||||
if self.rope_scaling["type"] != "linear":
|
||||
raise ValueError("rope_scaling 'type' currently only supports 'linear'")
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
@@ -101,12 +56,10 @@ class Attention(nn.Module):
|
||||
rope_scale = (
|
||||
1 / args.rope_scaling["factor"]
|
||||
if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
|
||||
else 2.0
|
||||
else 1
|
||||
)
|
||||
|
||||
self.rope = DynamicNTKScalingRoPE(
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
traditional=args.rope_traditional,
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
@@ -116,7 +69,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -141,8 +94,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.wo(output)
|
||||
@@ -171,7 +124,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.attention(self.attention_norm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -193,13 +146,14 @@ class InternLM2Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.tok_embeddings(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -222,20 +176,23 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.tok_embeddings.as_linear(out)
|
||||
else:
|
||||
out = self.output(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
return {k: v for k, v in weights.items() if "attention.rope.inv_freq" not in k}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,241 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
bias: bool = False
|
||||
qkv_bias: bool = False
|
||||
max_position_embeddings: int = 32768
|
||||
num_key_value_heads: int = None
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
if self.num_key_value_heads is None:
|
||||
self.num_key_value_heads = self.num_attention_heads
|
||||
|
||||
if self.rope_scaling:
|
||||
required_keys = {"factor", "rope_type"}
|
||||
if not all(key in self.rope_scaling for key in required_keys):
|
||||
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
||||
|
||||
if self.rope_scaling["rope_type"] not in ["linear", "dynamic"]:
|
||||
raise ValueError(
|
||||
"rope_scaling 'rope_type' currently only supports 'linear' or 'dynamic"
|
||||
)
|
||||
|
||||
|
||||
class DynamicNTKScalingRoPE(nn.Module):
|
||||
"""Implements the rotary positional encoding with Dynamic NTK scaling."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
max_position_embeddings: int = 2048,
|
||||
traditional: bool = False,
|
||||
base: float = 10000,
|
||||
scale: float = 1.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.original_base = base
|
||||
self.dims = dims
|
||||
self.traditional = traditional
|
||||
self.scale = scale
|
||||
|
||||
def extra_repr(self):
|
||||
return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, scaling_factor={self.scaling_factor}"
|
||||
|
||||
def __call__(self, x, offset: int = 0):
|
||||
seq_len = x.shape[1] + offset
|
||||
if seq_len > self.max_position_embeddings:
|
||||
base = self.original_base * (
|
||||
(self.scale * seq_len / self.max_position_embeddings) - (self.scale - 1)
|
||||
) ** (self.dims / (self.dims - 2))
|
||||
else:
|
||||
base = self.original_base
|
||||
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
self.dims,
|
||||
traditional=self.traditional,
|
||||
base=base,
|
||||
scale=self.scale,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
qkv_bias = args.qkv_bias
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
self.n_kv_groups = n_heads // args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=qkv_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=qkv_bias)
|
||||
|
||||
rope_scale = (
|
||||
1 / args.rope_scaling["factor"]
|
||||
if args.rope_scaling is not None
|
||||
and args.rope_scaling["rope_type"] == "linear"
|
||||
else 2.0
|
||||
)
|
||||
|
||||
self.rope = DynamicNTKScalingRoPE(
|
||||
head_dim,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
traditional=args.rope_traditional,
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim, bias):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size, args.bias)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class InternLM2Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
assert args.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = InternLM2Model(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
return {k: v for k, v in weights.items() if "attention.rope.inv_freq" not in k}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,118 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .deepseek_v3 import DeepseekV3Model
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextArgs(BaseModelArgs):
|
||||
vocab_size: int = 102400
|
||||
hidden_size: int = 4096
|
||||
intermediate_size: int = 11008
|
||||
moe_intermediate_size: int = 1407
|
||||
num_hidden_layers: int = 30
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 32
|
||||
n_shared_experts: Optional[int] = None
|
||||
n_routed_experts: Optional[int] = None
|
||||
routed_scaling_factor: float = 1.0
|
||||
kv_lora_rank: int = 512
|
||||
q_lora_rank: int = 1536
|
||||
qk_rope_head_dim: int = 64
|
||||
v_head_dim: int = 128
|
||||
qk_nope_head_dim: int = 128
|
||||
topk_method: str = "noaux_tc"
|
||||
scoring_func: str = "sigmoid"
|
||||
norm_topk_prob: bool = True
|
||||
n_group: Optional[int] = None
|
||||
topk_group: Optional[int] = None
|
||||
num_experts_per_tok: Optional[int] = None
|
||||
moe_layer_freq: int = 1
|
||||
first_k_dense_replace: int = 0
|
||||
max_position_embeddings: int = 2048
|
||||
rms_norm_eps: float = 1e-6
|
||||
rope_theta: float = 10000.0
|
||||
rope_scaling: Dict = None
|
||||
attention_bias: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
text_config: Union[TextArgs, dict]
|
||||
model_type: str
|
||||
|
||||
def __post_init__(self):
|
||||
self.text_config = TextArgs.from_dict(self.text_config)
|
||||
|
||||
|
||||
class LanguageModel(nn.Module):
|
||||
def __init__(self, config: TextArgs):
|
||||
super().__init__()
|
||||
self.args = config
|
||||
self.model = DeepseekV3Model(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, cache, mask)
|
||||
return self.lm_head(out)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = config
|
||||
self.model_type = config.model_type
|
||||
self.language_model = LanguageModel(config.text_config)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
):
|
||||
return self.language_model(inputs, cache, mask)
|
||||
|
||||
def sanitize(self, weights):
|
||||
def keep(key):
|
||||
return (
|
||||
"vision_tower" not in key
|
||||
and "rotary_emb" not in key
|
||||
and "multi_modal_projector" not in key
|
||||
)
|
||||
|
||||
weights = {k: v for k, v in weights.items() if keep(k)}
|
||||
# Stack experts
|
||||
for l in range(self.args.text_config.num_hidden_layers):
|
||||
prefix = f"language_model.model.layers.{l}"
|
||||
for m in [("gate_proj"), ("down_proj"), ("up_proj")]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
|
||||
for e in range(self.args.text_config.n_routed_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
|
||||
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.language_model.model.layers
|
||||
|
||||
@property
|
||||
def cast_predicate(self):
|
||||
def predicate(k):
|
||||
return "e_score_correction_bias" not in k
|
||||
|
||||
return predicate
|
||||
+42
-35
@@ -1,13 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
from .base import BaseModelArgs, KVCache, create_additive_causal_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -19,8 +16,6 @@ class ModelArgs(BaseModelArgs):
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
head_dim: Optional[int] = None
|
||||
max_position_embeddings: Optional[int] = None
|
||||
num_key_value_heads: Optional[int] = None
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
@@ -33,6 +28,14 @@ class ModelArgs(BaseModelArgs):
|
||||
if self.num_key_value_heads is None:
|
||||
self.num_key_value_heads = self.num_attention_heads
|
||||
|
||||
if self.rope_scaling:
|
||||
required_keys = {"factor", "type"}
|
||||
if not all(key in self.rope_scaling for key in required_keys):
|
||||
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
||||
|
||||
if self.rope_scaling["type"] != "linear":
|
||||
raise ValueError("rope_scaling 'type' currently only supports 'linear'")
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
@@ -42,8 +45,7 @@ class Attention(nn.Module):
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
if hasattr(args, "attention_bias"):
|
||||
attention_bias = args.attention_bias
|
||||
@@ -55,19 +57,23 @@ class Attention(nn.Module):
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
args.rope_traditional,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
rope_scale = (
|
||||
1 / args.rope_scaling["factor"]
|
||||
if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
|
||||
else 1
|
||||
)
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
traditional=args.rope_traditional,
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[KVCache] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -86,10 +92,9 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
@@ -130,7 +135,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[KVCache] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -155,17 +160,16 @@ class LlamaModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
if input_embeddings is not None:
|
||||
h = input_embeddings
|
||||
else:
|
||||
h = self.embed_tokens(inputs)
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = create_additive_causal_mask(
|
||||
h.shape[1], cache[0].offset if cache is not None else 0
|
||||
)
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -188,11 +192,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache, input_embeddings)
|
||||
out = self.model(inputs, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
@@ -201,13 +203,18 @@ class Model(nn.Module):
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
weights = {
|
||||
return {
|
||||
k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
|
||||
}
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,333 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import ChunkedKVCache, KVCache
|
||||
from .rope_utils import initialize_rope
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextArgs(BaseModelArgs):
|
||||
attention_bias: bool
|
||||
attention_chunk_size: int
|
||||
head_dim: int
|
||||
hidden_act: str
|
||||
hidden_size: int
|
||||
interleave_moe_layer_step: int
|
||||
intermediate_size: int
|
||||
intermediate_size_mlp: int
|
||||
max_position_embeddings: int
|
||||
model_type: str
|
||||
num_attention_heads: int
|
||||
num_experts_per_tok: int
|
||||
num_hidden_layers: int
|
||||
num_key_value_heads: int
|
||||
num_local_experts: int
|
||||
rms_norm_eps: float
|
||||
rope_scaling: Any
|
||||
rope_theta: float
|
||||
use_qk_norm: bool
|
||||
vocab_size: int
|
||||
attn_temperature_tuning: int = 4
|
||||
floor_scale: int = 8192
|
||||
attn_scale: float = 0.1
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
text_config: Union[TextArgs, dict]
|
||||
model_type: str
|
||||
|
||||
def __post_init__(self):
|
||||
self.text_config = TextArgs.from_dict(self.text_config)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: TextArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.use_rope = int((layer_idx + 1) % 4 != 0) # rope unused for dense layers
|
||||
self.attn_temperature_tuning = args.attn_temperature_tuning
|
||||
self.floor_scale = args.floor_scale
|
||||
self.attn_scale = args.attn_scale
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
|
||||
self.scale = head_dim**-0.5
|
||||
if hasattr(args, "attention_bias"):
|
||||
attention_bias = args.attention_bias
|
||||
else:
|
||||
attention_bias = False
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
|
||||
|
||||
self.use_qk_norm = args.use_qk_norm and self.use_rope
|
||||
|
||||
if self.use_rope:
|
||||
self.rope = initialize_rope(
|
||||
head_dim,
|
||||
args.rope_theta,
|
||||
traditional=True,
|
||||
scaling_config=args.rope_scaling,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
offset = cache.offset
|
||||
else:
|
||||
offset = 0
|
||||
|
||||
if self.use_rope:
|
||||
queries = self.rope(queries, offset=offset)
|
||||
keys = self.rope(keys, offset=offset)
|
||||
|
||||
if self.use_qk_norm:
|
||||
queries = mx.fast.rms_norm(queries, weight=None, eps=1e-6)
|
||||
keys = mx.fast.rms_norm(keys, weight=None, eps=1e-6)
|
||||
|
||||
if self.attn_temperature_tuning and not self.use_rope:
|
||||
attn_scales = (
|
||||
mx.log(
|
||||
mx.floor(mx.arange(offset + 1, offset + L + 1) / self.floor_scale)
|
||||
+ 1.0
|
||||
)
|
||||
* self.attn_scale
|
||||
+ 1.0
|
||||
)
|
||||
attn_scales = attn_scales[:, None]
|
||||
queries = (queries * attn_scales).astype(queries.dtype)
|
||||
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs, intermediate_size: int = None):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
hidden_dim = intermediate_size or args.intermediate_size
|
||||
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class MoE(nn.Module):
|
||||
def __init__(self, args):
|
||||
super().__init__()
|
||||
self.top_k = args.num_experts_per_tok
|
||||
self.num_experts = args.num_local_experts
|
||||
self.experts = SwitchGLU(
|
||||
args.hidden_size, args.intermediate_size, self.num_experts
|
||||
)
|
||||
self.router = nn.Linear(args.hidden_size, args.num_local_experts, bias=False)
|
||||
self.shared_expert = MLP(args)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
logits = self.router(x)
|
||||
k = self.top_k
|
||||
indices = mx.argpartition(-logits, kth=k - 1, axis=-1)[..., :k]
|
||||
scores = mx.take_along_axis(logits, indices, axis=-1)
|
||||
scores = mx.sigmoid(scores.astype(mx.float32)).astype(x.dtype)
|
||||
|
||||
out = self.experts(x * scores, indices).squeeze(2)
|
||||
return out + self.shared_expert(x)
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: TextArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args, layer_idx)
|
||||
self.is_moe_layer = (layer_idx % args.interleave_moe_layer_step) == (
|
||||
args.interleave_moe_layer_step - 1
|
||||
)
|
||||
if self.is_moe_layer:
|
||||
self.feed_forward = MoE(args)
|
||||
else:
|
||||
self.feed_forward = MLP(args, args.intermediate_size_mlp)
|
||||
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.feed_forward(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class LlamaModel(nn.Module):
|
||||
def __init__(self, args: TextArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [TransformerBlock(args, i) for i in range(args.num_hidden_layers)]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.attention_chunk_size = args.attention_chunk_size
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if cache is not None:
|
||||
for idx, c in enumerate(cache):
|
||||
if (idx + 1) % 4 != 0:
|
||||
c.maybe_trim_front()
|
||||
start = cache[0].start_position
|
||||
offset = cache[0].offset
|
||||
else:
|
||||
start = 0
|
||||
offset = 0
|
||||
end = offset + h.shape[1]
|
||||
linds = mx.arange(start, end)
|
||||
rinds = mx.arange(offset, end)[:, None]
|
||||
block_pos = mx.abs(
|
||||
(linds // self.attention_chunk_size) - (rinds // self.attention_chunk_size)
|
||||
)
|
||||
token_pos = linds <= rinds
|
||||
chunk_mask = (block_pos == 0) & token_pos
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
else:
|
||||
chunk_mask &= mask
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for idx, (layer, c) in enumerate(zip(self.layers, cache)):
|
||||
use_chunked_attention = (idx + 1) % 4 != 0
|
||||
if use_chunked_attention:
|
||||
local_mask = chunk_mask
|
||||
else:
|
||||
local_mask = mask
|
||||
h = layer(h, local_mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class LanguageModel(nn.Module):
|
||||
def __init__(self, args: TextArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = LlamaModel(self.args)
|
||||
self.lm_head = nn.Linear(
|
||||
self.args.hidden_size, self.args.vocab_size, bias=False
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.language_model = LanguageModel(args.text_config)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
return self.language_model(inputs, mask, cache)
|
||||
|
||||
def sanitize(self, weights):
|
||||
def to_remove(k):
|
||||
return "vision_model" in k or "multi_modal_projector" in k
|
||||
|
||||
# Remove vision weights
|
||||
weights = {k: v for k, v in weights.items() if not to_remove(k)}
|
||||
|
||||
# Rename expert weights for SwitchGLU
|
||||
for l in range(self.args.text_config.num_hidden_layers):
|
||||
prefix = f"language_model.model.layers.{l}.feed_forward.experts"
|
||||
if f"{prefix}.gate_up_proj" in weights:
|
||||
v = weights.pop(f"{prefix}.gate_up_proj")
|
||||
gate_k = f"{prefix}.gate_proj.weight"
|
||||
up_k = f"{prefix}.up_proj.weight"
|
||||
gate_proj, up_proj = mx.split(v, 2, axis=-1)
|
||||
weights[gate_k] = mx.swapaxes(gate_proj, 1, 2)
|
||||
weights[up_k] = mx.swapaxes(up_proj, 1, 2)
|
||||
if f"{prefix}.down_proj" in weights:
|
||||
down_proj = weights.pop(f"{prefix}.down_proj")
|
||||
weights[f"{prefix}.down_proj.weight"] = mx.swapaxes(down_proj, 1, 2)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.language_model.model.layers
|
||||
|
||||
def make_cache(self):
|
||||
chunk_size = self.args.text_config.attention_chunk_size
|
||||
caches = []
|
||||
for i in range(len(self.layers)):
|
||||
if (i + 1) % 4 != 0:
|
||||
caches.append(ChunkedKVCache(chunk_size))
|
||||
else:
|
||||
caches.append(KVCache())
|
||||
return caches
|
||||
@@ -1,242 +0,0 @@
|
||||
# Copyright © 2024-2025 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .cache import MambaCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
vocab_size: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
state_size: int
|
||||
num_hidden_layers: int
|
||||
conv_kernel: int
|
||||
use_bias: bool
|
||||
use_conv_bias: bool
|
||||
time_step_rank: int
|
||||
tie_word_embeddings: bool = True
|
||||
use_bcdt_rms: bool = False
|
||||
mixer_rms_eps: float = 1e-6
|
||||
|
||||
def __post_init__(self):
|
||||
if not hasattr(self, "hidden_size") and hasattr(self, "d_model"):
|
||||
self.hidden_size = self.d_model
|
||||
if not hasattr(self, "intermediate_size") and hasattr(self, "d_inner"):
|
||||
self.intermediate_size = self.d_inner
|
||||
if not hasattr(self, "state_size") and hasattr(self, "d_state"):
|
||||
self.state_size = self.d_state
|
||||
if not hasattr(self, "num_hidden_layers") and hasattr(self, "n_layer"):
|
||||
self.num_hidden_layers = self.n_layer
|
||||
if not hasattr(self, "num_hidden_layers") and hasattr(self, "n_layers"):
|
||||
self.num_hidden_layers = self.n_layers
|
||||
if not hasattr(self, "conv_kernel") and hasattr(self, "d_conv"):
|
||||
self.conv_kernel = self.d_conv
|
||||
if not hasattr(self, "use_bias") and hasattr(self, "bias"):
|
||||
self.use_bias = self.bias
|
||||
if not hasattr(self, "use_conv_bias") and hasattr(self, "conv_bias"):
|
||||
self.use_conv_bias = self.conv_bias
|
||||
|
||||
if self.time_step_rank == "auto":
|
||||
self.time_step_rank = math.ceil(self.hidden_size / 16)
|
||||
if self.model_type == "falcon_mamba":
|
||||
self.use_bcdt_rms = True
|
||||
|
||||
|
||||
class DepthWiseConv1d(nn.Module):
|
||||
def __init__(self, channels, kernel_size, bias=True, padding=0):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.kernel_size = kernel_size
|
||||
self.padding = padding
|
||||
self.weight = mx.random.normal((self.channels, kernel_size, 1))
|
||||
self.bias = mx.zeros((channels,)) if bias else None
|
||||
|
||||
def __call__(self, x, cache=None):
|
||||
B, L, C = x.shape
|
||||
groups, K, _ = self.weight.shape
|
||||
|
||||
if cache is not None:
|
||||
x = mx.concatenate([cache, x], axis=1)
|
||||
else:
|
||||
x = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
|
||||
|
||||
y = mx.conv_general(x, self.weight, groups=groups)
|
||||
|
||||
if self.bias is not None:
|
||||
y = y + self.bias
|
||||
|
||||
return y, x[:, -K + 1 :, :]
|
||||
|
||||
|
||||
class MambaBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.hidden_size = args.hidden_size
|
||||
self.ssm_state_size = args.state_size
|
||||
self.conv_kernel_size = args.conv_kernel
|
||||
self.intermediate_size = args.intermediate_size
|
||||
self.time_step_rank = int(args.time_step_rank)
|
||||
self.use_conv_bias = args.use_conv_bias
|
||||
self.use_bcdt_rms = args.use_bcdt_rms
|
||||
if self.use_bcdt_rms:
|
||||
self.mixer_norm = lambda x: mx.fast.rms_norm(
|
||||
x, mx.ones(x.shape[-1], x.dtype), eps=args.mixer_rms_eps
|
||||
)
|
||||
|
||||
self.in_proj = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size * 2, bias=args.use_bias
|
||||
)
|
||||
|
||||
self.conv1d = DepthWiseConv1d(
|
||||
channels=self.intermediate_size,
|
||||
kernel_size=self.conv_kernel_size,
|
||||
bias=self.use_conv_bias,
|
||||
padding=self.conv_kernel_size - 1,
|
||||
)
|
||||
|
||||
self.x_proj = nn.Linear(
|
||||
self.intermediate_size,
|
||||
self.time_step_rank + 2 * self.ssm_state_size,
|
||||
bias=False,
|
||||
)
|
||||
self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)
|
||||
|
||||
A = mx.repeat(
|
||||
mx.arange(1.0, self.ssm_state_size + 1.0).reshape([1, self.ssm_state_size]),
|
||||
repeats=self.intermediate_size,
|
||||
axis=0,
|
||||
)
|
||||
self.A_log = mx.log(A)
|
||||
self.D = mx.ones([self.intermediate_size])
|
||||
|
||||
self.out_proj = nn.Linear(
|
||||
self.intermediate_size, self.hidden_size, bias=args.use_bias
|
||||
)
|
||||
|
||||
def ssm_step(self, x, A, state=None):
|
||||
D = self.D
|
||||
deltaBC = self.x_proj(x)
|
||||
delta, B, C = map(
|
||||
self.mixer_norm if self.use_bcdt_rms else lambda x: x,
|
||||
mx.split(
|
||||
deltaBC,
|
||||
[self.time_step_rank, self.time_step_rank + self.ssm_state_size],
|
||||
axis=-1,
|
||||
),
|
||||
)
|
||||
if self.use_bcdt_rms:
|
||||
delta, B, C = map(self.mixer_norm, (delta, B, C))
|
||||
delta = nn.softplus(self.dt_proj(delta))
|
||||
new_state = mx.expand_dims(delta * x, -1) * mx.expand_dims(B, 1)
|
||||
if state is not None:
|
||||
new_state += state * mx.exp(mx.expand_dims(delta, -1) * A)
|
||||
y = (new_state @ mx.expand_dims(C, -1)).squeeze(2)
|
||||
y = y + D * x
|
||||
return y, new_state
|
||||
|
||||
def _process_sequence(self, x, conv_cache, state_cache):
|
||||
B, T, D = x.shape
|
||||
xz = self.in_proj(x)
|
||||
x, z = xz.split(indices_or_sections=2, axis=-1)
|
||||
|
||||
conv_out, new_conv_cache = self.conv1d(x, conv_cache)
|
||||
x = nn.silu(conv_out)
|
||||
|
||||
A = -mx.exp(self.A_log)
|
||||
|
||||
outputs = []
|
||||
current_state = state_cache
|
||||
y = []
|
||||
for t in range(T):
|
||||
y_t, current_state = self.ssm_step(x[:, t], A, current_state)
|
||||
y.append(y_t)
|
||||
y = mx.stack(y, axis=1)
|
||||
z = self.out_proj(nn.silu(z) * y)
|
||||
return z, (new_conv_cache, current_state)
|
||||
|
||||
def __call__(self, x, cache):
|
||||
if cache is None:
|
||||
conv_cache, state_cache = None, None
|
||||
else:
|
||||
conv_cache, state_cache = cache[0], cache[1]
|
||||
|
||||
output, (new_conv_cache, new_state_cache) = self._process_sequence(
|
||||
x, conv_cache, state_cache
|
||||
)
|
||||
|
||||
if isinstance(cache, MambaCache):
|
||||
cache[0] = new_conv_cache
|
||||
cache[1] = new_state_cache
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.mixer = MambaBlock(args)
|
||||
self.norm = nn.RMSNorm(args.hidden_size)
|
||||
|
||||
def __call__(self, x: mx.array, cache):
|
||||
return self.mixer(self.norm(x), cache) + x
|
||||
|
||||
|
||||
class Mamba(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [ResidualBlock(args) for _ in range(args.num_hidden_layers)]
|
||||
self.norm_f = nn.RMSNorm(args.hidden_size)
|
||||
|
||||
def __call__(self, x: mx.array, cache):
|
||||
x = self.embeddings(x)
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
for layer, c in zip(self.layers, cache):
|
||||
x = layer(x, c)
|
||||
return self.norm_f(x)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.backbone = Mamba(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(self, inputs: mx.array, cache=None):
|
||||
B, T = inputs.shape
|
||||
|
||||
x = self.backbone(inputs, cache)
|
||||
|
||||
if self.args.tie_word_embeddings:
|
||||
logits = self.backbone.embeddings.as_linear(x)
|
||||
else:
|
||||
logits = self.lm_head(x)
|
||||
|
||||
return logits
|
||||
|
||||
def sanitize(self, weights):
|
||||
for k, v in weights.items():
|
||||
if "conv1d.weight" in k and v.shape[-1] != 1:
|
||||
weights[k] = v.moveaxis(2, 1)
|
||||
return weights
|
||||
|
||||
def make_cache(self):
|
||||
return [MambaCache() for _ in range(len(self.layers))]
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.backbone.layers
|
||||
@@ -1,196 +0,0 @@
|
||||
# Copyright © 2023-2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
max_position_embeddings: int = 32768
|
||||
rope_theta: float = 10000.0
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = False
|
||||
num_nextn_predict_layers: int = 2
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
traditional=args.rope_traditional,
|
||||
scaling_config=args.rope_scaling,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class MiMoModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
self.num_nextn_predict_layers = args.num_nextn_predict_layers
|
||||
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
h = self.norm(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = MiMoModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
|
||||
return {
|
||||
k: v
|
||||
for k, v in weights.items()
|
||||
if "self_attn.rotary_emb.inv_freq" not in k
|
||||
and not k.startswith("model.mtp_layers.")
|
||||
}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+33
-23
@@ -1,13 +1,11 @@
|
||||
# Copyright © 2023-2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -23,7 +21,6 @@ class ModelArgs(BaseModelArgs):
|
||||
num_key_value_heads: int
|
||||
scale_depth: float
|
||||
scale_emb: float
|
||||
max_position_embeddings: Optional[int] = None
|
||||
rope_theta: float = 1000000.0
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[str, float]]] = None
|
||||
@@ -69,19 +66,24 @@ class Attention(nn.Module):
|
||||
self.num_heads * self.head_dim, self.hidden_size, bias=False
|
||||
)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
args.rope_traditional,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
rope_scale = (
|
||||
1 / args.rope_scaling["factor"]
|
||||
if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
|
||||
else 1
|
||||
)
|
||||
|
||||
self.rope = nn.RoPE(
|
||||
dims=self.head_dim,
|
||||
traditional=args.rope_traditional,
|
||||
base=self.rope_theta,
|
||||
scale=rope_scale,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
):
|
||||
B, L, _ = x.shape
|
||||
|
||||
@@ -101,8 +103,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
attn_output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
attn_output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
@@ -131,12 +133,12 @@ class DecoderLayer(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r * (self.scale_depth / self.num_hidden_layers**0.5)
|
||||
h = x + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r * (self.scale_depth / self.num_hidden_layers**0.5)
|
||||
out = h + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
|
||||
return out
|
||||
|
||||
|
||||
@@ -154,13 +156,14 @@ class MiniCPMModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs) * self.args.scale_emb
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -184,10 +187,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
|
||||
if not self.args.tie_word_embeddings:
|
||||
out = self.lm_head(out / (self.args.hidden_size / self.args.dim_model_base))
|
||||
@@ -204,3 +206,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,250 +0,0 @@
|
||||
# Copyright © 2023-2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import SuScaledRoPE
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
dim_model_base: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
q_lora_rank: int
|
||||
qk_nope_head_dim: int
|
||||
qk_rope_head_dim: int
|
||||
kv_lora_rank: int
|
||||
scale_depth: float
|
||||
scale_emb: float
|
||||
max_position_embeddings: int
|
||||
attention_bias: bool = False
|
||||
rope_theta: float = 1000000.0
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[str, float]]] = None
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.qk_rope_head_dim = self.args.qk_rope_head_dim
|
||||
self.qk_nope_head_dim = self.args.qk_nope_head_dim
|
||||
self.attention_bias = self.args.attention_bias
|
||||
self.kv_lora_rank = self.args.kv_lora_rank
|
||||
self.num_heads = self.args.num_attention_heads
|
||||
self.q_lora_rank = self.args.q_lora_rank
|
||||
self.hidden_size = self.args.hidden_size
|
||||
|
||||
self.v_head_dim = self.hidden_size // self.args.num_attention_heads
|
||||
self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
|
||||
self.softmax_scale = self.q_head_dim ** (-0.5)
|
||||
|
||||
self.q_a_proj = nn.Linear(
|
||||
self.hidden_size, self.q_lora_rank, bias=self.attention_bias
|
||||
)
|
||||
self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
|
||||
|
||||
self.q_b_proj = nn.Linear(
|
||||
self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
|
||||
)
|
||||
|
||||
self.kv_a_proj_with_mqa = nn.Linear(
|
||||
self.hidden_size,
|
||||
self.kv_lora_rank + self.qk_rope_head_dim,
|
||||
bias=self.attention_bias,
|
||||
)
|
||||
|
||||
self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
|
||||
|
||||
self.kv_b_proj = nn.Linear(
|
||||
self.kv_lora_rank,
|
||||
self.num_heads
|
||||
* (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.o_proj = nn.Linear(
|
||||
self.num_heads * self.v_head_dim,
|
||||
self.hidden_size,
|
||||
bias=self.attention_bias,
|
||||
)
|
||||
|
||||
self.rope = SuScaledRoPE(
|
||||
dims=args.qk_rope_head_dim,
|
||||
base=args.rope_theta,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
original_max_position_embeddings=args.rope_scaling.get(
|
||||
"original_max_position_embeddings", 4096
|
||||
),
|
||||
short_factor=args.rope_scaling.get("short_factor", 1.0),
|
||||
long_factor=args.rope_scaling.get("long_factor", 1.0),
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Dict[str, mx.array]] = None,
|
||||
):
|
||||
B, L, _ = x.shape
|
||||
|
||||
# Project query
|
||||
q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
|
||||
q = q.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
|
||||
q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
|
||||
|
||||
# Project key and value
|
||||
compressed_kv = self.kv_a_proj_with_mqa(x)
|
||||
compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
|
||||
k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
|
||||
|
||||
kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
|
||||
kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
|
||||
|
||||
# Apply RoPE to the query and key parts that need position embedding
|
||||
if cache is not None:
|
||||
q_pe = self.rope(q_pe, offset=cache.offset)
|
||||
k_pe = self.rope(k_pe, offset=cache.offset)
|
||||
else:
|
||||
q_pe = self.rope(q_pe)
|
||||
k_pe = self.rope(k_pe)
|
||||
|
||||
# Create the full query and key tensors by combining the parts
|
||||
# Broadcast k_pe to all heads
|
||||
k_pe_broadcasted = mx.broadcast_to(
|
||||
k_pe, (B, self.num_heads, L, self.qk_rope_head_dim)
|
||||
)
|
||||
|
||||
# Use concatenate for queries
|
||||
queries = mx.concatenate([q_nope, q_pe], axis=-1)
|
||||
|
||||
# Use concatenate for keys
|
||||
keys = mx.concatenate([k_nope, k_pe_broadcasted], axis=-1)
|
||||
|
||||
# Update cache if needed
|
||||
if cache is not None:
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
|
||||
# Perform attention
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.softmax_scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.hidden_size = args.hidden_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
self.scale_depth = args.scale_depth
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r * (self.scale_depth / (self.num_hidden_layers**0.5))
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r * (self.scale_depth / (self.num_hidden_layers**0.5))
|
||||
return out
|
||||
|
||||
|
||||
class MiniCPM3Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
assert self.vocab_size > 0
|
||||
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [DecoderLayer(args) for _ in range(args.num_hidden_layers)]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs) * self.args.scale_emb
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = MiniCPM3Model(args)
|
||||
|
||||
if not self.args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
|
||||
if not self.args.tie_word_embeddings:
|
||||
out = self.lm_head(out / (self.args.hidden_size / self.args.dim_model_base))
|
||||
else:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,49 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from . import llama
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
text_config: dict
|
||||
|
||||
def __post_init__(self):
|
||||
self.text_config["tie_word_embeddings"] = False
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask: Optional[mx.array] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
return self.language_model(
|
||||
inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
|
||||
)
|
||||
|
||||
def sanitize(self, weights):
|
||||
weights = tree_unflatten(list(weights.items()))
|
||||
weights.pop("vision_tower", None)
|
||||
weights.pop("multi_modal_projector", None)
|
||||
return dict(tree_flatten(weights))
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.language_model.model.layers
|
||||
+21
-13
@@ -1,12 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@@ -65,7 +64,7 @@ class MixtralAttention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -86,8 +85,8 @@ class MixtralAttention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
@@ -137,7 +136,7 @@ class MixtralDecoderLayer(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -161,13 +160,15 @@ class MixtralModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
T = h.shape[1]
|
||||
if T > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(T)
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -189,10 +190,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
@@ -217,3 +217,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,385 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AttentionConfig:
|
||||
no_op: bool = False
|
||||
replace_with_linear: bool = False
|
||||
sparsify: Optional[list[str]] = None
|
||||
n_heads_in_group: Optional[int] = None # GQA group size
|
||||
window_length: Optional[int] = None # Not directly used here, placeholder
|
||||
num_sink_tokens: Optional[int] = None # Not directly used here, placeholder
|
||||
use_prefill_window_in_sink_attention: bool = (
|
||||
False # Not directly used here, placeholder
|
||||
)
|
||||
unshifted_sink: bool = False # Not directly used here, placeholder
|
||||
|
||||
def __post_init__(self):
|
||||
# Ensure consistency: If no-op or linear, other attn params are irrelevant
|
||||
if self.no_op or self.replace_with_linear:
|
||||
# Use object.__setattr__ because the dataclass is frozen
|
||||
object.__setattr__(self, "n_heads_in_group", None)
|
||||
object.__setattr__(self, "window_length", None)
|
||||
object.__setattr__(self, "num_sink_tokens", None)
|
||||
# If it's a standard attention block, n_heads_in_group must be provided
|
||||
elif not self.no_op:
|
||||
if self.n_heads_in_group is None:
|
||||
raise ValueError(
|
||||
"n_heads_in_group must be specified for active attention blocks"
|
||||
)
|
||||
if self.n_heads_in_group <= 0:
|
||||
raise ValueError(
|
||||
f"n_heads_in_group must be positive, got {self.n_heads_in_group}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FFNConfig:
|
||||
no_op: bool = False
|
||||
replace_with_linear: bool = False
|
||||
sparsify: Optional[list[str]] = None
|
||||
ffn_mult: Optional[float] = None
|
||||
|
||||
def __post_init__(self):
|
||||
# Ensure consistency: If no-op or linear, ffn_mult is irrelevant
|
||||
if self.no_op or self.replace_with_linear:
|
||||
object.__setattr__(self, "ffn_mult", None)
|
||||
# If it's a standard FFN block, ffn_mult must be provided
|
||||
elif not self.no_op:
|
||||
if self.ffn_mult is None:
|
||||
raise ValueError("ffn_mult must be specified for active FFN blocks")
|
||||
# Round to prevent potential floating point inconsistencies if needed
|
||||
object.__setattr__(self, "ffn_mult", round(self.ffn_mult, 6))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BlockConfig:
|
||||
attention: AttentionConfig
|
||||
ffn: FFNConfig
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict):
|
||||
# Helper to create BlockConfig from a dictionary (e.g., loaded from JSON)
|
||||
attn_conf = AttentionConfig(**data.get("attention", {}))
|
||||
ffn_conf = FFNConfig(**data.get("ffn", {}))
|
||||
return cls(attention=attn_conf, ffn=ffn_conf)
|
||||
|
||||
|
||||
def _find_multiple(n: int, k: int) -> int:
|
||||
"""Finds the smallest multiple of k greater than or equal to n."""
|
||||
if n % k == 0:
|
||||
return n
|
||||
return n + k - (n % k)
|
||||
|
||||
|
||||
def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int:
|
||||
"""Calculates intermediate size based on multiplier, rounding up to multiple of 256."""
|
||||
intermediate_size = int(2 * ffn_mult * n_embd / 3)
|
||||
return _find_multiple(intermediate_size, 256)
|
||||
|
||||
|
||||
# Activation function mapping
|
||||
_ACT2FN = {
|
||||
"silu": nn.silu,
|
||||
"relu": nn.relu,
|
||||
"gelu": nn.gelu,
|
||||
"gelu_new": nn.gelu_approx,
|
||||
"gelu_fast": nn.gelu_approx,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str = "nemotron-nas"
|
||||
hidden_size: int = 8192
|
||||
num_hidden_layers: int = 80
|
||||
num_attention_heads: int = 64
|
||||
rms_norm_eps: float = 1e-5
|
||||
vocab_size: int = 128256
|
||||
block_configs: list = field(default_factory=list) # List of BlockConfig or dicts
|
||||
hidden_act: str = "silu"
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
rope_theta: float = 500000.0
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
max_position_embeddings: int = 131072
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
# Automatically parse block_configs if they are loaded as dicts
|
||||
if self.block_configs and isinstance(self.block_configs[0], dict):
|
||||
self.block_configs = [
|
||||
BlockConfig.from_dict(conf) for conf in self.block_configs
|
||||
]
|
||||
|
||||
if len(self.block_configs) != self.num_hidden_layers:
|
||||
raise ValueError(
|
||||
f"Number of block_configs ({len(self.block_configs)}) must match "
|
||||
f"num_hidden_layers ({self.num_hidden_layers})"
|
||||
)
|
||||
|
||||
# Basic validation for RoPE scaling if provided
|
||||
if self.rope_scaling:
|
||||
if "factor" not in self.rope_scaling:
|
||||
raise ValueError("rope_scaling must contain 'factor'")
|
||||
rope_type = self.rope_scaling.get("rope_type")
|
||||
if rope_type is None:
|
||||
raise ValueError("rope_scaling must contain 'rope_type'")
|
||||
|
||||
# Validate individual block configs (post_init in dataclasses already does some)
|
||||
for i, block_conf in enumerate(self.block_configs):
|
||||
attn_conf = block_conf.attention
|
||||
if not attn_conf.no_op and not attn_conf.replace_with_linear:
|
||||
if self.num_attention_heads % attn_conf.n_heads_in_group != 0:
|
||||
raise ValueError(
|
||||
f"Layer {i}: num_attention_heads ({self.num_attention_heads}) "
|
||||
f"must be divisible by n_heads_in_group ({attn_conf.n_heads_in_group})"
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""Standard GQA Attention mechanism for layers that use it."""
|
||||
|
||||
def __init__(self, args: ModelArgs, attention_config: AttentionConfig):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = n_heads // attention_config.n_heads_in_group
|
||||
|
||||
self.head_dim = head_dim = args.hidden_size // n_heads
|
||||
if (self.head_dim * n_heads) != dim:
|
||||
raise ValueError(
|
||||
f"hidden_size ({dim}) must be divisible by num_attention_heads ({n_heads})"
|
||||
)
|
||||
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
|
||||
|
||||
# Initialize RoPE based on global config
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
False, # Llama uses traditional=False
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, self.head_dim).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""Standard Feed-Forward Network for layers that use it."""
|
||||
|
||||
def __init__(self, args: ModelArgs, ffn_config: FFNConfig):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
# Calculate intermediate dim based on layer's specific config
|
||||
hidden_dim = _ffn_mult_to_intermediate_size(ffn_config.ffn_mult, dim)
|
||||
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
|
||||
|
||||
try:
|
||||
self.act_fn = _ACT2FN[args.hidden_act]
|
||||
except KeyError:
|
||||
raise ValueError(f"Unknown activation function: {args.hidden_act}")
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class LinearSubblockReplacement(nn.Module):
|
||||
"""A simple linear layer used to replace Attention or MLP blocks."""
|
||||
|
||||
def __init__(self, hidden_size: int, bias: bool):
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(hidden_size, hidden_size, bias=bias)
|
||||
|
||||
def __call__(self, x: mx.array, *args, **kwargs) -> mx.array:
|
||||
# Accepts potential extra args (like mask, cache) but ignores them
|
||||
return self.linear(x)
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
"""A single transformer block, potentially heterogeneous based on config."""
|
||||
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
# Get the specific configuration for this layer
|
||||
block_config = args.block_configs[layer_idx]
|
||||
self.attention_config = block_config.attention
|
||||
self.ffn_config = block_config.ffn
|
||||
|
||||
# Conditionally initialize Input LayerNorm (needed unless Attention is no-op)
|
||||
if not self.attention_config.no_op:
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
else:
|
||||
self.input_layernorm = None
|
||||
|
||||
# Conditionally initialize Attention block
|
||||
if self.attention_config.no_op:
|
||||
self.self_attn = None
|
||||
elif self.attention_config.replace_with_linear:
|
||||
self.self_attn = LinearSubblockReplacement(
|
||||
args.hidden_size, args.attention_bias
|
||||
)
|
||||
else:
|
||||
# Standard attention for this layer
|
||||
self.self_attn = Attention(args, self.attention_config)
|
||||
|
||||
# Conditionally initialize Post-Attention LayerNorm (needed unless FFN is no-op)
|
||||
if not self.ffn_config.no_op:
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
else:
|
||||
self.post_attention_layernorm = None
|
||||
|
||||
# Conditionally initialize MLP block
|
||||
if self.ffn_config.no_op:
|
||||
self.mlp = None
|
||||
elif self.ffn_config.replace_with_linear:
|
||||
self.mlp = LinearSubblockReplacement(args.hidden_size, args.mlp_bias)
|
||||
else:
|
||||
# Standard MLP for this layer
|
||||
self.mlp = MLP(args, self.ffn_config)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
|
||||
# Attention part (Input Norm -> Attention -> Residual)
|
||||
if self.self_attn is not None:
|
||||
residual = x
|
||||
h = self.input_layernorm(x)
|
||||
attn_out = self.self_attn(h, mask=mask, cache=cache)
|
||||
x = residual + attn_out
|
||||
|
||||
# MLP part (Post-Attention Norm -> MLP -> Residual)
|
||||
if self.mlp is not None:
|
||||
residual = x
|
||||
h = self.post_attention_layernorm(x)
|
||||
mlp_out = self.mlp(h)
|
||||
x = residual + mlp_out
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class NemotronNASModel(nn.Module):
|
||||
"""The core Nemotron-NAS style transformer model."""
|
||||
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args, layer_idx=i)
|
||||
for i in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[List[Any]] = None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
h = layer(h, mask, cache=cache[i])
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = NemotronNASModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
else:
|
||||
self.lm_head = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask=None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask=mask, cache=cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,220 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
hidden_act: str
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
head_dim: Optional[int] = None
|
||||
max_position_embeddings: Optional[int] = None
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
partial_rotary_factor: float = 0.5
|
||||
rope_theta: float = 10000.0
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
if self.rope_scaling:
|
||||
if not "factor" in self.rope_scaling:
|
||||
raise ValueError(f"rope_scaling must contain 'factor'")
|
||||
rope_type = self.rope_scaling.get("type") or self.rope_scaling.get(
|
||||
"rope_type"
|
||||
)
|
||||
if rope_type is None:
|
||||
raise ValueError(
|
||||
f"rope_scaling must contain either 'type' or 'rope_type'"
|
||||
)
|
||||
if rope_type not in ["linear"]:
|
||||
raise ValueError("rope_scaling 'type' currently only supports 'linear'")
|
||||
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def relu_squared(x):
|
||||
return nn.relu(x).square()
|
||||
|
||||
|
||||
class NemotronLayerNorm1P(nn.LayerNorm):
|
||||
def __call__(self, x):
|
||||
weight = self.weight + 1 if "weight" in self else None
|
||||
bias = self.bias if "bias" in self else None
|
||||
return mx.fast.layer_norm(x, weight, bias, self.eps)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
self.partial_rotary_factor = args.partial_rotary_factor
|
||||
|
||||
self.scale = head_dim**-0.5
|
||||
if hasattr(args, "attention_bias"):
|
||||
attention_bias = args.attention_bias
|
||||
else:
|
||||
attention_bias = False
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
|
||||
|
||||
rope_scale = 1.0
|
||||
if args.rope_scaling and args.rope_scaling["type"] == "linear":
|
||||
assert isinstance(args.rope_scaling["factor"], float)
|
||||
rope_scale = 1 / args.rope_scaling["factor"]
|
||||
self.rope = nn.RoPE(
|
||||
int(self.partial_rotary_factor * self.head_dim),
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, _ = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
hidden_dim = args.intermediate_size
|
||||
mlp_bias = args.mlp_bias
|
||||
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(relu_squared(self.up_proj(x)))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.input_layernorm = NemotronLayerNorm1P(args.hidden_size, eps=args.norm_eps)
|
||||
self.post_attention_layernorm = NemotronLayerNorm1P(
|
||||
args.hidden_size, eps=args.norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class NemotronModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = NemotronLayerNorm1P(args.hidden_size, eps=args.norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = NemotronModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+20
-15
@@ -1,19 +1,17 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from sys import exit
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask
|
||||
from .base import BaseModelArgs
|
||||
|
||||
try:
|
||||
import hf_olmo
|
||||
except ImportError:
|
||||
print("To run olmo install ai2-olmo: pip install ai2-olmo")
|
||||
sys.exit(1)
|
||||
exit(1)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -68,7 +66,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -98,7 +96,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.attend(self.att_norm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -124,13 +122,14 @@ class Transformer(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.wte(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.blocks)
|
||||
@@ -154,10 +153,9 @@ class OlmoModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
return self.transformer(inputs, mask, cache)
|
||||
return self.transformer(inputs, cache)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
@@ -170,11 +168,18 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
return self.model(inputs, mask, cache)
|
||||
return self.model(inputs, cache)
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.transformer.blocks
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.d_model // self.args.n_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.n_heads
|
||||
|
||||
@@ -1,212 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
head_dim: Optional[int] = None
|
||||
max_position_embeddings: Optional[int] = None
|
||||
num_key_value_heads: Optional[int] = None
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
if self.num_key_value_heads is None:
|
||||
self.num_key_value_heads = self.num_attention_heads
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
|
||||
self.scale = head_dim**-0.5
|
||||
if hasattr(args, "attention_bias"):
|
||||
attention_bias = args.attention_bias
|
||||
else:
|
||||
attention_bias = False
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
args.rope_traditional,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
)
|
||||
|
||||
self.q_norm = nn.RMSNorm(n_heads * head_dim, args.rms_norm_eps)
|
||||
self.k_norm = nn.RMSNorm(n_kv_heads * head_dim, args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
queries = self.q_norm(queries)
|
||||
keys = self.k_norm(keys)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
hidden_dim = args.intermediate_size
|
||||
if hasattr(args, "mlp_bias"):
|
||||
mlp_bias = args.mlp_bias
|
||||
else:
|
||||
mlp_bias = False
|
||||
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.post_feedforward_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.post_attention_layernorm(self.self_attn(x, mask, cache))
|
||||
h = x + r
|
||||
r = self.post_feedforward_layernorm(self.mlp(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class LlamaModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = LlamaModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
out = self.model(inputs, cache, mask)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
return {
|
||||
k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
|
||||
}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,217 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_experts: int
|
||||
num_experts_per_tok: int
|
||||
norm_topk_prob: bool = False
|
||||
head_dim: Optional[int] = None
|
||||
max_position_embeddings: Optional[int] = None
|
||||
num_key_value_heads: Optional[int] = None
|
||||
attention_bias: bool = False
|
||||
mlp_bias: bool = False
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
if self.num_key_value_heads is None:
|
||||
self.num_key_value_heads = self.num_attention_heads
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
|
||||
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
self.head_dim,
|
||||
args.rope_theta,
|
||||
args.rope_traditional,
|
||||
args.rope_scaling,
|
||||
args.max_position_embeddings,
|
||||
)
|
||||
|
||||
self.q_norm = nn.RMSNorm(n_heads * head_dim, args.rms_norm_eps)
|
||||
self.k_norm = nn.RMSNorm(n_kv_heads * head_dim, args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
queries = self.q_norm(queries)
|
||||
keys = self.k_norm(keys)
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class OlmoeSparseMoeBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_experts = args.num_experts
|
||||
self.top_k = args.num_experts_per_tok
|
||||
self.norm_topk_prob = args.norm_topk_prob
|
||||
|
||||
self.gate = nn.Linear(args.hidden_size, self.num_experts, bias=False)
|
||||
self.switch_mlp = SwitchGLU(
|
||||
args.hidden_size,
|
||||
args.intermediate_size,
|
||||
self.num_experts,
|
||||
bias=args.mlp_bias,
|
||||
)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
x_flat = x.reshape(-1, D)
|
||||
router_logits = self.gate(x_flat)
|
||||
routing_weights = mx.softmax(router_logits, axis=1, precise=True)
|
||||
k = self.top_k
|
||||
indices = mx.stop_gradient(
|
||||
mx.argpartition(-routing_weights, kth=k - 1, axis=-1)[..., :k]
|
||||
)
|
||||
scores = mx.take_along_axis(routing_weights, indices, axis=-1)
|
||||
if self.norm_topk_prob:
|
||||
scores = scores / scores.sum(axis=-1, keepdims=True)
|
||||
y = self.switch_mlp(x_flat, indices)
|
||||
y = (y * scores[..., None]).sum(axis=-2)
|
||||
return y.reshape(B, L, D)
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = OlmoeSparseMoeBlock(args)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
x = x + self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
x = x + self.mlp(self.post_attention_layernorm(x))
|
||||
return x
|
||||
|
||||
|
||||
class OlmoeModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = OlmoeModel(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
out = self.model(inputs, cache, mask)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
|
||||
return weights
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for n in ["up_proj", "down_proj", "gate_proj"]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
|
||||
for e in range(self.args.num_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{n}.{k}"] = mx.stack(to_join)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+19
-13
@@ -1,12 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -80,7 +78,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -107,8 +105,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
@@ -152,7 +150,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.attn(self.attn_norm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -178,13 +176,14 @@ class OpenELMModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.token_embeddings(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -207,10 +206,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.transformer(inputs, mask, cache)
|
||||
out = self.transformer(inputs, cache)
|
||||
if self.args.share_input_output_layers:
|
||||
out = self.transformer.token_embeddings.as_linear(out)
|
||||
else:
|
||||
@@ -221,3 +219,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.transformer.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.head_dim
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_kv_heads
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_additive_causal_mask
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParamsArgs(BaseModelArgs):
|
||||
dim: int
|
||||
ffn_type: str
|
||||
n_heads: int
|
||||
n_layers: int
|
||||
norm_eps: float
|
||||
positional_embedding_type: str
|
||||
post_embed_norm: bool
|
||||
qk_norm: bool
|
||||
vocab_size: int
|
||||
weight_tying: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
params_args_dict: ParamsArgs
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
self.dim = args.dim
|
||||
self.n_heads = args.n_heads
|
||||
self.head_dim = self.dim // self.n_heads
|
||||
self.qk_norm = args.qk_norm
|
||||
self.scale = self.head_dim**-0.5
|
||||
|
||||
self.in_proj = nn.Linear(self.dim, 3 * self.dim, bias=False)
|
||||
self.out_proj = nn.Linear(self.dim, self.dim, bias=False)
|
||||
if self.qk_norm:
|
||||
self.q_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
|
||||
self.k_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
|
||||
self.rope = nn.RoPE(
|
||||
self.head_dim,
|
||||
traditional=False,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.in_proj(x).split(3, axis=-1)
|
||||
|
||||
if self.qk_norm:
|
||||
queries = self.q_norm(queries)
|
||||
keys = self.q_norm(keys)
|
||||
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.out_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
# https://github.com/mlfoundations/open_lm/blob/c65b43042ff31c0fe26f930decf1ccab1b03ab4b/open_lm/model.py#L254C2-L254C3
|
||||
hidden_dim = 256 * ((int(2 * 4 * args.dim / 3) + 256 - 1) // 256)
|
||||
self.w12 = nn.Linear(args.dim, 2 * hidden_dim, bias=False)
|
||||
self.w3 = nn.Linear(hidden_dim, args.dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
gate, x = self.w12(x).split(2, axis=-1)
|
||||
return self.w3(nn.silu(gate) * x)
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.attention = Attention(args)
|
||||
self.feed_forward = MLP(args)
|
||||
self.ffn_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
|
||||
self.attention_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.attention(self.attention_norm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.feed_forward(self.ffn_norm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class OpenLM(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)
|
||||
self.layers = [TransformerBlock(args=args) for _ in range(args.n_layers)]
|
||||
self.norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
|
||||
self.output = nn.Linear(args.dim, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
):
|
||||
_, L = inputs.shape
|
||||
|
||||
h = self.tok_embeddings(inputs)
|
||||
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = create_additive_causal_mask(
|
||||
h.shape[1], cache[0].offset if cache is not None else 0
|
||||
)
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, cache=c)
|
||||
|
||||
return self.output(self.norm(h))
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
args.params_args_dict = ParamsArgs.from_dict(args.params_args_dict)
|
||||
self.args = args.params_args_dict
|
||||
self.model_type = args.model_type
|
||||
self.model = OpenLM(self.args)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, cache)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
return {k: v for k, v in weights.items() if "inv_freq" not in k}
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.dim // self.args.n_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.n_heads
|
||||
+23
-20
@@ -1,12 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -92,13 +91,8 @@ class PhiAttention(nn.Module):
|
||||
keys = self.rope(keys)
|
||||
|
||||
scale = math.sqrt(1 / queries.shape[-1])
|
||||
output = scaled_dot_product_attention(
|
||||
queries.astype(mx.float32),
|
||||
keys,
|
||||
values,
|
||||
cache=cache,
|
||||
scale=scale,
|
||||
mask=mask,
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries.astype(mx.float32), keys, values, scale=scale, mask=mask
|
||||
).astype(values.dtype)
|
||||
|
||||
output = output.moveaxis(2, 1).reshape(B, L, -1)
|
||||
@@ -111,9 +105,10 @@ class PhiMLP(nn.Module):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
|
||||
self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
|
||||
self.act = nn.GELU(approx="precise")
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.fc2(nn.gelu_approx(self.fc1(x)))
|
||||
return self.fc2(self.act(self.fc1(x)))
|
||||
|
||||
|
||||
class PhiDecoderLayer(nn.Module):
|
||||
@@ -141,15 +136,16 @@ class PhiModel(nn.Module):
|
||||
config.hidden_size, eps=config.layer_norm_eps
|
||||
)
|
||||
|
||||
def __call__(self, x, mask, cache):
|
||||
def __call__(self, x, cache):
|
||||
x = self.embed_tokens(x)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(x, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
mask = None
|
||||
if x.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
|
||||
mask = mask.astype(x.dtype)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
x = layer(x, mask, c)
|
||||
return self.final_layernorm(x)
|
||||
@@ -166,12 +162,19 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
y = self.model(x, mask, cache)
|
||||
cache: mx.array = None,
|
||||
) -> Tuple[mx.array, mx.array]:
|
||||
y = self.model(x, cache)
|
||||
return self.lm_head(y)
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
+33
-35
@@ -1,13 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import SuScaledRoPE
|
||||
from .base import BaseModelArgs
|
||||
from .su_rope import SuScaledRotaryEmbedding
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -19,14 +17,12 @@ class ModelArgs(BaseModelArgs):
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: Optional[int] = None
|
||||
num_key_value_heads: int = None
|
||||
rope_theta: float = 10000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, List[float]]]] = None
|
||||
partial_rotary_factor: float = 1.0
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
max_position_embeddings: int = 131072
|
||||
original_max_position_embeddings: int = 4096
|
||||
tie_word_embeddings: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
if self.num_key_value_heads is None:
|
||||
@@ -37,9 +33,9 @@ class ModelArgs(BaseModelArgs):
|
||||
if not all(key in self.rope_scaling for key in required_keys):
|
||||
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
||||
|
||||
if self.rope_scaling["type"] not in ["longrope", "su", "linear"]:
|
||||
if self.rope_scaling["type"] not in ["su", "linear"]:
|
||||
print(
|
||||
"[WARNING] rope_scaling 'type' currently only supports 'linear', 'su', and 'longrope'; setting rope scaling to false."
|
||||
"[WARNING] rope_scaling 'type' currently only supports 'linear' and 'su'; setting rope scaling to false."
|
||||
)
|
||||
self.rope_scaling = None
|
||||
|
||||
@@ -50,7 +46,6 @@ class Attention(nn.Module):
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
|
||||
@@ -61,23 +56,23 @@ class Attention(nn.Module):
|
||||
self.qkv_proj = nn.Linear(dim, op_size, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
rope_dim = int(head_dim * args.partial_rotary_factor)
|
||||
if args.rope_scaling and args.rope_scaling["type"] in ["longrope", "su"]:
|
||||
self.rope = SuScaledRoPE(
|
||||
rope_dim,
|
||||
rope_scale = 1.0
|
||||
if args.rope_scaling and args.rope_scaling["type"] == "su":
|
||||
self.rope = SuScaledRotaryEmbedding(
|
||||
head_dim,
|
||||
traditional=False,
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
original_max_position_embeddings=args.original_max_position_embeddings,
|
||||
short_factor=args.rope_scaling["short_factor"],
|
||||
long_factor=args.rope_scaling["long_factor"],
|
||||
)
|
||||
else:
|
||||
rope_scale = 1.0
|
||||
if args.rope_scaling and args.rope_scaling["type"] == "linear":
|
||||
assert isinstance(args.rope_scaling["factor"], float)
|
||||
rope_scale = 1 / args.rope_scaling["factor"]
|
||||
self.rope = nn.RoPE(
|
||||
rope_dim,
|
||||
head_dim,
|
||||
traditional=args.rope_traditional,
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
@@ -87,7 +82,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -110,8 +105,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
@@ -146,7 +141,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -171,13 +166,14 @@ class Phi3Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -193,23 +189,25 @@ class Model(nn.Module):
|
||||
super().__init__()
|
||||
self.model_type = args.model_type
|
||||
self.model = Phi3Model(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
out = self.model(inputs, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
+21
-16
@@ -1,14 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, Optional
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -22,14 +19,14 @@ class ModelArgs(BaseModelArgs):
|
||||
num_attention_heads: int
|
||||
layer_norm_epsilon: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
num_key_value_heads: int = None
|
||||
mup_attn_multiplier: float = 1.0
|
||||
mup_use_scaling: bool = True
|
||||
mup_embedding_multiplier: float = 10.0
|
||||
mup_width_multiplier: float = 8.0
|
||||
rope_embedding_base: float = 1000000
|
||||
rope_position_scale: float = 1.0
|
||||
blocksparse_block_size: int = 64
|
||||
blocksparse_block_size: int = (64,)
|
||||
blocksparse_num_local_blocks: int = 16
|
||||
blocksparse_vert_stride: int = 8
|
||||
|
||||
@@ -160,7 +157,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -188,8 +185,8 @@ class Attention(nn.Module):
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
else:
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.dense(output)
|
||||
@@ -229,7 +226,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -258,15 +255,16 @@ class Phi3Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
if self.mup_embedding_multiplier:
|
||||
h = self.mup_embedding_multiplier * h
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache, return_array=True)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -292,10 +290,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
if self.mup_width_multiplier:
|
||||
out = out / self.mup_width_multiplier
|
||||
@@ -306,8 +303,16 @@ class Model(nn.Module):
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
def sanitize(self, weights):
|
||||
# Remove unused precomputed rotary freqs
|
||||
return {
|
||||
k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
|
||||
}
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,214 +0,0 @@
|
||||
# Copyright © 2024 Apple Inc.
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import SuScaledRoPE
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str = "phimoe"
|
||||
vocab_size: int = 32064
|
||||
hidden_size: int = 4096
|
||||
intermediate_size: int = 6400
|
||||
num_hidden_layers: int = 32
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 8
|
||||
max_position_embeddings: int = 131072
|
||||
original_max_position_embeddings: int = 4096
|
||||
rms_norm_eps: float = 1e-6
|
||||
rope_scaling: Dict[str, Union[float, List[float]]] = None
|
||||
num_local_experts: int = 16
|
||||
num_experts_per_tok: int = 2
|
||||
rope_theta: float = 10000.0
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=True)
|
||||
|
||||
self.rope = SuScaledRoPE(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
original_max_position_embeddings=args.original_max_position_embeddings,
|
||||
short_factor=args.rope_scaling["short_factor"],
|
||||
long_factor=args.rope_scaling["long_factor"],
|
||||
short_mscale=args.rope_scaling["short_mscale"],
|
||||
long_mscale=args.rope_scaling["long_mscale"],
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class PhiMoESparseMoeBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.hidden_dim = args.hidden_size
|
||||
self.ffn_dim = args.intermediate_size
|
||||
self.num_experts = args.num_local_experts
|
||||
self.top_k = args.num_experts_per_tok
|
||||
|
||||
self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
|
||||
self.switch_mlp = SwitchGLU(self.hidden_dim, self.ffn_dim, self.num_experts)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
gates = self.gate(x)
|
||||
|
||||
k = self.top_k
|
||||
inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
|
||||
scores = mx.take_along_axis(gates, inds, axis=-1)
|
||||
scores = mx.softmax(scores, axis=-1, precise=True)
|
||||
|
||||
y = self.switch_mlp(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class PhiMoEDecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
|
||||
self.self_attn = Attention(args)
|
||||
self.block_sparse_moe = PhiMoESparseMoeBlock(args)
|
||||
self.input_layernorm = nn.LayerNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.LayerNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
residual = x
|
||||
hidden_states = self.input_layernorm(x)
|
||||
hidden_states = self.self_attn(hidden_states, mask=mask, cache=cache)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.block_sparse_moe(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class PhiMoEModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [PhiMoEDecoderLayer(args) for _ in range(args.num_hidden_layers)]
|
||||
self.norm = nn.LayerNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.model_type = args.model_type
|
||||
self.args = args
|
||||
self.model = PhiMoEModel(args)
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=True)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
if "model.layers.0.block_sparse_moe.experts.0.w1.weight" not in weights:
|
||||
return weights
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
|
||||
for k in ["weight", "scales", "biases"]:
|
||||
if f"{prefix}.block_sparse_moe.experts.0.{n}.{k}" in weights:
|
||||
to_join = [
|
||||
weights.pop(
|
||||
f"{prefix}.block_sparse_moe.experts.{e}.{n}.{k}"
|
||||
)
|
||||
for e in range(self.args.num_local_experts)
|
||||
]
|
||||
weights[f"{prefix}.block_sparse_moe.switch_mlp.{m}.{k}"] = (
|
||||
mx.stack(to_join)
|
||||
)
|
||||
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
+16
-15
@@ -1,5 +1,3 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import inspect
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
@@ -8,7 +6,6 @@ from typing import Tuple
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import create_attention_mask, scaled_dot_product_attention
|
||||
from .switch_layers import SwitchMLP
|
||||
|
||||
|
||||
@@ -71,13 +68,8 @@ class RoPEAttention(nn.Module):
|
||||
# Finally perform the attention computation
|
||||
scale = math.sqrt(1 / queries.shape[-1])
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries.astype(mx.float32),
|
||||
keys,
|
||||
values,
|
||||
cache=cache,
|
||||
scale=scale,
|
||||
mask=mask,
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries.astype(mx.float32), keys, values, scale=scale, mask=mask
|
||||
).astype(values.dtype)
|
||||
output = output.moveaxis(2, 1).reshape(B, L, -1)
|
||||
|
||||
@@ -173,11 +165,12 @@ class Model(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(x, cache)
|
||||
cache: mx.array = None,
|
||||
) -> Tuple[mx.array, mx.array]:
|
||||
mask = None
|
||||
if x.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
|
||||
mask = mask.astype(x.dtype)
|
||||
|
||||
y = self.transformer(x, mask, cache)
|
||||
return self.lm_head(y)
|
||||
@@ -200,3 +193,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.transformer.h
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.model_dim // self.args.num_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_heads
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_unflatten
|
||||
|
||||
from . import llama
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
text_config: dict
|
||||
|
||||
def __post_init__(self):
|
||||
self.text_config["tie_word_embeddings"] = False
|
||||
self.text_config["num_attention_heads"] = self.text_config.get(
|
||||
"num_attention_heads", 32
|
||||
)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache=None,
|
||||
mask: Optional[mx.array] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
return self.language_model(
|
||||
inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
|
||||
)
|
||||
|
||||
def sanitize(self, weights):
|
||||
weights = tree_unflatten(list(weights.items()))
|
||||
weights.pop("vision_tower", None)
|
||||
weights.pop("multi_modal_projector", None)
|
||||
return dict(tree_flatten(weights))
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.language_model.model.layers
|
||||
+24
-22
@@ -1,13 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Any, List, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -62,8 +60,8 @@ class Attention(nn.Module):
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
attention_mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> Tuple[mx.array, Tuple[mx.array, mx.array]]:
|
||||
bsz, q_len, _ = hidden_states.shape
|
||||
|
||||
queries = self.q_proj(hidden_states)
|
||||
@@ -89,14 +87,10 @@ class Attention(nn.Module):
|
||||
queries = self.rotary_emb(queries)
|
||||
keys = self.rotary_emb(keys)
|
||||
|
||||
keys = mx.tile(keys, [1, self.config.n_shared_head, 1, 1])
|
||||
values = mx.tile(values, [1, self.config.n_shared_head, 1, 1])
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
cache=cache,
|
||||
scale=self.scale,
|
||||
mask=attention_mask,
|
||||
)
|
||||
@@ -131,8 +125,8 @@ class PlamoDecoderLayer(nn.Module):
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
attention_mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
):
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> Tuple[Any, ...]:
|
||||
# from LlamaDecoder
|
||||
residual = hidden_states
|
||||
|
||||
@@ -173,13 +167,14 @@ class PlamoModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
cache: Optional[List[Union[Tuple[mx.array, mx.array], None]]] = None,
|
||||
) -> Tuple[mx.array, Optional[List[Union[Tuple[mx.array, mx.array], None]]]]:
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(self.embed_tokens.weight.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None for _ in range(len(self.layers.layers))]
|
||||
@@ -203,12 +198,19 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> mx.array:
|
||||
out = self.model(inputs, cache, mask)
|
||||
cache: Optional[List[Tuple[mx.array, mx.array]]] = None,
|
||||
) -> Tuple[mx.array, mx.array]:
|
||||
out = self.model(inputs, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_attention_heads // self.args.n_shared_head
|
||||
|
||||
@@ -1,599 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from mlx_lm.models.base import BaseModelArgs, create_attention_mask
|
||||
|
||||
from .cache import KVCache, MambaCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str = "plamo2"
|
||||
hidden_size: int = 4096
|
||||
num_hidden_layers: int = 32
|
||||
rms_norm_eps: float = 1e-6
|
||||
tie_word_embeddings: bool = True
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 4
|
||||
hidden_size_per_head: int = 128
|
||||
max_position_embeddings: int = 2048
|
||||
attention_window_size: int = 2048
|
||||
full_attention_idx: Optional[list[int]] = None
|
||||
mamba_d_state: int = 64
|
||||
mamba_d_conv: int = 4
|
||||
mamba_num_heads: int = 64
|
||||
mamba_step: int = 2
|
||||
mamba_chunk_size: int = 256
|
||||
mamba_enabled: bool = True
|
||||
intermediate_size: int = 13312
|
||||
vocab_size: int = 32000
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
eps: float = 1e-6,
|
||||
offset: float = 1.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.weight = mx.zeros(hidden_size)
|
||||
self.variance_epsilon = eps
|
||||
self.offset = offset
|
||||
|
||||
def __call__(self, hidden_states: mx.array) -> mx.array:
|
||||
return mx.fast.rms_norm(
|
||||
hidden_states, self.weight + self.offset, self.variance_epsilon
|
||||
)
|
||||
|
||||
|
||||
def get_initial_dt_bias(num_heads: int) -> mx.array:
|
||||
dt_min = 0.001
|
||||
dt_max = 0.1
|
||||
dt = mx.exp(
|
||||
mx.random.uniform(shape=(num_heads,)) * (math.log(dt_max) - math.log(dt_min))
|
||||
+ math.log(dt_min)
|
||||
)
|
||||
dt = mx.clip(dt, a_min=1e-4, a_max=None)
|
||||
inv_dt = dt + mx.log(-mx.expm1(-dt))
|
||||
return inv_dt
|
||||
|
||||
|
||||
def get_initial_A(num_heads: int) -> mx.array:
|
||||
A = mx.arange(1, num_heads + 1, dtype=mx.float32)
|
||||
return mx.log(A)
|
||||
|
||||
|
||||
# From: https://github.com/state-spaces/mamba/blob/0cce0fa645f100f00620ddf2333c2b7712abfdec/mamba_ssm/ops/triton/selective_state_update.py#L219
|
||||
def selective_state_update_ref(
|
||||
state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
|
||||
) -> tuple[mx.array, mx.array]:
|
||||
"""
|
||||
Argument:
|
||||
state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
|
||||
x: (batch, dim) or (batch, nheads, dim)
|
||||
dt: (batch, dim) or (batch, nheads, dim)
|
||||
A: (dim, dstate) or (nheads, dim, dstate)
|
||||
B: (batch, dstate) or (batch, ngroups, dstate)
|
||||
C: (batch, dstate) or (batch, ngroups, dstate)
|
||||
D: (dim,) or (nheads, dim)
|
||||
z: (batch, dim) or (batch, nheads, dim)
|
||||
dt_bias: (dim,) or (nheads, dim)
|
||||
Return:
|
||||
out: (batch, dim) or (batch, nheads, dim)
|
||||
"""
|
||||
has_heads = state.ndim > 3
|
||||
if state.ndim == 3:
|
||||
state = mx.expand_dims(state, 1)
|
||||
if x.ndim == 2:
|
||||
x = mx.expand_dims(x, 1)
|
||||
if dt.ndim == 2:
|
||||
dt = mx.expand_dims(dt, 1)
|
||||
if A.ndim == 2:
|
||||
A = mx.expand_dims(A, 0)
|
||||
if B.ndim == 2:
|
||||
B = mx.expand_dims(B, 1)
|
||||
if C.ndim == 2:
|
||||
C = mx.expand_dims(C, 1)
|
||||
if D is not None and D.ndim == 1:
|
||||
D = mx.expand_dims(D, 0)
|
||||
if z is not None and z.ndim == 2:
|
||||
z = mx.expand_dims(z, 1)
|
||||
if dt_bias is not None and dt_bias.ndim == 1:
|
||||
dt_bias = mx.expand_dims(dt_bias, 0)
|
||||
batch, nheads, dim, dstate = state.shape
|
||||
assert x.shape == (batch, nheads, dim)
|
||||
assert dt.shape == x.shape
|
||||
assert A.shape == (nheads, dim, dstate)
|
||||
ngroups = B.shape[1]
|
||||
assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
|
||||
assert B.shape == (batch, ngroups, dstate)
|
||||
assert C.shape == B.shape
|
||||
if D is not None:
|
||||
assert D.shape == (nheads, dim)
|
||||
if z is not None:
|
||||
assert z.shape == x.shape
|
||||
if dt_bias is not None:
|
||||
assert dt_bias.shape == (nheads, dim)
|
||||
dt = dt + dt_bias
|
||||
dt = nn.softplus(dt) if dt_softplus else dt
|
||||
dA = mx.exp(mx.expand_dims(dt, axis=-1) * A) # (batch, nheads, dim, dstate)
|
||||
B = mx.reshape(
|
||||
mx.repeat(mx.expand_dims(B, axis=2), nheads // ngroups, 2),
|
||||
(batch, nheads, dstate),
|
||||
) # (batch, nheads, dstate)
|
||||
C = mx.reshape(
|
||||
mx.repeat(mx.expand_dims(C, axis=2), nheads // ngroups, 2),
|
||||
(batch, nheads, dstate),
|
||||
) # (batch, nheads, dstate)
|
||||
dB = mx.expand_dims(dt, axis=-1) * mx.expand_dims(
|
||||
B, axis=-2
|
||||
) # (batch, nheads, dim, dstate)
|
||||
state = state * dA + dB * mx.expand_dims(x, axis=-1) # (batch, dim, dstate)
|
||||
out = mx.einsum("bhdn,bhn->bhd", state.astype(C.dtype), C)
|
||||
if D is not None:
|
||||
out += (x * D).astype(out.dtype)
|
||||
out = (out if z is None else out * nn.silu(z)).astype(x.dtype)
|
||||
if not has_heads:
|
||||
out = out.squeeze(1)
|
||||
return out, state
|
||||
|
||||
|
||||
def ssd_update_state(
|
||||
ssm_state: mx.array,
|
||||
x: mx.array,
|
||||
dt: mx.array,
|
||||
A: mx.array,
|
||||
B: mx.array,
|
||||
C: mx.array,
|
||||
D: mx.array,
|
||||
z: mx.array,
|
||||
dt_bias: mx.array,
|
||||
dt_softplus: bool,
|
||||
) -> tuple[mx.array, mx.array]:
|
||||
assert ssm_state.dtype == mx.float32
|
||||
dtype = x.dtype
|
||||
|
||||
hidden_size_per_head = x.shape[-1]
|
||||
d_state = B.shape[-1]
|
||||
A = mx.broadcast_to(
|
||||
A[:, None, None], (A.shape[0], hidden_size_per_head, d_state)
|
||||
).astype(mx.float32)
|
||||
dt = mx.broadcast_to(
|
||||
dt[..., None], (dt.shape[0], dt.shape[1], hidden_size_per_head)
|
||||
)
|
||||
dt_bias = mx.broadcast_to(
|
||||
dt_bias[:, None], (dt_bias.shape[0], hidden_size_per_head)
|
||||
)
|
||||
D = mx.broadcast_to(D[:, None], (D.shape[0], hidden_size_per_head))
|
||||
out, ssm_state = selective_state_update_ref(
|
||||
ssm_state,
|
||||
x.astype(dtype),
|
||||
dt.astype(dtype),
|
||||
A.astype(mx.float32),
|
||||
B.astype(dtype),
|
||||
C.astype(dtype),
|
||||
D.astype(mx.float32),
|
||||
z.astype(dtype),
|
||||
dt_bias.astype(mx.float32),
|
||||
dt_softplus=dt_softplus,
|
||||
)
|
||||
return out[:, None], ssm_state
|
||||
|
||||
|
||||
def ssd_chunk_scan_combined(
|
||||
x: mx.array,
|
||||
dt: mx.array,
|
||||
A: mx.array,
|
||||
B: mx.array,
|
||||
C: mx.array,
|
||||
D: mx.array,
|
||||
z: mx.array,
|
||||
dt_bias: mx.array,
|
||||
dt_softplus: bool,
|
||||
ssm_state: mx.array,
|
||||
) -> tuple[mx.array, mx.array]:
|
||||
assert ssm_state.dtype == mx.float32
|
||||
length = x.shape[1]
|
||||
ys = []
|
||||
for i in range(length):
|
||||
y, ssm_state = ssd_update_state(
|
||||
ssm_state,
|
||||
x[:, i],
|
||||
dt[:, i],
|
||||
A,
|
||||
B[:, i],
|
||||
C[:, i],
|
||||
D if D.ndim == 1 else D[:, i],
|
||||
z=z[:, i],
|
||||
dt_bias=dt_bias,
|
||||
dt_softplus=dt_softplus,
|
||||
)
|
||||
ys.append(y)
|
||||
return mx.concatenate(ys, axis=1), ssm_state
|
||||
|
||||
|
||||
def causal_conv1d_update(conv_state, x, weight) -> tuple[mx.array, mx.array]:
|
||||
_, seqlen, dim = x.shape
|
||||
state_len = conv_state.shape[-2]
|
||||
x = mx.concatenate([conv_state, x], axis=-2)
|
||||
conv_state = x[:, -state_len:]
|
||||
out = mx.conv1d(
|
||||
x,
|
||||
weight,
|
||||
padding=0,
|
||||
groups=dim,
|
||||
)[:, -seqlen:]
|
||||
return nn.silu(out), conv_state
|
||||
|
||||
|
||||
class Mamba(nn.Module):
|
||||
def __init__(self, config: ModelArgs) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.d_state = config.mamba_d_state
|
||||
self.d_conv = config.mamba_d_conv
|
||||
self.chunk_size = config.mamba_chunk_size
|
||||
self.num_heads = config.mamba_num_heads
|
||||
self.hidden_size_per_head = config.hidden_size_per_head
|
||||
|
||||
self.intermediate_size = self.num_heads * self.hidden_size_per_head
|
||||
|
||||
self.in_proj = nn.Linear(
|
||||
self.hidden_size, 2 * self.intermediate_size, bias=False
|
||||
)
|
||||
self.conv1d = nn.Conv1d(
|
||||
in_channels=self.intermediate_size,
|
||||
out_channels=self.intermediate_size,
|
||||
bias=False,
|
||||
kernel_size=self.d_conv,
|
||||
groups=self.intermediate_size,
|
||||
padding=0,
|
||||
)
|
||||
self.dt_dim = max(64, self.hidden_size // 16)
|
||||
self.bcdt_proj = nn.Linear(
|
||||
self.intermediate_size,
|
||||
self.dt_dim + 2 * self.d_state,
|
||||
bias=False,
|
||||
)
|
||||
self.dt_proj = nn.Linear(self.dt_dim, self.num_heads, bias=False)
|
||||
|
||||
self.dt_bias = get_initial_dt_bias(self.num_heads)
|
||||
self.A_log = get_initial_A(self.num_heads)
|
||||
self.D = mx.ones(self.num_heads, dtype=mx.float32)
|
||||
|
||||
self.dt_norm_weight = mx.ones(self.dt_dim)
|
||||
self.B_norm_weight = mx.ones(self.d_state)
|
||||
self.C_norm_weight = mx.ones(self.d_state)
|
||||
|
||||
self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache=None,
|
||||
):
|
||||
bsize, length, _ = hidden_states.shape
|
||||
|
||||
if cache is not None and cache[0] is not None:
|
||||
conv_state = cache[0]
|
||||
ssm_state = cache[1]
|
||||
else:
|
||||
conv_state = mx.zeros(
|
||||
(bsize, self.d_conv - 1, self.intermediate_size),
|
||||
dtype=hidden_states.dtype,
|
||||
)
|
||||
ssm_state = mx.zeros(
|
||||
(bsize, self.num_heads, self.hidden_size_per_head, self.d_state),
|
||||
dtype=mx.float32,
|
||||
)
|
||||
|
||||
zx = self.in_proj(hidden_states)
|
||||
zx = zx.reshape(bsize, length, self.num_heads, -1)
|
||||
# z: (bsize, length, num_heads, hidden_size_per_head)
|
||||
# x: (bsize, length, num_heads, hidden_size_per_head)
|
||||
z, x = mx.split(
|
||||
zx,
|
||||
[
|
||||
self.hidden_size_per_head,
|
||||
],
|
||||
axis=-1,
|
||||
)
|
||||
|
||||
x = x.reshape(bsize, -1, self.num_heads * self.hidden_size_per_head)
|
||||
x, conv_state = causal_conv1d_update(conv_state, x, self.conv1d.weight)
|
||||
BCdt = self.bcdt_proj(x)
|
||||
x = x.reshape(bsize, length, self.num_heads, -1)
|
||||
B, C, dt = mx.split(BCdt, [self.d_state, self.d_state * 2], axis=-1)
|
||||
|
||||
A = -mx.exp(self.A_log.astype(mx.float32)) # (num_heads,)
|
||||
dt = mx.fast.rms_norm(dt, self.dt_norm_weight, self.config.rms_norm_eps)
|
||||
B = mx.fast.rms_norm(B, self.B_norm_weight, self.config.rms_norm_eps)
|
||||
C = mx.fast.rms_norm(C, self.C_norm_weight, self.config.rms_norm_eps)
|
||||
|
||||
# (bsize, length, num_heads, 1)
|
||||
dt = self.dt_proj(dt)[..., None]
|
||||
|
||||
out, ssm_state = ssd_chunk_scan_combined(
|
||||
x,
|
||||
dt.reshape(bsize, length, -1),
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
D=self.D,
|
||||
z=z,
|
||||
dt_bias=self.dt_bias,
|
||||
dt_softplus=True,
|
||||
ssm_state=ssm_state,
|
||||
)
|
||||
|
||||
if cache is not None:
|
||||
cache[0] = conv_state
|
||||
cache[1] = ssm_state
|
||||
y = self.out_proj(out.reshape(bsize, length, -1))
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, config: ModelArgs) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
head_dim = config.hidden_size_per_head
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_num_heads = config.num_attention_heads
|
||||
self.qk_dim = self.v_dim = head_dim
|
||||
self.k_num_heads = self.v_num_heads = config.num_key_value_heads
|
||||
assert self.q_num_heads % self.k_num_heads == 0
|
||||
self.n_group = self.q_num_heads // self.k_num_heads
|
||||
|
||||
self.q_proj_dim = self.q_num_heads * self.qk_dim
|
||||
self.k_proj_dim = self.k_num_heads * self.qk_dim
|
||||
self.v_proj_dim = self.k_num_heads * self.v_dim
|
||||
self.qkv_proj = nn.Linear(
|
||||
self.hidden_size,
|
||||
self.q_proj_dim + self.k_proj_dim + self.v_proj_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.o_proj = nn.Linear(
|
||||
self.q_num_heads * self.v_dim, self.hidden_size, bias=False
|
||||
)
|
||||
|
||||
self.q_weight = mx.ones((self.q_num_heads, self.qk_dim))
|
||||
self.k_weight = mx.ones((self.k_num_heads, self.qk_dim))
|
||||
|
||||
self.rope = nn.RoPE(self.qk_dim)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache=None,
|
||||
):
|
||||
B, T, _ = hidden_states.shape
|
||||
|
||||
qkv = self.qkv_proj(hidden_states)
|
||||
q, k, v = mx.split(
|
||||
qkv, [self.q_proj_dim, self.q_proj_dim + self.k_proj_dim], axis=-1
|
||||
)
|
||||
q = q.reshape(B, T, self.q_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
|
||||
k = k.reshape(B, T, self.k_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
|
||||
v = v.reshape(B, T, self.v_num_heads, self.v_dim).transpose(0, 2, 1, 3)
|
||||
|
||||
q = mx.fast.rms_norm(q, weight=None, eps=1e-6) * self.q_weight[:, None]
|
||||
k = mx.fast.rms_norm(k, weight=None, eps=1e-6) * self.k_weight[:, None]
|
||||
|
||||
if cache is not None:
|
||||
q = self.rope(q, offset=cache.offset)
|
||||
k = self.rope(k, offset=cache.offset)
|
||||
k, v = cache.update_and_fetch(k, v)
|
||||
else:
|
||||
q = self.rope(q)
|
||||
k = self.rope(k)
|
||||
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
scale=self.scale,
|
||||
mask=mask,
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(
|
||||
B, T, self.q_num_heads * self.v_dim
|
||||
)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: ModelArgs) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
self.gate_up_proj = nn.Linear(
|
||||
self.hidden_size, self.intermediate_size * 2, bias=False
|
||||
)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
|
||||
def __call__(self, x: mx.array) -> mx.array:
|
||||
h = self.gate_up_proj(x)
|
||||
hs = mx.split(h, 2, axis=-1)
|
||||
return self.down_proj(nn.silu(hs[0]) * hs[1])
|
||||
|
||||
|
||||
class PlamoDecoderLayer(nn.Module):
|
||||
def __init__(self, config: ModelArgs, is_mamba: bool) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.is_mamba = is_mamba
|
||||
self.mixer: nn.Module
|
||||
if is_mamba:
|
||||
self.mixer = Mamba(config)
|
||||
else:
|
||||
self.mixer = Attention(config)
|
||||
self.mlp = MLP(config)
|
||||
self.pre_mixer_norm = RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps, offset=1.0
|
||||
)
|
||||
self.post_mixer_norm = RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / 5
|
||||
)
|
||||
self.pre_mlp_norm = RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps, offset=1.0
|
||||
)
|
||||
self.post_mlp_norm = RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / (5**1.5)
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache=None,
|
||||
):
|
||||
residual = hidden_states
|
||||
hidden_states = self.pre_mixer_norm(hidden_states)
|
||||
|
||||
hidden_states_sa = self.mixer(
|
||||
hidden_states=hidden_states,
|
||||
mask=mask,
|
||||
cache=cache,
|
||||
)
|
||||
|
||||
hidden_states_sa = self.post_mixer_norm(hidden_states_sa)
|
||||
hidden_states = residual + hidden_states_sa
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.pre_mlp_norm(hidden_states)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states_mlp = self.mlp(hidden_states)
|
||||
|
||||
# Residual
|
||||
hidden_states_mlp = self.post_mlp_norm(hidden_states_mlp)
|
||||
return residual + hidden_states_mlp
|
||||
|
||||
|
||||
def is_mamba(config: ModelArgs, i: int) -> bool:
|
||||
if not config.mamba_enabled:
|
||||
return False
|
||||
assert config.mamba_step > 1
|
||||
assert i < config.num_hidden_layers
|
||||
|
||||
if config.num_hidden_layers <= (config.mamba_step // 2):
|
||||
# use attention in last layer
|
||||
return i != config.num_hidden_layers - 1
|
||||
return (i % config.mamba_step) != (config.mamba_step // 2)
|
||||
|
||||
|
||||
class PlamoDecoder(nn.Module):
|
||||
def __init__(self, config: ModelArgs) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.layers = [
|
||||
PlamoDecoderLayer(config, is_mamba=is_mamba(config, i))
|
||||
for i in range(config.num_hidden_layers)
|
||||
]
|
||||
|
||||
def __call__(self, x: mx.array, mask: mx.array, cache):
|
||||
for i, decoder_layer in enumerate(self.layers):
|
||||
x = decoder_layer(
|
||||
x,
|
||||
mask=mask,
|
||||
cache=cache[i],
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
class PlamoModel(nn.Module):
|
||||
def __init__(self, config: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = PlamoDecoder(config) # type: ignore
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache=None,
|
||||
):
|
||||
batch_size, seq_length = inputs.shape
|
||||
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, [cache[1]] if cache is not None else None)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers.layers)
|
||||
|
||||
# decoder layers
|
||||
out = self.layers(
|
||||
h,
|
||||
mask,
|
||||
cache,
|
||||
)
|
||||
|
||||
return self.norm(out)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, config: ModelArgs) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.model_type = config.model_type
|
||||
self.model = PlamoModel(config)
|
||||
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
if not config.tie_word_embeddings:
|
||||
self.lm_head: nn.Module = nn.Linear(
|
||||
config.hidden_size, self.vocab_size, bias=False
|
||||
)
|
||||
|
||||
def sanitize(self, weights: dict[Any, Any]) -> dict[Any, Any]:
|
||||
for k, v in weights.items():
|
||||
if "conv1d.weight" in k and v.shape[-1] != 1:
|
||||
weights[k] = v.moveaxis(2, 1)
|
||||
return weights
|
||||
|
||||
def make_cache(self):
|
||||
# TODO use RotatingKVCache is not full_attn
|
||||
# full_attn = self.layer_idx in self.config.full_attention_idx
|
||||
return [MambaCache() if l.is_mamba else KVCache() for l in self.layers]
|
||||
|
||||
def __call__(
|
||||
self, inputs: mx.array, mask: Optional[mx.array] = None, cache=None
|
||||
) -> mx.array:
|
||||
outputs = self.model(
|
||||
inputs=inputs,
|
||||
mask=None,
|
||||
cache=cache,
|
||||
)
|
||||
if self.config.tie_word_embeddings:
|
||||
logits = self.model.embed_tokens.as_linear(outputs)
|
||||
else:
|
||||
logits = self.lm_head(outputs)
|
||||
|
||||
return logits
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers.layers
|
||||
+19
-9
@@ -1,11 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -64,8 +63,8 @@ class Attention(nn.Module):
|
||||
queries = self.rotary_emb(queries)
|
||||
keys = self.rotary_emb(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
|
||||
@@ -123,8 +122,11 @@ class QwenModel(nn.Module):
|
||||
def __call__(self, inputs, mask=None, cache=None):
|
||||
x = self.wte(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(x, cache)
|
||||
mask = None
|
||||
T = x.shape[1]
|
||||
if T > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(T)
|
||||
mask = mask.astype(x.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.h)
|
||||
@@ -149,11 +151,19 @@ class Model(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
cache: mx.array = None,
|
||||
) -> Tuple[mx.array, mx.array]:
|
||||
y = self.transformer(x, mask, cache)
|
||||
return self.lm_head(y)
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.transformer.h
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_attention_heads
|
||||
|
||||
+41
-27
@@ -1,13 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -19,13 +16,24 @@ class ModelArgs(BaseModelArgs):
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
max_position_embeddings: int = 32768
|
||||
num_key_value_heads: int = None
|
||||
rope_theta: float = 1000000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
tie_word_embeddings: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
if self.num_key_value_heads is None:
|
||||
self.num_key_value_heads = self.num_attention_heads
|
||||
|
||||
if self.rope_scaling:
|
||||
required_keys = {"factor", "type"}
|
||||
if not all(key in self.rope_scaling for key in required_keys):
|
||||
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
||||
|
||||
if self.rope_scaling["type"] != "linear":
|
||||
raise ValueError("rope_scaling 'type' currently only supports 'linear'")
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
@@ -33,7 +41,6 @@ class Attention(nn.Module):
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
@@ -44,19 +51,23 @@ class Attention(nn.Module):
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.rope = initialize_rope(
|
||||
rope_scale = (
|
||||
1 / args.rope_scaling["factor"]
|
||||
if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
|
||||
else 1
|
||||
)
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
traditional=args.rope_traditional,
|
||||
scaling_config=args.rope_scaling,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
base=args.rope_theta,
|
||||
scale=rope_scale,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -75,8 +86,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
@@ -110,7 +121,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -135,17 +146,14 @@ class Qwen2Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
if input_embeddings is not None:
|
||||
h = input_embeddings
|
||||
else:
|
||||
h = self.embed_tokens(inputs)
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -168,11 +176,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache, input_embeddings)
|
||||
out = self.model(inputs, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
@@ -190,3 +196,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
+21
-15
@@ -1,12 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@@ -23,7 +22,7 @@ class ModelArgs(BaseModelArgs):
|
||||
shared_expert_intermediate_size: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: Optional[int] = None
|
||||
num_key_value_heads: int = None
|
||||
rope_theta: float = 1000000
|
||||
rope_traditional: bool = False
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
@@ -48,7 +47,6 @@ class Attention(nn.Module):
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.hidden_size // n_heads
|
||||
@@ -69,7 +67,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -88,8 +86,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
@@ -161,7 +159,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -186,13 +184,14 @@ class Qwen2MoeModel(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -214,10 +213,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
@@ -238,3 +236,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -1,189 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .rope_utils import initialize_rope
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
max_position_embeddings: int
|
||||
rope_theta: float
|
||||
head_dim: int
|
||||
tie_word_embeddings: bool
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = args.head_dim
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
|
||||
self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
|
||||
self.rope = initialize_rope(
|
||||
head_dim,
|
||||
base=args.rope_theta,
|
||||
traditional=False,
|
||||
scaling_config=args.rope_scaling,
|
||||
max_position_embeddings=args.max_position_embeddings,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.num_attention_heads = args.num_attention_heads
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args)
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.args = args
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class Qwen3Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = Qwen3Model(args)
|
||||
if not args.tie_word_embeddings:
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
out = self.lm_head(out)
|
||||
return out
|
||||
|
||||
def sanitize(self, weights):
|
||||
if self.args.tie_word_embeddings:
|
||||
weights.pop("lm_head.weight", None)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,240 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
num_experts: int
|
||||
num_experts_per_tok: int
|
||||
decoder_sparse_step: int
|
||||
mlp_only_layers: List[int]
|
||||
moe_intermediate_size: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
head_dim: int
|
||||
rope_theta: float
|
||||
tie_word_embeddings: bool
|
||||
max_position_embeddings: int
|
||||
norm_topk_prob: bool
|
||||
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
|
||||
dim = args.hidden_size
|
||||
self.n_heads = n_heads = args.num_attention_heads
|
||||
assert args.num_key_value_heads is not None
|
||||
self.n_kv_heads = n_kv_heads = args.num_key_value_heads
|
||||
|
||||
head_dim = getattr(
|
||||
args, "head_dim", args.hidden_size // args.num_attention_heads
|
||||
)
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
|
||||
self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
|
||||
self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
|
||||
|
||||
self.rope = nn.RoPE(
|
||||
head_dim,
|
||||
traditional=False,
|
||||
base=args.rope_theta,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
# Prepare the queries, keys and values for the attention computation
|
||||
queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim, hidden_dim):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def __call__(self, x) -> mx.array:
|
||||
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class Qwen3MoeSparseMoeBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
dim = args.hidden_size
|
||||
intermediate_size = args.moe_intermediate_size
|
||||
|
||||
self.num_experts = num_experts = args.num_experts
|
||||
self.top_k = args.num_experts_per_tok
|
||||
self.norm_topk_prob = args.norm_topk_prob
|
||||
|
||||
self.gate = nn.Linear(dim, num_experts, bias=False)
|
||||
self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
):
|
||||
gates = self.gate(x)
|
||||
gates = mx.softmax(gates, axis=-1, precise=True)
|
||||
|
||||
k = self.top_k
|
||||
inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
|
||||
scores = mx.take_along_axis(gates, inds, axis=-1)
|
||||
if self.norm_topk_prob:
|
||||
scores /= mx.sum(scores, axis=-1, keepdims=True)
|
||||
|
||||
y = self.switch_mlp(x, inds)
|
||||
y = (y * scores[..., None]).sum(axis=-2)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class Qwen3MoeDecoderLayer(nn.Module):
|
||||
def __init__(self, args: ModelArgs, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = args.hidden_size
|
||||
self.self_attn = Attention(args, layer_idx)
|
||||
|
||||
self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(
|
||||
args.hidden_size, eps=args.rms_norm_eps
|
||||
)
|
||||
self.args = args
|
||||
|
||||
if (layer_idx not in args.mlp_only_layers) and (
|
||||
args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0
|
||||
):
|
||||
self.mlp = Qwen3MoeSparseMoeBlock(args)
|
||||
else:
|
||||
self.mlp = MLP(args.hidden_size, args.intermediate_size)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
r = self.mlp(self.post_attention_layernorm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class Qwen3MoeModel(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.num_hidden_layers = args.num_hidden_layers
|
||||
assert self.vocab_size > 0
|
||||
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
||||
self.layers = [
|
||||
Qwen3MoeDecoderLayer(args=args, layer_idx=i)
|
||||
for i in range(args.num_hidden_layers)
|
||||
]
|
||||
self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for layer, c in zip(self.layers, cache):
|
||||
h = layer(h, mask, c)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.model_type = args.model_type
|
||||
self.model = Qwen3MoeModel(args)
|
||||
self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
return self.lm_head(out)
|
||||
|
||||
def sanitize(self, weights):
|
||||
if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
|
||||
return weights
|
||||
for l in range(self.args.num_hidden_layers):
|
||||
prefix = f"model.layers.{l}"
|
||||
for n in ["up_proj", "down_proj", "gate_proj"]:
|
||||
if f"{prefix}.mlp.experts.0.{n}.weight" in weights:
|
||||
to_join = [
|
||||
weights.pop(f"{prefix}.mlp.experts.{e}.{n}.weight")
|
||||
for e in range(self.args.num_experts)
|
||||
]
|
||||
weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(to_join)
|
||||
return weights
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
@@ -1,458 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .cache import MambaCache, RotatingKVCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
attention_bias: bool
|
||||
conv1d_width: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
logits_soft_cap: float
|
||||
num_attention_heads: int
|
||||
num_hidden_layers: int
|
||||
num_key_value_heads: int
|
||||
rms_norm_eps: float
|
||||
rope_theta: float
|
||||
attention_window_size: int
|
||||
vocab_size: int
|
||||
embeddings_scale_by_sqrt_dim: bool = True
|
||||
block_types: Optional[List[str]] = None
|
||||
_block_types: Optional[List[str]] = None
|
||||
|
||||
def __post_init__(self):
|
||||
# For some reason these have different names in 2B and 9B
|
||||
if self.block_types is None:
|
||||
self.block_types = self._block_types
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dims: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.weight = mx.ones((dims,))
|
||||
self.eps = eps
|
||||
|
||||
def __call__(self, x):
|
||||
return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
|
||||
|
||||
|
||||
def rnn_scan(x, a, h0):
|
||||
assert x.ndim == 3
|
||||
assert a.shape == x.shape[-a.ndim :]
|
||||
assert a.dtype == x.dtype
|
||||
|
||||
if x.shape[1] == 1:
|
||||
# Using scan in sampling mode.
|
||||
if h0 is None:
|
||||
return x, x[:, 0]
|
||||
|
||||
else:
|
||||
y = a * h0[:, None] + x
|
||||
return y, y[:, -1]
|
||||
|
||||
else:
|
||||
# Using scan in linear mode.
|
||||
if h0 is not None:
|
||||
h_t = h0
|
||||
else:
|
||||
B, _, D = x.shape
|
||||
h_t = mx.zeros((B, D), dtype=x.dtype)
|
||||
|
||||
y = mx.zeros_like(x)
|
||||
for t in range(x.shape[1]):
|
||||
h_t = a[:, t] * h_t + x[:, t]
|
||||
y[:, t] = h_t
|
||||
|
||||
return y, h_t
|
||||
|
||||
|
||||
class Conv1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
kernel_size: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.weight = mx.zeros((channels, kernel_size, 1))
|
||||
self.bias = mx.zeros((channels,))
|
||||
|
||||
def __call__(self, x, cache=None):
|
||||
B, L, C = x.shape
|
||||
groups, K, _ = self.weight.shape
|
||||
|
||||
if cache is not None:
|
||||
x = mx.concatenate([cache, x], axis=1)
|
||||
else:
|
||||
x = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
|
||||
|
||||
y = mx.conv_general(x, self.weight, groups=groups)
|
||||
y = y + self.bias
|
||||
|
||||
return y, x[:, -K + 1 :, :]
|
||||
|
||||
|
||||
class RGLRU(nn.Module):
|
||||
"""A Real-Gated Linear Recurrent Unit (RG-LRU) layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width: int,
|
||||
num_heads: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = self.width // self.num_heads
|
||||
|
||||
self.recurrent_param = mx.zeros((self.width,))
|
||||
|
||||
self.input_gate_weight = mx.zeros(
|
||||
(self.num_heads, self.head_dim, self.head_dim),
|
||||
)
|
||||
self.input_gate_bias = mx.zeros((self.num_heads, self.head_dim))
|
||||
|
||||
self.recurrent_gate_weight = mx.zeros(
|
||||
(self.num_heads, self.head_dim, self.head_dim),
|
||||
)
|
||||
self.recurrent_gate_bias = mx.zeros((self.num_heads, self.head_dim))
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache=None,
|
||||
):
|
||||
B, L, _ = x.shape
|
||||
|
||||
def apply_block_linear(h, w, b):
|
||||
h = h.reshape((B, L, self.num_heads, self.head_dim))
|
||||
h = (h.swapaxes(1, 2) @ w).swapaxes(1, 2) + b
|
||||
return mx.sigmoid(h.flatten(2, 3))
|
||||
|
||||
# Gates for x and a.
|
||||
gate_x = apply_block_linear(x, self.input_gate_weight, self.input_gate_bias)
|
||||
gate_a = apply_block_linear(
|
||||
x, self.recurrent_gate_weight, self.recurrent_gate_bias
|
||||
)
|
||||
|
||||
# Compute the parameter `A` of the recurrence.
|
||||
log_a = -8.0 * gate_a * nn.softplus(self.recurrent_param)
|
||||
a = mx.exp(log_a)
|
||||
a_square = mx.exp(2 * log_a)
|
||||
|
||||
# Gate the input.
|
||||
gated_x = x * gate_x
|
||||
|
||||
# Apply gamma normalization to the input.
|
||||
multiplier = mx.sqrt(1 - a_square)
|
||||
if cache is None:
|
||||
multiplier[:, 0, :] = 1.0
|
||||
normalized_x = gated_x * multiplier.astype(x.dtype)
|
||||
|
||||
y, last_h = rnn_scan(
|
||||
x=normalized_x,
|
||||
a=a,
|
||||
h0=cache,
|
||||
)
|
||||
|
||||
return y, last_h
|
||||
|
||||
|
||||
class RecurrentBlock(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width: int,
|
||||
num_heads: int,
|
||||
lru_width: int = None,
|
||||
conv1d_temporal_width: int = 4,
|
||||
):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.num_heads = num_heads
|
||||
self.lru_width = lru_width or width
|
||||
self.conv1d_temporal_width = conv1d_temporal_width
|
||||
|
||||
self.linear_y = nn.Linear(width, self.lru_width)
|
||||
self.linear_x = nn.Linear(width, self.lru_width)
|
||||
self.linear_out = nn.Linear(self.lru_width, width)
|
||||
self.conv_1d = Conv1d(
|
||||
channels=self.lru_width,
|
||||
kernel_size=self.conv1d_temporal_width,
|
||||
)
|
||||
self.rg_lru = RGLRU(
|
||||
width=self.lru_width,
|
||||
num_heads=self.num_heads,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
# y branch.
|
||||
y = self.linear_y(x)
|
||||
y = nn.gelu_approx(y)
|
||||
|
||||
# x branch.
|
||||
x = self.linear_x(x)
|
||||
if cache is None:
|
||||
cache = [None, None]
|
||||
x, cache[0] = self.conv_1d(x=x, cache=cache[0])
|
||||
x, cache[1] = self.rg_lru(x=x, cache=cache[1])
|
||||
|
||||
x = x * y
|
||||
x = self.linear_out(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class LocalAttentionBlock(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width: int,
|
||||
num_heads: int,
|
||||
window_size: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.num_heads = num_heads
|
||||
self.window_size = window_size
|
||||
self.scale = (width // num_heads) ** (-0.5)
|
||||
|
||||
self.head_dim = self.width // self.num_heads
|
||||
self.q_proj = nn.Linear(self.width, self.width, bias=False)
|
||||
self.k_proj = nn.Linear(self.width, self.head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(self.width, self.head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(self.width, self.width, bias=True)
|
||||
self.rope = nn.RoPE(
|
||||
self.head_dim // 2,
|
||||
traditional=False,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
B, L, D = x.shape
|
||||
|
||||
queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
|
||||
queries = queries.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
|
||||
keys = keys.reshape(B, L, 1, -1).transpose(0, 2, 1, 3)
|
||||
values = values.reshape(B, L, 1, -1).transpose(0, 2, 1, 3)
|
||||
|
||||
if cache is not None:
|
||||
queries = self.rope(queries, offset=cache.offset)
|
||||
keys = self.rope(keys, offset=cache.offset)
|
||||
keys, values = cache.update_and_fetch(keys, values)
|
||||
else:
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
class MLPBlock(nn.Module):
|
||||
|
||||
def __init__(self, width: int, expanded_width: int):
|
||||
super().__init__()
|
||||
self.up_proj = nn.Linear(width, expanded_width // 2)
|
||||
self.gate_proj = nn.Linear(width, expanded_width // 2)
|
||||
self.down_proj = nn.Linear(expanded_width // 2, width)
|
||||
|
||||
def __call__(self, x: mx.array):
|
||||
gate = self.gate_proj(x)
|
||||
x = self.up_proj(x)
|
||||
return self.down_proj(nn.gelu_approx(gate) * x)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width: int,
|
||||
mlp_expanded_width: int,
|
||||
num_heads: int,
|
||||
attention_window_size: int,
|
||||
temporal_block_type: str,
|
||||
lru_width: Optional[int] = None,
|
||||
conv1d_temporal_width: int = 4,
|
||||
):
|
||||
"""Initializes the residual block.
|
||||
|
||||
Args:
|
||||
width: The width of the block.
|
||||
mlp_expanded_width: The width of the expansion inside the MLP block.
|
||||
num_heads: The number of heads for the Attention or the RG-LRU.
|
||||
attention_window_size: The window size for the local attention block.
|
||||
temporal_block_type: Either "recurrent" or "attention", specifying the
|
||||
type of recurrent block to use.
|
||||
lru_width: The width of the RG-LRU if different from `width`.
|
||||
conv1d_temporal_width: The width of the temporal convolution.
|
||||
"""
|
||||
super().__init__()
|
||||
self.width = width
|
||||
self.mlp_expanded_width = mlp_expanded_width
|
||||
self.num_heads = num_heads
|
||||
self.attention_window_size = attention_window_size
|
||||
self.temporal_block_type = temporal_block_type
|
||||
self.lru_width = lru_width
|
||||
self.conv1d_temporal_width = conv1d_temporal_width
|
||||
|
||||
self.temporal_pre_norm = RMSNorm(width)
|
||||
if self.temporal_block_type == "recurrent":
|
||||
self.temporal_block = RecurrentBlock(
|
||||
width=self.width,
|
||||
num_heads=self.num_heads,
|
||||
lru_width=self.lru_width,
|
||||
conv1d_temporal_width=self.conv1d_temporal_width,
|
||||
)
|
||||
|
||||
else:
|
||||
self.temporal_block = LocalAttentionBlock(
|
||||
width=self.width,
|
||||
num_heads=self.num_heads,
|
||||
window_size=self.attention_window_size,
|
||||
)
|
||||
|
||||
self.channel_pre_norm = RMSNorm(width)
|
||||
self.mlp_block = MLPBlock(
|
||||
width=self.width,
|
||||
expanded_width=self.mlp_expanded_width,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache=None,
|
||||
mask=None,
|
||||
):
|
||||
raw_x = x
|
||||
|
||||
inputs_normalized = self.temporal_pre_norm(raw_x)
|
||||
|
||||
x = self.temporal_block(inputs_normalized, cache=cache, mask=mask)
|
||||
residual = x + raw_x
|
||||
|
||||
x = self.channel_pre_norm(residual)
|
||||
x = self.mlp_block(x)
|
||||
|
||||
x = x + residual
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Griffin(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
self.embed_tokens = nn.Embedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
)
|
||||
|
||||
self.scale_by_sqrt_dim = config.embeddings_scale_by_sqrt_dim
|
||||
block_types = config.block_types
|
||||
|
||||
self.layers = [
|
||||
ResidualBlock(
|
||||
width=config.hidden_size,
|
||||
mlp_expanded_width=config.intermediate_size,
|
||||
num_heads=config.num_attention_heads,
|
||||
attention_window_size=config.attention_window_size,
|
||||
temporal_block_type=block_types[i % len(block_types)],
|
||||
lru_width=None,
|
||||
)
|
||||
for i in range(config.num_hidden_layers)
|
||||
]
|
||||
self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
tokens,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
x = self.embed_tokens(tokens)
|
||||
if self.scale_by_sqrt_dim:
|
||||
x = x * math.sqrt(x.shape[-1])
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
|
||||
for i, block in enumerate(self.layers):
|
||||
if block.temporal_block_type != "recurrent":
|
||||
mask_cache = [cache[i]]
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(x, mask_cache)
|
||||
|
||||
for i, block in enumerate(self.layers):
|
||||
x = block(x, mask=mask, cache=cache[i])
|
||||
|
||||
return self.final_norm(x)
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
|
||||
def __init__(self, config):
|
||||
self.args = config
|
||||
self.model = Griffin(config)
|
||||
self.model_type = config.model_type
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
|
||||
def __call__(self, tokens: mx.array, mask: mx.array = None, cache=None) -> mx.array:
|
||||
"""
|
||||
Args:
|
||||
tokens: Sequence of input tokens.
|
||||
"""
|
||||
logits = self.model(tokens, mask=mask, cache=cache)
|
||||
if "lm_head" in self:
|
||||
logits = self.lm_head(logits)
|
||||
else:
|
||||
logits = self.model.embed_tokens.as_linear(logits)
|
||||
|
||||
c = self.args.logits_soft_cap
|
||||
if c:
|
||||
logits = mx.tanh(logits / c) * c
|
||||
return logits
|
||||
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
def sanitize(self, weights):
|
||||
for k, v in weights.items():
|
||||
if "conv_1d.weight" in k and v.shape[-1] != 1:
|
||||
weights[k] = v.moveaxis(2, 1)
|
||||
if "lm_head.weight" not in weights:
|
||||
self.pop("lm_head")
|
||||
return weights
|
||||
|
||||
def make_cache(self):
|
||||
cache = []
|
||||
for layer in self.layers:
|
||||
if layer.temporal_block_type == "recurrent":
|
||||
cache.append(MambaCache())
|
||||
else:
|
||||
cache.append(RotatingKVCache(max_size=self.args.attention_window_size))
|
||||
return cache
|
||||
@@ -1,255 +0,0 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
|
||||
class SuScaledRoPE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
base: float = 10000.0,
|
||||
max_position_embeddings: int = 131072,
|
||||
original_max_position_embeddings: int = 4096,
|
||||
short_factor: Union[List[float], float] = 1.0,
|
||||
long_factor: Union[List[float], float] = 1.0,
|
||||
short_mscale: float = None,
|
||||
long_mscale: float = None,
|
||||
):
|
||||
"""
|
||||
Su Scaled Rotary Embedding layer.
|
||||
|
||||
Args:
|
||||
dims (int): The feature dimensions to be rotated.
|
||||
base (int, optional): Base for the exponential scaling.
|
||||
max_position_embeddings (int, optional): The maximum sequence
|
||||
length that this model was trained with. This is used to determine
|
||||
the size of the original RoPE embeddings when using long scaling.
|
||||
Default: ``131072``.
|
||||
original_max_position_embeddings (int, optional): The maximum
|
||||
sequence length that this model was trained with. This is used to
|
||||
determine the size of the original RoPE embeddings when using long
|
||||
scaling. Default: ``4096``.
|
||||
short_factor (float or list[float], optional): List of scaling
|
||||
factors for sequences of length lesser than
|
||||
``original_max_position_embeddings``. Default: ``1.0``.
|
||||
long_factor (float or list[float], optional): List of scaling
|
||||
factors for sequences of length greater than
|
||||
``original_max_position_embeddings``. Default: ``1.0``.
|
||||
short_mscale (float, optional): Scale the input prior to embedding.
|
||||
long_mscale (float, optional): Scale the input prior to embedding.
|
||||
"""
|
||||
super().__init__()
|
||||
freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
|
||||
self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
|
||||
self.original_max_position_embeddings = original_max_position_embeddings
|
||||
self.scale = long_mscale or math.sqrt(
|
||||
1
|
||||
+ math.log(max_position_embeddings / original_max_position_embeddings)
|
||||
/ math.log(original_max_position_embeddings)
|
||||
)
|
||||
self.dim = dims
|
||||
|
||||
def __call__(self, x, offset: int = 0):
|
||||
x[..., : self.dim] = self.scale * x[..., : self.dim]
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
self.dim,
|
||||
traditional=False,
|
||||
base=None,
|
||||
scale=1.0,
|
||||
offset=offset,
|
||||
freqs=self._freqs,
|
||||
)
|
||||
|
||||
|
||||
class Llama3RoPE(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
max_position_embeddings: int = 2048,
|
||||
traditional: bool = False,
|
||||
base: float = 10000,
|
||||
scaling_config: dict = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.dims = dims
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.traditional = traditional
|
||||
|
||||
factor = scaling_config["factor"]
|
||||
low_freq_factor = scaling_config.get("low_freq_factor", 1.0)
|
||||
high_freq_factor = scaling_config.get("high_freq_factor", 4.0)
|
||||
old_context_len = scaling_config.get(
|
||||
"original_max_position_embeddings",
|
||||
8192,
|
||||
)
|
||||
|
||||
low_freq_wavelen = old_context_len / low_freq_factor
|
||||
high_freq_wavelen = old_context_len / high_freq_factor
|
||||
|
||||
freqs = base ** (mx.arange(0, dims, 2) / dims)
|
||||
wavelens = 2 * mx.pi * freqs
|
||||
|
||||
freqs = mx.where(wavelens > low_freq_wavelen, freqs * factor, freqs)
|
||||
is_medium_freq = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
|
||||
smooth_factors = (old_context_len / wavelens - low_freq_factor) / (
|
||||
high_freq_factor - low_freq_factor
|
||||
)
|
||||
smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
|
||||
self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs)
|
||||
|
||||
def extra_repr(self):
|
||||
return (
|
||||
f"{self.dims}, traditional={self.traditional}, "
|
||||
f"max_position_embeddings={self.max_position_embeddings}"
|
||||
)
|
||||
|
||||
def __call__(self, x, offset: int = 0):
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
self.dims,
|
||||
traditional=self.traditional,
|
||||
base=None,
|
||||
scale=1.0,
|
||||
offset=offset,
|
||||
freqs=self._freqs,
|
||||
)
|
||||
|
||||
|
||||
class YarnRoPE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims,
|
||||
traditional=False,
|
||||
max_position_embeddings=2048,
|
||||
base=10000,
|
||||
scaling_factor=1.0,
|
||||
original_max_position_embeddings=4096,
|
||||
beta_fast=32,
|
||||
beta_slow=1,
|
||||
mscale=1,
|
||||
mscale_all_dim=0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
def yarn_find_correction_dim(num_rotations):
|
||||
return (
|
||||
dims
|
||||
* math.log(
|
||||
original_max_position_embeddings / (num_rotations * 2 * math.pi)
|
||||
)
|
||||
) / (2 * math.log(base))
|
||||
|
||||
def yarn_find_correction_range():
|
||||
low = math.floor(yarn_find_correction_dim(beta_fast))
|
||||
high = math.ceil(yarn_find_correction_dim(beta_slow))
|
||||
return max(low, 0), min(high, dims - 1)
|
||||
|
||||
def yarn_get_mscale(scale=1, mscale=1):
|
||||
if scale <= 1:
|
||||
return 1.0
|
||||
return 0.1 * mscale * math.log(scale) + 1.0
|
||||
|
||||
def yarn_linear_ramp_mask(min_val, max_val, dim):
|
||||
if min_val == max_val:
|
||||
max_val += 0.001 # Prevent singularity
|
||||
|
||||
linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (
|
||||
max_val - min_val
|
||||
)
|
||||
return mx.clip(linear_func, 0, 1)
|
||||
|
||||
self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
|
||||
scaling_factor, mscale_all_dim
|
||||
)
|
||||
freq_extra = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
|
||||
freq_inter = scaling_factor * base ** (
|
||||
mx.arange(0, dims, 2, dtype=mx.float32) / dims
|
||||
)
|
||||
low, high = yarn_find_correction_range()
|
||||
freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
|
||||
self._freqs = (freq_inter * freq_extra) / (
|
||||
freq_inter * freq_mask + freq_extra * (1 - freq_mask)
|
||||
)
|
||||
self.dims = dims
|
||||
self.traditional = traditional
|
||||
|
||||
def __call__(self, x, offset=0):
|
||||
if self.mscale != 1.0:
|
||||
x[..., : self.dims] = self.mscale * x[..., : self.dims]
|
||||
return mx.fast.rope(
|
||||
x,
|
||||
self.dims,
|
||||
traditional=self.traditional,
|
||||
base=None,
|
||||
scale=1.0,
|
||||
offset=offset,
|
||||
freqs=self._freqs,
|
||||
)
|
||||
|
||||
|
||||
def initialize_rope(
|
||||
dims,
|
||||
base,
|
||||
traditional,
|
||||
scaling_config: Optional[dict] = None,
|
||||
max_position_embeddings: Optional[int] = None,
|
||||
):
|
||||
if scaling_config is not None:
|
||||
rope_type = scaling_config.get("type") or scaling_config.get(
|
||||
"rope_type", "default"
|
||||
)
|
||||
else:
|
||||
rope_type = "default"
|
||||
|
||||
if rope_type in ["default", "linear"]:
|
||||
scale = 1 / scaling_config["factor"] if rope_type == "linear" else 1.0
|
||||
return nn.RoPE(dims, traditional=traditional, base=base, scale=scale)
|
||||
|
||||
elif rope_type == "llama3":
|
||||
return Llama3RoPE(
|
||||
dims=dims,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
traditional=traditional,
|
||||
base=base,
|
||||
scaling_config=scaling_config,
|
||||
)
|
||||
elif rope_type == "yarn":
|
||||
scaling_factor = scaling_config["factor"]
|
||||
rope_kwargs = {
|
||||
key: scaling_config[key]
|
||||
for key in [
|
||||
"original_max_position_embeddings",
|
||||
"beta_fast",
|
||||
"beta_slow",
|
||||
"mscale",
|
||||
"mscale_all_dim",
|
||||
]
|
||||
if key in scaling_config
|
||||
}
|
||||
return YarnRoPE(
|
||||
dims=dims,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
traditional=traditional,
|
||||
base=base,
|
||||
**rope_kwargs,
|
||||
)
|
||||
elif rope_type == "longrope":
|
||||
return SuScaledRoPE(
|
||||
dims=dims,
|
||||
base=base,
|
||||
max_position_embeddings=max_position_embeddings,
|
||||
original_max_position_embeddings=scaling_config[
|
||||
"original_max_position_embeddings"
|
||||
],
|
||||
short_factor=scaling_config["short_factor"],
|
||||
long_factor=scaling_config["long_factor"],
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported RoPE type {rope_type}")
|
||||
+18
-10
@@ -1,12 +1,11 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -120,8 +119,8 @@ class Attention(nn.Module):
|
||||
|
||||
# Finally perform the attention computation
|
||||
scale = math.sqrt(1 / queries.shape[-1])
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=scale, mask=mask
|
||||
).astype(values.dtype)
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
return self.o_proj(output)
|
||||
@@ -197,11 +196,12 @@ class Model(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
) -> mx.array:
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(x, cache)
|
||||
cache: mx.array = None,
|
||||
) -> Tuple[mx.array, mx.array]:
|
||||
mask = None
|
||||
if x.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
|
||||
mask = mask.astype(x.dtype)
|
||||
|
||||
y = self.model(x, mask, cache)
|
||||
return self.lm_head(y)
|
||||
@@ -209,3 +209,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
+19
-13
@@ -1,12 +1,10 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
|
||||
from .base import BaseModelArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -45,7 +43,7 @@ class Attention(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
B, L, D = x.shape
|
||||
|
||||
@@ -64,8 +62,8 @@ class Attention(nn.Module):
|
||||
queries = self.rope(queries)
|
||||
keys = self.rope(keys)
|
||||
|
||||
output = scaled_dot_product_attention(
|
||||
queries, keys, values, cache=cache, scale=self.scale, mask=mask
|
||||
output = mx.fast.scaled_dot_product_attention(
|
||||
queries, keys, values, scale=self.scale, mask=mask
|
||||
)
|
||||
|
||||
output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
|
||||
@@ -100,7 +98,7 @@ class TransformerBlock(nn.Module):
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
cache: Optional[Tuple[mx.array, mx.array]] = None,
|
||||
) -> mx.array:
|
||||
r = self.self_attn(self.input_layernorm(x), mask, cache)
|
||||
h = x + r
|
||||
@@ -125,13 +123,14 @@ class Starcoder2Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
h = self.embed_tokens(inputs)
|
||||
|
||||
if mask is None:
|
||||
mask = create_attention_mask(h, cache)
|
||||
mask = None
|
||||
if h.shape[1] > 1:
|
||||
mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
|
||||
mask = mask.astype(h.dtype)
|
||||
|
||||
if cache is None:
|
||||
cache = [None] * len(self.layers)
|
||||
@@ -154,10 +153,9 @@ class Model(nn.Module):
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: mx.array = None,
|
||||
cache=None,
|
||||
):
|
||||
out = self.model(inputs, mask, cache)
|
||||
out = self.model(inputs, cache)
|
||||
if self.args.tie_word_embeddings:
|
||||
out = self.model.embed_tokens.as_linear(out)
|
||||
else:
|
||||
@@ -167,3 +165,11 @@ class Model(nn.Module):
|
||||
@property
|
||||
def layers(self):
|
||||
return self.model.layers
|
||||
|
||||
@property
|
||||
def head_dim(self):
|
||||
return self.args.hidden_size // self.args.num_attention_heads
|
||||
|
||||
@property
|
||||
def n_kv_heads(self):
|
||||
return self.args.num_key_value_heads
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
import math
|
||||
from typing import List, Union
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
class SuScaledRotaryEmbedding:
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
traditional: bool = False,
|
||||
base: float = 10000.0,
|
||||
scale: float = 1.0,
|
||||
max_position_embeddings: int = 131072,
|
||||
original_max_position_embeddings: int = 4096,
|
||||
short_factor: Union[List[float], float] = 1.0,
|
||||
long_factor: Union[List[float], float] = 1.0,
|
||||
):
|
||||
"""
|
||||
Phi3Su Scaled Rotary Embedding layer for Phi-3 models.
|
||||
|
||||
Args:
|
||||
dims (int): The feature dimensions to be rotated.
|
||||
traditional (bool, optional): Unused. Default: ``False``.
|
||||
base (int, optional): Base for the exponential scaling.
|
||||
scale (float, optional): The scale used to scale the positions.
|
||||
Default: ``1.0``.
|
||||
max_position_embeddings (int, optional): The maximum sequence
|
||||
length that this model was trained with. This is used to determine
|
||||
the size of the original RoPE embeddings when using long scaling.
|
||||
Default: ``131072``.
|
||||
original_max_position_embeddings (int, optional): The maximum
|
||||
sequence length that this model was trained with. This is used to
|
||||
determine the size of the original RoPE embeddings when using long
|
||||
scaling. Default: ``4096``.
|
||||
short_factor (float or list[float], optional): List of scaling
|
||||
factors for sequences of length lesser than
|
||||
``original_max_position_embeddings``. Default: ``1.0``.
|
||||
long_factor (float or list[float], optional): List of scaling
|
||||
factors for sequences of length greater than
|
||||
``original_max_position_embeddings``. Default: ``1.0``.
|
||||
"""
|
||||
self.inv_freq_short = 1.0 / (
|
||||
mx.array(short_factor, dtype=mx.float32)
|
||||
* base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
|
||||
)
|
||||
self.inv_freq_long = 1.0 / (
|
||||
scale
|
||||
* mx.array(long_factor, dtype=mx.float32)
|
||||
* base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
|
||||
)
|
||||
self.original_max_position_embeddings = original_max_position_embeddings
|
||||
self.scaling_factor = math.sqrt(
|
||||
1
|
||||
+ math.log(max_position_embeddings / original_max_position_embeddings)
|
||||
/ math.log(original_max_position_embeddings)
|
||||
)
|
||||
|
||||
def _get_cos_sin(self, offset, L):
|
||||
position_ids = mx.arange(offset, offset + L, dtype=mx.float32)
|
||||
inv_freq = (
|
||||
self.inv_freq_long
|
||||
if (offset + L) > self.original_max_position_embeddings
|
||||
else self.inv_freq_short
|
||||
)
|
||||
freqs = position_ids[:, None] * inv_freq[None, :]
|
||||
emb = mx.concatenate([freqs, freqs], axis=-1)
|
||||
cos = mx.cos(emb) * self.scaling_factor
|
||||
sin = mx.sin(emb) * self.scaling_factor
|
||||
return cos, sin
|
||||
|
||||
def __call__(self, x, offset: int = 0):
|
||||
def _rotate_half(_x):
|
||||
midpoint = _x.shape[-1] // 2
|
||||
x1, x2 = _x[..., :midpoint], _x[..., midpoint:]
|
||||
return mx.concatenate([-x2, x1], axis=-1)
|
||||
|
||||
cos, sin = self._get_cos_sin(offset, x.shape[2])
|
||||
return (x * cos) + (_rotate_half(x) * sin)
|
||||
@@ -1,26 +1,9 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
|
||||
def _gather_sort(x, indices):
|
||||
*_, M = indices.shape
|
||||
indices = indices.flatten()
|
||||
order = mx.argsort(indices)
|
||||
inv_order = mx.argsort(order)
|
||||
return x.flatten(0, -3)[order // M], indices[order], inv_order
|
||||
|
||||
|
||||
def _scatter_unsort(x, inv_order, shape=None):
|
||||
x = x[inv_order]
|
||||
if shape is not None:
|
||||
x = mx.unflatten(x, 0, shape)
|
||||
return x
|
||||
|
||||
|
||||
class QuantizedSwitchLinear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -53,6 +36,12 @@ class QuantizedSwitchLinear(nn.Module):
|
||||
# Freeze this model's parameters
|
||||
self.freeze()
|
||||
|
||||
def unfreeze(self, *args, **kwargs):
|
||||
"""Wrap unfreeze so that we unfreeze any layers we might contain but
|
||||
our parameters will remain frozen."""
|
||||
super().unfreeze(*args, **kwargs)
|
||||
self.freeze(recurse=False)
|
||||
|
||||
@property
|
||||
def input_dims(self):
|
||||
return self.scales.shape[2] * self.group_size
|
||||
@@ -65,7 +54,7 @@ class QuantizedSwitchLinear(nn.Module):
|
||||
def num_experts(self):
|
||||
return self.weight.shape[0]
|
||||
|
||||
def __call__(self, x, indices, sorted_indices=False):
|
||||
def __call__(self, x, indices):
|
||||
x = mx.gather_qmm(
|
||||
x,
|
||||
self["weight"],
|
||||
@@ -75,7 +64,6 @@ class QuantizedSwitchLinear(nn.Module):
|
||||
transpose=True,
|
||||
group_size=self.group_size,
|
||||
bits=self.bits,
|
||||
sorted_indices=sorted_indices,
|
||||
)
|
||||
if "bias" in self:
|
||||
x = x + mx.expand_dims(self["bias"][indices], -2)
|
||||
@@ -109,13 +97,8 @@ class SwitchLinear(nn.Module):
|
||||
def num_experts(self):
|
||||
return self.weight.shape[0]
|
||||
|
||||
def __call__(self, x, indices, sorted_indices=False):
|
||||
x = mx.gather_mm(
|
||||
x,
|
||||
self["weight"].swapaxes(-1, -2),
|
||||
rhs_indices=indices,
|
||||
sorted_indices=sorted_indices,
|
||||
)
|
||||
def __call__(self, x, indices):
|
||||
x = mx.gather_mm(x, self["weight"].swapaxes(-1, -2), rhs_indices=indices)
|
||||
if "bias" in self:
|
||||
x = x + mx.expand_dims(self["bias"][indices], -2)
|
||||
return x
|
||||
@@ -137,7 +120,7 @@ class SwitchGLU(nn.Module):
|
||||
input_dims: int,
|
||||
hidden_dims: int,
|
||||
num_experts: int,
|
||||
activation=nn.SiLU(),
|
||||
activation=nn.silu,
|
||||
bias: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
@@ -150,24 +133,9 @@ class SwitchGLU(nn.Module):
|
||||
def __call__(self, x, indices) -> mx.array:
|
||||
x = mx.expand_dims(x, (-2, -3))
|
||||
|
||||
# When we have many tokens, then sort them to make sure that the access
|
||||
# of different experts is in order.
|
||||
do_sort = indices.size >= 64
|
||||
idx = indices
|
||||
inv_order = None
|
||||
if do_sort:
|
||||
x, idx, inv_order = _gather_sort(x, indices)
|
||||
|
||||
x_up = self.up_proj(x, idx, sorted_indices=do_sort)
|
||||
x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
|
||||
x = self.down_proj(
|
||||
self.activation(x_gate) * x_up,
|
||||
idx,
|
||||
sorted_indices=do_sort,
|
||||
)
|
||||
|
||||
if do_sort:
|
||||
x = _scatter_unsort(x, inv_order, indices.shape)
|
||||
x_up = self.up_proj(x, indices)
|
||||
x_gate = self.gate_proj(x, indices)
|
||||
x = self.down_proj(self.activation(x_gate) * x_up, indices)
|
||||
|
||||
return x.squeeze(-2)
|
||||
|
||||
@@ -178,7 +146,7 @@ class SwitchMLP(nn.Module):
|
||||
input_dims: int,
|
||||
hidden_dims: int,
|
||||
num_experts: int,
|
||||
activation=nn.GELU(approx="precise"),
|
||||
activation=nn.gelu_approx,
|
||||
bias: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
@@ -190,19 +158,8 @@ class SwitchMLP(nn.Module):
|
||||
def __call__(self, x, indices) -> mx.array:
|
||||
x = mx.expand_dims(x, (-2, -3))
|
||||
|
||||
# When we have many tokens, then sort them to make sure that the access
|
||||
# of different experts is in order.
|
||||
do_sort = indices.size >= 64
|
||||
idx = indices
|
||||
inv_order = None
|
||||
if do_sort:
|
||||
x, idx, inv_order = _gather_sort(x, indices)
|
||||
|
||||
x = self.fc1(x, idx, sorted_indices=do_sort)
|
||||
x = self.fc1(x, indices)
|
||||
x = self.activation(x)
|
||||
x = self.fc2(x, idx, sorted_indices=do_sort)
|
||||
|
||||
if do_sort:
|
||||
x = _scatter_unsort(x, inv_order, indices.shape)
|
||||
x = self.fc2(x, indices)
|
||||
|
||||
return x.squeeze(-2)
|
||||
|
||||
@@ -1,588 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict
|
||||
from urllib import request
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from mlx.utils import tree_flatten, tree_map, tree_map_with_path
|
||||
from tqdm import tqdm
|
||||
|
||||
from mlx_lm.models.base import create_attention_mask
|
||||
from mlx_lm.models.switch_layers import SwitchLinear
|
||||
from mlx_lm.quant.utils import load_data
|
||||
from mlx_lm.utils import (
|
||||
fetch_from_hub,
|
||||
get_model_path,
|
||||
save,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScaleConfig:
|
||||
prev: nn.Module
|
||||
layers: list[nn.Module]
|
||||
block: nn.Module | None = None
|
||||
kwargs: list = field(default_factory=list)
|
||||
use_config: Callable[[nn.Module], bool] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AWQConfig:
|
||||
embed: str
|
||||
lm_head: str
|
||||
no_clip: list[str]
|
||||
scale_configs: list[ScaleConfig]
|
||||
lm_key: str | None = None
|
||||
|
||||
|
||||
def update(cfg, **kwargs):
|
||||
cfg = copy.deepcopy(cfg)
|
||||
for k, v in kwargs.items():
|
||||
setattr(cfg, k, v)
|
||||
return cfg
|
||||
|
||||
|
||||
llama_awq = AWQConfig(
|
||||
embed="embed_tokens",
|
||||
lm_head="lm_head",
|
||||
no_clip=["q_proj", "k_proj"],
|
||||
scale_configs=[
|
||||
ScaleConfig(
|
||||
block="self_attn",
|
||||
prev="input_layernorm",
|
||||
layers=["q_proj", "k_proj", "v_proj"],
|
||||
kwargs=["mask"],
|
||||
),
|
||||
ScaleConfig(prev="mlp.up_proj", layers=["mlp.down_proj"]),
|
||||
ScaleConfig(
|
||||
block="mlp",
|
||||
prev="post_attention_layernorm",
|
||||
layers=["gate_proj", "up_proj"],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
gemma3_text_awq = AWQConfig(
|
||||
embed="embed_tokens",
|
||||
lm_head="lm_head",
|
||||
no_clip=["q_proj", "k_proj"],
|
||||
scale_configs=[
|
||||
ScaleConfig(
|
||||
block="self_attn",
|
||||
prev="input_layernorm",
|
||||
layers=["q_proj", "k_proj", "v_proj"],
|
||||
kwargs=["mask"],
|
||||
),
|
||||
ScaleConfig(prev="mlp.up_proj", layers=["mlp.down_proj"]),
|
||||
ScaleConfig(
|
||||
block="mlp",
|
||||
prev="pre_feedforward_layernorm",
|
||||
layers=["gate_proj", "up_proj"],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
gemma3_awq = update(gemma3_text_awq, lm_key="language_model")
|
||||
|
||||
deepseek_v2_awq = AWQConfig(
|
||||
embed="embed_tokens",
|
||||
lm_head="lm_head",
|
||||
no_clip=["q_proj", "q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"],
|
||||
scale_configs=[
|
||||
ScaleConfig(
|
||||
block="self_attn",
|
||||
prev="input_layernorm",
|
||||
layers=["q_proj", "kv_a_proj_with_mqa"],
|
||||
kwargs=["mask"],
|
||||
),
|
||||
ScaleConfig(
|
||||
prev="self_attn.kv_a_layernorm",
|
||||
layers=["self_attn.kv_b_proj"],
|
||||
),
|
||||
ScaleConfig(
|
||||
prev="mlp.up_proj",
|
||||
layers=["mlp.down_proj"],
|
||||
use_config=lambda block: not "switch_mlp" in block.mlp,
|
||||
),
|
||||
ScaleConfig(
|
||||
prev="mlp.shared_experts.up_proj",
|
||||
layers=["mlp.shared_experts.down_proj"],
|
||||
use_config=lambda block: "switch_mlp" in block.mlp,
|
||||
),
|
||||
ScaleConfig(
|
||||
prev="mlp.switch_mlp.up_proj",
|
||||
layers=["mlp.switch_mlp.down_proj"],
|
||||
use_config=lambda block: "switch_mlp" in block.mlp,
|
||||
kwargs=["indices"],
|
||||
),
|
||||
ScaleConfig(
|
||||
block="mlp",
|
||||
prev="post_attention_layernorm",
|
||||
layers=["gate_proj", "up_proj"],
|
||||
use_config=lambda block: not "switch_mlp" in block.mlp,
|
||||
),
|
||||
ScaleConfig(
|
||||
block="mlp",
|
||||
prev="post_attention_layernorm",
|
||||
layers=[
|
||||
"switch_mlp.gate_proj",
|
||||
"switch_mlp.up_proj",
|
||||
"shared_experts.gate_proj",
|
||||
"shared_experts.up_proj",
|
||||
"gate", # not quantized, just scaled
|
||||
],
|
||||
use_config=lambda block: "switch_mlp" in block.mlp,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
AWQ_MODEL_CONFIGS = {
|
||||
"llama": llama_awq,
|
||||
"mistral": llama_awq,
|
||||
"qwen2": llama_awq,
|
||||
"qwen3": llama_awq,
|
||||
"gemma3_text": gemma3_text_awq,
|
||||
"gemma3": update(gemma3_text_awq, lm_key="language_model"),
|
||||
"deepseek_v2": deepseek_v2_awq,
|
||||
}
|
||||
|
||||
|
||||
def mse(x, y):
|
||||
return ((x - y).astype(mx.float32)) ** 2
|
||||
|
||||
|
||||
def submodule_from_key(module, key):
|
||||
keys = key.split(".")
|
||||
for k in keys:
|
||||
module = module[k]
|
||||
return module
|
||||
|
||||
|
||||
def run_layer(
|
||||
layer: nn.Module,
|
||||
x: mx.array,
|
||||
indices: mx.array | None = None,
|
||||
batch_size: int = 32,
|
||||
**kwargs,
|
||||
):
|
||||
y = []
|
||||
for i in range(0, x.shape[0], batch_size):
|
||||
if indices is not None:
|
||||
y.append(
|
||||
layer(x[i : i + batch_size], indices[i : i + batch_size], **kwargs)
|
||||
)
|
||||
else:
|
||||
y.append(layer(x[i : i + batch_size], **kwargs))
|
||||
mx.eval(y)
|
||||
y = mx.concatenate(y, axis=0)
|
||||
return y
|
||||
|
||||
|
||||
def dist_split(x: mx.array, group: mx.distributed.Group):
|
||||
N = group.size()
|
||||
if N == 1:
|
||||
return x
|
||||
B = x.shape[0]
|
||||
assert B % N == 0
|
||||
r = group.rank()
|
||||
local_B = (B + N - 1) // N
|
||||
return x[r * local_B : (r + 1) * local_B]
|
||||
|
||||
|
||||
def search_best_scale(
|
||||
layers: list[nn.Module],
|
||||
quantize_func: Callable,
|
||||
block: nn.Module | None,
|
||||
layer_kwargs: dict,
|
||||
n_grid: int,
|
||||
):
|
||||
group = mx.distributed.init()
|
||||
|
||||
layer_kwargs = layer_kwargs or {}
|
||||
|
||||
x = layers[0].input_feat
|
||||
|
||||
block = block or layers[0]
|
||||
out = block(x, **layer_kwargs)
|
||||
|
||||
x_max = x.abs().mean(axis=(0, 1))
|
||||
|
||||
best_error = float("inf")
|
||||
best_scales = None
|
||||
|
||||
weights = tree_flatten(block.parameters())
|
||||
|
||||
# Search across different scaling ratios
|
||||
# and take the best loss.
|
||||
for ratio in range(n_grid):
|
||||
ratio = ratio / n_grid
|
||||
scales = mx.maximum(x_max**ratio, 1e-4).reshape(-1)
|
||||
scales = scales / (scales.max() * scales.min()).sqrt()
|
||||
for layer in layers:
|
||||
if isinstance(layer, (nn.Linear, SwitchLinear)):
|
||||
layer.weight = quantize_func(layer.weight * scales) / scales
|
||||
|
||||
out_q = run_layer(block, x, **layer_kwargs)
|
||||
loss = mse(out, out_q).sum()
|
||||
if group is not None:
|
||||
loss = mx.distributed.all_sum(loss) / group.size()
|
||||
loss /= out.size
|
||||
mx.eval(loss)
|
||||
if loss.item() < best_error:
|
||||
best_error = loss.item()
|
||||
best_scales = scales
|
||||
|
||||
# reload the original weights
|
||||
block.load_weights(weights)
|
||||
|
||||
best_scales = best_scales.reshape(-1)
|
||||
mx.eval(best_scales)
|
||||
return best_scales
|
||||
|
||||
|
||||
def apply_scale(prev_op, layers, scales):
|
||||
# Fuse the scales into the previous op
|
||||
if isinstance(prev_op, (nn.Linear, SwitchLinear)):
|
||||
assert len(layers) == 1
|
||||
prev_op.weight = prev_op.weight / scales[:, mx.newaxis]
|
||||
if hasattr(prev_op, "bias"):
|
||||
prev_op.bias = prev_op.bias / scales
|
||||
layers[0].weight = layers[0].weight * scales
|
||||
elif isinstance(prev_op, (nn.LayerNorm, nn.RMSNorm)):
|
||||
prev_op.weight = prev_op.weight / scales
|
||||
if hasattr(prev_op, "bias"):
|
||||
prev_op.bias = prev_op.bias / scales
|
||||
for layer in layers:
|
||||
layer.weight = layer.weight * scales
|
||||
elif prev_op.__class__.__name__ == "RMSNorm": # For gemma models
|
||||
dt = prev_op.weight.dtype
|
||||
prev_op.weight = (
|
||||
(1.0 + prev_op.weight.astype(mx.float32)) / scales - 1.0
|
||||
).astype(dt)
|
||||
for layer in layers:
|
||||
layer.weight = layer.weight * scales
|
||||
else:
|
||||
raise NotImplementedError(f"Could not apply scale to prev_op: {prev_op}")
|
||||
|
||||
for layer in layers:
|
||||
if hasattr(layer, "input_feat"):
|
||||
layer.input_feat = layer.input_feat / scales
|
||||
|
||||
|
||||
def scale_block(
|
||||
block: nn.Module,
|
||||
configs: list[ScaleConfig],
|
||||
quantize_func: Callable,
|
||||
layer_kwargs: dict,
|
||||
n_grid: int,
|
||||
):
|
||||
for conf in configs:
|
||||
if conf.use_config is not None and not conf.use_config(block):
|
||||
continue
|
||||
if conf.block is not None:
|
||||
local_block = block[conf.block]
|
||||
layers = [submodule_from_key(local_block, l) for l in conf.layers]
|
||||
else:
|
||||
local_block = None
|
||||
layers = [submodule_from_key(block, l) for l in conf.layers]
|
||||
local_kwargs = {k: layer_kwargs[k] for k in conf.kwargs if k in layer_kwargs}
|
||||
for k in conf.kwargs:
|
||||
if hasattr(layers[0], k):
|
||||
local_kwargs[k] = getattr(layers[0], k)
|
||||
|
||||
scales = search_best_scale(
|
||||
layers=layers,
|
||||
block=local_block,
|
||||
layer_kwargs=local_kwargs,
|
||||
quantize_func=quantize_func,
|
||||
n_grid=n_grid,
|
||||
)
|
||||
apply_scale(submodule_from_key(block, conf.prev), layers, scales)
|
||||
|
||||
|
||||
def search_best_clip(
|
||||
module: nn.Module,
|
||||
quantize_func: Callable,
|
||||
group_size: int,
|
||||
n_grid: int,
|
||||
max_shrink: float = 0.5,
|
||||
batch_size: int = 64,
|
||||
n_frames: int = 512,
|
||||
):
|
||||
group = mx.distributed.init()
|
||||
|
||||
# subsample the input features
|
||||
x = module.input_feat.flatten(0, 1)
|
||||
stride = (x.shape[0] + n_frames - 1) // n_frames
|
||||
x = x[::stride]
|
||||
|
||||
w = module.weight
|
||||
x = x.reshape(x.shape[0], -1, group_size)
|
||||
|
||||
w_init_shape = w.shape
|
||||
w_all = mx.flatten(w, 0, w.ndim - 2)
|
||||
w_max_all = []
|
||||
|
||||
# batch across W to save memory
|
||||
for b in range(0, w_all.shape[0], batch_size):
|
||||
w = w_all[b : b + batch_size]
|
||||
|
||||
group_shape = (w.shape[0], w.shape[-1] // group_size)
|
||||
best_error = mx.full(group_shape, float("inf"))
|
||||
best_w_max = mx.zeros((*group_shape, 1), dtype=x.dtype)
|
||||
|
||||
w_shape = w.shape
|
||||
|
||||
w = w.reshape(*w.shape[:-1], -1, group_size)
|
||||
out = mx.einsum("bdg,odg->bod", x, w)
|
||||
init_max = w.abs().max(axis=-1, keepdims=True)
|
||||
|
||||
# try a range of clips and pick the one with the smallest loss
|
||||
for i in range(int(max_shrink * n_grid)):
|
||||
p = 1 - i / n_grid
|
||||
w_max = p * init_max
|
||||
w_m = mx.clip(w, -w_max, w_max).reshape(w_shape)
|
||||
|
||||
w_q = quantize_func(w_m)
|
||||
|
||||
w_q = w_q.reshape(*w_q.shape[:-1], -1, group_size)
|
||||
out_q = mx.einsum("bdg,odg->bod", x, w_q)
|
||||
|
||||
# Take the mean across the input batch
|
||||
loss = mse(out, out_q).sum(axis=0)
|
||||
if group is not None:
|
||||
loss = mx.distributed.all_sum(loss) / group.size()
|
||||
loss /= out.shape[0]
|
||||
best_indices = loss < best_error
|
||||
best_error = mx.where(best_indices, loss, best_error)
|
||||
best_w_max = mx.where(best_indices[..., mx.newaxis], w_max, best_w_max)
|
||||
mx.eval(best_w_max, best_error)
|
||||
|
||||
w_max_all.append(best_w_max)
|
||||
|
||||
best_w_max = mx.concatenate(w_max_all, axis=0)
|
||||
|
||||
w_r = w_all.reshape(*w_all.shape[:-1], -1, group_size)
|
||||
best_w = mx.clip(w_r, -best_w_max, best_w_max)
|
||||
best_w = best_w.reshape(w_init_shape)
|
||||
|
||||
mx.eval(best_w)
|
||||
return best_w
|
||||
|
||||
|
||||
def clip_block(
|
||||
block: nn.Module,
|
||||
no_clip_keys: list[str],
|
||||
quantize_func: Callable,
|
||||
group_size: int,
|
||||
n_grid: int = 20,
|
||||
):
|
||||
def apply_clip(path, module):
|
||||
if isinstance(module, (nn.Linear, SwitchLinear)) and all(
|
||||
k not in path for k in no_clip_keys
|
||||
):
|
||||
best_weight = search_best_clip(
|
||||
module,
|
||||
quantize_func=quantize_func,
|
||||
group_size=group_size,
|
||||
n_grid=n_grid,
|
||||
)
|
||||
module.weight = best_weight
|
||||
|
||||
tree_map_with_path(apply_clip, block.leaf_modules(), is_leaf=nn.Module.is_module)
|
||||
|
||||
|
||||
def awq_quantize(
|
||||
model,
|
||||
inputs: mx.array,
|
||||
awq_config: AWQConfig,
|
||||
group_size: int = 64,
|
||||
bits: int = 3,
|
||||
embed_group_size: int = 32,
|
||||
embed_bits: int = 4,
|
||||
n_grid: int = 20,
|
||||
):
|
||||
if awq_config.lm_key is not None:
|
||||
model = model[awq_config.lm_key]
|
||||
|
||||
group = mx.distributed.init()
|
||||
|
||||
def quantize_func(w):
|
||||
wq = mx.quantize(w, bits=bits, group_size=group_size)
|
||||
return mx.dequantize(*wq, bits=bits, group_size=group_size)
|
||||
|
||||
mask = create_attention_mask(inputs)
|
||||
|
||||
embed_key = awq_config.embed
|
||||
model.model[embed_key] = model.model[embed_key].to_quantized(
|
||||
group_size=embed_group_size, bits=embed_bits
|
||||
)
|
||||
inputs = model.model[embed_key](inputs)
|
||||
|
||||
def capture(module):
|
||||
if not isinstance(module, (nn.Linear, SwitchLinear)):
|
||||
return module
|
||||
|
||||
class Catcher(nn.Module):
|
||||
def __call__(self, x: mx.array, *args, **kwargs):
|
||||
# Store the input features on the original modules.
|
||||
if hasattr(module, "input_feat"):
|
||||
module.input_feat = mx.concatenate([module.input_feat, x], axis=0)
|
||||
else:
|
||||
module.input_feat = x
|
||||
|
||||
# Also store the MOE indices if applicabale
|
||||
if isinstance(module, SwitchLinear):
|
||||
indices = args[0]
|
||||
if hasattr(module, "indices"):
|
||||
module.indices = mx.concatenate(
|
||||
[module.indices, indices], axis=0
|
||||
)
|
||||
else:
|
||||
module.indices = indices
|
||||
|
||||
return module(x, *args, **kwargs)
|
||||
|
||||
return Catcher()
|
||||
|
||||
for e, block in enumerate(tqdm(model.layers)):
|
||||
# Capture the input features for each of the layers in the transformer block
|
||||
orig_leaves = block.leaf_modules()
|
||||
capture_leaves = tree_map(capture, orig_leaves, is_leaf=nn.Module.is_module)
|
||||
block.update_modules(capture_leaves)
|
||||
outputs = run_layer(block, inputs, mask=mask)
|
||||
block.update_modules(orig_leaves)
|
||||
del capture_leaves
|
||||
|
||||
# Quantize the block without AWQ to obtain a reference loss
|
||||
nn.quantize(block, group_size=group_size, bits=bits)
|
||||
outputs_q = run_layer(block, inputs, mask=mask)
|
||||
before_loss = mse(outputs, outputs_q).sum()
|
||||
if group is not None:
|
||||
before_loss = mx.distributed.all_sum(before_loss) / group.size()
|
||||
before_loss /= outputs.size
|
||||
block.update_modules(orig_leaves)
|
||||
orig_params = block.parameters()
|
||||
|
||||
scale_block(
|
||||
block=block,
|
||||
configs=awq_config.scale_configs,
|
||||
quantize_func=quantize_func,
|
||||
n_grid=n_grid,
|
||||
layer_kwargs={"mask": mask},
|
||||
)
|
||||
|
||||
clip_block(
|
||||
block=block,
|
||||
no_clip_keys=awq_config.no_clip,
|
||||
quantize_func=quantize_func,
|
||||
group_size=group_size,
|
||||
n_grid=n_grid,
|
||||
)
|
||||
|
||||
# Quantize the scaled and clipped block
|
||||
nn.quantize(block, group_size=group_size, bits=bits)
|
||||
outputs_q = run_layer(block, inputs, mask=mask)
|
||||
after_loss = mse(outputs, outputs_q).sum()
|
||||
if group is not None:
|
||||
after_loss = mx.distributed.all_sum(after_loss) / group.size()
|
||||
after_loss /= outputs.size
|
||||
tqdm.write(f"Loss reduction: {after_loss / before_loss}")
|
||||
if after_loss > before_loss:
|
||||
# Reload original weights and quantize
|
||||
block.update_modules(orig_leaves)
|
||||
block.update(orig_params)
|
||||
nn.quantize(block, group_size=group_size, bits=bits)
|
||||
tqdm.write("Loss is not reduced, falling back to original weights.")
|
||||
|
||||
inputs = outputs
|
||||
|
||||
mx.eval(block)
|
||||
mx.clear_cache()
|
||||
|
||||
if (lm_head := awq_config.lm_head) in model:
|
||||
model[lm_head] = model[lm_head].to_quantized(
|
||||
group_size=embed_group_size, bits=embed_bits
|
||||
)
|
||||
|
||||
|
||||
def update_config(
|
||||
model: nn.Module,
|
||||
config: Dict[str, Any],
|
||||
):
|
||||
# dummy
|
||||
config["quantization"] = {"group_size": 64, "bits": 4}
|
||||
|
||||
def update_config(path, module):
|
||||
if hasattr(module, "bits"):
|
||||
config["quantization"][path] = {
|
||||
"group_size": module.group_size,
|
||||
"bits": module.bits,
|
||||
}
|
||||
else:
|
||||
config["quantization"][path] = False
|
||||
|
||||
tree_map_with_path(update_config, model.leaf_modules(), is_leaf=nn.Module.is_module)
|
||||
return config
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--model", "-m", default="mlx-community/Qwen2.5-7B-Instruct-bf16"
|
||||
)
|
||||
parser.add_argument("--mlx-path", default="mlx_model")
|
||||
parser.add_argument("--bits", type=int, default=4)
|
||||
parser.add_argument("--group-size", type=int, default=64)
|
||||
parser.add_argument("--embed-bits", type=int, default=4)
|
||||
parser.add_argument("--embed-group-size", type=int, default=32)
|
||||
parser.add_argument("--num-samples", type=int, default=128)
|
||||
parser.add_argument("--sequence-length", type=int, default=512)
|
||||
parser.add_argument("--n-grid", type=int, default=20)
|
||||
parser.add_argument("--seed", type=int, default=123)
|
||||
args = parser.parse_args()
|
||||
|
||||
group = mx.distributed.init()
|
||||
|
||||
num_samples = args.num_samples
|
||||
if group is not None and num_samples % group.size() > 0:
|
||||
num_samples += group.size() - num_samples % group.size()
|
||||
|
||||
mx.random.seed(args.seed)
|
||||
|
||||
model_path, hf_repo = get_model_path(args.model, revision=None)
|
||||
model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
|
||||
|
||||
model_type = config["model_type"]
|
||||
if (awq_config := AWQ_MODEL_CONFIGS.get(model_type, None)) is None:
|
||||
raise NotImplementedError(f"AWQ support for {model_type} models NYI.")
|
||||
|
||||
calibration_data = load_data(tokenizer, args.num_samples, args.sequence_length)
|
||||
|
||||
calibration_data = dist_split(calibration_data, group)
|
||||
|
||||
awq_quantize(
|
||||
model,
|
||||
calibration_data,
|
||||
awq_config,
|
||||
bits=args.bits,
|
||||
group_size=args.group_size,
|
||||
embed_bits=args.embed_bits,
|
||||
embed_group_size=args.embed_group_size,
|
||||
n_grid=args.n_grid,
|
||||
)
|
||||
|
||||
config = update_config(model, config)
|
||||
save(
|
||||
args.mlx_path,
|
||||
model_path,
|
||||
model,
|
||||
tokenizer,
|
||||
config,
|
||||
hf_repo=hf_repo,
|
||||
)
|
||||
@@ -1,251 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import time
|
||||
import types
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import mlx.optimizers as optimizers
|
||||
import numpy as np
|
||||
from mlx.utils import tree_flatten, tree_map
|
||||
from tqdm import tqdm
|
||||
|
||||
from mlx_lm.tuner.datasets import load_dataset
|
||||
from mlx_lm.tuner.losses import kl_div_loss
|
||||
from mlx_lm.tuner.trainer import grad_checkpoint, iterate_batches
|
||||
from mlx_lm.tuner.utils import print_trainable_parameters
|
||||
from mlx_lm.utils import (
|
||||
fetch_from_hub,
|
||||
get_model_path,
|
||||
quantize_model,
|
||||
save,
|
||||
)
|
||||
|
||||
|
||||
class Catcher(nn.Module):
|
||||
def __init__(self, module):
|
||||
super().__init__()
|
||||
self.module = module
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
self.outputs = self.module(*args, **kwargs)
|
||||
return self.outputs
|
||||
|
||||
|
||||
def dwq_quantize(
|
||||
model,
|
||||
q_model,
|
||||
opt,
|
||||
data,
|
||||
batch_size: int = 2,
|
||||
max_seq_length: int = 2048,
|
||||
activation_layer_step: float = 0.25,
|
||||
activation_loss_weight: float = 1.0,
|
||||
dtype: mx.Dtype = mx.bfloat16,
|
||||
gradient_checkpoint: bool = False,
|
||||
):
|
||||
group = mx.distributed.init()
|
||||
world_size = group.size()
|
||||
rank = group.rank()
|
||||
|
||||
def unfreeze(_, m):
|
||||
if hasattr(m, "bits") and hasattr(m, "group_size"):
|
||||
m.unfreeze(keys=["scales", "biases"], recurse=False)
|
||||
|
||||
q_model.apply_to_modules(unfreeze)
|
||||
print_trainable_parameters(q_model)
|
||||
|
||||
layer_id_step = max(int(activation_layer_step * len(model.layers)), 1)
|
||||
layer_ids = list(range(len(model.layers)))[layer_id_step::layer_id_step]
|
||||
|
||||
for lid in layer_ids:
|
||||
model.layers[lid] = Catcher(model.layers[lid])
|
||||
q_model.layers[lid] = Catcher(q_model.layers[lid])
|
||||
|
||||
if gradient_checkpoint:
|
||||
grad_checkpoint(q_model.layers[0])
|
||||
|
||||
def forward(model, inputs):
|
||||
logits = model(inputs)
|
||||
extra_targets = [
|
||||
model.layers[lid].outputs.astype(mx.float32) for lid in layer_ids
|
||||
]
|
||||
for lid in layer_ids:
|
||||
model.layers[lid].outputs = None
|
||||
return logits, extra_targets
|
||||
|
||||
def loss_fn(params, x, targets, extra_targets, lengths):
|
||||
q_model.update(tree_map(lambda x: x.astype(dtype), params))
|
||||
logits, q_extra_targets = forward(q_model, x)
|
||||
losses = kl_div_loss(logits, targets)
|
||||
mask = mx.arange(1, 1 + targets.shape[1]) < lengths[:, 1:]
|
||||
ntoks = mask.sum()
|
||||
kl_loss = (mask * losses).sum() / ntoks
|
||||
act_loss = mx.stack(
|
||||
[
|
||||
(mask * (qe - e).abs().mean(axis=-1)).sum() / ntoks
|
||||
for qe, e in zip(q_extra_targets, extra_targets)
|
||||
]
|
||||
)
|
||||
loss = kl_loss + activation_loss_weight * act_loss.mean()
|
||||
return loss, ntoks
|
||||
|
||||
def step(inputs, targets, extra_targets, lengths, params):
|
||||
(loss, ntoks), grads = mx.value_and_grad(loss_fn)(
|
||||
params, inputs, targets, extra_targets, lengths
|
||||
)
|
||||
grads = nn.average_gradients(grads)
|
||||
params = opt.apply_gradients(grads, params)
|
||||
return loss, ntoks, params
|
||||
|
||||
# Accumulate learned weights in higher precision
|
||||
params = tree_map(
|
||||
lambda x: x.astype(mx.float32),
|
||||
q_model.trainable_parameters(),
|
||||
)
|
||||
|
||||
total_loss = 0.0
|
||||
total_tokens = 0
|
||||
tokens = 0
|
||||
tic = time.time()
|
||||
for it, (batch, lengths) in (
|
||||
pbar := tqdm(
|
||||
enumerate(iterate_batches(data, batch_size, max_seq_length)),
|
||||
total=len(data) // batch_size,
|
||||
)
|
||||
):
|
||||
batch = batch[:, :-1]
|
||||
targets, extra_targets = forward(model, batch)
|
||||
mx.eval(targets, extra_targets)
|
||||
loss, ntoks, params = step(batch, targets, extra_targets, lengths, params)
|
||||
mx.eval(loss, params)
|
||||
loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
|
||||
ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
|
||||
tokens += ntoks
|
||||
total_loss += loss * ntoks
|
||||
if rank == 0:
|
||||
pbar.set_description(desc=f"{loss=:.4f}")
|
||||
if (it + 1) % 20 == 0:
|
||||
toks_per_sec = tokens / (time.time() - tic)
|
||||
peak_memory_gb = mx.get_peak_memory() / 1e9
|
||||
avg_loss = total_loss / tokens
|
||||
total_tokens += tokens
|
||||
tqdm.write(
|
||||
f"{it=}, {avg_loss=:.4f}, {total_tokens=},"
|
||||
f" {toks_per_sec=:.3f}, {peak_memory_gb=:.3f}",
|
||||
)
|
||||
tic = time.time()
|
||||
tokens = 0
|
||||
total_loss = 0
|
||||
q_model.update(tree_map(lambda x: x.astype(dtype), params))
|
||||
for lid in layer_ids:
|
||||
q_model.layers[lid] = q_model.layers[lid].module
|
||||
|
||||
|
||||
def load_data(tokenizer, data_path: str, num_samples: int, max_seq_length: int):
|
||||
args = types.SimpleNamespace(
|
||||
hf_dataset={
|
||||
"path": data_path,
|
||||
"train_split": f"train",
|
||||
"valid_split": "train[:1]",
|
||||
},
|
||||
train=True,
|
||||
test=False,
|
||||
)
|
||||
dataset = load_dataset(args, tokenizer)[0]
|
||||
perm = np.random.permutation(len(dataset))[:num_samples].tolist()
|
||||
|
||||
def process(idx):
|
||||
tokens, offset = dataset.process(dataset[idx])
|
||||
return (tokens[:max_seq_length], offset)
|
||||
|
||||
return [process(i) for i in perm]
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", "-m", default="Qwen/Qwen3-4B")
|
||||
parser.add_argument("--quantized-model", default=None)
|
||||
parser.add_argument(
|
||||
"--mlx-path", default="mlx_model", help="Path to save the quantized model."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bits",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Bits per weight for quantization.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--group-size", type=int, default=64, help="Group size for quantization."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-samples",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Number of samples to use for training.",
|
||||
)
|
||||
parser.add_argument("--max-seq-length", type=int, default=2049)
|
||||
parser.add_argument("--seed", type=int, default=123)
|
||||
parser.add_argument("--learning-rate", type=float, default=1e-6)
|
||||
parser.add_argument("--batch-size", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--data-path",
|
||||
type=str,
|
||||
default="allenai/tulu-3-sft-mixture",
|
||||
help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--grad-checkpoint",
|
||||
action="store_true",
|
||||
help="Use gradient checkpointing to reduce memory use.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
group = mx.distributed.init()
|
||||
|
||||
num_samples = args.num_samples
|
||||
if num_samples % group.size() > 0:
|
||||
num_samples += group.size() - num_samples % group.size()
|
||||
|
||||
np.random.seed(args.seed)
|
||||
mx.random.seed(args.seed)
|
||||
|
||||
model_path, hf_repo = get_model_path(args.model, revision=None)
|
||||
model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
|
||||
|
||||
calibration_data = load_data(
|
||||
tokenizer, args.data_path, args.num_samples, args.max_seq_length
|
||||
)
|
||||
|
||||
if args.quantized_model is not None:
|
||||
q_model_path = get_model_path(args.quantized_model, revision=None)
|
||||
q_model, config, _ = fetch_from_hub(q_model_path, lazy=True)
|
||||
else:
|
||||
q_model = copy.deepcopy(model)
|
||||
_, config = quantize_model(
|
||||
q_model,
|
||||
config,
|
||||
q_group_size=args.group_size,
|
||||
q_bits=args.bits,
|
||||
)
|
||||
|
||||
opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
|
||||
dwq_quantize(
|
||||
model,
|
||||
q_model,
|
||||
opt,
|
||||
calibration_data,
|
||||
batch_size=args.batch_size,
|
||||
max_seq_length=args.max_seq_length,
|
||||
gradient_checkpoint=args.grad_checkpoint,
|
||||
)
|
||||
save(
|
||||
args.mlx_path,
|
||||
model_path,
|
||||
q_model,
|
||||
tokenizer,
|
||||
config,
|
||||
hf_repo=hf_repo,
|
||||
)
|
||||
@@ -1,349 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import math
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
|
||||
from tqdm import tqdm
|
||||
|
||||
from mlx_lm.quant.utils import load_data
|
||||
from mlx_lm.tuner.losses import kl_div_loss
|
||||
from mlx_lm.tuner.trainer import grad_checkpoint
|
||||
from mlx_lm.tuner.utils import get_total_parameters
|
||||
from mlx_lm.utils import (
|
||||
compute_bits_per_weight,
|
||||
fetch_from_hub,
|
||||
get_model_path,
|
||||
load,
|
||||
quantize_model,
|
||||
save,
|
||||
)
|
||||
|
||||
|
||||
def make_quant_predicate(config):
|
||||
def quant_predicate(p, m, _):
|
||||
if not hasattr(m, "to_quantized"):
|
||||
return False
|
||||
return config.get(p, True)
|
||||
|
||||
return quant_predicate
|
||||
|
||||
|
||||
def eval_ppl(model, data, batch_size=8):
|
||||
all_loss = 0.0
|
||||
ntoks = 0
|
||||
for s in range(0, len(data), batch_size):
|
||||
batch = data[s : s + batch_size]
|
||||
logits = model(batch[:, :-1]).astype(mx.float32)
|
||||
losses = nn.losses.cross_entropy(logits, batch[:, 1:])
|
||||
all_loss += losses.sum().item()
|
||||
ntoks += losses.size
|
||||
ppl = math.exp(all_loss / ntoks)
|
||||
return ppl
|
||||
|
||||
|
||||
def make_options(
|
||||
low_bits, low_group_size, high_bits, high_group_size, include_bpw=True
|
||||
):
|
||||
options = []
|
||||
min_bpw = low_bits + 32 / low_group_size
|
||||
max_bpw = high_bits + 32 / high_group_size
|
||||
for b in range(low_bits, high_bits + 1):
|
||||
for g in [32, 64, 128]:
|
||||
cbpw = b + 32 / g
|
||||
if b == 7 or not (min_bpw <= cbpw <= max_bpw):
|
||||
continue
|
||||
options.append({"bits": b, "group_size": g, "bpw": cbpw})
|
||||
options.sort(key=lambda x: x["bpw"])
|
||||
if not include_bpw:
|
||||
for o in options:
|
||||
o.pop("bpw")
|
||||
|
||||
return options
|
||||
|
||||
|
||||
def estimate_sensitivities(
|
||||
model,
|
||||
data,
|
||||
low_bits,
|
||||
low_group_size,
|
||||
high_bits,
|
||||
high_group_size,
|
||||
batch_size: int = 4,
|
||||
gradient_accum_dtype: mx.Dtype = mx.float32,
|
||||
gradient_checkpoint: bool = False,
|
||||
):
|
||||
def qdq(w, bits, group_size):
|
||||
w, s, b = mx.quantize(w, bits=bits, group_size=group_size)
|
||||
return mx.dequantize(w, scales=s, biases=b, bits=bits, group_size=group_size)
|
||||
|
||||
layers = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module)
|
||||
layers = {k: l for k, l in layers if hasattr(l, "to_quantized")}
|
||||
q_model = copy.deepcopy(model)
|
||||
q_layers = copy.deepcopy(layers)
|
||||
for l in q_layers.values():
|
||||
l.weight = qdq(l.weight, low_bits, low_group_size)
|
||||
# Freeze everything but the quantizable weight
|
||||
l.freeze()
|
||||
l.unfreeze(keys=["weight"])
|
||||
q_model.freeze()
|
||||
q_model.update_modules(tree_unflatten(list(q_layers.items())))
|
||||
|
||||
def loss_fn(batch, targets):
|
||||
return kl_div_loss(q_model(batch), targets).mean()
|
||||
|
||||
if gradient_checkpoint:
|
||||
grad_checkpoint(q_model.layers[0])
|
||||
|
||||
grad_accum = tree_map(
|
||||
lambda x: mx.zeros(x.shape, dtype=gradient_accum_dtype),
|
||||
q_model.trainable_parameters(),
|
||||
)
|
||||
for e, s in tqdm(
|
||||
enumerate(range(0, len(data), batch_size)),
|
||||
total=len(data) // batch_size,
|
||||
desc="Estimating sensitivities",
|
||||
):
|
||||
batch = data[s : s + batch_size]
|
||||
targets = model(batch)
|
||||
mx.eval(targets)
|
||||
_, grads = nn.value_and_grad(q_model, loss_fn)(batch, targets)
|
||||
grad_accum = tree_map(lambda x, y: x + y, grad_accum, grads)
|
||||
del grads
|
||||
mx.eval(grad_accum)
|
||||
|
||||
options = make_options(low_bits, low_group_size, high_bits, high_group_size)
|
||||
current_bpw = options[0]["bpw"]
|
||||
|
||||
def compute_sensitivity(gradient, low_q_weight, original_weight):
|
||||
n_batches = (len(data) + batch_size - 1) // batch_size
|
||||
gradient = gradient / n_batches
|
||||
scores = [{"loss_change": 0, "extra_bits": 0}]
|
||||
for opt in options[1:]:
|
||||
extra_bits = (opt["bpw"] - current_bpw) * original_weight.size
|
||||
other_weight = qdq(original_weight, opt["bits"], opt["group_size"])
|
||||
loss_change = (gradient * (low_q_weight - other_weight)).sum()
|
||||
scores.append({"loss_change": loss_change, "extra_bits": extra_bits})
|
||||
return scores
|
||||
|
||||
sensitivities = tree_map(
|
||||
compute_sensitivity,
|
||||
grad_accum,
|
||||
q_model.parameters(),
|
||||
model.parameters(),
|
||||
)
|
||||
mx.eval(sensitivities)
|
||||
|
||||
sensitivities = [
|
||||
(k.replace(".weight", ""), s.item() if isinstance(s, mx.array) else s)
|
||||
for k, s in tree_flatten(sensitivities)
|
||||
]
|
||||
|
||||
return sensitivities
|
||||
|
||||
|
||||
def compute_bit_budget(model, target_bpw):
|
||||
model_bytes = tree_reduce(
|
||||
lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
|
||||
)
|
||||
model_params = get_total_parameters(model)
|
||||
|
||||
return model_params * target_bpw - model_bytes * 8
|
||||
|
||||
|
||||
def estimate_threshold(
|
||||
model,
|
||||
sensitivities,
|
||||
target_bpw,
|
||||
low_bits,
|
||||
low_group_size,
|
||||
high_bits,
|
||||
high_group_size,
|
||||
):
|
||||
options = make_options(
|
||||
low_bits, low_group_size, high_bits, high_group_size, include_bpw=False
|
||||
)
|
||||
sensitivities = tree_flatten(
|
||||
tree_unflatten(list(sensitivities.items())),
|
||||
is_leaf=lambda x: isinstance(x, list) and "loss_change" in x[0],
|
||||
)
|
||||
|
||||
q_model = copy.deepcopy(model)
|
||||
nn.quantize(q_model, group_size=low_group_size, bits=low_bits)
|
||||
budget = int(compute_bit_budget(q_model, target_bpw))
|
||||
benefit_map = {}
|
||||
|
||||
def benefit(layer, option, budget):
|
||||
if (layer, option, budget) in benefit_map:
|
||||
return benefit_map[layer, option, budget]
|
||||
|
||||
stack = [(layer, option, budget)]
|
||||
while stack:
|
||||
layer, option, budget = stack[-1]
|
||||
|
||||
if budget <= 0 or layer < 0 or option < 0:
|
||||
benefit_map[layer, option, budget] = 0
|
||||
stack.pop()
|
||||
continue
|
||||
|
||||
# We either not use this option
|
||||
prev_layer = layer if option > 0 else layer - 1
|
||||
prev_option = (option if option > 0 else len(options)) - 1
|
||||
if (prev_layer, prev_option, budget) not in benefit_map:
|
||||
stack.append((prev_layer, prev_option, budget))
|
||||
continue
|
||||
a = benefit_map[prev_layer, prev_option, budget]
|
||||
|
||||
# Or we use it so we have less budget for before
|
||||
b = float("-inf")
|
||||
info = sensitivities[layer][1][option]
|
||||
prev_layer = layer - 1
|
||||
prev_option = len(options) - 1
|
||||
prev_budget = budget - info["extra_bits"]
|
||||
if (
|
||||
prev_layer,
|
||||
prev_option,
|
||||
prev_budget,
|
||||
) not in benefit_map and prev_budget >= 0:
|
||||
stack.append((prev_layer, prev_option, prev_budget))
|
||||
continue
|
||||
if prev_budget >= 0:
|
||||
b = benefit_map[prev_layer, prev_option, prev_budget]
|
||||
b += info["loss_change"]
|
||||
|
||||
benefit_map[layer, option, budget] = max(a, b)
|
||||
stack.pop()
|
||||
|
||||
return benefit_map[layer, option, budget]
|
||||
|
||||
def backtrack(layer, budget):
|
||||
selected = []
|
||||
while layer >= 0:
|
||||
prev_benefit = benefit(layer - 1, len(options) - 1, budget)
|
||||
option_benefits = [benefit(layer, i, budget) for i in range(len(options))]
|
||||
idx, v = max(enumerate(option_benefits), key=lambda x: x[1] - prev_benefit)
|
||||
info = sensitivities[layer][1][idx]
|
||||
if v != 0:
|
||||
budget -= info["extra_bits"]
|
||||
selected.append((layer, idx))
|
||||
layer -= 1
|
||||
return selected[::-1]
|
||||
|
||||
selected = backtrack(len(sensitivities) - 1, budget)
|
||||
config = {sensitivities[l][0]: options[i] for l, i in selected}
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", "-m", default="Qwen/Qwen3-0.6B-base")
|
||||
parser.add_argument(
|
||||
"--mlx-path", default="mlx_model", help="Path to save the model"
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=123)
|
||||
parser.add_argument(
|
||||
"--sensitivities",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a pre-computed sensitivity JSON file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-bpw", type=float, default=5.0, help="Target bits per weight."
|
||||
)
|
||||
parser.add_argument("--low-bits", type=int, default=4)
|
||||
parser.add_argument("--low-group-size", type=int, default=128)
|
||||
parser.add_argument("--high-bits", type=int, default=5)
|
||||
parser.add_argument("--high-group-size", type=int, default=32)
|
||||
parser.add_argument(
|
||||
"--report-ppl",
|
||||
action="store_true",
|
||||
help="Compute the perplexity of the base and quantized models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--grad-checkpoint",
|
||||
action="store_true",
|
||||
help="Use gradient checkpointing to reduce memory use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--accumulation-dtype",
|
||||
default="float32",
|
||||
choices=["float32", "bfloat16"],
|
||||
help="What type to use to accumulate the gradients for the sensitivities",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
group = mx.distributed.init()
|
||||
|
||||
if args.sensitivities is None:
|
||||
model, tokenizer = load(args.model)
|
||||
mx.random.seed(args.seed)
|
||||
data = load_data(tokenizer, num_samples=-1, sequence_length=512)
|
||||
|
||||
sensitivities = estimate_sensitivities(
|
||||
model,
|
||||
data,
|
||||
args.low_bits,
|
||||
args.low_group_size,
|
||||
args.high_bits,
|
||||
args.high_group_size,
|
||||
gradient_accum_dtype=getattr(mx, args.accumulation_dtype),
|
||||
gradient_checkpoint=args.grad_checkpoint,
|
||||
)
|
||||
model_name = args.model.replace("/", "_")
|
||||
with open(f"{model_name}_sensitivities.json", "w") as fid:
|
||||
json.dump(sensitivities, fid)
|
||||
else:
|
||||
with open(args.sensitivities, "r") as fid:
|
||||
sensitivities = json.load(fid)
|
||||
|
||||
sensitivities = dict(sensitivities)
|
||||
model_path, hf_repo = get_model_path(args.model, revision=None)
|
||||
model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
|
||||
mx.random.seed(args.seed)
|
||||
data = load_data(tokenizer, num_samples=-1, sequence_length=512)
|
||||
|
||||
if args.report_ppl:
|
||||
ppl = eval_ppl(model, data)
|
||||
print(f"Original PPL: {ppl:.3f}")
|
||||
|
||||
quant_config = estimate_threshold(
|
||||
model,
|
||||
sensitivities,
|
||||
target_bpw=args.target_bpw,
|
||||
low_bits=args.low_bits,
|
||||
low_group_size=args.low_group_size,
|
||||
high_bits=args.high_bits,
|
||||
high_group_size=args.high_group_size,
|
||||
)
|
||||
|
||||
model, config = quantize_model(
|
||||
model,
|
||||
config,
|
||||
q_group_size=args.low_group_size,
|
||||
q_bits=args.low_bits,
|
||||
quant_predicate=make_quant_predicate(quant_config),
|
||||
)
|
||||
|
||||
if args.report_ppl:
|
||||
ppl = eval_ppl(model, data)
|
||||
print(f"Quantized PPL: {ppl:.3f}")
|
||||
|
||||
save(
|
||||
args.mlx_path,
|
||||
model_path,
|
||||
model,
|
||||
tokenizer,
|
||||
config,
|
||||
hf_repo=hf_repo,
|
||||
)
|
||||
print(f"Peak memory used: {mx.get_peak_memory() / 1000**3:.3f}GB")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,26 +0,0 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
def load_data(tokenizer, num_samples: int, sequence_length: int) -> mx.array:
|
||||
save_dir = Path.home() / ".cache/mlx-lm/calibration_v5.txt"
|
||||
if not save_dir.exists():
|
||||
from urllib import request
|
||||
|
||||
save_dir.parent.mkdir(parents=True, exist_ok=True)
|
||||
url = "https://gist.githubusercontent.com/tristandruyen/9e207a95c7d75ddf37525d353e00659c/raw/571fda718462de863e5a0171078c175420c7649a/calibration_data_v5_rc.txt"
|
||||
request.urlretrieve(url, save_dir)
|
||||
with open(save_dir) as fid:
|
||||
texts = fid.read()
|
||||
tokens = tokenizer.encode(texts, return_tensors="mlx")[0]
|
||||
|
||||
# select random non-overlapping chunks
|
||||
tokens = tokens[: (tokens.size // sequence_length) * sequence_length]
|
||||
tokens = tokens.reshape(-1, sequence_length)
|
||||
segments = mx.random.permutation(tokens.shape[0])
|
||||
if num_samples > 0:
|
||||
segments = segments[:num_samples]
|
||||
return tokens[segments]
|
||||
@@ -0,0 +1,6 @@
|
||||
mlx>=0.14.1
|
||||
numpy
|
||||
transformers>=4.39.3
|
||||
protobuf
|
||||
pyyaml
|
||||
jinja2
|
||||
+14
-290
@@ -1,310 +1,34 @@
|
||||
# Copyright © 2023-2024 Apple Inc.
|
||||
|
||||
import math
|
||||
from functools import partial
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
def make_sampler(
|
||||
temp: float = 0.0,
|
||||
top_p: float = 0.0,
|
||||
min_p: float = 0.0,
|
||||
min_tokens_to_keep: int = 1,
|
||||
top_k: int = 0,
|
||||
xtc_probability: float = 0.0,
|
||||
xtc_threshold: float = 0.0,
|
||||
xtc_special_tokens: List[int] = [],
|
||||
) -> Callable[mx.array, mx.array]:
|
||||
"""
|
||||
Make a sampler function for use with ``generate_step``.
|
||||
|
||||
Args:
|
||||
temp (float): The temperature for sampling, if 0 the argmax is used.
|
||||
Default: ``0``.
|
||||
top_p (float, optional): Nulceus sampling, higher means model considers
|
||||
more less likely words.
|
||||
min_p (float, optional): The minimum value (scaled by the top token's
|
||||
probability) that a token probability must have to be considered.
|
||||
min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
|
||||
be filtered by min_p sampling.
|
||||
top_k (int, optional): The top k tokens ranked by probability to constrain
|
||||
the sampling to.
|
||||
xtc_probability (float, optional): The probability of applying XTC
|
||||
sampling.
|
||||
xtc_threshold (float, optional): The threshold the probs need to reach
|
||||
for being sampled.
|
||||
xtc_special_tokens (list(int), optional): List of special tokens IDs to
|
||||
be excluded from XTC sampling.
|
||||
|
||||
|
||||
Returns:
|
||||
Callable[mx.array, mx.array]:
|
||||
A sampler which takes log-probabilities and returns tokens.
|
||||
"""
|
||||
if temp == 0:
|
||||
return lambda x: mx.argmax(x, axis=-1)
|
||||
|
||||
# Create sampler chain
|
||||
sampling_methods = []
|
||||
if top_k > 0:
|
||||
sampling_methods.append(lambda x: apply_top_k(x, top_k))
|
||||
if top_p > 0 and top_p < 1.0:
|
||||
sampling_methods.append(lambda x: apply_top_p(x, top_p))
|
||||
if min_p != 0.0:
|
||||
sampling_methods.append(lambda x: apply_min_p(x, min_p, min_tokens_to_keep))
|
||||
if xtc_probability > 0.0:
|
||||
sampling_methods.append(
|
||||
lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens)
|
||||
)
|
||||
|
||||
# Apply the sampling methods
|
||||
def sampler(logits):
|
||||
for method in sampling_methods:
|
||||
logits = method(logits)
|
||||
|
||||
# Return the sampled token
|
||||
return categorical_sampling(logits, temp)
|
||||
|
||||
return sampler
|
||||
|
||||
|
||||
def make_logits_processors(
|
||||
logit_bias: Optional[Dict[int, float]] = None,
|
||||
repetition_penalty: Optional[float] = None,
|
||||
repetition_context_size: Optional[int] = 20,
|
||||
):
|
||||
"""
|
||||
Make logits processors for use with ``generate_step``.
|
||||
|
||||
Args:
|
||||
repetition_penalty (float, optional): The penalty factor for repeating
|
||||
tokens.
|
||||
repetition_context_size (int, optional): The number of tokens to
|
||||
consider for repetition penalty. Default: ``20``.
|
||||
logit_bias (dictionary, optional): Additive logit bias.
|
||||
|
||||
Returns:
|
||||
List[Callable[[mx.array, mx.array], mx.array]]:
|
||||
A list of logits processors. Each processor in the list is a
|
||||
callable which takes an array of tokens and an array of logits
|
||||
and returns the updated logits.
|
||||
"""
|
||||
logits_processors = []
|
||||
if logit_bias:
|
||||
indices = mx.array(list(logit_bias.keys()))
|
||||
values = mx.array(list(logit_bias.values()))
|
||||
|
||||
def logit_bias_processor(_, logits):
|
||||
logits[:, indices] += values
|
||||
return logits
|
||||
|
||||
logits_processors.append(logit_bias_processor)
|
||||
|
||||
if repetition_penalty and repetition_penalty != 0.0:
|
||||
logits_processors.append(
|
||||
make_repetition_penalty(repetition_penalty, repetition_context_size)
|
||||
)
|
||||
return logits_processors
|
||||
|
||||
|
||||
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
|
||||
def apply_top_k(
|
||||
logprobs: mx.array,
|
||||
top_k: int,
|
||||
) -> mx.array:
|
||||
"""
|
||||
Sample from only the top K tokens ranked by probability.
|
||||
|
||||
Args:
|
||||
logprobs: A vector of log probabilities.
|
||||
top_k (int): Top k tokens to sample from.
|
||||
"""
|
||||
vocab_size = logprobs.shape[-1]
|
||||
if not isinstance(top_k, int) or not (0 < top_k < vocab_size):
|
||||
raise ValueError(
|
||||
f"`top_k` has to be an integer in the (0, {vocab_size}] interval,"
|
||||
f" but is {top_k}."
|
||||
)
|
||||
mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]
|
||||
masked_logprobs = mx.put_along_axis(
|
||||
logprobs, mask_idx, mx.array(-float("inf"), logprobs.dtype), axis=-1
|
||||
)
|
||||
return masked_logprobs
|
||||
|
||||
|
||||
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
|
||||
def apply_min_p(
|
||||
logprobs: mx.array,
|
||||
min_p: float,
|
||||
min_tokens_to_keep: int = 1,
|
||||
) -> mx.array:
|
||||
"""
|
||||
Apply min-p sampling to the logprobs.
|
||||
|
||||
Min-p keeps all tokens that are above a minimum probability, scaled by the
|
||||
probability of the most likely token. As a result, the filter is more
|
||||
aggressive given a very high-probability token.
|
||||
|
||||
Args:
|
||||
logprobs: A vector of log probabilities.
|
||||
min_p (float): Minimum token probability. Typical values are in the
|
||||
0.01-0.2 range, comparably selective as setting `top_p` in the
|
||||
0.99-0.8 range.
|
||||
min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
|
||||
be filtered. Default: ``1``.
|
||||
|
||||
"""
|
||||
if not (0 <= min_p <= 1.0):
|
||||
raise ValueError(
|
||||
f"`min_p` has to be a float in the [0, 1] interval, but is {min_p}"
|
||||
)
|
||||
if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
|
||||
raise ValueError(
|
||||
f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"
|
||||
)
|
||||
# reference implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L531-L605
|
||||
|
||||
# Indices sorted in decreasing order
|
||||
sorted_indices = mx.argsort(-logprobs, axis=-1)
|
||||
sorted_logprobs = mx.take_along_axis(logprobs, sorted_indices, axis=-1)
|
||||
|
||||
# Top probability
|
||||
top_logprobs = sorted_logprobs[:, 0:1]
|
||||
|
||||
# Calculate the min_p threshold
|
||||
scaled_min_p = top_logprobs + math.log(min_p)
|
||||
|
||||
# Mask tokens that have a probability less than the scaled min_p
|
||||
tokens_to_remove = sorted_logprobs < scaled_min_p
|
||||
tokens_to_remove[..., :min_tokens_to_keep] = False
|
||||
|
||||
# Create pool of tokens with probability less than scaled min_p
|
||||
selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)
|
||||
|
||||
# Create a mapping to rearrange back to original indices
|
||||
inverse_indices = mx.put_along_axis(
|
||||
mx.zeros_like(sorted_indices),
|
||||
sorted_indices,
|
||||
mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
|
||||
axis=-1,
|
||||
)
|
||||
|
||||
# Rearrange selected_logprobs back to original order
|
||||
original_order_logprobs = mx.take_along_axis(
|
||||
selected_logprobs, inverse_indices, axis=-1
|
||||
)
|
||||
|
||||
return original_order_logprobs
|
||||
|
||||
|
||||
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
|
||||
def apply_top_p(logprobs: mx.array, top_p: float) -> mx.array:
|
||||
def top_p_sampling(logits: mx.array, top_p: float, temperature: float) -> mx.array:
|
||||
"""
|
||||
Apply top-p (nucleus) sampling to logits.
|
||||
|
||||
Args:
|
||||
logprobs: A vector of log probabilities.
|
||||
logits: The logits from the model's output.
|
||||
top_p: The cumulative probability threshold for top-p filtering.
|
||||
temperature: Temperature parameter for softmax distribution reshaping.
|
||||
Returns:
|
||||
token selected based on the top-p criterion.
|
||||
"""
|
||||
# referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460
|
||||
probs = mx.exp(logprobs)
|
||||
# sort in ascending order
|
||||
sorted_indices = mx.argsort(logprobs, axis=-1)
|
||||
sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
|
||||
probs = mx.softmax(logits / temperature, axis=-1)
|
||||
|
||||
# sort probs in ascending order
|
||||
sorted_indices = mx.argsort(probs, axis=-1)
|
||||
sorted_probs = probs[..., sorted_indices.squeeze(0)]
|
||||
|
||||
cumulative_probs = mx.cumsum(sorted_probs, axis=-1)
|
||||
|
||||
# Rearrange cumulative probs back to original order
|
||||
inverse_indices = mx.put_along_axis(
|
||||
mx.zeros_like(sorted_indices),
|
||||
sorted_indices,
|
||||
mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
|
||||
axis=-1,
|
||||
)
|
||||
cumulative_probs = mx.take_along_axis(cumulative_probs, inverse_indices, axis=-1)
|
||||
|
||||
# select tokens with cumulative probs below threshold
|
||||
return mx.where(
|
||||
top_probs = mx.where(
|
||||
cumulative_probs > 1 - top_p,
|
||||
logprobs,
|
||||
-float("inf"),
|
||||
sorted_probs,
|
||||
mx.zeros_like(sorted_probs),
|
||||
)
|
||||
|
||||
sorted_token = mx.random.categorical(mx.log(top_probs))
|
||||
token = sorted_indices.squeeze(0)[sorted_token]
|
||||
|
||||
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
|
||||
def apply_xtc(
|
||||
logits: mx.array,
|
||||
xtc_probability: float,
|
||||
xtc_threshold: float,
|
||||
xtc_special_tokens: List[int],
|
||||
) -> mx.array:
|
||||
"""
|
||||
Apply XTC sampling to the logits.
|
||||
|
||||
Args:
|
||||
logits: The logits from the model's output.
|
||||
xtc_probability (float): Probability of XTC sampling to happen for each token
|
||||
xtc_threshold (float): The threshold the probs need to reach for being sampled.
|
||||
special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling.
|
||||
"""
|
||||
if not (0 <= xtc_threshold <= 0.5):
|
||||
raise ValueError(
|
||||
f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}"
|
||||
)
|
||||
if not (0 <= xtc_probability <= 1.0):
|
||||
raise ValueError(
|
||||
f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}"
|
||||
)
|
||||
|
||||
probs = mx.softmax(logits, -1)
|
||||
mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
|
||||
if xtc_special_tokens:
|
||||
mask[..., xtc_special_tokens] = False
|
||||
|
||||
return mx.where(
|
||||
mx.random.uniform(0, 1) > xtc_probability,
|
||||
logits,
|
||||
mx.where(mask, -mx.inf, logits),
|
||||
)
|
||||
|
||||
|
||||
@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
|
||||
def categorical_sampling(logits, temp):
|
||||
return mx.random.categorical(logits * (1 / temp))
|
||||
|
||||
|
||||
def make_repetition_penalty(penalty: float, context_size: int = 20):
|
||||
"""
|
||||
Make repetition penalty processor.
|
||||
|
||||
Paper: https://arxiv.org/abs/1909.05858
|
||||
|
||||
Args:
|
||||
penalty (float): The repetition penalty factor to be applied.
|
||||
context_size (int): The number of previous tokens to use.
|
||||
Default: ``20``.
|
||||
|
||||
Returns:
|
||||
Callable[[mx.array, List[int]], mx.array]:
|
||||
The repetition penalty processor.
|
||||
"""
|
||||
if penalty < 0 or not isinstance(penalty, (int, float)):
|
||||
raise ValueError(f"penalty must be a non-negative float, got {penalty}")
|
||||
|
||||
def repetition_penalty_processor(tokens, logits):
|
||||
if len(tokens) > 0:
|
||||
tokens = tokens[-context_size:]
|
||||
selected_logits = logits[:, tokens]
|
||||
selected_logits = mx.where(
|
||||
selected_logits < 0,
|
||||
selected_logits * penalty,
|
||||
selected_logits / penalty,
|
||||
)
|
||||
logits[:, tokens] = selected_logits
|
||||
return logits
|
||||
|
||||
return repetition_penalty_processor
|
||||
return token
|
||||
|
||||
+187
-678
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user