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