fix

2024-08-21 03:51:25 +00:00 · 2024-08-21 03:51:25 +00:00 · 6aface9316
parent 193030f696
commit 6aface9316
1 changed files with 6 additions and 1 deletions
--- a/colossalai/shardformer/policies/llama.py
+++ b/colossalai/shardformer/policies/llama.py
@ -391,7 +391,12 @@ class LlamaForSequenceClassificationPolicy(LlamaPolicy):
                LlamaForSequenceClassification: ModulePolicyDescription(
                    sub_module_replacement=[
                        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,
+                            ),
                        )
                    ]
                )