Fix FSDP parameter counts and HF generation compatibility

docs/Checkpoints.md

@@ -48,6 +48,20 @@ python scripts/convert_olmo2_to_hf.py --input_dir /path/to/olmo/checkpoint --out
 
 *Warning*: As we continue to develop and improve OLMo, our implementation in this repo may become incompatible with the implementation in the Transformer library. During these periods, OLMo checkpoints may not be convertible to Transformers checkpoint. At present, all OLMo checkpoints of our officially released models are convertible to Transformers checkpoints.
 
+Once a checkpoint has been converted to Transformers format, it can be evaluated with
+[OLMES](https://github.com/allenai/olmes) by passing the local or remote checkpoint path directly to
+`--model`. For example:
+
+```bash
+olmes --model /path/to/hf/checkpoint --task gsm8k::olmes --output-dir /path/to/eval-results
+```
+
+The in-training evaluation code in this repository is optimized for quick perplexity-style measurements over
+prepared continuations, so tasks such as GSM8K may appear there as multiple-choice or continuation-scoring tasks.
+Canonical generative evaluations, including GSM8K answer generation and extraction, should be run offline with
+OLMES. Custom OLMo-based checkpoints need to be converted first so the converted directory contains the
+Transformers files expected by OLMES, including a `config.json` with a supported `model_type`.
+
 HF OLMo checkpoints
 ---
 

hf_olmo/modeling_olmo.py

@@ -4,7 +4,6 @@
 
 import torch
 from transformers import GenerationMixin, PreTrainedModel
-from transformers.cache_utils import Cache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.auto import AutoModelForCausalLM
 
@@ -48,6 +47,7 @@ class OLMoForCausalLM(PreTrainedModel, GenerationMixin):
     _no_split_modules = ["OLMoBlock"]
     _supports_flash_attn_2 = True
     _supports_sdpa = True
+    _supports_cache_class = False
     supports_gradient_checkpointing = True
 
     def __init__(self, config: OLMoConfig, model: Optional[OLMo] = None, init_params: bool = False):
@@ -93,10 +93,11 @@ def forward(
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        cache_position: Optional[
-            Cache
-        ] = None,  # This is a hack mitigation of an issue in transformers `4.39.x` https://github.com/huggingface/transformers/issues/29426
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
+        del cache_position, position_ids
+
         if use_cache is None:
             use_cache = self.config.use_cache
 

olmo/model.py

@@ -1569,7 +1569,14 @@ def num_params(self, include_embedding: bool = True) -> int:
                 lambda np: ".wte." not in np[0] and ".wpe." not in np[0],
                 params,
             )
-        return sum(p.numel() for _, p in params)
+        total = 0
+        for _, p in params:
+            unpadded_size = getattr(p, "_unpadded_unsharded_size", None)
+            if unpadded_size is not None:
+                total += unpadded_size.numel()
+            else:
+                total += p.numel()
+        return total
 
     @property
     def num_fwd_flops(self):

tests/model_test.py

@@ -473,6 +473,25 @@ def test_layer_norm(train_config: TrainConfig, elementwise_affine: bool, include
     torch.testing.assert_close(y_actual, y_expected)
 
 
+def test_num_params_uses_fsdp_unpadded_size(train_config: TrainConfig):
+    train_config.model.d_model = 16
+    train_config.model.n_heads = 2
+    train_config.model.n_layers = 1
+    train_config.model.vocab_size = 32
+    train_config.model.embedding_size = 32
+    train_config.model.max_sequence_length = 16
+
+    model = OLMo(train_config.model, init_params=False)
+    expected_num_params = model.num_params()
+
+    param = next(model.parameters())
+    padded_numel = param.numel() + 8
+    param.data = torch.empty(padded_numel)
+    param._unpadded_unsharded_size = torch.Size([padded_numel - 8])  # type: ignore[attr-defined]
+
+    assert model.num_params() == expected_num_params
+
+
 def test_block_groups():
     model_with_block_groups = OLMo(ModelConfig(d_model=128, n_heads=2, n_layers=9, block_group_size=3)).eval()
     model_without_block_groups = OLMo(ModelConfig(d_model=128, n_heads=2, n_layers=9, block_group_size=1)).eval()