Perform a training step on a batch of data. Args: data: The input data for training, typically containing tensors and metadata. loss_function: A function that computes the loss and metrics given a batch and predictions. Returns: dict[str
(self, data: TensorDict, loss_function: Callable)
| 110 | raise NotImplementedError |
| 111 | |
| 112 | def train_batch(self, data: TensorDict, loss_function: Callable) -> Any: |
| 113 | """ |
| 114 | Perform a training step on a batch of data. |
| 115 | |
| 116 | Args: |
| 117 | data: The input data for training, typically containing tensors and metadata. |
| 118 | loss_function: A function that computes the loss and metrics given a batch and predictions. |
| 119 | |
| 120 | Returns: |
| 121 | dict[str, torch.Tensor]: A dictionary containing the aggregated training metrics for the batch. |
| 122 | """ |
| 123 | maybe_fix_3d_position_ids(data) |
| 124 | |
| 125 | self.optimizer_zero_grad() |
| 126 | outputs = self.forward_backward_batch(data, loss_function, forward_only=False) |
| 127 | grad_norm = self.optimizer_step() |
| 128 | if self.is_mp_src_rank_with_outputs(): |
| 129 | assert "grad_norm" not in outputs["metrics"] |
| 130 | outputs["metrics"]["grad_norm"] = grad_norm |
| 131 | return outputs |
| 132 | |
| 133 | def infer_batch(self, data: TensorDict, loss_function: Optional[Callable] = None) -> Any: |
| 134 | """ |