Home/AI/ML Notes/GenAI and LLMs Summary/Advanced Training: Efficiency & Parallelism
⌂ Main Page
GenAI and LLMs Summary

Advanced Training: Efficiency & Parallelism

Imitation and proxy learning, continual pretraining, RLAIF, FlashAttention, gradient checkpointing, mixed precision, distributed parallelism, and ZeRO.

Imitation learning - fine-tune LLM w/outputs from more powerful LLM (~distillation) - Alpaca, Vicuna, GPT4ALL.

Proxy tuning: a) fine-tune Llama 7B, b) compute differences between weights of fine-tuned and original Llama 7B, c) add differences to Llama 70B + normalized model outputs and generate the desired response.

Continual Pre-Training - common scenarios new data arrives:
1) Regular pretraining: initialize model w/random weights and pretrain on dataset D1
2) Continued pretraining: adopt pretrained model from 1) and pretrain it on dataset D2
3) Retrain: same as 1), but train on datasets D1 + D2 – 2x more expensive than continual pretraining.

Third option commonly used, but continual pre-train saves significant compute. Concern w/2) distribution shift caused by new data may degrade performance on prev. data (catastrophic forgetting) or poor adaptation to new data. To combat: 1) Re-warm (gradually increase) and re-decay (gradually decrease) LR (re-apply typical LR schedule), 2) add small portion (5%) of original dataset D1 to new dataset (D2) to prevent catastrophic forgetting.

A diagram of a diagram Description automatically generated

RLAIF - like RLHF, but AI ranks the quality of several generated responses for each prompt (reward model used).

Online AI Feedback (OAIF) - no reward model - uses an annotator LM in the loop during training.

Achieving alignment (SFT + RLHF) w/greater efficiency and less human effort – it collects real-time preference judgments from a separate annotator LM evaluating pairs of responses to determine which responses better align with goals. No reward model. Easy customization through flexible annotator prompting. Used when loss of precision from humans tolerated. Risk of bias or integrity issues => need human oversight.

Training efficiencies

Flash Attention: complexity of self-attention O(N^2), f.a. breaks Q, K, V matrices into small blocks that fit into fast GPU memory - SRAM, then it optimizes Softmax computations, drastically reducing read / writes to GPU’s slow memory HBM (High Bandwidth Memory).

SFTTrainer(…, gradient_checkpointing=True) - saves memory, but slower backward pass. G.ch. selectively stores only a subset of forward pass’s intermediate activations for the backward pass (to compute gradients) by discarding most of them. If gradients are needed for an activation that was not stored, the algo recomputes the forward pass for that part of the network - significantly reduces memory usage, allows to train large models or use larger batch sizes within the same hardware constraints. Why do we need this: despite PEFT, the huge model can still pose challenges due to the high memory requirements for storing activations during training and this method allows to fine-tune models on available hardware by mitigating memory bottlenecks.

Mixed Precision Training: model maintains a "Master Copy" of weights in full precision (FP32), forward and backward passes (activations and gradients) are conducted in half precision (FP16), and final weight update is back to FP32 master copy to preserve numerical stability.

Distributed parallelism - when 1 GPU not enough - split along 1 of 4 axes: data splits the batch, tensor splits the matrix, pipeline splits the stack, context splits the sequence. Other methods - optimizations on top:

  1. Data Parallelism (DP) - if batch too large for a GPU, split batch across GPUs, gradients are synchronized in the end (synchr. or asynchr.). Simple, scales well, but each GPU holds full model copy (memory-heavy),