Merge branch 'lm_workload' of github.com:mlcommons/algorithmic-efficiency into lm_workload

Author: rka97

algoperf/workloads/lm/input_pipeline.py

@@ -126,7 +126,7 @@ def get_lm_dataset(
     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
+    ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
   elif split == 'validation':
     ds = batch_with_padding(
       sequences_ds,
@@ -139,6 +139,6 @@ def get_lm_dataset(
     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
+    ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
 
   return ds

algoperf/workloads/lm/workload.py

@@ -61,11 +61,11 @@ def num_train_examples(self) -> int:
 
   @property
   def num_eval_train_examples(self) -> int:
-    return 500 # Subset for evaluation. # TODO(kasimbeg): update
+    return 10_000 # Subset for evaluation.
 
   @property
   def num_validation_examples(self) -> int:
-    return 500  # TODO(kasimbeg update)
+    return 100_000 # sequences
 
   @property
   def num_test_examples(self) -> int:
@@ -85,11 +85,11 @@ def train_stddev(self):
 
   @property
   def max_allowed_runtime_sec(self) -> int:
-    return 3600 * 5  # 4 hours  TODO(kasimbeg): update
+    return 3600 * 14 # 14 hours  TODO(kasimbeg): update
 
   @property
   def eval_period_time_sec(self) -> int:
-    return 600  # 10 minutes  TODO(kasimbeg): update
+    return 1200  # 20 minutes  TODO(kasimbeg): update
 
   @property
   def step_hint(self) -> int:

algorithms/target_setting_algorithms/fineweb_edu_lm/jax_nadamw_target_setting.py

@@ -170,8 +170,8 @@ def init_optimizer_state(
   del rng
 
   def jax_cosine_warmup(step_hint: int, hyperparameters):
-    # Create learning rate schedule.
     step_hint = 0.75 * step_hint
+    # Create learning rate schedule.
     warmup_steps = int(hyperparameters.warmup_factor * step_hint)
     warmup_fn = optax.linear_schedule(
       init_value=0.0,
@@ -343,7 +343,7 @@ def update_params(
   )
 
   # Log loss, grad_norm.
-  if global_step % 1 == 0 and workload.metrics_logger is not None:
+  if global_step % 100 == 0 and workload.metrics_logger is not None:
     workload.metrics_logger.append_scalar_metrics(
       {'loss': loss.item(), 'grad_norm': grad_norm.item()}, global_step
     )

algorithms/target_setting_algorithms/fineweb_edu_lm/tuning_search_space.json

@@ -1,11 +1,11 @@
 [
   {
     "dropout_rate": 0.0,
-    "label_smoothing": 0.1,
-    "learning_rate": 0.0003955553491092581,
-    "one_minus_beta1": 0.06124602712,
-    "beta2": 0.9535169492059872,
-    "weight_decay": 0.03268700808664715,
-    "warmup_factor": 0.0375
+    "label_smoothing": 0.0,
+    "learning_rate":  0.00038418421332238876,
+    "one_minus_beta1": 0.01564758865,
+    "beta2": 0.992362328914093,
+    "weight_decay": 0.25551270901641954,
+    "warmup_factor": 0.05
   }
 ]