lm workload dataset integration in jax

Author: rka97

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -71,17 +71,6 @@ def _build_input_queue(
                                      cache,
                                      repeat_final_dataset)
 
-  def sync_batch_stats(
-      self, model_state: spec.ModelAuxiliaryState) -> spec.ModelAuxiliaryState:
-    """Sync the batch statistics across replicas."""
-    # An axis_name is passed to pmap which can then be used by pmean.
-    # In this case each device has its own version of the batch statistics
-    # and we average them.
-    avg_fn = jax.pmap(lambda x: lax.pmean(x, 'x'), 'x')
-    new_model_state = model_state.copy()
-    new_model_state['batch_stats'] = avg_fn(model_state['batch_stats'])
-    return new_model_state
-
   def init_model_fn(
       self,
       rng: spec.RandomState,

algoperf/workloads/lm/input_pipeline.py

@@ -87,19 +87,19 @@ def batch_iterator():
         tokens = jax.nn.one_hot(token_ids, num_classes=vocab_size)
         inputs, targets = tokens[:, :-1], tokens[:, 1:]
         inputs, targets = jax.device_put(inputs), jax.device_put(targets)
-      yield inputs, targets
-
+      batch = {
+          "inputs": inputs,
+          "targets": targets,
+      }
+      yield batch
   return batch_iterator()
 
 
 def get_lm_dataset(data_rng: jax.random.PRNGKey,
                    split: str,
                    data_dir: str,
-                   vocab_size: int,
                    global_batch_size: int,
-                   num_batches: Optional[int] = None,
-                   repeat_final_dataset: bool = False,
-                   vocab_path: Optional[str] = None):
+                   num_batches: Optional[int] = None):
   """Load HF dataset and return a TF dataset."""
 
   dataset_path = os.path.join(data_dir, split)

algoperf/workloads/lm/lm_jax/models.py

@@ -14,5 +14,6 @@ def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
         return nn.Dense(
             self.vocab_size,
             kernel_init=nn.initializers.normal(0.02),
-            bias_init=nn.initializers.zeros
+            bias_init=nn.initializers.zeros,
+            name="output"
         )(x)

algoperf/workloads/lm/lm_jax/workload.py

@@ -2,16 +2,36 @@
 
 from typing import Dict, Optional, Tuple
 
+import jax
 import jax.numpy as jnp
+import optax
 from flax import jax_utils
 from algoperf import param_utils
+from algoperf import sharding_utils
 from algoperf import spec
 from algoperf.workloads.lm.workload import BaseLmWorkload
 from algoperf.workloads.lm.lm_jax.models import LinearModel
+from algoperf.workloads.lm.input_pipeline import get_hf_dataloader, get_lm_dataset
 
 
 class LmWorkload(BaseLmWorkload):
   """LM JAX workload."""
