github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
ColossalAI/colossalai/shardformer/shard/__init__.py

7 lines
320 B
Python
Raw Normal View History

[shardformer, pipeline] add `gradient_checkpointing_ratio` and heterogenous shard policy for llama (#5508) * 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
2024-04-01 03:34:58 +00:00
from .grad_ckpt_config import GradientCheckpointConfig, PipelineGradientCheckpointConfig
[shardformer] refactored the user api (#3828) * [shardformer] refactored the user api * polish code
2023-05-24 08:01:26 +00:00
from .shard_config import ShardConfig
[shardformer] Refactor shardformer api (#4001) * fix an error in readme * simplify code * refactor shardformer * add todo * remove slicer * resolve code review
2023-06-15 09:55:42 +00:00
from .sharder import ModelSharder
from .shardformer import ShardFormer
[shardformer] refactored the user api (#3828) * [shardformer] refactored the user api * polish code
2023-05-24 08:01:26 +00:00
[shardformer, pipeline] add `gradient_checkpointing_ratio` and heterogenous shard policy for llama (#5508) * 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
2024-04-01 03:34:58 +00:00
__all__ = ["ShardConfig", "ModelSharder", "ShardFormer", "PipelineGradientCheckpointConfig", "GradientCheckpointConfig"]