Generate a batch.
(data_iterator)
| 158 | |
| 159 | |
| 160 | def get_batch(data_iterator): |
| 161 | """Generate a batch.""" |
| 162 | |
| 163 | # TODO: this is pretty hacky, find a better way |
| 164 | if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()): |
| 165 | return None, None, None, None, None |
| 166 | |
| 167 | args = get_args() |
| 168 | |
| 169 | if "-Raw" in args.dataset: |
| 170 | # get batches based on the TP rank you are on |
| 171 | batch = get_batch_on_this_tp_rank_original(data_iterator) |
| 172 | # slice batch along sequence dimension for context parallelism |
| 173 | batch = get_batch_on_this_cp_rank(batch) |
| 174 | |
| 175 | elif "-Idxmap" in args.dataset: |
| 176 | # get batches based on the TP rank you are on |
| 177 | batch = get_batch_on_this_tp_rank(data_iterator) |
| 178 | # slice batch along sequence dimension for context parallelism |
| 179 | batch = get_batch_on_this_cp_rank(batch) |
| 180 | |
| 181 | else: |
| 182 | raise ValueError("please set correct --dataset ") |
| 183 | |
| 184 | return batch.values() |
| 185 | |
| 186 | def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor): |
| 187 | """Loss function. |
no test coverage detected