| 341 | return idxs |
| 342 | |
| 343 | def forward(self, forward_input): |
| 344 | # We need to offset the seed by the microbatch ID. Save it in a local var to |
| 345 | # ensure it is preserved in the closure. Otherwise checkpointed forward funcs |
| 346 | # will see a different offset. |
| 347 | self.micro_offset += 1 |
| 348 | |
| 349 | def exec_range_func(start, end): |
| 350 | ''' Helper function to be used with checkpoint() |
| 351 | Adapted from torch.utils.checkpoint:checkpoint_sequential() |
| 352 | ''' |
| 353 | local_micro_offset = self.micro_offset + 1 |
| 354 | |
| 355 | def exec_func(*inputs): |
| 356 | # Single tensor inputs need to be unwrapped |
| 357 | if len(inputs) == 1: |
| 358 | inputs = inputs[0] |
| 359 | for idx, layer in enumerate(self.forward_funcs[start:end]): |
| 360 | self.curr_layer = idx + self._local_start |
| 361 | if self.seed_layers: |
| 362 | new_seed = (self.base_seed * local_micro_offset) + self.curr_layer |
| 363 | if self.seed_fn: |
| 364 | self.seed_fn(new_seed) |
| 365 | else: |
| 366 | ds_utils.set_random_seed(new_seed) |
| 367 | |
| 368 | inputs = layer(inputs) |
| 369 | return inputs |
| 370 | |
| 371 | return exec_func |
| 372 | |
| 373 | if self.activation_checkpoint_interval == 0: |
| 374 | func = exec_range_func(0, len(self.forward_funcs)) |
| 375 | x = func(forward_input) |
| 376 | else: |
| 377 | num_layers = len(self.forward_funcs) |
| 378 | x = forward_input |
| 379 | for start_idx, is_checkpointable_result in \ |
| 380 | zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results): |
| 381 | |
| 382 | end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) |
| 383 | |
| 384 | funcs = self.forward_funcs[start_idx:end_idx] |
| 385 | # Since we either pass tensors or tuples of tensors without unpacking, we |
| 386 | # need to be careful not to double-wrap tensors with tuple. |
| 387 | if not isinstance(x, tuple): |
| 388 | x = (x, ) |
| 389 | |
| 390 | if is_checkpointable_result: |
| 391 | x = self.activation_checkpoint_func(exec_range_func(start_idx, end_idx), *x) |
| 392 | else: |
| 393 | x = exec_range_func(start_idx, end_idx)(*x) |
| 394 | return x |
| 395 | |
| 396 | def _partition_layers(self, method='uniform'): |
| 397 | num_stages = self._topo.get_dim('pipe') |