OLMo 2

[ysjprojects]

https://huggingface.co/collections/allenai/olmo-2-674117b93ab84e98afc72edc
https://arxiv.org/abs/2501.00656

Version 2 of OLMo released by Ai2.

Comes in 7B and 13B Base, Instruct, and additional SFT and DPO models.

First, we find that OLMo 2 7B and 13B are the best fully-open models to-date, often outperforming open weight models of equivalent size. Not only do we observe a dramatic improvement in performance across all tasks compared to our earlier OLMo 0424 model but, notably, OLMo 2 7B outperforms LLama-3.1 8B and OLMo 2 13B outperforms Qwen 2.5 7B despite its lower total training FLOPs. The OLMo 2 models sit at the Pareto frontier of training FLOPs vs model average performance (see figure above).

[rasbt]

Hi there,
just wanted to say thanks for taking on this PR (I know this is a lot of work)! The OLMo models are awesome, and I'd be great to have OLMo 2 in LitGPT.

[ysjprojects]

Hi there, just wanted to say thanks for taking on this PR (I know this is a lot of work)! The OLMo models are awesome, and I'd be great to have OLMo 2 in LitGPT.

Thanks mate!

Currently on vacation, will resume working on this PR once I'm back.

[ysjprojects]

Performed some fixes today, now test_model passes for Olmo2.

[Borda]

seems like almost all tests are failing on:

FAILED tests/test_model.py::test_sdpa_choice_kv_cache[SmolLM2-1.7B-Instruct] - RuntimeError: Deterministic behavior was enabled with either `torch.use_deterministic_algorithms(True)` or `at::Context::setDeterministicAlgorithms(true)`, but this operation is not deterministic because it uses CuBLAS and you have CUDA >= 10.2. To enable deterministic behavior in this case, you must set an environment variable before running your PyTorch application: CUBLAS_WORKSPACE_CONFIG=:4096:8 or CUBLAS_WORKSPACE_CONFIG=:16:8. For more information, go to https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility

[lantiga]

Looks good! Left a few comments

[lantiga]

maybe a less special-casey way of doing this could be to avoid the introduction of the boolean norm_1 and norm_2 configs, but rather just have Identity as the norm class itself

[ysjprojects]

self.norm_1 = nn.Identity() if not config.norm_1 else config.norm_class(config.n_embd, eps=config.norm_eps)
        self.attn = CausalSelfAttention(config, block_idx)
        self.post_attention_norm = (
            config.norm_class(config.n_embd, eps=config.norm_eps) if config.post_attention_norm else nn.Identity()
        )
        self.norm_2 = (
            nn.Identity()
            if not config.norm_2
            else (None if config.shared_attention_norm else config.norm_class(config.n_embd, eps=config.norm_eps))
        )
        self.mlp = config.mlp_class(config)
        self.post_mlp_norm = (
            config.norm_class(config.n_embd, eps=config.norm_eps) if config.post_mlp_norm else nn.Identity()
        )

The issue is that olmo2 selectively use RMSNorm for post_attention_norm and post_mlp_norm but Identity for norm_1 and norm_2

Perhaps a way to get rid of the booleans would be to specify it as a special case for olmo2:

self.norm_1 = nn.Identity() if config.name.lower().startswith(("olmo-2-")) else...

IMO that's the easiest workaround to getting rid of norm_1 and nom_2 booleans.

[t-vi]

How about norm_1_class and norm_2_class as overrides to norm_class in the config file?
Then reading the config could set up norm_1_class, norm_2_class either from the config names or from norm_class.
This would move the cases from the model to the config. Ideally, we'd also subsume the shared_attention_norm, which seems to be very similar to not norm_2.

[lantiga]

that's a good idea, it will be advantageous in the future, wdyt @ysjprojects?

[lantiga]

we can also do it in a follow up PR, doesn't need to be this one

[ysjprojects]

@lantiga I have addressed the issues and left some follow-up comments.

[ysjprojects]

@Borda Changes should be ready to merge

[Borda]

@Borda Changes should be ready to merge

@t-vi mind have a look as codeowner, pls ^^