ColossalAI/colossalai/shardformer/modeling
Wenhao Chen e614aa34f3
[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 11:34:58 +08:00
..
chatglm2_6b [hotfix] fix typo change enabel to enable under colossalai/shardformer/ (#5317) 2024-03-05 21:48:46 +08:00
__init__.py [shardformer] added development protocol for standardization (#4149) 2023-07-04 16:05:01 +08:00
bert.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
blip2.py [shardformer] update colo attention to support custom mask (#5510) 2024-03-27 11:19:32 +08:00
bloom.py [hotfix] fix typo change enabel to enable under colossalai/shardformer/ (#5317) 2024-03-05 21:48:46 +08:00
chatglm2.py [shardformer] update colo attention to support custom mask (#5510) 2024-03-27 11:19:32 +08:00
falcon.py [shardformer]: support gpt-j, falcon, Mistral and add interleaved pipeline for bert (#5088) 2023-11-28 16:54:42 +08:00
gpt2.py [format] applied code formatting on changed files in pull request 5510 (#5517) 2024-03-27 11:21:03 +08:00
gptj.py [shardformer] update colo attention to support custom mask (#5510) 2024-03-27 11:19:32 +08:00
jit.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
llama.py [shardformer, pipeline] add `gradient_checkpointing_ratio` and heterogenous shard policy for llama (#5508) 2024-04-01 11:34:58 +08:00
mistral.py [feat] refactored extension module (#5298) 2024-01-25 17:01:48 +08:00
opt.py [shardformer] update colo attention to support custom mask (#5510) 2024-03-27 11:19:32 +08:00
sam.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
t5.py [hotfix] fix typo change enabel to enable under colossalai/shardformer/ (#5317) 2024-03-05 21:48:46 +08:00
vit.py [shardformer] update colo attention to support custom mask (#5510) 2024-03-27 11:19:32 +08:00
whisper.py [shardformer] update colo attention to support custom mask (#5510) 2024-03-27 11:19:32 +08:00