Repository: EleutherAI/lm-evaluation-harness
Issue: Unexpected effect of batch size parameters on observed batch dimension for wikitext task
Root cause: Documentation Deficiency

I'm trying to evaluate a HuggingFace language model on the `wikitext` task, and I'm struggling to understand the relationship between:
1. The `batch_size` argument to `lm_eval.evaluator.simple_evaluate`
2. The `max_batch_size` argument to `lm_eval.evaluator.simple_evaluate`
3. The `batch_size` attribute of the `lm_eval.models.gpt2.HFLM` class
4. The actual batch dimension of the input tensor passed in by the harness

I'm not sure if the behavior I'm seeing is a bug or not, and I haven't been able to find much documentation on the intended meaning of these parameters. I'd be grateful for any pointers to relevant documentation I might have missed.

## Steps to Reproduce

Run the following script, varying the values of `simple_evaluate_batch_size`, `simple_evaluate_max_batch_size`, and `HFLM_batch_size`:

```python
#!/usr/bin/env python3

from transformers import GPTNeoXForCausalLM, GPTNeoXTokenizerFast
from lm_eval import evaluator
from lm_eval.models.gpt2 import HFLM

################################################################################
# Try varying these parameters:
simple_evaluate_batch_size = # ...
simple_evaluate_max_batch_size = # ...
HFLM_batch_size = # ...
################################################################################

model = GPTNeoXForCausalLM.from_pretrained("EleutherAI/pythia-70m-deduped")
tokenizer = GPTNeoXTokenizerFast.from_pretrained("EleutherAI/pythia-70m-deduped")

# Monkey-patch model to report batch dimension
orig_forward = model.forward
def patched_forward(input_ids, *args, **kwargs):
    print("batch dimension:", input_ids.size(0))
    return orig_forward(input_ids, *args, **kwargs)
model.forward = patched_forward

lm = HFLM(
    device="cuda:0",
    pretrained=model,
    tokenizer=tokenizer,
    batch_size=HFLM_batch_size,
)

results = evaluator.simple_evaluate(
    model=lm,
    tasks=("wikitext",),
    no_cache=True,
    batch_size=simple_evaluate_batch_size,
    max_batch_size=simple_evaluate_max_batch_size,
    device="cuda:0",
)
```

## Observed Behavior

I observe the following batch dimensions under different settings of the parameters in the script given above:

| `simple_evaluate_batch_size` | `simple_evaluate_max_batch_size` | `HFLM_batch_size`| Observed Batch Dimension |
|--|--|--|--|
| `None` | `None` | `1` | `1` |
| `None` | `None` | `2` | between `1` and `2` |
| `None` | `None` | `6` | between `1` and `6` |
| `None` | `None` | `128` | between `1` and `7` |
| `None` | `None` | `4096` | between `1` and `7` |
| `None` | `None` | `"auto"` | `512` |
| `128` | `128` | `1` | `1` |
| `128` | `128` | `128` | between `1` and `7` |
| `128` | `128` | `"auto"` | `512` |

My main conclusions from this are:
- **The `simple_evaluate_batch_size` and `simple_evaluate_max_batch_size` parameters appear to have no effect.**
  - It seems that **only `HFLM_batch_size` matters.**
- When `HFLM_batch_size` is a numerical value, the value **only seems to matter up to about `7`**, beyond which the batch size cannot increase.
- When `HFLM_batch_size` is `"auto"`, the batch size seems **pinned at `512`**, even if I try to restrict it by setting `simple_evaluate_max_batch_size` to some smaller value.

I wouldn't expect any of this behavior. I would like to set my batch size to around `128`, but I can't find a way to achieve this using these three parameters -- any tips for doing so would be appreciated.

- `simple_evaluate` will use `batch_size` and `max_batch_size` when constructing from `str` or `transformers.PreTrainedModel`, otherwise `HFLM`/etc. will take precedence, so `batch_size` and `max_batch_size` can be set in `HFLM`/etc. and omitted in `simple_evaluate`.
- Those small variable batch sizes are specific to perplexity tasks.
- Observing `512` in `patched_forward` is due to maximal batch size detection trying `512` -> `256` -> ... (since it runs `forward` to see if there is a memory exception). It will settle on a smaller maximal batch size.

As @gakada mentioned, just to elaborate on the perplexity task point: 

https://github.com/EleutherAI/lm-evaluation-harness/blob/d1537059b515511801ae9b742f8e949f1bfcd010/lm_eval/base.py#L257C13-L257C13

Currently, we don't batch beyond ~7 inputs for PerplexityTasks such as `wikitext`, because we dispatch the chunks of our string to calculate perplexity on that string all together for only one string at a time. This is definitely not desirable, and we'll see about doing this more efficiently as well as allowing for varied sizes of rolling windows or chunks in future!

For non-PerplexityTasks,  any situation where you're observing "between `1` and ~`7`" you should instead observe `HFLM_batch_size` consistently. Confirmed this was the case using your script for me with `"arc_challenge"`! Sorry for the confusion, we'll strive to document the usage of these three parameters better in future.

@gakada @haileyschoelkopf That all makes sense -- thanks so much for your help!