github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-26 20:23:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer/fix overlap bug] fix overlap bug, add overlap as an option in shardco… (#4516)

Browse Source 
* fix overlap bug and support bert, add overlap as an option in shardconfig

* support overlap for chatglm and bloom
This commit is contained in:
Bin Jia
2023-08-28 17:16:40 +08:00
committed by GitHub
parent 376533a564
commit c554b7f559
7 changed files with 63 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/shardformer/policies/chatglm2.py
View File

@@ -50,6 +50,7 @@ class ChatGLMPolicy(Policy):
policy = {}
use_sequence_parallel = self.shard_config.enable_sequence_parallelism
overlap = self.shard_config.enable_sequence_overlap
if self.shard_config.enable_tensor_parallelism:
policy[ChatGLMModel] = ModulePolicyDescription(attribute_replacement={},
sub_module_replacement=[
@@ -81,7 +82,8 @@ class ChatGLMPolicy(Policy):
target_module=col_nn.Linear1D_Col,
kwargs={
'seq_parallel': use_sequence_parallel,
'seq_parallel_dim': 0
'seq_parallel_dim': 0,
'overlap': overlap
}),
SubModuleReplacementDescription(suffix="self_attention.dense",
target_module=col_nn.Linear1D_Row,