finish merge

Author: priyakasimbeg

algoperf/workloads/lm/input_pipeline.py

@@ -54,14 +54,14 @@ def batch_with_padding(
 def get_data_iter(data_rng: jax.random.PRNGKey,
   split: str,
   data_dir: str,
-  global_batch_size: int,
+  batch_size: int,
   num_batches: Optional[int] = None,):
 
-  ds = get_lm_dataset(data_rng, split, data_dir, global_batch_size, num_batches)
+  ds = get_lm_dataset(data_rng, split, data_dir, batch_size, num_batches)
 
   it = map(
     functools.partial(
-      data_utils.shard_and_maybe_pad_np, global_batch_size=global_batch_size
+      data_utils.shard_and_maybe_pad_np, global_batch_size=batch_size
     ),
     ds,
   )
@@ -72,7 +72,7 @@ def get_lm_dataset(
   data_rng: jax.random.PRNGKey,
   split: str,
   data_dir: str,
-  global_batch_size: int,
+  batch_size: int,
   num_batches: Optional[int] = None,
 ):
   """Load preprocessed TF dataset."""
@@ -98,14 +98,15 @@ def get_lm_dataset(
     },
     num_parallel_calls=AUTOTUNE,
   )
-  sequences_ds = sequences_ds.repeat()
   if split == 'train':
     ds = sequences_ds.shuffle(
       SHUFFLE_BUFFER_SIZE, seed=shuffle_seed
     )
     ds = ds.batch(
-      global_batch_size, drop_remainder=False
+      batch_size, drop_remainder=False
     )
+    ds = ds.take(num_batches) if num_batches is not None else ds
+    ds = ds.repeat()
     ds = ds.map(lambda x: {
          'inputs': x['inputs'],
          'targets': x['targets'],
@@ -115,12 +116,14 @@ def get_lm_dataset(
   elif split == 'eval_train':
     ds = batch_with_padding(
       sequences_ds,
-      global_batch_size,
+      batch_size,
       padded_shapes={
-        'inputs': (global_batch_size, None),
-        'targets': (global_batch_size, None),
+        'inputs': (batch_size, None),
+        'targets': (batch_size, None),
       },
     )
+    ds = ds.take(num_batches) if num_batches is not None else ds
+    ds = ds.repeat()
     ds = ds.map(lambda x: {'inputs': x['inputs'],
                           'targets': x['targets'],
                           'weights': tf.where(tf.equal(x['inputs'], PAD_ID), 0.0, 1.0)
@@ -129,12 +132,14 @@ def get_lm_dataset(
   elif split == 'validation':
     ds = batch_with_padding(
       sequences_ds,
-      global_batch_size,
+      batch_size,
       padded_shapes={
-        'inputs': (global_batch_size, None),
-        'targets': (global_batch_size, None),
+        'inputs': (batch_size, None),
+        'targets': (batch_size, None),
       },
     )
+    ds = ds.take(num_batches) if num_batches is not None else ds
+    ds = ds.repeat()
     ds = ds.map(lambda x: {'inputs': x['inputs'],
                           'targets': x['targets'],
                           'weights': tf.where(tf.equal(x['inputs'], PAD_ID), 0.0, 1.0)

algoperf/workloads/lm/lm_jax/workload.py

@@ -21,16 +21,17 @@ def _build_input_queue(self,
                          split: str,
                          data_dir: str,
                          global_batch_size: int,
-                         num_batches: Optional[int] = None,
-                         repeat_final_dataset: bool = False):
+                         cache: Optional[bool] = None,
+                         repeat_final_dataset: Optional[bool] = None,
+                         num_batches: Optional[int] = None):
     """Build an input queue using pre-cached FineWeb dataset."""
-    del num_batches
-    del repeat_final_dataset
+    del cache, repeat_final_dataset
     ds = get_data_iter(
         data_rng=data_rng,
         split=split,
         data_dir=data_dir,
-        global_batch_size=global_batch_size)
+        batch_size=global_batch_size,
+        num_batches=num_batches)
     ds = map(jax_sharding_utils.shard_along_batch_dim, ds)
     return ds
 
@@ -71,7 +72,7 @@ def model_fn(
       mode: spec.ForwardPassMode,
       rng: spec.RandomState,
       update_batch_norm: bool,
-      dropout_rate: float = None) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
+      dropout_rate: float = 0.0) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     del mode, rng, update_batch_norm, model_state, dropout_rate
     inputs = batch['inputs']
     # Convert one-hot inputs to token IDs if needed

algoperf/workloads/lm/lm_pytorch/workload.py

@@ -98,17 +98,17 @@ def _build_input_queue(
       split: str,
       data_dir: str,
       global_batch_size: int,
-      num_batches: Optional[int] = None,
-      repeat_final_dataset: bool = False) -> Iterator[Dict[str, spec.Tensor]]:
+      cache: Optional[bool] = None,
+      repeat_final_dataset: Optional[bool] = None,
+      num_batches: Optional[int] = None) -> Iterator[Dict[str, spec.Tensor]]:
     """Build an input queue for the given split."""
+    del cache, repeat_final_dataset
     local_batch_size = global_batch_size // N_GPUS
-    # In DDP mode, pass local_device_count=1 to prevent shard_and_maybe_pad_np
-    # from seeing all GPUs via torch.cuda.device_count()
     loader = get_data_iter(
         data_rng=data_rng,
         split=split,
         data_dir=data_dir,
-        global_batch_size=local_batch_size,
+        batch_size=local_batch_size,
         num_batches=num_batches,
     )
     if USE_PYTORCH_DDP:

algoperf/workloads/lm/workload.py

@@ -2,11 +2,11 @@
 
 import abc
 import math
+import numpy as np
 import os
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Iterator
 
 import jax
-import numpy as np
 from absl import flags
 
 from algoperf import spec
@@ -85,11 +85,11 @@ def train_stddev(self):
 
   @property
   def max_allowed_runtime_sec(self) -> int:
-    return 3600 * 14  # 14 hours
+    return 3600 * 14 # 14 hours  TODO(kasimbeg): update
 
   @property
   def eval_period_time_sec(self) -> int:
-    return 1200  # 20 minutes
+    return 1200  # 20 minutes  TODO(kasimbeg): update
 
   @property
   def step_hint(self) -> int:
@@ -119,9 +119,10 @@ def _build_input_queue(
     split: str,
     data_dir: str,
     global_batch_size: int,
+    cache: Optional[bool] = None,
+    repeat_final_dataset: Optional[bool] = None,
     num_batches: Optional[int] = None,
-    repeat_final_dataset: bool = False,
-  ):
+  ) -> Iterator[Dict[str, Any]]:
     """Build an input queue for the given split."""
 
 
@@ -150,8 +151,7 @@ def _eval_model_on_split(
         split,
         data_dir,
         global_batch_size,
-        num_batches,
-        repeat_final_dataset=True,
+        num_batches=num_batches
       )
 
     eval_metrics = {}