Testing with linear model

Author: rka97

algoperf/workloads/lm/input_pipeline.py

@@ -86,7 +86,6 @@ def batch_iterator():
         token_ids = jnp.stack([_tokenize(x) for x in doc['text']])
         tokens = jax.nn.one_hot(token_ids, num_classes=vocab_size)
         inputs, targets = tokens[:, :-1], tokens[:, 1:]
-        devices = jax.devices("gpu")
         inputs, targets = jax.device_put(inputs), jax.device_put(targets)
       yield inputs, targets
 

algoperf/workloads/lm/lm_jax/models.py

@@ -0,0 +1,18 @@
+from flax import linen as nn
+import jax.numpy as jnp
+
+class LinearModel(nn.Module):
+    vocab_size: int
+    
+    @nn.compact
+    def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
+        x = nn.Dense(
+            512,
+            kernel_init=nn.initializers.normal(0.02),
+            bias_init=nn.initializers.zeros
+        )(inputs)
+        return nn.Dense(
+            self.vocab_size,
+            kernel_init=nn.initializers.normal(0.02),
+            bias_init=nn.initializers.zeros
+        )(x)

algoperf/workloads/lm/lm_jax/workload.py

@@ -2,8 +2,12 @@
 
 from typing import Dict, Optional, Tuple
 
+import jax.numpy as jnp
+from flax import jax_utils
+from algoperf import param_utils
 from algoperf import spec
 from algoperf.workloads.lm.workload import BaseLmWorkload
+from algoperf.workloads.lm.lm_jax.models import LinearModel
 
 
 class LmWorkload(BaseLmWorkload):
@@ -14,18 +18,32 @@ def init_model_fn(
       rng: spec.RandomState,
       dropout_rate: Optional[float] = None,
       aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
-    """aux_dropout_rate is used as attention_dropout_rate."""
-    pass
+    
+    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')
+    
+    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)
+    
+    return params, model_state
 
   def model_fn(
       self,
       params: spec.ParameterContainer,
-      augmented_and_preprocessed_input_batch: Dict[str, spec.Tensor],
+      batch: Dict[str, spec.Tensor],
       model_state: spec.ModelAuxiliaryState,
       mode: spec.ForwardPassMode,
       rng: spec.RandomState,
       update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
-    pass
+    
+    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
 
   def _eval_batch(self,
                   params: spec.ParameterContainer,

algoperf/workloads/lm/lm_pytorch/models.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+
+class LinearLayer(nn.Module):
+    def __init__(self, vocab_size: int):
+        super().__init__()
+        self.bottleneck = nn.Linear(vocab_size, 512)
+        self.output = nn.Linear(512, vocab_size)
+        self.reset_parameters()
+        
+    def reset_parameters(self):
+        nn.init.normal_(self.bottleneck.weight, std=0.02)
+        nn.init.zeros_(self.bottleneck.bias)
+        nn.init.normal_(self.output.weight, std=0.02)
+        nn.init.zeros_(self.output.bias)
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.output(self.bottleneck(x))

algoperf/workloads/lm/lm_pytorch/workload.py

@@ -5,10 +5,13 @@
 import jax
 import torch
 import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
 
+from algoperf import param_utils
 from algoperf import pytorch_utils
 from algoperf import spec
 from algoperf.workloads.lm.workload import BaseLmWorkload
+from algoperf.workloads.lm.lm_pytorch.models import LinearLayer
 
 USE_PYTORCH_DDP, RANK, DEVICE, N_GPUS = pytorch_utils.pytorch_setup()
 
@@ -21,18 +24,39 @@ def init_model_fn(
       rng: spec.RandomState,
       dropout_rate: Optional[float] = None,
       aux_dropout_rate: Optional[float] = None) -> spec.ModelInitState:
-    """aux_dropout_rate is used as attention_dropout_rate."""
-    pass
+    
+    if hasattr(self, '_model'):
+        self._model.reset_parameters()
+        return self._model, None
+
+    torch.manual_seed(rng[0])
+    self._model = LinearLayer(vocab_size=self._vocab_size)
+    self._param_shapes = param_utils.pytorch_param_shapes(self._model)
+    self._param_types = param_utils.pytorch_param_types(self._param_shapes)
+    self._model.to(DEVICE)
+    
+    if N_GPUS > 1:
+        if USE_PYTORCH_DDP:
+            self._model = DDP(self._model, device_ids=[RANK], output_device=RANK)
+        else:
+            self._model = torch.nn.DataParallel(self._model)
+            
+    return self._model, None
 
   def model_fn(
       self,
       params: spec.ParameterContainer,
-      augmented_and_preprocessed_input_batch: Dict[str, spec.Tensor],
+      batch: Dict[str, spec.Tensor],
       model_state: spec.ModelAuxiliaryState,
       mode: spec.ForwardPassMode,
       rng: spec.RandomState,
       update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
-    pass
+    
+    del model_state, rng, update_batch_norm  # Not used for linear model
+    model = params
+    inputs = batch['inputs'].float()  # Convert one-hot to float
+    logits = model(inputs)
+    return logits, None
 
   def _build_input_queue(
       self,

algoperf/workloads/lm/tests/test_linear_model.py

@@ -0,0 +1,39 @@
+import jax
+import jax.numpy as jnp
+import torch
+
+TEST_SEQ_LEN = 512
+
+def test_pytorch_linear():
+    from algoperf.workloads.lm.lm_pytorch.models import LinearLayer
+    vocab_size = 32000
+    model = LinearLayer(vocab_size)
+    
+    batch_size = 8
+    seq_len = TEST_SEQ_LEN
+    inputs = torch.randn(batch_size, seq_len, vocab_size)
+    outputs = model(inputs)
+    
+    assert outputs.shape == (batch_size, seq_len, vocab_size)
+    assert not torch.isnan(outputs).any()
+
+def test_jax_linear():
+    from algoperf.workloads.lm.lm_jax.models import LinearModel
+
+    vocab_size = 32000
+    seq_len = TEST_SEQ_LEN
+    batch_size = 8
+    model = LinearModel(vocab_size)
+    rng = jax.random.PRNGKey(0)
+    params = model.init(rng, jnp.ones((1, seq_len, vocab_size)))
+    
+    inputs = jax.random.normal(rng, (batch_size, seq_len, vocab_size))
+    outputs = model.apply(params, inputs)
+    
+    assert outputs.shape == (batch_size, seq_len, vocab_size)
+    assert not jnp.isnan(outputs).any()
+
+if __name__ == '__main__':
+    test_pytorch_linear()
+    test_jax_linear()
+    print("All tests passed!")

algoperf/workloads/lm/workload.py

@@ -20,8 +20,8 @@
 class BaseLmWorkload(spec.Workload):
   """LM workload."""
 
-  _vocab_size: int = 32000
-  _seq_len: int = 2048
+  _vocab_size: int = 50257
+  _seq_len: int = 512
 
   def __init__(self) -> None:
     pass
@@ -106,24 +106,15 @@ def activation(self) -> str:
   def glu(self) -> bool:
     return True
 
+  @abc.abstractmethod
   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):
-    ds = input_pipeline.get_lm_dataset(
-        data_rng,
-        split,
-        data_dir,
-        vocab_size=self._vocab_size,
-        global_batch_size=global_batch_size,
-        num_batches=num_batches,
-        repeat_final_dataset=repeat_final_dataset)
-
-    for batch in iter(ds):
-      yield batch
+    """Build an input queue for the given split."""
 
   @abc.abstractmethod
   def _eval_batch(self,