Loads preprocessed data for training.
| 28 | |
| 29 | |
| 30 | class DataLoader: |
| 31 | """ |
| 32 | Loads preprocessed data for training. |
| 33 | """ |
| 34 | |
| 35 | def __init__(self, config, mesh, data_iterator, goodput_recorder): |
| 36 | self.config = config |
| 37 | self.goodput_recorder = goodput_recorder |
| 38 | if isinstance(data_iterator, list): |
| 39 | self.data_iterator_list = data_iterator |
| 40 | self.data_iterator_index = 0 |
| 41 | self.data_iterator = self.data_iterator_list[self.data_iterator_index] |
| 42 | else: |
| 43 | self.data_iterator = data_iterator |
| 44 | self.last_batch = None |
| 45 | self.input_data_shardings = get_input_data_sharding(config, mesh) |
| 46 | |
| 47 | def update_data_iterator(self): |
| 48 | """Update to the next data iterator in the list, if applicable.""" |
| 49 | if hasattr(self, "data_iterator_list"): |
| 50 | self.data_iterator_index = (self.data_iterator_index + 1) % len(self.data_iterator_list) |
| 51 | self.data_iterator = self.data_iterator_list[self.data_iterator_index] |
| 52 | |
| 53 | def load_next_batch_pre_sharding(self): |
| 54 | """Loads the next batch w/o sharding. Can keep reusing the same batch for performance reasons.""" |
| 55 | with maybe_record_goodput(self.goodput_recorder, GoodputEvent.DATA_LOADING): |
| 56 | try: |
| 57 | if self.config.reuse_example_batch and self.last_batch: |
| 58 | example_batch = self.last_batch |
| 59 | else: |
| 60 | example_batch = next(self.data_iterator) |
| 61 | self.update_data_iterator() |
| 62 | self.last_batch = example_batch |
| 63 | self.check_example_batch() |
| 64 | except Exception as e: # pylint: disable=broad-except |
| 65 | if isinstance(e, StopIteration): |
| 66 | raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)} exception: ({e})") |
| 67 | else: |
| 68 | raise exceptions.StopTraining(f"`load_next_batch()` failed with {type(e)} exception: ({e}).") |
| 69 | return self.last_batch |
| 70 | |
| 71 | def load_next_batch(self, *args, **kwargs): |
| 72 | """Loads the next batch with sharding hint""" |
| 73 | return maybe_shard_with_name( |
| 74 | self.load_next_batch_pre_sharding(), |
| 75 | self.input_data_shardings, |
| 76 | self.config.shard_mode, |
| 77 | ) |
| 78 | |
| 79 | def check_example_batch(self): |
| 80 | if self.config.max_checkify: |
| 81 | jittable_f = checkify.checkify(lambda x: checkify.check(jnp.any(x > -1), "Batch contains bad synthetic data!")) |
| 82 | # Check if inputs in batch contains bad synthetic data. |
| 83 | # pylint: disable=not-callable |
| 84 | err, _ = jax.jit(jittable_f)(self.last_batch["inputs"][: self.config.global_batch_size_to_train_on, :]) |
| 85 | err.throw() |
| 86 | |
| 87 |
no outgoing calls