| 266 | |
| 267 | |
| 268 | def build_dataloader( |
| 269 | cfg, |
| 270 | tokenizer, |
| 271 | device_batch_size, |
| 272 | count_padding_tokens=True, |
| 273 | device_microbatch_size: int | None = None, |
| 274 | ): |
| 275 | split_batch_fn = None |
| 276 | num_samples_in_batch_fn = None |
| 277 | num_tokens_in_batch_fn = None |
| 278 | |
| 279 | if cfg.name == "text": |
| 280 | data_loader = text_data_module.build_text_dataloader( |
| 281 | cfg, |
| 282 | tokenizer, |
| 283 | device_batch_size, |
| 284 | device_microbatch_size=device_microbatch_size, |
| 285 | ) |
| 286 | else: |
| 287 | raise ValueError(f"Not sure how to build dataloader with config: {cfg}") |
| 288 | |
| 289 | if not count_padding_tokens: |
| 290 | num_tokens_in_batch_fn = get_num_tokens_in_batch_unpadded |
| 291 | if cfg.get("sequence_packing", False): |
| 292 | split_batch_fn = split_packed_batch |
| 293 | num_samples_in_batch_fn = get_num_samples_in_packed_batch |
| 294 | |
| 295 | data_loader = DataSpec( |
| 296 | data_loader, |
| 297 | get_num_tokens_in_batch=num_tokens_in_batch_fn, |
| 298 | split_batch=split_batch_fn, |
| 299 | get_num_samples_in_batch=num_samples_in_batch_fn, |
| 300 | ) |
| 301 | return data_loader |
| 302 | |
| 303 | |
| 304 | def build_model(cfg: DictConfig): |