A DataLoader that implements batch size ramp-up. It dynamically increases the 'global_batch_size_current' in the config object based on the training step. The rest of the training pipeline (including the parent's `check_example_batch` and the training step itself) is assumed to read this
| 86 | |
| 87 | |
| 88 | class RampUpDataLoader(DataLoader): |
| 89 | """ |
| 90 | A DataLoader that implements batch size ramp-up. |
| 91 | |
| 92 | It dynamically increases the 'global_batch_size_current' in the config |
| 93 | object based on the training step. The rest of the training pipeline |
| 94 | (including the parent's `check_example_batch` and the training step itself) |
| 95 | is assumed to read this config value to determine the logical batch size. |
| 96 | """ |
| 97 | |
| 98 | def __init__(self, config, mesh, data_iterator, goodput_recorder): |
| 99 | # Call parent constructor |
| 100 | super().__init__(config, mesh, data_iterator, goodput_recorder) |
| 101 | |
| 102 | self.rampup_active = True |
| 103 | self.batch_buffer = None |
| 104 | self.buffer_start = 0 |
| 105 | |
| 106 | def load_next_batch(self, *args, rampup_manager=None, **kwargs): |
| 107 | """ |
| 108 | Updates the batch size based on the schedule and then loads the next |
| 109 | batch using the parent method. |
| 110 | """ |
| 111 | # If ramp-up is not active, just behave like the parent |
| 112 | if not self.rampup_active: |
| 113 | return super().load_next_batch() |
| 114 | |
| 115 | slice_start, slice_end = self.buffer_start, self.buffer_start + rampup_manager.global_batch_size_current |
| 116 | |
| 117 | # Load new batch if batch_buffer is None |
| 118 | if self.batch_buffer is None: |
| 119 | self.batch_buffer = super().load_next_batch_pre_sharding() |
| 120 | slice_start, slice_end = 0, rampup_manager.global_batch_size_current |
| 121 | |
| 122 | # If the slice end overpast batch end we collect new batch data |
| 123 | if slice_end > rampup_manager.global_batch_size_end: |
| 124 | old_buffer, self.batch_buffer = self.batch_buffer, super().load_next_batch_pre_sharding() |
| 125 | |
| 126 | # self.global_batch_size_end is batch_buffer size |
| 127 | def _slice_and_concat(old_data, new_data): |
| 128 | sliced_old_data = jax.lax.dynamic_slice_in_dim( |
| 129 | old_data, |
| 130 | slice_start, |
| 131 | rampup_manager.global_batch_size_end - slice_start, |
| 132 | axis=0, |
| 133 | ) |
| 134 | sliced_new_data = jax.lax.dynamic_slice_in_dim( |
| 135 | new_data, |
| 136 | 0, |
| 137 | slice_end - rampup_manager.global_batch_size_end, |
| 138 | axis=0, |
| 139 | ) |
| 140 | return jax.lax.concatenate((sliced_old_data, sliced_new_data), dimension=0) |
| 141 | |
| 142 | self.buffer_start = slice_end - rampup_manager.global_batch_size_end |
| 143 | output = jax.tree.map(_slice_and_concat, old_buffer, self.batch_buffer) |
| 144 | else: |
| 145 |
no outgoing calls