Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pass. :param func: the function to evaluate. :param inputs: the argument sequence to pass to `func`. :param params: a sequence of parameters
(func, inputs, params, flag)
| 84 | #ckpt = deepspeed.checkpointing.checkpoint |
| 85 | ckpt = torch.utils.checkpoint.checkpoint |
| 86 | def checkpoint(func, inputs, params, flag): |
| 87 | """ |
| 88 | Evaluate a function without caching intermediate activations, allowing for |
| 89 | reduced memory at the expense of extra compute in the backward pass. |
| 90 | :param func: the function to evaluate. |
| 91 | :param inputs: the argument sequence to pass to `func`. |
| 92 | :param params: a sequence of parameters `func` depends on but does not |
| 93 | explicitly take as arguments. |
| 94 | :param flag: if False, disable gradient checkpointing. |
| 95 | """ |
| 96 | if flag: |
| 97 | try: |
| 98 | return ckpt(func, *inputs) |
| 99 | except: |
| 100 | args = tuple(inputs) + tuple(params) |
| 101 | return CheckpointFunction.apply(func, len(inputs), *args) |
| 102 | else: |
| 103 | return func(*inputs) |
| 104 | |
| 105 | |
| 106 | class CheckpointFunction(torch.autograd.Function): |
no outgoing calls
no test coverage detected