refactor evaluation pipeline for lm

Author: priyakasimbeg

algoperf/workloads/lm/input_pipeline.py

@@ -5,6 +5,7 @@
 from typing import Optional
 
 import jax
+import numpy as np
 import tensorflow as tf
 
 from algoperf import data_utils
@@ -106,7 +107,7 @@ def get_lm_dataset(
     repeated_sequences_dataset = shuffled_sequences_ds.repeat()
     ds = repeated_sequences_dataset.batch(
       global_batch_size, drop_remainder=False
-    ).take(100).prefetch(tf.data.experimental.AUTOTUNE)
+    ).prefetch(tf.data.experimental.AUTOTUNE)
   elif split == 'eval_train':
     ds = batch_with_padding(
       sequences_ds,
@@ -115,7 +116,11 @@ def get_lm_dataset(
         'inputs': (global_batch_size, None),
         'targets': (global_batch_size, None),
       },
-    ).take(100).prefetch(tf.data.experimental.AUTOTUNE)  # todo(kasimbeg): set final size of validation
+    )
+    ds = ds.map(lambda x: {'inputs': x['inputs'],
+                          'targets': x['targets'],
+                          'weights': tf.where(tf.equal(x['inputs'], PAD_ID), 0.0, 1.0)})
+    ds = ds.take(1000).prefetch(tf.data.experimental.AUTOTUNE)  # todo(kasimbeg): set final size of validation
   elif split == 'validation':
     ds = batch_with_padding(
       sequences_ds,
@@ -124,6 +129,10 @@ def get_lm_dataset(
         'inputs': (global_batch_size, None),
         'targets': (global_batch_size, None),
       },
-    ).take(100).prefetch(tf.data.experimental.AUTOTUNE)  # todo(kasimbeg): set final size
+    )
+    ds = ds.map(lambda x: {'inputs': x['inputs'],
+                          'targets': x['targets'],
+                          'weights': tf.where(tf.equal(x['inputs'], PAD_ID), 0.0, 1.0)})
+    ds = ds.take(1000).prefetch(tf.data.experimental.AUTOTUNE)  # todo(kasimbeg): set final size
 
   return ds

algoperf/workloads/lm/lm_jax/workload.py

@@ -28,26 +28,13 @@ def _build_input_queue(self,
     """Build an input queue using pre-cached FineWeb dataset."""
     del num_batches
     del repeat_final_dataset
-    loader = get_data_iter(
+    ds = get_data_iter(
         data_rng=data_rng,
         split=split,
         data_dir=data_dir,
         global_batch_size=global_batch_size)
-    loader = map(jax_sharding_utils.shard_along_batch_dim, loader)
-    return loader
-
-  def _build_hf_input_queue(self,
-                         data_rng: jax.random.PRNGKey,
-                         split: str,
-                         data_dir: str,
-                         global_batch_size: int,
-                         num_batches: Optional[int] = None,
-                         repeat_final_dataset: bool = False):
-    """Build an input queue using HuggingFace FineWeb dataset."""
-    del num_batches
-    del repeat_final_dataset
-    iter = get_data_iter(data_rng, split, data_dir, global_batch_size)
-    return iter
+    ds = map(jax_sharding_utils.shard_along_batch_dim, ds)
+    return ds
 
   def init_model_fn(
       self,
@@ -156,9 +143,10 @@ def _eval_batch(self,
     """Evaluate the model on a single batch."""
     logits, _ = self.model_fn(
         params, batch, model_state, spec.ForwardPassMode.EVAL, rng, False)
-    targets = batch['targets']
-
     # Calculate cross-entropy loss
     # TODO(kasimbeg): add weights?
-    loss_metrics = self.compute_weighted_cross_entropy(logits, targets)
-    return loss_metrics
+    metrics = self.compute_weighted_cross_entropy(logits, batch['targets'], batch['weights'])
+    return {
+      'loss': metrics['summed'],
+      'denominator': metrics['n_valid_examples'],
+    }

algoperf/workloads/lm/workload.py

@@ -2,6 +2,7 @@
 
 import abc
 import math
+import numpy as np
 import os
 from typing import Any, Dict, Optional
 
@@ -44,11 +45,11 @@ def validation_target_value(self) -> float:
     return 20.0  # Target perplexity
 
   def has_reached_test_target(self, eval_result: Dict[str, float]) -> bool:
-    return eval_result['test/ppl'] <= self.test_target_value
+    return True # No test targets
 
   @property
   def test_target_value(self) -> float:
-    return 20.0  # Target perplexity
+    return None # No test targets
 
   @property
   def loss_type(self) -> spec.LossType:
@@ -60,19 +61,19 @@ def num_train_examples(self) -> int:
 
   @property
   def num_eval_train_examples(self) -> int:
-    return 10000  # Subset for evaluation
+    return 500 # Subset for evaluation. # TODO(kasimbeg): update
 
   @property
   def num_validation_examples(self) -> int:
-    return 50000
+    return 500  # TODO(kasimbeg update)
 
   @property
   def num_test_examples(self) -> int:
-    return 50000
+    return 0
 
   @property
   def eval_batch_size(self) -> int:
-    return 8
+    return 32
 
   @property
   def train_mean(self):
@@ -84,7 +85,7 @@ def train_stddev(self):
 
   @property
   def max_allowed_runtime_sec(self) -> int:
-    return 3600 * 4  # 4 hours
+    return 3600 * 5  # 4 hours
 
   @property
   def eval_period_time_sec(self) -> int:
@@ -93,7 +94,7 @@ def eval_period_time_sec(self) -> int:
   @property
   def step_hint(self) -> int:
     """Approx. steps the baseline can do in the allowed runtime budget."""
-    return 7000
+    return 54000
 
   @property
   def pre_ln(self) -> bool:
@@ -141,7 +142,7 @@ def _eval_batch(
     )
 
     loss_dict = self.loss_fn(batch['targets'], logits)
-    return loss_dict['summed']
+    return loss_dict
 
   def _eval_model_on_split(
     self,
@@ -170,12 +171,15 @@ def _eval_model_on_split(
     eval_metrics = {}
     for _ in range(num_batches):
       eval_batch = next(self._eval_iters[split])
-      metrics = self._eval_batch(params, eval_batch)
+      metrics = self._eval_batch(params, eval_batch, model_state, rng)
       for metric_name, metric_value in metrics.items():
         if metric_name not in eval_metrics:
           eval_metrics[metric_name] = 0.0
         eval_metrics[metric_name] += metric_value
-      eval_results = self._normalize_eval_metrics(num_examples, eval_metrics)
+
+    eval_results = self._normalize_eval_metrics(num_examples, eval_metrics)
+    eval_results['ppl'] = np.exp(eval_results['loss'])
+    print(eval_results)
 
     return eval_results