fine tuning the updated Phi-2 with flash-attn-2 produces very high loss > 2 #28488
Comments
|
I experienced the same thing! Over 3 epochs same set up just updated code and flash attention, the loss went from 6 to 2. And on the old code without flash attention it was .60 to ~.29 . Very strange. |
|
Hi @abacaj, as per @pacman100 guidelines in #28142 / #28142 (comment) you need to make sure to load your model in full-precision and train with autocast (bf16=True). Also can you share more insights on how you train your model? (do you load the model in bf16/fp16, do you use PEFT, packing, etc.) ? |
|
Hi @younesbelkada, this is a full fine tune using HF trainer. Padding only. Model is loaded in bf16. I try loading in "fp32" but get error: model = AutoModelForCausalLM.from_pretrained(
model_args.model_name_or_path,
trust_remote_code=True,
config=config,
attn_implementation="flash_attention_2",
torch_dtype=torch.float32,
cache_dir=training_args.cache_dir,
) |
|
Ok thanks @abacaj for getting back ! I think you get that error because the patch #28142 has not been released on pypi - can you try to build transformers from source? pip install -U git+https://github.com/huggingface/transformers.gitThat should hopefully solve it, let me know if you face into more issues! |
|
Ok so I remove the explicit |
|
@abacaj which padding side are you using for training? |
|
I use
|
|
FWIW changing padding side doesn't do anything to the loss, it's the same |
|
I see, as a sanity check, can you share your |
|
|
During my testing, I used bf16, trust remote code, no gradient ckpt, for SFT, with flshattn. The resulting model was terrible, I knew off the of (6 to 2)loss it wasn’t going to preform, but during testing it was worse than expected, very mangled answers. However when I trained the model, same arguments; just using the old phi repo code, and no flshattnt I got a great model. The loss from .6 to .29. Both were full fine tunes. Flash attention is critical for 24gb cards as without it it’s training off shared memory. I can help out more with testing when it’s done training in ~30 hours off shared mem 😭. The script I used is on #28381 . (Keep in mind the script doesn’t reflect me using bf16, however both times I trained the model I did have compute dtype set to bf16.) |
|
Hello everyone! Could you all please test using the latest revision on Regards, |
|
FWIW - the model still comes out significantly worse using FA2. If anyone wants to fine-tune this model, I recommend you use it without FA2 currently. Running code benchmarks with FA2 < 50% on heval. Without FA2 (and all other hparams are identical, including seed) > 60% heval. |
|
The first graph is a comparison between using and not using flash attention 2. It seems that the loss doesn't change much with fa2 (yellowish curve). |
|
@abacaj could you please provide a minimal snippet to reproduce your fine-tuning? We want to investigate it further more and attempt to find the root of the problem. We are doing a line-by-line comparison between the new model's code and the previous one. |
I second this, just woke up and checked training after 3 epochs with FA2 I went from .54 to ~.40, meanwhile, with no FA2 I went from .60 to .30. Both full fine tunes. I’m gonna train the fa2 check point on an additional epoch to see if it gives me the same loss as with out FA2. Or to see if it over fits. EDIT:
The loss is acting quite random, in comparison to no FA2. Where the loss consistently went down. SECOND UPDATE: I restarted the training with the same checkpoint and upped the learning rate by a order of 1, so from 2e5 to 2e6 and now the loss is more consistent, confusing why this hyper parameter differs in training when using fa2 and not using fa2. Not perfect but better. THIRD UPDATE: I tried retraining the base model with fa2 and the loss isnt going anywhere. After 1.5 epochs. Its almost as if the weights aren’t being updated at all, and if so very marginally. Just consistently staying between .5 and .4 but random at every logging step. |
|
sorry to comment on this closed issue but I still have issues with FA2
|
|
I do another two runs using FA2 and without FA2 (blue line). Testing the models out using vLLM, the model without FA2 (blue line) scores 58% on humaneval, the model with FA2 scores 49%. I basically stopped using FA2 because the model comes out so much worse and I can't pinpoint why (the runs are identical with exception of use_flash_2)
|
|
Hi @ArthurZucker, |
|
Just wanted to acknowledge I have the same issue with using Fast Attention 2 with phi-2, the training loss hardly decreases with FA2 turned on, and works pretty well with it turned off. |
|
same question... |
@gugarosa is there any update on fixing FA2 for this amazing model? |
|
There is a PR that should fix it but is hard to merge #28673 |
|
I am also seeing similar issue where loss is trending downwards but quite unstable and it seems to learn very slowly. I am running full fine-tuning of latest Phi2 model on my dataset. @ArthurZucker I just started another run after reinstalling transformers with changes from #28673 to see if it fixes the issue (still using FA2). will post loss curve in next few hours. [Incorrect] Update-1: loss curve after reinstalling transformers with changes from [#28673]. Looks like there is no change.. Update-2: Looks like my new transformer installation didn't include changes from #28673 so essentially both plot should be same. I tried reinstalling transformers again with PR changes and now training is failing: |
|
Update, it was a torch error, its training now, but the loss is the same as before, I lowered my dataset to 1k examples over 3 epochs with a lr of 2e6 and still the loss is random. Never consistently going down. |
|
How are you guys testing this? It does seem to matter when doing a full fine tune, and a lora fine tune. Using FA2 I could never get the loss to even be consistent with a full fine tune ( with SFT). Right now I am doing a DPO of phi2 with QLORA, and the loss is not only consistent, it’s consistently going down; from .69 to .27 at just a single epoch. I have not tried SFT with a lora, but maybe if we wanna use FA2 its just better to stick with using lora. |
|
hi there, now that SDPA has been merged #29108 you can use FA-2 through pytorch interface: 0- Install pytorch 2.2 - trainer.train()
+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+ trainer.train()Perhaps this will fix all instability issues with respect to FA2 ! |
|
@younesbelkada Hey! Thanks for your response!(before starting my training run) I got pytorch2.2, and I pulled the latest commits from transformers and installed from source. I’m using the DPO.py, from TRL, and I saw the commit, so I tried to pass “—attn_implementation SPDA” but it gave me a SPDA not currently supported error, I wish I still had the error up, I’ll try it out again, once my training run ends in a little less than an hour. However I had only tried and pass it as a flag, not how you are just now telling me. |
|
Hi @NickWithBotronics ! |
|
@younesbelkada sorry, I don’t think I said what I meant to say properly. I had used ‘git pull’ to get the latest commits onto where I have my transformer saved, and I installed with ‘pip uninstall transformers’ , and then ‘pip install .’ I did that after I saw the SPDA commit. So I thought I did it right, however I’ll try it out again with ‘ pip install -U git+https://github.com/huggingface/transformers’ And try adding the Python’’’
|
|
Ah yes sorry, I overead that part, can you share your full snippet? |
|
I’m currently using this script, and when I tried using spda I used the argument “—attn_implementation spda” |
|
@NickWithBotronics I think it failed because you passed |
|
CMD.txt Actually it ran after the line edit, without using the --attn flag, idk if its using it as its using the same memory as before, so its going into my shared memory; meanwhile, with FA2 it fits with dpo with my current arguments. I think i tried a variation of everything, using "--attn_implementation flash_attention_2" does work though |
|
Hi @NickWithBotronics - thanks for getting back! import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "microsoft/phi-2"
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="sdpa")
print(model)And I can confirm that I am able to succesfully load phi-2 with SDPA. Can you try that script? If it works on your end I'll check out the dpo script and see if it fails there for me 🙏 |
|
@younesbelkada I was able to test out the script you provided, and it worked. (I'm assuming it has something to do with the DPO trainer.) |
|
OK thanks for getting back , I will try that on my end and get back to you |
Does this work? I'm on main and have torch==2.2.0 |
|
Hi @abacaj PhiForCausalLM(
(model): PhiModel(
(embed_tokens): Embedding(51200, 2560)
(embed_dropout): Dropout(p=0.0, inplace=False)
(layers): ModuleList(
(0-31): 32 x PhiDecoderLayer(
(self_attn): PhiSdpaAttention(
(q_proj): Linear(in_features=2560, out_features=2560, bias=True)
(k_proj): Linear(in_features=2560, out_features=2560, bias=True)
(v_proj): Linear(in_features=2560, out_features=2560, bias=True)
(dense): Linear(in_features=2560, out_features=2560, bias=True)
(rotary_emb): PhiRotaryEmbedding()
)
(mlp): PhiMLP(
(activation_fn): NewGELUActivation()
(fc1): Linear(in_features=2560, out_features=10240, bias=True)
(fc2): Linear(in_features=10240, out_features=2560, bias=True)
)
(input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
(resid_dropout): Dropout(p=0.1, inplace=False)
)
)
(final_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)
)
(lm_head): Linear(in_features=2560, out_features=51200, bias=True)
) |
I tried again same error: I'm using: |
|
Ah I see issue it seems to throw that error when I create model with AutoConfig, it is working now will give it a try |
|
Got this error during a training step: |
I dont get this error (or any error) but the loss is completely different than with eager attention still |
|
Has there been any update? |
Same here. |
|
@Becomebright Thanks for looking into this. I'm curious, you're saying it actually works fine (no high loss) for you with SDPA + FP16? Reading through #28673, I'm not sure there's a way to get around the attention overflow issues. It seems Phi-2 might only be compatible with FP32 (short of retraining the model), and the only way to use the model for now is to run it in FP32. (Or maybe someone can try running it on FP32 with fine-tuning on FP16/BF16 LoRa to see if that works?) |
|
@hackyon I'm not certain whether FP16 is effective since I've only tried over-fitting one sample. Update: Finetuning Phi-2 using DeepSpeed Zero2, SDPA, and FP16 has encountered an overflow issue: |
Update: I just tried finetuning again on Phi2 with the same dataset and settings, once with and without FA2. Can confirm the differences now were marginally, while FA2 was a lot faster. Here are the results: FA2: TrainOutput(global_step=356, training_loss=1.24568430895216, metrics={'train_runtime': 3220.3242, 'train_samples_per_second': 3.538, 'train_steps_per_second': 0.111, 'total_flos': 1.862385933484032e+17, 'train_loss': 1.24568430895216, 'epoch': 1.0}) Without FA2: TrainOutput(global_step=356, training_loss=1.1397285970409265, metrics={'train_runtime': 4415.1268, 'train_samples_per_second': 2.581, 'train_steps_per_second': 0.081, 'total_flos': 1.862385933484032e+17, 'train_loss': 1.1397285970409265, 'epoch': 1.0}). The only thing is, I finetuned Phi-2 for a summarisation task. I also compared the results with and without FA2, and here are the Rouge differences (higher is better). Without FA2: Rouge1:66.5, RougeL: 49.3, WITH FA2: Rouge1:62.9 RougeL: 54.8. Does this mean with FA2 one has to finetune for a longer period of time to get the same results. Which doesn't really stroke with the whole idea of getting faster training with FA2. What am I missing here? |
|
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread. Please note that issues that do not follow the contributing guidelines are likely to be ignored. |
System Info
The updated code of phi-2 produces a high loss, I have tried fp16, bf16, deepspeed and fsdp the result is the same -> loss starts at 2 and keeps going higher. Setting
use_flash_attention_2=Falsefixes this or using the old phi-2 modeling file.torch==2.1.2
flash-attn==2.4.2
transformers==4.37.0.dev0
Who can help?
No response
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
Fine-tune the updated phi-2 model using transformers trainer
Expected behavior
Loss go down
The text was updated successfully, but these errors were encountered: