| 212 | """ |
| 213 | |
| 214 | def __init__( |
| 215 | self, |
| 216 | flops_available: float, |
| 217 | log_dict: Callable[[Dict, int], None], |
| 218 | window_size: int = 100, |
| 219 | time_unit: str = 'hours', |
| 220 | log_iter_interval: int = 1, |
| 221 | ): |
| 222 | self.flops_available = flops_available |
| 223 | self.log_dict = log_dict |
| 224 | self.log_iter_interval = log_iter_interval |
| 225 | # Track the batch num samples and wct to compute throughput over a window of batches |
| 226 | self.history_samples: Deque[int] = deque(maxlen=window_size + 1) |
| 227 | self.history_training_loss: Deque[int] = deque(maxlen=log_iter_interval) |
| 228 | self.history_wct: Deque[float] = deque(maxlen=window_size + 1) |
| 229 | self.history_lengths: Deque[int] = deque(maxlen=window_size + 1) |
| 230 | self.history_flops: Deque[int] = deque(maxlen=window_size + 1) |
| 231 | |
| 232 | self.divider = 1 |
| 233 | if time_unit == 'seconds': |
| 234 | self.divider = 1 |
| 235 | elif time_unit == 'minutes': |
| 236 | self.divider = 60 |
| 237 | elif time_unit == 'hours': |
| 238 | self.divider = 60 * 60 |
| 239 | elif time_unit == 'days': |
| 240 | self.divider = 60 * 60 * 24 |
| 241 | else: |
| 242 | raise ValueError( |
| 243 | f'Invalid time_unit: {time_unit}. Must be one of "seconds", "minutes", "hours", or "days".' |
| 244 | ) |
| 245 | |
| 246 | # Keep track of time spent evaluating |
| 247 | self.total_eval_wct = 0.0 |
| 248 | self.iter = -1 |
| 249 | |
| 250 | def on_train_batch_end( |
| 251 | self, |