LnL-AI

/

dbrx-base-converted-v2

like 6

Follow

LnL AI 11

what happened if flash attention enable?

@fahadh4ilyas

Flash Attention 2 barfs when applying padding

@fahadh4ilyas the padding issue may be caused by my custom training dataset code. I am going to double check.

@fahadh4ilyas confirmed. there is nothing wrong with fa2. It was my custom train code that was breaking it. Remove note about fa2 compat.

@Qubitium what is your resource specification to use for training? Does 8×A100 enough?

@fahadh4ilyas I am currently testing train on it with

1. fa2 enabled
2. trl/sft_trainer
3. batch 1
4. max seq len 2048
5. adam 8bit optimizer
6. bfloat16

I am using 767.44GB of vram right now. @winglian bf16 test shows he is using only 8x80 640GB so not sure what magic he is doing to use so much lower vram. Though I am using trl and he is using axotol.

I have a slight hunch that databricks purposely made the model just big enough to be out of normal 8xA100/H100 range.

Galore may be an option to get it back to normal 8xA100 range but our early tests show galore a lot slower so there may be a trade-off, or we are not testing galore correctly. I haven't personally validated the galore tests.

@Qubitium what kind of parallelization engine that you use for using multiple gpu? Deepspeed? or else? And did you do full finetune or lora?

@fahadh4ilyas Zero parallelization at the moment. Just dumb accelerate/trl integration where model layers spread across multiple gpus but only 1 gpu is particpating in training at any given moment so extremely inefficient. This is our first attempt to train on something that requires more than 1 or 2 gpu for full-finetuning so have not tested out deepspeed yet (it should help).

We only do full-finetuning and not lora/qlora at the moment.

@fahadh4ilyas With optimizer set to paged_adam_8bit memory usage went down to ~670GB of our setup. However, we still reverted back to adam_8bit as the paged_adam_8bit memory pattern was triggering a CUDA/Nvidia issue where UVM process (started by nvidia) which controls gpu/cpu memory sharing is started. This caused training speed to slow down 3x. UVM appears to be unified memory sharing for gpu/cpu that is designed reduce OOM. Not sure how to disable this on linux and paged_adam_8bit triggers this 100% in our setup.