1
0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-09 17:08:23 +00:00
ColossalAI/colossalai
Zhongkai Zhao 8e412a548e
[shardformer] Sequence Parallelism Optimization ()
* sequence parallel optimization

* validate sequence parallel in llama (code to be polished)

* shardformer api writing

* integrate sequence parallel in ShardFormer

* fix pp bugs and sp bugs for LlaMa model

* integrating ring-based sequence parallelism into ShardFormer

* [sequence parallelism]: Add fused megatron function

* integrating ring-based sequence parallelism into ShardFormer

---------

Co-authored-by: linsj20 <linsj20@mails.tsinghua.edu.cn>

* fix bugs when useing sp and flashattention together

* fix operation function name

* support flash attention for ulysses-style sp

* clarify sp process group

* fix compatibility bugs in moe plugin

* fix fused linear bugs

* fix linear layer test

* support gpt model all-to-all sp

* modify shard data dimension (meant to be dim=-1)

* support megtron-style sp and distributed attn for llama model

* [shardformer] add megatron sp to llama

* support llama7B 128k with distributed attention

* [shardformer] robustness enhancement

* add block attn

* sp mode 1: keep input as a complete sequence

* fix sp compatability

* finish sp mode 3 support for gpt

* using all_to_all_single when batch size is 1

* support mode 2 sp in gpt2 ()

* [shardformer] add megatron sp to llama

* support llama7B 128k with distributed attention

* [shardformer] robustness enhancement

* add block attn

* sp mode 1: keep input as a complete sequence

* fix sp compatability

* refactor ring implementation

* support mode 2 sp in gpt2

* polish code

* enable distributed attn mask when using sp mode 2 and 3 in llama

* automatically enable flash attn when using sp mode 2 and 3 in llama

* inplace attn mask

* add zero2 support for sequence parallel

* polish code

* fix bugs

* fix gemini checkpoint io

* loose tensor checking atol and rtol

* add comment

* fix llama layernorm grad

* fix zero grad

* fix zero grad

* fix conflict

* update split and gather auto grad func

* sequence parallel: inside text split ()

* polish code (part 1)

* polish code (part 2)

* polish code (part 2.5)

* polish code (part 3)

* sequence parallel: inside text split

* miscellaneous minor fixes

* polish code

* fix ulysses style ZeRO

* sequence parallel: inside text split

* miscellaneous minor fixes

* disaggregate sp group and dp group for  sp

* fix llama and gpt sp

* polish code

* move ulysses grad sync to ddp ()

* remove zero_stage and unbind the grad sync for alltoall sp

* add 2d group creation test

* move ulysses grad sync to ddp

* add 2d group creation test

* remove useless code

* change shard config not to enable sp when enable_all_optimizations

* add sp warnings for several model

* remove useless code

---------

Co-authored-by: linsj20 <linsj20@mails.tsinghua.edu.cn>
2024-04-03 17:15:47 +08:00
..
_analyzer [misc] update pre-commit and run all files () 2023-09-19 14:20:26 +08:00
_C [setup] support pre-build and jit-build of cuda kernels () 2023-01-06 20:50:26 +08:00
accelerator [hotfix] fix typo change MoECheckpintIO to MoECheckpointIO () 2024-03-05 21:52:30 +08:00
amp [npu] change device to accelerator api () 2024-01-09 10:20:05 +08:00
auto_parallel [hotfix] Fix wrong import in meta_registry () 2024-02-20 19:24:43 +08:00
autochunk [misc] update pre-commit and run all files () 2023-09-19 14:20:26 +08:00
booster [shardformer] Sequence Parallelism Optimization () 2024-04-03 17:15:47 +08:00
checkpoint_io [hotfix] fix typo change MoECheckpintIO to MoECheckpointIO () 2024-03-05 21:52:30 +08:00
cli [devops] fix extention building () 2024-03-05 15:35:54 +08:00
cluster [shardformer] Sequence Parallelism Optimization () 2024-04-03 17:15:47 +08:00
context [moe] merge moe into main () 2023-11-02 02:21:24 +00:00
device [npu] add npu support for hybrid plugin and llama () 2023-11-22 19:23:21 +08:00
fx [misc] update pre-commit and run all files () 2023-09-19 14:20:26 +08:00
inference [shardformer, pipeline] add gradient_checkpointing_ratio and heterogenous shard policy for llama () 2024-04-01 11:34:58 +08:00
interface [lazy] support from_pretrained () 2023-09-26 11:04:11 +08:00
kernel [shardformer] update colo attention to support custom mask () 2024-03-27 11:19:32 +08:00
lazy [doc] add lazy init docs () 2023-09-27 10:24:04 +08:00
legacy Fix ColoTensorSpec for py11 () 2024-03-26 15:56:49 +08:00
logging [misc] update pre-commit and run all files () 2023-09-19 14:20:26 +08:00
moe [hotfix] fix typo change MoECheckpintIO to MoECheckpointIO () 2024-03-05 21:52:30 +08:00
nn [shardformer] update colo attention to support custom mask () 2024-03-27 11:19:32 +08:00
pipeline [shardformer, pipeline] add gradient_checkpointing_ratio and heterogenous shard policy for llama () 2024-04-01 11:34:58 +08:00
shardformer [shardformer] Sequence Parallelism Optimization () 2024-04-03 17:15:47 +08:00
tensor fixed layout converter caching and updated tester 2024-03-26 17:22:27 +08:00
testing [shardformer] update colo attention to support custom mask () 2024-03-27 11:19:32 +08:00
utils Merge pull request from hpcaitech/feature/npu 2024-01-29 13:49:39 +08:00
zero [shardformer] Sequence Parallelism Optimization () 2024-04-03 17:15:47 +08:00
__init__.py [accelerator] init the accelerator module () 2023-11-30 13:25:17 +08:00
initialize.py [npu] change device to accelerator api () 2024-01-09 10:20:05 +08:00