Fix torch sharding issue, update input pipeline and workload classes to use int32 for tensor types and add dropout rate parameter

Author: rka97

algoperf/workloads/lm/input_pipeline.py

@@ -119,8 +119,8 @@ def tf_generator():
   ds = tf.data.Dataset.from_generator(
       tf_generator,
       output_signature={
-          "inputs": tf.TensorSpec(shape=(None,), dtype=tf.int64),
-          "targets": tf.TensorSpec(shape=(None,), dtype=tf.int64),
+          "inputs": tf.TensorSpec(shape=(None,), dtype=tf.int32),
+          "targets": tf.TensorSpec(shape=(None,), dtype=tf.int32),
       })
 
   # Avoid creating too many threads when using PyTorch DDP.

algoperf/workloads/lm/lm_jax/workload.py

@@ -90,8 +90,9 @@ def model_fn(
       model_state: spec.ModelAuxiliaryState,
       mode: spec.ForwardPassMode,
       rng: spec.RandomState,
-      update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
-    del mode, rng, update_batch_norm, model_state
+      update_batch_norm: bool,
+      dropout_rate: float) -> 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
     if inputs.ndim == 3:  # one-hot encoded

algoperf/workloads/lm/lm_pytorch/workload.py

@@ -6,7 +6,8 @@
 import torch
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
-
+from itertools import islice
+from algoperf import data_utils
 from algoperf import param_utils
 from algoperf import pytorch_utils
 from algoperf import spec
@@ -84,19 +85,22 @@ def _build_input_queue(
       num_batches: Optional[int] = None,
       repeat_final_dataset: bool = False) -> Iterator[Dict[str, spec.Tensor]]:
     """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,
+    from algoperf.workloads.lm.input_pipeline import get_lm_dataset
+    local_batch_size = global_batch_size // N_GPUS
+    
+    loader = get_lm_dataset(
         data_rng=data_rng,
-        batch_size=global_batch_size,
-        seq_len=self._seq_len,
-        framework="torch",
-        split=split)
+        split=split,
+        data_dir=data_dir,
+        global_batch_size=local_batch_size,
+        num_batches=num_batches
+    )
+    if USE_PYTORCH_DDP:
+       loader = islice(loader, RANK, None, N_GPUS)
     seq_len = self._seq_len
     weights = None
 
-    dtype = torch.long
+    dtype = torch.int32
     is_train = split == 'train'
 
     for batch in loader:
@@ -109,17 +113,16 @@ def _build_input_queue(
           per_device_batch_size = torch.tensor(
               targets.shape[0], dtype=dtype, device=DEVICE)
           dist.broadcast(per_device_batch_size, src=0)
-
+          local_batch_size = per_device_batch_size.item()
         # Broadcast to all devices
-        dist.broadcast(inputs, src=0)
-        dist.broadcast(targets, src=0)
+        #dist.broadcast(inputs, src=0)
+        #dist.broadcast(targets, src=0)
 
       if weights is None:
-        batch_size = targets.shape[0] if not USE_PYTORCH_DDP else per_device_batch_size.item()
-        weights = torch.ones((batch_size, seq_len), device=DEVICE)
+        weights = torch.ones((local_batch_size, seq_len), device=DEVICE)
       batch = {
-          'inputs': inputs,
-          'targets': targets,
+          'inputs': torch.tensor(inputs, device=DEVICE, dtype=dtype),
+          'targets': torch.tensor(targets, device=DEVICE, dtype=dtype),
           'weights': weights,
       }
       yield batch

algoperf/workloads/lm/tests/test_build_input_queue_torch.py

@@ -17,9 +17,9 @@ def sync_ddp():
 def test_dataloader_torch():
   # Test config.
   rng_seed = 1996
-  data_dir = '/fast/najroldi/data/finewebedu'
+  data_dir = '/home/ak4605/data/finewebedu/'
   split = 'train'
-  global_batch_size = 8
+  global_batch_size = 64
   dtype = torch.int32
   seq_len = 2048
 
@@ -44,35 +44,40 @@ def test_dataloader_torch():
   # print(f"inputs: {inputs}")
 
   # Start test.
-  for _ in range(100):
+  for _ in range(1):
 
     batch = next(input_queue)
+    print(f"RANK {RANK} got batch")
 
     assert type(batch) == dict
     assert 'inputs' in batch
     assert 'targets' in batch
 
     inputs, targets = batch['inputs'], batch['targets']
-
+    print(f"RANK {RANK} inputs.shape: {inputs.shape}")
+    print(f"RANK {RANK} targets.shape: {targets.shape}")
+    print(f"RANK {RANK} type(inputs): {type(inputs)}")
     assert type(inputs) == torch.Tensor
     assert type(targets) == torch.Tensor
 
     assert inputs.device == DEVICE
     assert targets.device == DEVICE
-
     assert inputs.dtype == dtype
     assert targets.dtype == dtype
 
+    print(local_batch_size, seq_len)
     assert inputs.shape == (local_batch_size, seq_len)
     assert targets.shape == (local_batch_size, seq_len)
 
     assert torch.equal(inputs[:, 1:], targets[:, :-1])
+    print(f"RANK {RANK} inputs[0, :10]: {inputs[0, :10]}")
 
   print(f"=== ALL TEST PASSED ===")
 
 
 def main():
   profiler = PassThroughProfiler()
+  print(USE_PYTORCH_DDP, RANK, DEVICE, N_GPUS)
   pytorch_init(USE_PYTORCH_DDP, RANK, profiler)
   test_dataloader_torch()
 

algoperf/workloads/lm/workload.py

@@ -132,7 +132,8 @@ def _eval_batch(self,
         model_state,
         spec.ForwardPassMode.EVAL,
         rng,
-        update_batch_norm=False)
+        update_batch_norm=False,
+        dropout_rate=None)
 
     loss_dict = self.loss_fn(batch['targets'], logits)
     return loss_dict['summed']