Lightweight HF integration

[AkshitaB]

• Creates a config.json file at the checkpoint location, which allows Olmo models to be loaded as HF models.

python hf_olmo/add_hf_config_to_olmo_checkpoint.py --checkpoint-dir <olmo-checkpoint-location>

• The model, config, and tokenizer classes are registered with the relevant HF auto classes. Importing them allows models to be loaded as HF-compatible models.

from hf_olmo import *
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(<olmo-checkpoint-location>)
tokenizer = AutoTokenizer.from_pretrained(<olmo-checkpoint-location>)

Using HF pipeline works:

from hf_olmo.modeling_olmo import OLMoConfig, OLMoForCausalLM  # noqa: F401

model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
output = pipeline("question: who wrote romeo and juliet? answer: ", max_new_tokens=30)
# [{'generated_text': 'question: who wrote romeo and juliet? answer: romeo and juliet is a play by william shakespeare. the play was first performed in 1605. the play is set in the city'}]

Instruct-eval tasks also work (adding from hf_olmo import * to run_eval.py should suffice):
Tested with mmlu and bbh.

To Do:

• Test bbh instruct eval task. @OyvindTafjord pointed out that bbh tasks from instruct eval may be a more comprehensive test, since it uses more of the HF api.
• Add implementation for AutoModelForCausalLM-specific methods, so that HF's .generate() works.
• Test with deepspeed code from @hamishivi .

[epwalsh]

Overall looks good. I'm curious what your thoughts are about how we'll ultimately make this integration available when we release the code? I was thinking we'd make this a submodule of olmo that people could install as an extra, like pip install olmo[hf] or something. But keeping it a separate package is probably fine too. Maybe we want to call it something more specific like olmo_hf though?

[epwalsh]

Do we need to implement this or can we get this for free using HF built-in generate functionality?

[AkshitaB]

We can, now. Needed a couple small methods implemented.

[AkshitaB]

Maybe we want to call it something more specific like olmo_hf though?

Makes sense. I've renamed it.

[epwalsh]

Just some comments about CI, otherwise looks good

[epwalsh]

Should we add another job to run the HF tests?

[AkshitaB]

Discussed offline: since this requires updating the beaker image on which we run the GPU tests, and since we expect to reconfigure this at some point, it's not worth the effort now.

I've confirmed using instruct-eval that the HF integration runs on GPU.