Method test_rampup_data_loader

tests/data_loader_test.py:129–153

Tests that RampUpLoader correctly slices and increment.

(self)

    self.assertTrue(str(e.exception).startswith("You may have run out of training data."))

  def test_rampup_data_loader(self):
    """Tests that RampUpLoader correctly slices and increment."""
    # Mock iterator returns a FULL batch (size 4)
    full_batch_size = int(self.config_rampup.per_device_batch_size * self.config_rampup.num_target_devices)
    full_shape = [full_batch_size, self.config_rampup.max_target_length]
    full_batch = {"inputs": np.ones(full_shape, dtype=int)}
    self.mock_data_iterator.__next__.return_value = full_batch

    # Create the RampUpDataLoader
    rampup_manager = RampupBatchManager(self.config_rampup, -1)
    data_loader = RampUpDataLoader(self.config_rampup, self.mesh, self.mock_data_iterator, None)

    # Expected batch sizes based on test config.
    # The end global batch size is self.num_devices * per_device_batch_size
    # The rampup should be: 5 steps of size 4, 3 steps of size 8, 2 steps of size 12, then size 16.
    expected_batch_sizes = [4, 4, 4, 4, 4, 8, 8, 8, 12, 12, 16, 16]
    for i, expected_size in enumerate(expected_batch_sizes):
      batch = data_loader.load_next_batch(rampup_manager=rampup_manager)
      expected_shape = (expected_size, self.config_rampup.max_target_length)
      self.assertEqual(
          batch["inputs"].shape,
          expected_shape,
          f"Mismatch at step {i+1}: expected {expected_shape}, got {batch['inputs'].shape}",
      )
      self.assertTrue((batch["inputs"] == 1).all())

  def test_rampup_data_loader_from_checkpointing(self):
    """Tests that RampUpLoader correctly slices and increment resumed from checkpointing."""