mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-09-01 04:14:38 +08:00
support internlm2 (#797)
* support internlm2 * only attention projections --------- Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
@@ -397,6 +397,23 @@ class TestModels(unittest.TestCase):
|
||||
len(args.ffn_multipliers),
|
||||
)
|
||||
|
||||
def test_internlm2(self):
|
||||
from mlx_lm.models import internlm2
|
||||
|
||||
args = internlm2.ModelArgs(
|
||||
model_type="internlm2",
|
||||
hidden_size=1024,
|
||||
num_hidden_layers=4,
|
||||
intermediate_size=2048,
|
||||
num_attention_heads=4,
|
||||
rms_norm_eps=1e-5,
|
||||
vocab_size=10000,
|
||||
)
|
||||
model = internlm2.Model(args)
|
||||
self.model_test_runner(
|
||||
model, args.model_type, args.vocab_size, args.num_hidden_layers
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
Reference in New Issue
Block a user