[fix] fix test_lora

2024-11-15 10:26:30 +00:00 · 2024-11-15 10:26:30 +00:00 · 2980da559f
parent 0fb500c7d4
commit 2980da559f
1 changed files with 2 additions and 1 deletions
--- a/colossalai/shardformer/policies/llama.py
+++ b/colossalai/shardformer/policies/llama.py
@ -506,8 +506,9 @@ class LlamaForSequenceClassificationPolicy(LlamaPolicy):
                )
            }
            policy.update(new_item)
+        # TODO: test lora bug here
        # enable tp, replace layer to LinearWithGradAccum
-        else:
+        elif use_zbv:
            # add a new item for sequence classification
            new_item = {
                LlamaForSequenceClassification: ModulePolicyDescription(