Method on_train_batch_end

lit_gpt/speed_monitor.py:250–338

(
        self,
        samples: int,  # total samples seen (per device)
        train_elapsed: float,  # total training time (seconds)
        world_size: int,
        step_count: int,
        flops_per_batch: Optional[int] = None,  # (per device)
        lengths: Optional[int] = None,  # total length of the samples seen (per device)
        train_loss: Optional[float] = None,
    )

        self.iter = -1

    def on_train_batch_end(
        self,
        samples: int,  # total samples seen (per device)
        train_elapsed: float,  # total training time (seconds)
        world_size: int,
        step_count: int,
        flops_per_batch: Optional[int] = None,  # (per device)
        lengths: Optional[int] = None,  # total length of the samples seen (per device)
        train_loss: Optional[float] = None,
    ):
        self.iter += 1
        metrics = {}

        self.history_samples.append(samples)
        self.history_training_loss.append(train_loss)
        if lengths is not None:
            self.history_lengths.append(lengths)
            # if lengths are passed, there should be as many values as samples
            assert len(self.history_samples) == len(self.history_lengths)
        self.history_wct.append(train_elapsed)
        if len(self.history_wct) == self.history_wct.maxlen:
            elapsed_batches = len(self.history_samples) - 1
            elapsed_samples = self.history_samples[-1] - self.history_samples[0]
            elapsed_wct = self.history_wct[-1] - self.history_wct[0]
            samples_per_sec = elapsed_samples * world_size / elapsed_wct
            dev_samples_per_sec = elapsed_samples / elapsed_wct
            metrics.update(
                {
                    'throughput/batches_per_sec': elapsed_batches
                    * world_size
                    / elapsed_wct,
                    'throughput/samples_per_sec': samples_per_sec,
                    'throughput/device/batches_per_sec': elapsed_batches / elapsed_wct,
                    'throughput/device/samples_per_sec': dev_samples_per_sec,
                }
            )
            if lengths is not None:
                elapsed_lengths = int(self.history_lengths[-1]) - int(
                    self.history_lengths[0]
                )
                avg_length = elapsed_lengths / elapsed_batches
                metrics.update(
                    {
                        'throughput/tokens_per_sec': samples_per_sec * avg_length,
                        'throughput/device/tokens_per_sec': dev_samples_per_sec
                        * avg_length,
                        'total_tokens': avg_length * world_size * samples,
                    }
                )
                if train_loss is not None:
                    avg_loss = sum(self.history_training_loss) / len(
                        self.history_training_loss
                    )
                    metrics.update(
                        {
                            'metric/train_loss': avg_loss,
                            'metric/train_ppl': math.exp(avg_loss),
                        }