Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,157 Bytes
d2c9b66 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
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
# when passing p as a fraction number (e.g. 1/3), it will be interpreted
# as a string in argv, thus we need eval("1/3") here for fractions.
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,
)
|