gi111f11cb7f354a74a8c123f5b155498e
/
OpenELM-1_1B-Instruct_a13593210708684800672017
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

first commit

This commit is contained in:
xxl 2024-11-14 15:47:40 +08:00
parent 3df7e8e54d
commit cfc523952e
8 changed files with 1934 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
LICENSE Normal file
View File

@ -0,0 +1,47 @@
Copyright (C) 2024 Apple Inc. All Rights Reserved.
Disclaimer: IMPORTANT: This Apple software is supplied to you by Apple
Inc. ("Apple") in consideration of your agreement to the following
terms, and your use, installation, modification or redistribution of
this Apple software constitutes acceptance of these terms. If you do
not agree with these terms, please do not use, install, modify or
redistribute this Apple software.
In consideration of your agreement to abide by the following terms, and
subject to these terms, Apple grants you a personal, non-exclusive
license, under Apple's copyrights in this original Apple software (the
"Apple Software"), to use, reproduce, modify and redistribute the Apple
Software, with or without modifications, in source and/or binary forms;
provided that if you redistribute the Apple Software in its entirety and
without modifications, you must retain this notice and the following
text and disclaimers in all such redistributions of the Apple Software.
Neither the name, trademarks, service marks or logos of Apple Inc. may
be used to endorse or promote products derived from the Apple Software
without specific prior written permission from Apple. Except as
expressly stated in this notice, no other rights or licenses, express or
implied, are granted by Apple herein, including but not limited to any
patent rights that may be infringed by your derivative works or by other
works in which the Apple Software may be incorporated.
The Apple Software is provided by Apple on an "AS IS" basis. APPLE
MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION
THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS
FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND
OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS.
IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL
OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION,
MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED
AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE),
STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
-------------------------------------------------------------------------------
SOFTWARE DISTRIBUTED IN THIS REPOSITORY:
This software includes a number of subcomponents with separate
copyright notices and license terms - please see the file ACKNOWLEDGEMENTS.
-------------------------------------------------------------------------------

190
README.md
View File

