ColossalAI/colossalai
oahzxl 5db3a5bf42
[fx] allow control of ckpt_codegen init (#2498)
* [fx] allow control of ckpt_codegen init

Currently in ColoGraphModule, ActivationCheckpointCodeGen will be set automatically in __init__. But other codegen can't be set if so. 
So I add an arg to control whether to set ActivationCheckpointCodeGen in __init__.

* code style
2023-01-18 17:02:46 +08:00
..
_C
amp
auto_parallel
autochunk
builder
cli
communication
context
device
engine
fx [fx] allow control of ckpt_codegen init (#2498) 2023-01-18 17:02:46 +08:00
gemini
kernel
logging
nn
pipeline
registry
tensor
testing
trainer
utils
zero
__init__.py
constants.py
core.py
global_variables.py
initialize.py