github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[format] applied code formatting on changed files in pull request 4441 (#4445) 

Co-authored-by: github-actions <github-actions@github.com>
pull/4426/merge
github-actions[bot] 2023-08-16 10:47:23 +08:00 committed by GitHub
parent 5d4efdf58f
commit d20dceb9a3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 41 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

80
colossalai/shardformer/policies/vit.py
View File

@ -40,7 +40,7 @@ class ViTPolicy(Policy):
suffix="dropout", suffix="dropout",
target_module=DropoutForReplicatedInput, target_module=DropoutForReplicatedInput,
) )
]) ])
policy[ViTLayer] = ModulePolicyDescription(attribute_replacement={ policy[ViTLayer] = ModulePolicyDescription(attribute_replacement={
"attention.attention.num_attention_heads": "attention.attention.num_attention_heads":
@ -48,45 +48,45 @@ class ViTPolicy(Policy):
"attention.attention.all_head_size": "attention.attention.all_head_size":
self.model.config.hidden_size // self.shard_config.tensor_parallel_size, self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
}, },
param_replacement=[], param_replacement=[],
sub_module_replacement=[ sub_module_replacement=[
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="attention.attention.query", suffix="attention.attention.query",
target_module=col_nn.Linear1D_Col, target_module=col_nn.Linear1D_Col,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="attention.attention.key", suffix="attention.attention.key",
target_module=col_nn.Linear1D_Col, target_module=col_nn.Linear1D_Col,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="attention.attention.value", suffix="attention.attention.value",
target_module=col_nn.Linear1D_Col, target_module=col_nn.Linear1D_Col,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="attention.attention.dropout", suffix="attention.attention.dropout",
target_module=col_nn.DropoutForParallelInput, target_module=col_nn.DropoutForParallelInput,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="attention.output.dense", suffix="attention.output.dense",
target_module=col_nn.Linear1D_Row, target_module=col_nn.Linear1D_Row,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="attention.output.dropout", suffix="attention.output.dropout",
target_module=col_nn.DropoutForReplicatedInput, target_module=col_nn.DropoutForReplicatedInput,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="intermediate.dense", suffix="intermediate.dense",
target_module=col_nn.Linear1D_Col, target_module=col_nn.Linear1D_Col,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="output.dense", suffix="output.dense",
target_module=col_nn.Linear1D_Row, target_module=col_nn.Linear1D_Row,
), ),
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="output.dropout", suffix="output.dropout",
target_module=col_nn.DropoutForReplicatedInput, target_module=col_nn.DropoutForReplicatedInput,
), ),
]) ])
# use flash attention # use flash attention
if self.shard_config.enable_flash_attention: if self.shard_config.enable_flash_attention:

2
tests/test_pipeline/test_stage_manager.py
View File

@ -21,7 +21,7 @@ def check_stage_manager():
1: [0, 1], 1: [0, 1],
2: [2, 3], 2: [2, 3],
3: [2, 3], 3: [2, 3],
} }
pg_mesh = ProcessGroupMesh(DP_SIZE, PP_SIZE) pg_mesh = ProcessGroupMesh(DP_SIZE, PP_SIZE)
stage_manager = PipelineStageManager(pg_mesh, PP_DIM) stage_manager = PipelineStageManager(pg_mesh, PP_DIM)
rank = dist.get_rank() rank = dist.get_rank()