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

fix

pull/6023/head
wangbluo 2024-08-21 03:51:25 +00:00
parent 193030f696
commit 6aface9316
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/shardformer/policies/llama.py
View File

@ -391,7 +391,12 @@ class LlamaForSequenceClassificationPolicy(LlamaPolicy):
LlamaForSequenceClassification: ModulePolicyDescription( LlamaForSequenceClassification: ModulePolicyDescription(
sub_module_replacement=[ sub_module_replacement=[
SubModuleReplacementDescription( SubModuleReplacementDescription(
suffix="score", target_module=Linear1D_Col, kwargs=dict(gather_output=True) suffix="score",
target_module=Linear1D_Col,
kwargs=dict(
gather_output=True,
fp8_communication=self.shard_config.fp8_communication,
),
) )
] ]
) )