(ctx, run_function, *args)
| 272 | |
| 273 | @staticmethod |
| 274 | def forward(ctx, run_function, *args): |
| 275 | ctx.run_function = run_function |
| 276 | |
| 277 | # Copy the rng states. |
| 278 | ctx.fwd_cpu_rng_state = torch.get_rng_state() |
| 279 | ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state() |
| 280 | ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() |
| 281 | |
| 282 | with torch.no_grad(): |
| 283 | outputs = run_function(*args) |
| 284 | |
| 285 | # Divide hidden states across model parallel group and only keep |
| 286 | # the chunk corresponding to the current rank. |
| 287 | if _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER is not None: |
| 288 | ctx.input_0_shape = args[0].data.shape |
| 289 | args[0].data = split_tensor_into_1d_equal_chunks(args[0].data) |
| 290 | args[0].data = _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER.add(args[0].data) |
| 291 | |
| 292 | # Store everything. |
| 293 | ctx.save_for_backward(*args) |
| 294 | |
| 295 | return outputs |
| 296 | |
| 297 | @staticmethod |
| 298 | def backward(ctx, *args): |
nothing calls this directly
no test coverage detected