A stateful class tracking current batch size given train step
| 19 | |
| 20 | |
| 21 | class RampupBatchManager: |
| 22 | """ |
| 23 | A stateful class tracking current batch size given train step |
| 24 | """ |
| 25 | |
| 26 | def __init__(self, config, step_num): |
| 27 | self._verify_inputs(config) |
| 28 | self._init_values(config) |
| 29 | self.num_accum_samples = 0 |
| 30 | |
| 31 | # Compute the number of samples already used given recovered step num |
| 32 | self._recover_states(step_num) |
| 33 | |
| 34 | def _verify_inputs(self, config): |
| 35 | """Verify the rampup batch related inputs.""" |
| 36 | diff_batch_size = config.per_device_batch_size - config.per_device_batch_size_start |
| 37 | if diff_batch_size <= 0: |
| 38 | raise ValueError( |
| 39 | "per_device_batch_size must be greater than per_device_batch_size_start. " |
| 40 | f"get batch size is {config.per_device_batch_size} and " |
| 41 | f"batch size start is {config.per_device_batch_size_start}." |
| 42 | ) |
| 43 | if diff_batch_size % config.per_device_batch_size_increment: |
| 44 | raise ValueError( |
| 45 | "Expect rampup batch size change divisible by batch size increment." |
| 46 | f"Got per_device_batch_size={config.per_device_batch_size} and " |
| 47 | f"per_device_batch_size_start={config.per_device_batch_size_start}." |
| 48 | ) |
| 49 | |
| 50 | def _init_values(self, config): |
| 51 | """Initialize rampup batch related parameters""" |
| 52 | diff_batch_size = config.per_device_batch_size - config.per_device_batch_size_start |
| 53 | num_increments = diff_batch_size // config.per_device_batch_size_increment |
| 54 | self.samples_per_increment = config.global_rampup_samples / num_increments |
| 55 | num_devices = int(config.num_target_devices) |
| 56 | self.global_batch_size_end = int(num_devices * config.per_device_batch_size) |
| 57 | self.global_batch_size_start = int(num_devices * config.per_device_batch_size_start) |
| 58 | self.increment = int(num_devices * config.per_device_batch_size_increment) |
| 59 | self.global_rampup_samples = config.global_rampup_samples |
| 60 | self.global_batch_size_current = self.global_batch_size_start |
| 61 | self.total_rampup_steps = self._compute_total_rampup_steps(config) |
| 62 | self.total_used_samples = 0 |
| 63 | |
| 64 | def _compute_total_rampup_steps(self, config): |
| 65 | """Compute total number of rampup steps""" |
| 66 | batch_size_start = config.per_device_batch_size_start |
| 67 | batch_size_end = config.per_device_batch_size |
| 68 | batch_size_increment = config.per_device_batch_size_increment |
| 69 | diff_batch_size = batch_size_end - batch_size_start |
| 70 | num_increments = diff_batch_size // batch_size_increment |
| 71 | rampup_samples = config.global_rampup_samples / config.num_target_devices |
| 72 | rampup_samples_per_increment = rampup_samples / num_increments |
| 73 | total_rampup_steps = 0 |
| 74 | current_batch_size = batch_size_start |
| 75 | |
| 76 | while current_batch_size < batch_size_end: |
| 77 | steps_for_this_stage = math.ceil(rampup_samples_per_increment / current_batch_size) |
| 78 | total_rampup_steps += steps_for_this_stage |
no outgoing calls