Faster eval2

eval/eval_openlm_ckpt.py

@@ -46,8 +46,9 @@ def evaluate(model, tokenizer, cfg):
 
     composer_model = SimpleComposerOpenLMCausalLM(model, tokenizer)
 
+    cfg_icl_tasks = [dict(i) for i in cfg.icl_tasks]
     evaluators, logger_keys = build_icl_evaluators(
-        cfg.icl_tasks, tokenizer, cfg.max_seq_len, cfg.device_eval_batch_size
+        cfg_icl_tasks, tokenizer, cfg.max_seq_len, cfg.device_eval_batch_size
     )
 
     in_memory_logger = InMemoryLogger()  # track metrics in the in_memory_logger
@@ -123,6 +124,7 @@ def main():
 
     state_dict = checkpoint["state_dict"]
     state_dict = {x.replace("module.", ""): y for x, y in state_dict.items()}
+    state_dict = {x.replace("_orig_mod.", ""): y for x, y in state_dict.items()}
     open_lm.model.load_state_dict(state_dict)
     open_lm.model.eval()
 

open_lm/model_configs/open_lm_1b_swiglutorch.json

@@ -0,0 +1,14 @@
+{
+    "hidden_dim": 2048,
+    "n_layers": 24,
+    "n_heads": 16,
+    "seq_len": 2048,
+    "vocab_size": 50432,
+    "post_embed_norm": false,
+    "weight_tying": false,
+    "model_norm": "gain_only_lp_layer_norm",
+    "norm_type": "gain_only_lp_layer_norm",
+    "ffn_type": "swiglu_torch",
+    "qk_norm": true,
+    "positional_embedding_type": "rotary"
+}

open_lm/utils/llm_foundry_wrapper.py

@@ -3,7 +3,7 @@
 
 """Implements a Hugging Causal LM wrapped inside a :class:`.ComposerModel`."""
 
-from typing import Mapping, Union
+from typing import Union, Optional, Any
 from llmfoundry.eval.metrics.nlp import (
     InContextLearningLMAccuracy,
     InContextLearningLMExpectedCalibrationError,
@@ -16,6 +16,8 @@
     LanguagePerplexity,
 )
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+import torch
+from torch import dist
 
 from composer.models.huggingface import HuggingFaceModel
 
@@ -51,3 +53,34 @@ def __init__(self, model, tokenizer):
 
     def generate(self, input_ids=None, inputs_embeds=None, **kwargs):
         return super().generate(input_ids=input_ids, **kwargs)
+
+    def eval_forward(self, batch, outputs: Optional[Any] = None):
+        # If input_ids are all 0 after a certain point, we can skip that
+        num_zeros = 0
+        if batch["input_ids"] is not None:
+            # Find indices of the last non-zero element in each row
+            last_nonzero_indices = torch.max(torch.nonzero(batch["input_ids"], as_tuple=True)[1])
+            if last_nonzero_indices < batch["input_ids"].shape[1] - 1:
+                num_zeros = batch["input_ids"].shape[1] - last_nonzero_indices - 1
+                batch["input_ids"] = batch["input_ids"][:, : last_nonzero_indices + 1]
+                if batch["attention_mask"] is not None:
+                    batch["attention_mask"] = batch["attention_mask"][:, : last_nonzero_indices + 1]
+                if batch["labels"] is not None:
+                    batch["labels"] = batch["labels"][:, : last_nonzero_indices + 1]
+
+        output = super().eval_forward(batch, outputs)
+
+        # Add back the 0 that we removed
+        if num_zeros:
+            if hasattr(output, "logits") or isinstance(output, dict):
+                fake_logits = torch.zeros(
+                    output["logits"].shape[0], num_zeros, output["logits"].shape[2], device=output["logits"].device
+                )
+                fake_logits[:, :, 0] = 1
+                output["logits"] = torch.cat([output["logits"], fake_logits], dim=1)
+            else:
+                fake_logits = torch.zeros(output.shape[0], num_zeros, output.shape[2], device=output.device)
+                fake_logits[:, :, 0] = 1
+                output = torch.cat([output, fake_logits], dim=1)
+
+        return output

open_lm/utils/transformers/hf_model.py

@@ -98,13 +98,15 @@ def forward(
         ```"""
         assert position_ids is None, "Position IDs are not supported"
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+
         logits, _, past_key_values = self.model(
             input_ids=input_ids,
             inputs_embeds=inputs_embeds,
             past_key_values=past_key_values,
             use_cache=use_cache,
             attention_mask=attention_mask,
         )
+
         loss = None
         if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()
@@ -115,6 +117,7 @@ def forward(
             loss = loss_fct(shift_logits, shift_labels)
 
         output = CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values, loss=loss)
+
         return output
 
     def prepare_inputs_for_generation(