LocalLLaMA

3 readers

1 users here now

Community to discuss about Llama, the family of large language models created by Meta AI.

founded 1 year ago

MODERATORS

communick@poweruser.forum

Guidance on training 8K context size with 4bit (alien.top)

submitted 1 year ago by Bright-Question-6485@alien.top to c/localllama@poweruser.forum

2 comments fedilink hide all child comments

Hello all,

I got two NVIDIA P40 with in total 48GB vRAM and are trying to train a LLaMAv2 base instruction following model which has a base context size of 8.192 (8K). I successfully trained and had really strong results with the default SFTTrainer using 2.048 and 4.096 so, 2k and 4k. However when I switch it 8K I always hit the OOM wall. I set all to 4bit and the initial loading memory use is less then 2-3GB per GPU but the moment he starts training it dies. Does anyone have an idea or suggestion here?

I tried double quant but it is not compatible with the P40, same as for flash attention.

For my use case I need an 8K context. So far all my previous tests with 2-4K went really good with strong results so I am quite confident in my overall training setup.

With fastapi I managed to run even 60B and 34B models for inference using 4bit and a special split GPU switch where I could limit the GPU memory usage to 18GB:22GB (don’t know why but only this worked stable). I wonder if something similar can help here.

For batch size I tried all from 1 to 64 (which I used successfully with smaller context sizes).

Thanks a lot!

top 2 comments

sorted by: hot top controversial new old

[–] mcmoose1900@alien.top 1 points 1 year ago (1 children)

Are you doing a full finetune?

Try a LoRA, or better yet LongLora which is specifically optimized for long context: https://github.com/huggingface/peft/issues/958

[–] Bright-Question-6485@alien.top 1 points 1 year ago

Hi u/mcmoose1900 thanks a lot for the reply!

By my understanding, i already make use of peft and lora since starting this endeavour.

See excerpts of the code here (there is a chance that maybe it does not get used as intended due to the often weird ways Python works).

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    load_in_8bit=False,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16
)

base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

and here

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.2,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
)

max_seq_length = MAX_SEQ_LENGTH
trainer = SFTTrainer(
    model=base_model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    formatting_func=formatting_func,
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    args=training_args,
)

and the parameters

MAX_SEQ_LENGTH = 8192
LEARNING_RATE = 2e-5
PER_DEVICE_BATCH_SIZE = 1
GRADIENT_ACCUMULATION_STEPS = 1
USE_EVAL = True
QUANT_BIT_8 = False
QUANT_BIT_4 = not QUANT_BIT_8

The numbers above are very low as i tried lowering them to mitigate the OOM issue without success. Normally they would not make sense.