Adding Paloma Evaluation + Restructure of Model and Train Loops

[rdiehlmartinez]

This is a big PR - so please please read through it carefully.

Several things changed.

1. The main one is that now we support doing evaluation! The evaluation metric I decided on is Paloma. Why? Paloma was made by the same people that created the OLMO dataset which we use. This means that Paloma was designed so that the data it contains does not occur in OLMO.
2. To support ^ we need a more expansive setup script. I went a bit overboard but I hope that setting up the project is now as easy as source setup.sh and can be run both if you've never setup installation before and if you have.
3. In order to run the evaluation script, the model needs to be saved as a huggingface model. We've been saving it as a basic torch model. So we need a way to convert the model to huggingface. I wrote a wrapper script that does this. Let me know if the code is documented enough for users who aren't super familiar with this stuff to understand what is going on.
4. Still on point 3 - in checkpointing we now save a normal model AND a huggingface model. This has the nice benefit that when we upload these two models to the hugginface hub, huggingface recognizes the hugginface model and you can then do cool things. I'll stay vague on this point haha
5. As a result of 3 still, we also need to rewrite the model class a bit and decompose it so that it works nicely with the standard hugginface API - mostly in how the KV cache is processed, but also in how model weights are sent to different devices.
6. POSSIBLE BUGs might be introduced by 5 -- I've been testing on a single GPU. When we go to multi-GPU the way things get sharded to different devices might cause a bit of headaches. This is a future me/us problem.

Those are the main things. There are a bunch of smaller cosmetic stuff as well.

[DavidDemitriAfrica]

Had an issue with dependencies being incompatible but not sure if it's just my machine. I suspect it's the version of cuda compatible with mine, but if @Yu-val-weiss could test it also that would be cool

[DavidDemitriAfrica]

Difficulty installing jsonnet. Something about the wheels? I have Microsoft C++ Build Tools already installed so it shouldn't be that. Attached the install log.

install_log.txt

[rdiehlmartinez]

google/jsonnet#912

[rdiehlmartinez]

Ok might have to set a different derivative of jsonnet as the dependency

[DavidDemitriAfrica]

gojsonnet also doesn't work, just tried it

[rdiehlmartinez]

can you try pip install jsonnet-binary

[DavidDemitriAfrica]

Works now

[rdiehlmartinez]

was the fix jsonnet-binary?

[DavidDemitriAfrica]

Leads to a weird error down the line. Seems to be that there's a version mismatch, which only comes about since jsonnet-binary is in 0.17.0 and jsonnet is 0.20.0?

(.venv) C:\Users\David Africa\Cambridge Research\pico-live\pico>"C:/Users/David Africa/Cambridge Research/pico-live/pico/.venv/Scripts/python.exe" "c:/Users/David Africa/Cambridge Research/pico-live/pico/train.py" wandb: WARNING This integration is tested and supported for lightning Fabric 2.1.3. wandb: WARNING Please report any issues to https://github.com/wandb/wandb/issues with the tag lightning-fabric`.
Using 16-bit Automatic Mixed Precision (AMP)
Resolving data files: 100%|██████████████████████████████████████████████████████████████████████████| 100/100 [00:00<00:00, 390.40it/s]
C:\Users\David Africa\Cambridge Research\pico-live\pico.venv\Lib\site-packages\huggingface_hub\utils_deprecation.py:131: FutureWarning: 'Repository' (from 'huggingface_hub.repository') is deprecated and will be removed from version '1.0'. Please prefer the http-based alternatives instead. Given its large adoption in legacy code, the complete removal is only planned on next major release.
For more details, please read https://huggingface.co/docs/huggingface_hub/concepts/git_vs_http.
warnings.warn(warning_message, FutureWarning)
Cloning https://huggingface.co/pico-lm/demo into local empty directory.
Checked out 2024-12-03_14-41-23 from 2024-12-03_14-41-23.
branch '2024-12-03_14-41-23' set up to track 'origin/2024-12-03_14-41-23'.

model.safetensors: 100%|███████████████████████████████████████████████████████████████████████████| 2.40k/2.40k [00:00<00:00, 8.21kB/s]
model.pt: 100%|████████████████████████████████████████████████████████████████████████████████████| 7.13k/7.13k [00:00<00:00, 47.4kB/s]
Traceback (most recent call last):
File "c:\Users\David Africa\Cambridge Research\pico-live\pico\train.py", line 462, in
main()
File "C:\Users\David Africa\Cambridge Research\pico-live\pico.venv\Lib\site-packages\click\core.py", line 1157, in call
return self.main(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\David Africa\Cambridge Research\pico-live\pico.venv\Lib\site-packages\click\core.py", line 1078, in main
rv = self.invoke(ctx)
^^^^^^^^^^^^^^^^
File "C:\Users\David Africa\Cambridge Research\pico-live\pico.venv\Lib\site-packages\click\core.py", line 1434, in invoke
return ctx.invoke(self.callback, **ctx.params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\David Africa\Cambridge Research\pico-live\pico.venv\Lib\site-packages\click\core.py", line 783, in invoke
return __callback(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "c:\Users\David Africa\Cambridge Research\pico-live\pico\train.py", line 401, in main
evaluation_results = run_evaluation(evaluation_config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "c:\Users\David Africa\Cambridge Research\pico-live\pico\utils\evaluation.py", line 294, in run_evaluation
paloma_config_path = setup_paloma_config(model_path, evaluation_config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "c:\Users\David Africa\Cambridge Research\pico-live\pico\utils\evaluation.py", line 140, in setup_paloma_config
json_str = _jsonnet.evaluate_snippet("config", jsonnet_template, ext_vars=ext_vars)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: RUNTIME ERROR: field does not exist: fgets
lib/olmo-eval/configs/utils.libsonnet:44:14-23 thunk <main_string>
lib/olmo-eval/configs/utils.libsonnet:41:43-54 thunk
std.jsonnet:1373:27-30 thunk
std.jsonnet:28:26
std.jsonnet:28:17-28 thunk
std.jsonnet:28:17-40 function
std.jsonnet:28:17-40 function
std.jsonnet:1373:14-31 function
lib/olmo-eval/configs/utils.libsonnet:41:16-55
lib/olmo-eval/configs/utils.libsonnet:41:5-56 function
...
std.jsonnet:789:24-47 thunk
std.jsonnet:789:9-57 function
std.jsonnet:789:9-57 function
std.jsonnet:790:5-28 function
lib/olmo-eval/configs/utils.libsonnet:(51:45)-(64:2) function <create_model_location_steps>
lib/olmo-eval/configs/utils.libsonnet:236:34-69 thunk <model_location_steps>
lib/olmo-eval/configs/utils.libsonnet:255:9-29 thunk <all_steps>
lib/olmo-eval/configs/utils.libsonnet:263:5-14 function
config:28:16-90 object
During manifestation`

[rdiehlmartinez]

the easiest solution might just be to not use jsonnet - the only reason we use it is because the third-party evaluation pipeline has some util functions that we call on which use jsonnet. We can just re-write those to not use jsonnet.

[Yu-val-weiss]

lgtm!