mlx Model (loglikelihood & generate_until)

[chimezie]

This adds a new model type for mlx models. In particular, it implements the loglikelihood and generate_until interfaces. Works with the current versions of mlx and mlx-lm

The new model type is mlx, so the harness can be run this way to evaluate against a local mlx model:

lm_eval --model mlx --model_args model=.. model name or path ..   --tasks medqa_4options 

The expected model args are:

• model (huggingface model or local path to mlx model)
• adapter_path (path to a LoRa adapter to apply to the model)
• trust_remote_code
• eos_token
• top_p (defaults to 1)
• max_tokens (defaults to 2048)
• batch_size (defaults to 4)
• max_gen_tokens (defaults to 256)
• ensure_bos_token (defaults to False) : Whether or not to ensure the first token is a defined BOS token

[CLAassistant]

All committers have signed the CLA.

[chimezie]

I'm getting the following traceback running the evaluation this way (in an environment with mlx and mlx-lm):

lm_eval --model mlx --model_args model=internistai/base-7b-v0.2 \
    --tasks medqa_4options \
    --batch_size 64

Traceback:

2024-05-29:13:18:14,114 INFO     [__main__.py:254] Verbosity set to INFO
2024-05-29:13:18:16,354 INFO     [__main__.py:341] Selected Tasks: ['medqa_4options']
2024-05-29:13:18:16,355 INFO     [evaluator.py:141] Setting random seed to 0 | Setting numpy seed to 1234 | Setting torch manual seed to 1234
2024-05-29:13:18:16,355 INFO     [evaluator.py:178] Initializing mlx model, with arguments: {'model': 'internistai/base-7b-v0.2'}
Fetching 9 files: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9/9 [00:00<00:00, 32968.33it/s]
You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
2024-05-29:13:18:20,863 INFO     [mlx_llms.py:28] Model type is '<class 'mlx_lm.models.llama.Model'>
2024-05-29:13:18:22,781 INFO     [task.py:398] Building contexts for medqa_4options on rank 0...
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1273/1273 [00:00<00:00, 198223.53it/s]
2024-05-29:13:18:22,818 INFO     [evaluator.py:395] Running loglikelihood requests
Running loglikelihood requests (79 batches):  37%|███████████████████████████████████████▋                                                                    | 29/79 [10:13<15:22, 18.46s/it]Running loglikelihood requests (79 batches): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 79/79 [26:40<00:00, 20.26s/it]
[..snip..]
Traceback (most recent call last):
  File "/path/to/bin/lm_eval", line 8, in <module>
    sys.exit(cli_evaluate())
             ^^^^^^^^^^^^^^
  File "/path/to/lm_eval/__main__.py", line 347, in cli_evaluate
    results = evaluator.simple_evaluate(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/to/lm_eval/utils.py", line 321, in _wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/path/to/lm_eval/evaluator.py", line 256, in simple_evaluate
    results = evaluate(
              ^^^^^^^^^
  File "/path/to/lm_eval/utils.py", line 321, in _wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/path/to/lm_eval/evaluator.py", line 421, in evaluate
    task.apply_filters()
  File "/path/to/lm_eval/api/task.py", line 1000, in apply_filters
    f.apply(self._instances)
  File "/path/to/lm_eval/api/filter.py", line 55, in apply
    for inst, resp in zip(instances, resps):
  File "/path/to/lm_eval/filters/selection.py", line 23, in <lambda>
    return map(lambda r: r[0], resps)

The implemented loglikelihood function returns a list of 5,056 pairs of (log-likelihood, boolean). However, for some reason, the TakeFirstFilter.apply method receives a resps parameter with 5,092 resources, the last of which are empty lists, which seems to be causing the traceback.

Any help would be greatly appreciated.

[chimezie]

However, I was able to run it against mmlu_professional_medicine:

lm_eval --model mlx --model_args model=internistai/base-7b-v0.2 \
>     --tasks mmlu_professional_medicine \
>     --batch_size 64
[..snip..]
mlx (model=internistai/base-7b-v0.2), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 64
|        Tasks        |Version|Filter|n-shot|Metric|Value |   |Stderr|
|---------------------|------:|------|-----:|------|-----:|---|-----:|
|professional_medicine|      0|none  |     0|acc   |0.1838|±  |0.0235|

[chimezie]

Oddly enough, I can get a clean eval of internistai/base-7b-v0.2 against mmlu_professional_medicine tasks on MLX and then HF but still get the issue above when run against the medqa_4options task:

% time lm_eval --model mlx --model_args model=internistai/base-7b-v0.2 \                                              
    --tasks mmlu_professional_medicine \
    --batch_size 64 
2024-05-31:15:31:05,832 INFO     [evaluator.py:395] Running loglikelihood requests
Running loglikelihood requests (17 batches): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 17/17 [04:55<00:00, 17.36s/it]
mlx (model=internistai/base-7b-v0.2), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 64
|        Tasks        |Version|Filter|n-shot|Metric|Value |   |Stderr|
|---------------------|------:|------|-----:|------|-----:|---|-----:|
|professional_medicine|      0|none  |     0|acc   |0.7647|±  |0.0258|

lm_eval --model mlx --model_args model=internistai/base-7b-v0.2 --tasks   64  7.96s user 35.39s system 13% cpu 5:10.00 total

Hugging Face run on the same model:

% time lm_eval --model hf --model_args pretrained=internistai/base-7b-v0.2,dtype="float" --tasks mmlu_professional_medicine --device mps  --batch_size 64
hf (pretrained=internistai/base-7b-v0.2,dtype=float), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 64
|        Tasks        |Version|Filter|n-shot|Metric|Value |   |Stderr|
|---------------------|------:|------|-----:|------|-----:|---|-----:|
|professional_medicine|      0|none  |     0|acc   |0.7647|±  |0.0258|

lm_eval --model hf --model_args  --tasks mmlu_professional_medicine --device   28.83s user 117.90s system 63% cpu 3:49.41 total

[chimezie]

I fixed some handling of batch remainders, and it looks good; running comparisons against HF/MPS/Pytorch for medqa and some related subsets of MMLU

[lintangsutawika]

Could add installation dependancies (like lm_eval[mlx] see pyproject.toml) and a way to check if library is installed when called (see lm_eval/models/anthropic_llms.py)

[lintangsutawika]

@haileyschoelkopf bringing this to your attention as well.

[baberabb]

Hi! thanks for the substantial PR, and sorry it took so long. Left a couple of comments mainly about the indexing to extract the logprobs. A couple of other comments:

1. I think you can leave the tokenization to TemplateLM.loglikelihood (and move the loglikelihood logic to loglikelihood_tokens). This is mainly because we want to use encode_pair, which deals with a bug in some sentencepiece tokenizers.
2. Would also be great if you could add a test!

[chimezie]

Hi! thanks for the substantial PR, and sorry it took so long.

No worries

Left a couple of comments mainly about the indexing to extract the logprobs. A couple of other comments:

Thanks.

1. I think you can leave the tokenization to TemplateLM.loglikelihood (and move the loglikelihood logic to loglikelihood_tokens). This is mainly because we want to use encode_pair, which deals with a bug in some sentencepiece tokenizers.

Got it. Thanks. It wasn't always clear to me how to override this behavior in the least disruptive way, but this helps. I'll move this to loglikelihood_tokens.

2. Would also be great if you could add a test!

I will do that. Are there examples of other tests for lm_eval models that I can use to determine what conventions to follow?

[chimezie]

@baberabb I've removed all dependencies on the caching and I'm able to get similar answer log prob and greedy = continuation values for a handful of questions I probed. However, the final top-level figures still don't match, and I have run out of ideas why and wonder if the issue is at the level above _loglikelihood_tokens:

% lm_eval --model mlx --model_args model=internistai/base-7b-v0.2 --tasks mmlusr_question_and_answer_clinical_knowledge --batch_size 56
[..snip..]
mlx (model=internistai/base-7b-v0.2), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 56
|      Tasks       |Version|Filter|n-shot|Metric|   |Value |   |Stderr|
|------------------|------:|------|-----:|------|---|-----:|---|-----:|
|clinical knowledge|      1|none  |     0|acc   |↑  |0.2302|±  |0.0259|

% lm_eval --model hf --model_args pretrained=internistai/base-7b-v0.2,dtype="float32" --tasks mmlusr_question_and_answer_clinical_knowledge --batch_size 56 --device mps
[..snip..]
hf (pretrained=internistai/base-7b-v0.2,dtype=float32), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 56
|      Tasks       |Version|Filter|n-shot|Metric|   |Value |   |Stderr|
|------------------|------:|------|-----:|------|---|-----:|---|-----:|
|clinical knowledge|      1|none  |     0|acc   |↑  |0.5132|±  |0.0308|