@ -1,3 +1,189 @@
# OpenELM-1_1B-Instruct_a13593210708684800672017
---
license: other
license_name: apple-sample-code-license
license_link: LICENSE
---
OpenELM-1_1B-Instruct
# OpenELM
*Sachin Mehta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Iman Mirzadeh, Mahyar Najibi, Dmitry Belenko, Peter Zatloukal, Mohammad Rastegari*
We introduce **OpenELM**, a family of **Open** **E**fficient **L**anguage **M**odels. OpenELM uses a layer-wise scaling strategy to efficiently allocate parameters within each layer of the transformer model, leading to enhanced accuracy. We pretrained OpenELM models using the [CoreNet](https://github.com/apple/corenet) library. We release both pretrained and instruction tuned models with 270M, 450M, 1.1B and 3B parameters. We release the complete framework, encompassing data preparation, training, fine-tuning, and evaluation procedures, alongside multiple pre-trained checkpoints and training logs, to facilitate open research.
Our pre-training dataset contains RefinedWeb, deduplicated PILE, a subset of RedPajama, and a subset of Dolma v1.6, totaling approximately 1.8 trillion tokens. Please check license agreements and terms of these datasets before using them.
## Usage
We have provided an example function to generate output from OpenELM models loaded via [HuggingFace Hub](https://huggingface.co/docs/hub/) in `generate_openelm.py`.
You can try the model by running the following command:
```
python generate_openelm.py --model apple/OpenELM-1_1B-Instruct --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2
```
Please refer to [this link](https://huggingface.co/docs/hub/security-tokens) to obtain your hugging face access token.
Additional arguments to the hugging face generate function can be passed via `generate_kwargs`. As an example, to speedup the inference, you can try [lookup token speculative generation](https://huggingface.co/docs/transformers/generation_strategies) by passing the `prompt_lookup_num_tokens` argument as follows:
```
python generate_openelm.py --model apple/OpenELM-1_1B-Instruct --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 prompt_lookup_num_tokens=10
```
Alternatively, try model-wise speculative generation with an [assistive model](https://huggingface.co/blog/assisted-generation) by passing a smaller model through the `assistant_model` argument, for example:
```
python generate_openelm.py --model apple/OpenELM-1_1B-Instruct --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 --assistant_model [SMALLER_MODEL]
```
## Main Results
### Zero-Shot
| **Model Size** | **ARC-c** | **ARC-e** | **BoolQ** | **HellaSwag** | **PIQA** | **SciQ** | **WinoGrande** | **Average** |
|-----------------------------------------------------------------------------|-----------|-----------|-----------|---------------|-----------|-----------|----------------|-------------|
| [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M) | 26.45 | 45.08 | **53.98** | 46.71 | 69.75 | **84.70** | **53.91** | 54.37 |
| [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) | **30.55** | **46.68** | 48.56 | **52.07** | **70.78** | 84.40 | 52.72 | **55.11** |
| [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M) | 27.56 | 48.06 | 55.78 | 53.97 | 72.31 | 87.20 | 58.01 | 57.56 |
| [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) | **30.38** | **50.00** | **60.37** | **59.34** | **72.63** | **88.00** | **58.96** | **59.95** |
| [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B) | 32.34 | **55.43** | 63.58 | 64.81 | **75.57** | **90.60** | 61.72 | 63.44 |
| [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) | **37.97** | 52.23 | **70.00** | **71.20** | 75.03 | 89.30 | **62.75** | **65.50** |
| [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B) | 35.58 | 59.89 | 67.40 | 72.44 | 78.24 | **92.70** | 65.51 | 67.39 |
| [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct) | **39.42** | **61.74** | **68.17** | **76.36** | **79.00** | 92.50 | **66.85** | **69.15** |
### LLM360
| **Model Size** | **ARC-c** | **HellaSwag** | **MMLU** | **TruthfulQA** | **WinoGrande** | **Average** |
|-----------------------------------------------------------------------------|-----------|---------------|-----------|----------------|----------------|-------------|
| [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M) | 27.65 | 47.15 | 25.72 | **39.24** | **53.83** | 38.72 |
| [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) | **32.51** | **51.58** | **26.70** | 38.72 | 53.20 | **40.54** |
| [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M) | 30.20 | 53.86 | **26.01** | 40.18 | 57.22 | 41.50 |
| [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) | **33.53** | **59.31** | 25.41 | **40.48** | **58.33** | **43.41** |
| [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B) | 36.69 | 65.71 | **27.05** | 36.98 | 63.22 | 45.93 |
| [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) | **41.55** | **71.83** | 25.65 | **45.95** | **64.72** | **49.94** |
| [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B) | 42.24 | 73.28 | **26.76** | 34.98 | 67.25 | 48.90 |
| [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct) | **47.70** | **76.87** | 24.80 | **38.76** | **67.96** | **51.22** |
### OpenLLM Leaderboard
| **Model Size** | **ARC-c** | **CrowS-Pairs** | **HellaSwag** | **MMLU** | **PIQA** | **RACE** | **TruthfulQA** | **WinoGrande** | **Average** |
|-----------------------------------------------------------------------------|-----------|-----------------|---------------|-----------|-----------|-----------|----------------|----------------|-------------|
| [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M) | 27.65 | **66.79** | 47.15 | 25.72 | 69.75 | 30.91 | **39.24** | **53.83** | 45.13 |
| [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) | **32.51** | 66.01 | **51.58** | **26.70** | **70.78** | 33.78 | 38.72 | 53.20 | **46.66** |
| [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M) | 30.20 | **68.63** | 53.86 | **26.01** | 72.31 | 33.11 | 40.18 | 57.22 | 47.69 |
| [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) | **33.53** | 67.44 | **59.31** | 25.41 | **72.63** | **36.84** | **40.48** | **58.33** | **49.25** |
| [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B) | 36.69 | **71.74** | 65.71 | **27.05** | **75.57** | 36.46 | 36.98 | 63.22 | 51.68 |
| [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) | **41.55** | 71.02 | **71.83** | 25.65 | 75.03 | **39.43** | **45.95** | **64.72** | **54.40** |
| [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B) | 42.24 | **73.29** | 73.28 | **26.76** | 78.24 | **38.76** | 34.98 | 67.25 | 54.35 |
| [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct) | **47.70** | 72.33 | **76.87** | 24.80 | **79.00** | 38.47 | **38.76** | **67.96** | **55.73** |
See the technical report for more results and comparison.
## Evaluation
### Setup
Install the following dependencies:
```bash
# install public lm-eval-harness
harness_repo="public-lm-eval-harness"
git clone https://github.com/EleutherAI/lm-evaluation-harness ${harness_repo}
cd ${harness_repo}
# use main branch on 03-15-2024, SHA is dc90fec
git checkout dc90fec
pip install -e .
cd ..
# 66d6242 is the main branch on 2024-04-01
pip install datasets@git+https://github.com/huggingface/datasets.git@66d6242
pip install tokenizers>=0.15.2 transformers>=4.38.2 sentencepiece>=0.2.0
```
### Evaluate OpenELM
```bash
# OpenELM-1_1B-Instruct
hf_model=apple/OpenELM-1_1B-Instruct
# this flag is needed because lm-eval-harness set add_bos_token to False by default, but OpenELM uses LLaMA tokenizer which requires add_bos_token to be True
tokenizer=meta-llama/Llama-2-7b-hf
add_bos_token=True
batch_size=1
mkdir lm_eval_output
shot=0
task=arc_challenge,arc_easy,boolq,hellaswag,piqa,race,winogrande,sciq,truthfulqa_mc2
lm_eval --model hf \
--model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
--tasks ${task} \
--device cuda:0 \
--num_fewshot ${shot} \
--output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
--batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
shot=5
task=mmlu,winogrande
lm_eval --model hf \
--model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
--tasks ${task} \
--device cuda:0 \
--num_fewshot ${shot} \
--output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
--batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
shot=25
task=arc_challenge,crows_pairs_english
lm_eval --model hf \
--model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
--tasks ${task} \
--device cuda:0 \
--num_fewshot ${shot} \
--output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
--batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
shot=10
task=hellaswag
lm_eval --model hf \
--model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
--tasks ${task} \
--device cuda:0 \
--num_fewshot ${shot} \
--output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
--batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
```
## Bias, Risks, and Limitations
The release of OpenELM models aims to empower and enrich the open research community by providing access to state-of-the-art language models. Trained on publicly available datasets, these models are made available without any safety guarantees. Consequently, there exists the possibility of these models producing outputs that are inaccurate, harmful, biased, or objectionable in response to user prompts. Thus, it is imperative for users and developers to undertake thorough safety testing and implement appropriate filtering mechanisms tailored to their specific requirements.
## Citation
If you find our work useful, please cite:
```BibTex
@article{mehtaOpenELMEfficientLanguage2024,
title = {{OpenELM}: {An} {Efficient} {Language} {Model} {Family} with {Open} {Training} and {Inference} {Framework}},
shorttitle = {{OpenELM}},
url = {https://arxiv.org/abs/2404.14619v1},
language = {en},
urldate = {2024-04-24},
journal = {arXiv.org},
author = {Mehta, Sachin and Sekhavat, Mohammad Hossein and Cao, Qingqing and Horton, Maxwell and Jin, Yanzi and Sun, Chenfan and Mirzadeh, Iman and Najibi, Mahyar and Belenko, Dmitry and Zatloukal, Peter and Rastegari, Mohammad},
month = apr,
year = {2024},
}
@inproceedings{mehta2022cvnets,
author = {Mehta, Sachin and Abdolhosseini, Farzad and Rastegari, Mohammad},
title = {CVNets: High Performance Library for Computer Vision},
year = {2022},
booktitle = {Proceedings of the 30th ACM International Conference on Multimedia},
series = {MM '22}
}
```

124
config.json Normal file
View File

@ -0,0 +1,124 @@
{
"activation_fn_name": "swish",
"architectures": [
"OpenELMForCausalLM"
],
"auto_map": {
"AutoConfig": "configuration_openelm.OpenELMConfig",
"AutoModelForCausalLM": "modeling_openelm.OpenELMForCausalLM"
},
"bos_token_id": 1,
"eos_token_id": 2,
"ffn_dim_divisor": 256,
"ffn_multipliers": [
0.5,
0.63,
0.76,
0.89,
1.02,
1.15,
1.28,
1.41,
1.54,
1.67,
1.8,
1.93,
2.06,
2.19,
2.31,
2.44,
2.57,
2.7,
2.83,
2.96,
3.09,
3.22,
3.35,
3.48,
3.61,
3.74,
3.87,
4.0
],
"ffn_with_glu": true,
"head_dim": 64,
"initializer_range": 0.02,
"max_context_length": 2048,
"model_dim": 2048,
"model_type": "openelm",
"normalization_layer_name": "rms_norm",
"normalize_qk_projections": true,
"num_gqa_groups": 4,
"num_kv_heads": [
4,
4,
4,
5,
5,
5,
5,
5,
5,
5,
6,
6,
6,
6,
6,
6,
6,
6,
7,
7,
7,
7,
7,
7,
8,
8,
8,
8
],
"num_query_heads": [
16,
16,
16,
20,
20,
20,
20,
20,
20,
20,
24,
24,
24,
24,
24,
24,
24,
24,
28,
28,
28,
28,
28,
28,
32,
32,
32,
32
],
"num_transformer_layers": 28,
"qkv_multipliers": [
0.5,
1.0
],
"rope_freq_constant": 10000,
"rope_max_length": 4096,
"share_input_output_layers": true,
"torch_dtype": "bfloat16",
"transformers_version": "4.39.3",
"use_cache": true,
"vocab_size": 32000
}

318
configuration_openelm.py Normal file
View File

@ -0,0 +1,318 @@
#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
#
"""Implements HF OpenELMConfig based on PretrainedConfig"""
from numbers import Number
from typing import List, Optional, Union
import numpy as np
from transformers import PretrainedConfig
def make_divisible(
v: Union[float, int],
divisor: Optional[int] = 8,
min_value: Optional[Union[float, int]] = None,
) -> Union[float, int]:
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by the divisor
It can be seen at:
https://github.com/tensorflow/models/blob/2cfc99eff5e5eb729c6793d2f3d03aa1c9be2b15/research/slim/nets/mobilenet/mobilenet.py#L62
Args:
v: input value
divisor: default to 8
min_value: minimum divisor value
Returns:
new_v: new divisible value
"""
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_v < 0.9 * v:
new_v += divisor
return new_v
def compute_heads(model_dim: int, head_dim: int) -> int:
"""Compute the number of heads.
Args:
model_dim: Model dimension.
head_dim: Head dimension.
Returns:
An integer denoting number of heads in multi-head attention is returned.
Raises:
ValueError: if model dimension is not divisible by head dimension.
"""
if model_dim % head_dim == 0:
return model_dim // head_dim
else:
raise ValueError(
f"Model dimension should be divisible by head dimension. Got: {model_dim} and {head_dim}."
)
OpenELM_CONFIGS = {
"OpenELM-270M": dict(
num_transformer_layers=16,
model_dim=1280,
head_dim=64,
num_gqa_groups=4,
normalize_qk_projections=True,
share_input_output_layers=True,
# Vary the FFN and QKV multipliers to create variable FFN and attention layers respectively.
ffn_multipliers=(0.5, 4.0),
qkv_multipliers=(0.5, 1.0),
),
"OpenELM-450M": dict(
num_transformer_layers=20,
model_dim=1536,
head_dim=64,
num_gqa_groups=4,
normalize_qk_projections=True,
share_input_output_layers=True,
# Vary the FFN and QKV multipliers to create variable FFN and attention layers respectively.
ffn_multipliers=(0.5, 4.0),
qkv_multipliers=(0.5, 1.0),
),
"OpenELM-1_1B": dict(
num_transformer_layers=28,
model_dim=2048,
head_dim=64,
num_gqa_groups=4,
normalize_qk_projections=True,
share_input_output_layers=True,
# Vary the FFN and QKV multipliers to create variable FFN and attention layers respectively.
ffn_multipliers=(0.5, 4.0),
qkv_multipliers=(0.5, 1.0),
),
"OpenELM-3B": dict(
num_transformer_layers=36,
model_dim=3072,
head_dim=128,
num_gqa_groups=4,
normalize_qk_projections=True,
share_input_output_layers=True,
# Vary the FFN and QKV multipliers to create variable FFN and attention layers respectively.
ffn_multipliers=(0.5, 4.0),
qkv_multipliers=(0.5, 1.0),
),
}
class OpenELMConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`OpenELMModel`]. It is used to instantiate an OpenELM model according to the specified arguments, defining the model architecture.
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
vocab_size (`int`, *optional*, defaults to 32000):
Vocabulary size of the OpenELM model.
max_context_length (`int`, *optional*, defaults to 2048):
Maximum number of input tokens.
num_transformer_layers (`int`, *optional*, defaults to 12):
Number of hidden layers in the Transformer decoder.
model_dim (`int`, *optional*, defaults to 2048):
Dimension of the hidden representations.
head_dim (`int`, *optional*, defaults to 128):
The attention head dimension.
qkv_multipliers (`Union[Number, List[Number]]`, *optional*, defaults to 1.0):
If the qkv_multipliers is a Number, then all attention layers have the same latent dimensions,
resulting in uniform allocation of parameters.
If the qkv_multipliers is a List of Number, then each attention layer have different latent dimensions
assuming qkv_multipliers[0] != qkv_multipliers[1]. This results in variable allocation of parameters in attention layer.
This scaling is known as layer-wise or block-wise scaling: https://arxiv.org/abs/2008.00623
num_query_heads (`Union[int, None]`, *optional*, defaults to None):
The number of query heads, computed from `compute_heads(model_dim=model_dim, head_dim=head_dim)`.
num_gqa_groups (`int`, *optional*, defaults to 1):
This variable allows to switch between multi-head attention, group query attention, and multi-query attention.
When num_gqa_groups == 1, then it is multi-head attention.
When 1 < num_gqa_groups < num_heads and num_heads is divisible by num_gqa_groups, then it is group query attention
When num_gqa_groups == num_heads, then it is multi-query attention
ffn_multipliers (`Union[Number, List[Number]]`, *optional*, defaults to 4.0):
Feed-forward network (FFN) multipliers.
If the ffn_multipliers is a Number, then all FFN layers have the same latent dimensions,
resulting in uniform allocation of parameters.
If the ffn_multipliers is a List of Number, then each FFN layer have different latent dimensions
assuming ffn_multipliers[0] != ffn_multipliers[1]. This results in variable allocation of parameters in FFN layer.
This scaling is known as layer-wise or block-wise scaling: https://arxiv.org/abs/2008.00623
ffn_with_glu (`bool`, *optional*, defaults to True):
Whether to use FFN with Gated Linear Unit (GLU)
ffn_dim_divisor (`int`, *optional*, defaults to 256):
The ffn layer dimension divisor.
activation_fn_name (`str` or `function`, *optional*, defaults to `"swish"`):
The non-linear activation function (function or string) in the decoder.
normalization_layer_name (`str` or `function`, *optional*, defaults to `"rms_norm"`):
Type of normalization layer.
normalize_qk_projections (`bool`, *optional*, defaults to False):
Whether to normalize queries and keys after projections
share_input_output_layers (`bool`, *optional*, defaults to False):
Whether to share the embedding between input and output linear layer
rope_freq_constant (`int`, *optional*, defaults to 10000):
The base period of the RoPE embeddings.
rope_max_length (`int`, *optional*, defaults to 4096):
That rope_max_length is set to twice of max_context_length.
This allows flexibility in token lengths during training or fine-tuning.
initializer_range (`float`, *optional*, defaults to 0.02):
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
use_cache (`bool`, *optional*, defaults to `True`):
Whether or not the model should return the last key/values attentions (not used by all models). Only
relevant if `config.is_decoder=True`.
bos_token_id (`int`, *optional*, defaults to 2):
Beginning of stream token id.
eos_token_id (`int`, *optional*, defaults to 1):
End of stream token id.
"""
model_type = "openelm"
def __init__(
self,
vocab_size: int = 32000,
max_context_length: int = 2048,
num_transformer_layers: int = 12,
model_dim: int = 2048,
head_dim: int = 128,
qkv_multipliers: Union[Number, List[Number]] = 1.0,
num_query_heads: Union[int, None] = None,
num_gqa_groups: int = 1,
ffn_multipliers: Union[Number, List[Number]] = 4.0,
ffn_with_glu: bool = True,
ffn_dim_divisor: int = 256,
activation_fn_name: str = "swish",
normalization_layer_name: str = "rms_norm",
normalize_qk_projections: bool = False,
share_input_output_layers: bool = False,
rope_freq_constant: int = 10000,
rope_max_length: int = 4096,
initializer_range: float = 0.02,
use_cache: bool = True,
bos_token_id: int = 1,
eos_token_id: int = 2,
**kwargs,
) -> None:
self.vocab_size = vocab_size
self.max_context_length = max_context_length
self.num_transformer_layers = num_transformer_layers
self.model_dim = model_dim
self.head_dim = head_dim
self.qkv_multipliers = qkv_multipliers
self.num_query_heads = num_query_heads
self.num_gqa_groups = num_gqa_groups
self.ffn_multipliers = ffn_multipliers
self.ffn_with_glu = ffn_with_glu
self.ffn_dim_divisor = ffn_dim_divisor
self.activation_fn_name = activation_fn_name
self.normalization_layer_name = normalization_layer_name
self.normalize_qk_projections = normalize_qk_projections
self.share_input_output_layers = share_input_output_layers
self.rope_freq_constant = rope_freq_constant
self.rope_max_length = rope_max_length
self.num_query_heads = (
compute_heads(model_dim=model_dim, head_dim=head_dim)
if num_query_heads is None
else num_query_heads
)
self.initializer_range = initializer_range
self.__post_init__()
super().__init__(
use_cache=use_cache,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
**kwargs,
)
def __post_init__(self) -> None:
if self.num_gqa_groups is not None:
head_multiple_of = self.num_gqa_groups
else:
head_multiple_of = 2
if isinstance(self.qkv_multipliers, Number):
# All attention layers have the same latent dimensions, resulting in uniform allocation of parameters.
qkv_dim = make_divisible(
self.model_dim * self.qkv_multipliers,
divisor=self.head_dim * head_multiple_of,
)
query_dims = [int(qkv_dim)] * self.num_transformer_layers
elif (
isinstance(self.qkv_multipliers, (tuple, list))
and len(self.qkv_multipliers) == 2
):
# Each attention layer have different latent dimensions assuming qkv_multipliers[0] != qkv_multipliers[1].
# This results in variable allocation of parameters in attention layer.
# This scaling is known as layer-wise or block-wise scaling: https://arxiv.org/abs/2008.00623
qkv_multipliers = [
round(v, 2)
for v in np.linspace(
self.qkv_multipliers[0],
self.qkv_multipliers[1],
num=self.num_transformer_layers,
dtype=float,
)
]
# Make sure that scaled model dimension is divisible by scaled head dimension.
query_dims = [
int(
make_divisible(
self.model_dim * m, divisor=self.head_dim * head_multiple_of
)
)
for m in qkv_multipliers
]
else:
raise NotImplementedError(
f"QKV multipliers should be a single number or a list containing exactly two numbers. Got: {qkv_multipliers}."
)
# compute the number of query, key, and value heads
# For multi-head and multi-query attention, the number of heads for query, key, and value are the same.
# For group query attention, the number of key and value heads are the same.
self.num_query_heads = [
int(compute_heads(q_dim, self.head_dim)) for q_dim in query_dims
]
self.num_kv_heads = [
q_heads // self.num_gqa_groups for q_heads in self.num_query_heads
]
# Feed-forward network (FFN) multipliers
if isinstance(self.ffn_multipliers, Number):
# All FFN layers have the same latent dimensions, resulting in uniform allocation of parameters.
self.ffn_multipliers = [self.ffn_multipliers] * self.num_transformer_layers
elif isinstance(self.ffn_multipliers, (tuple, list)):
# Each FFN layer have different latent dimensions assuming ffn_multipliers[0] != ffn_multipliers[1].
# This results in variable allocation of parameters in FFN layer.
# This scaling is known as layer-wise or block-wise scaling: https://arxiv.org/abs/2008.00623
if len(self.ffn_multipliers) == 2:
self.ffn_multipliers = [
round(v, 2)
for v in np.linspace(
self.ffn_multipliers[0],
self.ffn_multipliers[1],
num=self.num_transformer_layers,
dtype=float,
)
]
else:
assert (
len(self.ffn_multipliers) == self.num_transformer_layers
), f"{len(self.ffn_multipliers)=}!={self.num_transformer_layers=}"
else:
raise NotImplementedError(
f"FFN multipliers should be a single number or a list containing exactly two numbers. Got: {qkv_multipliers}."
)
# check num_query_heads divisible by num_kv_heads for every layer
for layer_idx in range(len(query_dims)):
assert self.num_query_heads[layer_idx] % self.num_kv_heads[layer_idx] == 0

240
generate_openelm.py Normal file
View File

@ -0,0 +1,240 @@
#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
#
"""Module to generate OpenELM output given a model and an input prompt."""
import os
import logging
import time
import argparse
from typing import Optional, Union
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
def generate(
prompt: str,
model: Union[str, AutoModelForCausalLM],
hf_access_token: str = None,
tokenizer: Union[str, AutoTokenizer] = 'meta-llama/Llama-2-7b-hf',
device: Optional[str] = None,
max_length: int = 1024,
assistant_model: Optional[Union[str, AutoModelForCausalLM]] = None,
generate_kwargs: Optional[dict] = None,
) -> str:
""" Generates output given a prompt.
Args:
prompt: The string prompt.
model: The LLM Model. If a string is passed, it should be the path to
the hf converted checkpoint.
hf_access_token: Hugging face access token.
tokenizer: Tokenizer instance. If model is set as a string path,
the tokenizer will be loaded from the checkpoint.
device: String representation of device to run the model on. If None
and cuda available it would be set to cuda:0 else cpu.
max_length: Maximum length of tokens, input prompt + generated tokens.
assistant_model: If set, this model will be used for
speculative generation. If a string is passed, it should be the
path to the hf converted checkpoint.
generate_kwargs: Extra kwargs passed to the hf generate function.
Returns:
output_text: output generated as a string.
generation_time: generation time in seconds.
Raises:
ValueError: If device is set to CUDA but no CUDA device is detected.
ValueError: If tokenizer is not set.
ValueError: If hf_access_token is not specified.
"""
if not device:
if torch.cuda.is_available() and torch.cuda.device_count():
device = "cuda:0"
logging.warning(
'inference device is not set, using cuda:0, %s',
torch.cuda.get_device_name(0)
)
else:
device = 'cpu'
logging.warning(
(
'No CUDA device detected, using cpu, '
'expect slower speeds.'
)
)
if 'cuda' in device and not torch.cuda.is_available():
raise ValueError('CUDA device requested but no CUDA device detected.')
if not tokenizer:
raise ValueError('Tokenizer is not set in the generate function.')
if not hf_access_token:
raise ValueError((
'Hugging face access token needs to be specified. '
'Please refer to https://huggingface.co/docs/hub/security-tokens'
' to obtain one.'
)
)
if isinstance(model, str):
checkpoint_path = model
model = AutoModelForCausalLM.from_pretrained(
checkpoint_path,
trust_remote_code=True
)
model.to(device).eval()
if isinstance(tokenizer, str):
tokenizer = AutoTokenizer.from_pretrained(
tokenizer,
token=hf_access_token,
)
# Speculative mode
draft_model = None
if assistant_model:
draft_model = assistant_model
if isinstance(assistant_model, str):
draft_model = AutoModelForCausalLM.from_pretrained(
assistant_model,
trust_remote_code=True
)
draft_model.to(device).eval()
# Prepare the prompt
tokenized_prompt = tokenizer(prompt)
tokenized_prompt = torch.tensor(
tokenized_prompt['input_ids'],
device=device
)
tokenized_prompt = tokenized_prompt.unsqueeze(0)
# Generate
stime = time.time()
output_ids = model.generate(
tokenized_prompt,
max_length=max_length,
pad_token_id=0,
assistant_model=draft_model,
**(generate_kwargs if generate_kwargs else {}),
)
generation_time = time.time() - stime
output_text = tokenizer.decode(
output_ids[0].tolist(),
skip_special_tokens=True
)
return output_text, generation_time
def openelm_generate_parser():
"""Argument Parser"""
class KwargsParser(argparse.Action):
"""Parser action class to parse kwargs of form key=value"""
def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, dict())
for val in values:
if '=' not in val:
raise ValueError(
(
'Argument parsing error, kwargs are expected in'
' the form of key=value.'
)
)
kwarg_k, kwarg_v = val.split('=')
try:
converted_v = int(kwarg_v)
except ValueError:
try:
converted_v = float(kwarg_v)
except ValueError:
converted_v = kwarg_v
getattr(namespace, self.dest)[kwarg_k] = converted_v
parser = argparse.ArgumentParser('OpenELM Generate Module')
parser.add_argument(
'--model',
dest='model',
help='Path to the hf converted model.',
required=True,
type=str,
)
parser.add_argument(
'--hf_access_token',
dest='hf_access_token',
help='Hugging face access token, starting with "hf_".',
type=str,
)
parser.add_argument(
'--prompt',
dest='prompt',
help='Prompt for LLM call.',
default='',
type=str,
)
parser.add_argument(
'--device',
dest='device',
help='Device used for inference.',
type=str,
)
parser.add_argument(
'--max_length',
dest='max_length',
help='Maximum length of tokens.',
default=256,
type=int,
)
parser.add_argument(
'--assistant_model',
dest='assistant_model',
help=(
(
'If set, this is used as a draft model '
'for assisted speculative generation.'
)
),
type=str,
)
parser.add_argument(
'--generate_kwargs',
dest='generate_kwargs',
help='Additional kwargs passed to the HF generate function.',
type=str,
nargs='*',
action=KwargsParser,
)
return parser.parse_args()
if __name__ == '__main__':
args = openelm_generate_parser()
prompt = args.prompt
output_text, genertaion_time = generate(
prompt=prompt,
model=args.model,
device=args.device,
max_length=args.max_length,
assistant_model=args.assistant_model,
generate_kwargs=args.generate_kwargs,
hf_access_token=args.hf_access_token,
)
print_txt = (
f'\r\n{"=" * os.get_terminal_size().columns}\r\n'
'\033[1m Prompt + Generated Output\033[0m\r\n'
f'{"-" * os.get_terminal_size().columns}\r\n'
f'{output_text}\r\n'
f'{"-" * os.get_terminal_size().columns}\r\n'
'\r\nGeneration took'
f'\033[1m\033[92m {round(genertaion_time, 2)} \033[0m'
'seconds.\r\n'
)
print(print_txt)

6
generation_config.json Normal file
View File

@ -0,0 +1,6 @@
{
"_from_model_config": true,
"bos_token_id": 1,
"eos_token_id": 2,
"transformers_version": "4.39.3"
}

BIN
model.safetensors (Stored with Git LFS) Normal file
View File

Binary file not shown.

1008
modeling_openelm.py Normal file
View File

File diff suppressed because it is too large Load Diff