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

[fix] fix test_lora

pull/6114/head
duanjunwen 2024-11-15 10:26:30 +00:00
parent 0fb500c7d4
commit 2980da559f
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -506,8 +506,9 @@ class LlamaForSequenceClassificationPolicy(LlamaPolicy):
) )
} }
policy.update(new_item) policy.update(new_item)
# TODO: test lora bug here
# enable tp, replace layer to LinearWithGradAccum # enable tp, replace layer to LinearWithGradAccum
else: elif use_zbv:
# add a new item for sequence classification # add a new item for sequence classification
new_item = { new_item = {
LlamaForSequenceClassification: ModulePolicyDescription( LlamaForSequenceClassification: ModulePolicyDescription(