| from functools import partial |
|
|
| from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( |
| CheckpointImpl, |
| apply_activation_checkpointing, |
| checkpoint_wrapper, |
| ) |
|
|
|
|
| non_reentrant_wrapper = partial( |
| checkpoint_wrapper, |
| checkpoint_impl=CheckpointImpl.NO_REENTRANT, |
| ) |
|
|
|
|
| def apply_checkpointing(model, block, p): |
| """ |
| Apply selective activation checkpointing. |
| |
| Selectivity is defined as a percentage p, which means we apply ac |
| on p of the total blocks. p is a floating number in the range of |
| [0, 1]. |
| |
| Some examples: |
| p = 0: no ac for all blocks. same as `fsdp_activation_checkpointing=False` |
| p = 1: apply ac on every block. i.e. "full ac". |
| p = 1/2: [ac, no-ac, ac, no-ac, ...] |
| p = 1/3: [no-ac, ac, no-ac, no-ac, ac, no-ac, ...] |
| p = 2/3: [ac, no-ac, ac, ac, no-ac, ac, ...] |
| Since blocks are homogeneous, we make ac blocks evenly spaced among |
| all blocks. |
| |
| Implementation: |
| For a given ac ratio p, we should essentially apply ac on every "1/p" |
| blocks. The first ac block can be as early as the 0th block, or as |
| late as the "1/p"th block, and we pick the middle one: (0.5p)th block. |
| Therefore, we are essentially to apply ac on: |
| (0.5/p)th block, (1.5/p)th block, (2.5/p)th block, etc., and of course, |
| with these values rounding to integers. |
| Since ac is applied recursively, we can simply use the following math |
| in the code to apply ac on corresponding blocks. |
| """ |
| block_idx = 0 |
| cut_off = 1 / 2 |
| |
| |
| p = eval(p) if isinstance(p, str) else p |
|
|
| def selective_checkpointing(submodule): |
| nonlocal block_idx |
| nonlocal cut_off |
|
|
| if isinstance(submodule, block): |
| block_idx += 1 |
| if block_idx * p >= cut_off: |
| cut_off += 1 |
| return True |
| return False |
|
|
| apply_activation_checkpointing( |
| model, |
| checkpoint_wrapper_fn=non_reentrant_wrapper, |
| check_fn=selective_checkpointing, |
| ) |
|
|