+  def _build_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 pre-cached FineWeb dataset."""
+    del num_batches
+    del repeat_final_dataset
+    loader = get_lm_dataset(
+        data_rng=data_rng,
+        split=split,
+        data_dir=data_dir,
+        global_batch_size=global_batch_size)
+    return loader
 
   def init_model_fn(
       self,
@@ -21,14 +41,15 @@ def init_model_fn(
 
     model = LinearModel(vocab_size=self._vocab_size)
     input_shape = (1, self._seq_len, self._vocab_size)
-    variables = model.init(rng, jnp.ones(input_shape, jnp.float32))
-    model_state, params = variables.pop('params')
-    
+    params_rng, init_rng = jax.random.split(rng)
+    variables = jax.jit(model.init)({'params': params_rng},
+                                  jnp.ones(input_shape, jnp.float32))
+    params = variables['params'] 
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    model_state = jax_utils.replicate(model_state)
-    params = jax_utils.replicate(params)
-    
+    params = sharding_utils.shard_replicated(params)
+    model_state = None
+    self._model = model
     return params, model_state
 
   def model_fn(
@@ -40,15 +61,40 @@ def model_fn(
       rng: spec.RandomState,
       update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
 
-    del mode, rng, update_batch_norm  # Not used for linear model
-    inputs = batch['inputs']
-    logits = self._model.apply({'params': params, **model_state}, inputs)
-    return logits, model_state
+    del mode, rng, update_batch_norm, model_state 
+    inputs = jax.nn.one_hot(batch['inputs'], self._vocab_size, axis=-1)
+    logits = self._model.apply({'params': params}, inputs)
+    return logits, None
+
+  def loss_fn(
+      self,
+      label_batch: spec.Tensor,  # One-hot labels.
+      logits_batch: spec.Tensor, # Dense logits.
+      mask_batch: Optional[spec.Tensor] = None,
+      label_smoothing: Optional[float] = 0.0) -> Dict[str, spec.Tensor]: 
+    del mask_batch, label_smoothing
+    logits_flat = logits_batch.reshape(-1, self._vocab_size)
+    targets = jax.nn.one_hot(label_batch, self._vocab_size, axis=-1)
+    targets_flat = targets.reshape(-1, self._vocab_size)
+    # Cross-entropy loss
+    loss = -jnp.sum(targets_flat * jax.nn.log_softmax(logits_flat, axis=-1))
+    n_valid_examples = logits_flat.shape[0]
+    return {'summed': loss, 'n_valid_examples': n_valid_examples}
 
+  def is_output_params(self, param_name: str) -> bool:
+    """Return whether the given parameter is an output parameter."""
+    return param_name.contains('output') 
+    
   def _eval_batch(self,
                   params: spec.ParameterContainer,
                   batch: Dict[str, spec.Tensor],
                   model_state: spec.ModelAuxiliaryState,
                   rng: spec.RandomState) -> spec.Tensor:
     """Evaluate the model on a single batch."""
-    pass
+    logits, _ = self.model_fn(
+        params, batch, model_state, spec.ForwardPassMode.EVAL, rng, False)
+    targets = batch['targets']
+    
+    # Calculate cross-entropy loss
+    loss = -jnp.sum(targets * jax.nn.log_softmax(logits, axis=-1))
+    return loss

algoperf/workloads/lm/lm_pytorch/workload.py

@@ -66,35 +66,30 @@ def _build_input_queue(
       global_batch_size: int,
       num_batches: Optional[int] = None,
       repeat_final_dataset: bool = False) -> Iterator[Dict[str, spec.Tensor]]:
-    not_train = split != 'train'
-    per_device_batch_size = int(global_batch_size / N_GPUS)
-
-    seq_len = self._seq_len  # TODO: define it somewehere else?
-    dtype = torch.int32  # TODO: decide between int32 and int64.
-
-    # Only create and iterate over tf input pipeline in one Python process to
-    # avoid creating too many threads.
-    if RANK == 0:
-      np_iter = super()._build_input_queue(
-          data_rng=data_rng,
-          split=split,
-          data_dir=data_dir,
-          global_batch_size=global_batch_size,
-          num_batches=num_batches,
-          repeat_final_dataset=repeat_final_dataset)
+    """Build an input queue for the given split."""
+    from algoperf.workloads.lm.input_pipeline import get_hf_dataloader
+    
+    loader = get_hf_dataloader(
+        cache_dir=data_dir,
+        data_rng=data_rng,
+        batch_size=global_batch_size,
+        seq_len=self._seq_len,
+        framework="torch",
+        split=split)
+    seq_len = self._seq_len 
     weights = None
-
+    
     while True:
       # Only iterate over tf input pipeline in one Python process to
       # avoid creating too many threads.
       if RANK == 0:
-        batch = next(np_iter)  # pylint: disable=stop-iteration-return
+        batch = next(dataset_iter)  # pylint: disable=stop-iteration-return
         inputs = torch.as_tensor(
             batch['inputs'], dtype=dtype,
-            device=DEVICE)  # (N_GPUS, global_batch_size, seq_len)
+            device=DEVICE)  # (N_GPUS, per_device_batch_size, seq_len)
         targets = torch.as_tensor(
             batch['targets'], dtype=dtype,
-            device=DEVICE)  # (N_GPUS, global_batch_size, seq_len)
+            device=DEVICE)  # (N_GPUS, per_device_batch_size, seq_len)
 
         # Send batch to other devices when using DDP.
         if USE_PYTORCH_DDP:
@@ -138,10 +133,22 @@ def _build_input_queue(
       }
       yield batch
 
+  def is_output_params(self, param_name: str) -> bool:
+    """Return whether the given parameter is an output parameter."""
+    return 'output.weight' in param_name or 'output.bias' in param_name
+    
   def _eval_batch(self,
                   params: spec.ParameterContainer,
                   batch: Dict[str, spec.Tensor],
                   model_state: spec.ModelAuxiliaryState,
                   rng: spec.RandomState) -> spec.Tensor:
     """Evaluate the model on a single batch."""
-    pass
+    model = params
+    logits, _ = self.model_fn(
+        model, batch, model_state, spec.ForwardPassMode.EVAL, rng, False)
+    targets = batch['targets']
+    
+    # Calculate cross-entropy loss
+    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+    loss = -torch.sum(targets * log_probs)
+    return loss