Closed NouamaneTazi closed 7 months ago
Selective recomputation is handled by flash-attn now so there's no need to keep it in config. We still keep the handy @checkpoint_method decorator if we want to activate recomputation for some method
@checkpoint_method
class MyFancyModule(nn.Module): def __init__(self): ... self.do_checkpoint: bool = True @checkpoint_method(attr_name="do_checkpoint") def forward(self, x): ...
Selective recomputation is handled by flash-attn now so there's no need to keep it in config. We still keep the handy
@checkpoint_method
decorator if we want to activate recomputation for some method