mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-25 11:44:03 +00:00
e614aa34f3
* feat: add `GradientCheckpointConfig` and `PipelineGradientCheckpointConfig` * feat: apply `GradientCheckpointConfig` to policy and llama_forward * feat: move `distribute_layer` and `get_stage_index` to PipelineStageManager * fix: add optional args for `distribute_layer` and `get_stage_index` * fix: fix changed API calls * test: update llama tests * style: polish `GradientCheckpointConfig` * fix: fix pipeline utils tests
7 lines
320 B
Python
7 lines
320 B
Python
from .grad_ckpt_config import GradientCheckpointConfig, PipelineGradientCheckpointConfig
|
|
from .shard_config import ShardConfig
|
|
from .sharder import ModelSharder
|
|
from .shardformer import ShardFormer
|
|
|
|
__all__ = ["ShardConfig", "ModelSharder", "ShardFormer", "PipelineGradientCheckpointConfig", "GradientCheckpointConfig"]
|