Simple

config = Config(
  model_name_or_path="openchat/openchat_3.5",
  apply_lora=True,
)

Advanced

config = Config(
  model_name_or_path="openchat/openchat_3.5",
  apply_lora=True,
  lora_rank=8,
  lora_alpha=32,
  lora_dropout=0.05,
  raw_lora_target_modules="all",
  # Names of modules to apply LoRA. A comma-separated string, for example: "k,q,v" or "all".
)

To train the QLoRA model, we need to load the backbone model using bitsandbytes library and int4 (or int8) weights.

Simple

config = Config(
  model_name_or_path="01-ai/Yi-34B",
  apply_lora=True,
  load_in_4bit=True,
  prepare_model_for_kbit_training=True,
)

Advanced

config = Config(
  model_name_or_path="01-ai/Yi-34B",
  apply_lora=True,
  load_in_4bit=True,
  prepare_model_for_kbit_training=True,
  llm_int8_threshold=6.0,
  llm_int8_has_fp16_weight=True,
  bnb_4bit_use_double_quant=True,
  bnb_4bit_quant_type="nf4",
)

Before that, you must log in to Huggingface Hub or add an API Token to the environment variables.

config = Config(
  model_name_or_path="HuggingFaceH4/zephyr-7b-beta",
  push_to_hub=True,
  hub_private_repo=True,
  hub_model_id="BobaZooba/AntModel-7B-XLLM-Demo-LoRA",
  save_steps=25,
)

• Checkpoints will be saved locally and in Huggingface Hub each save_steps
• If you train a model with LoRA, then only LoRA weights will be saved

Before that, you must log in to W&B or add an API Token to the environment variables.

config = Config(
  model_name_or_path="HuggingFaceH4/zephyr-7b-beta",
  report_to_wandb=True,
  logging_steps=5,
  wandb_project="xllm-demo",
)

This will help to use less GPU memory during training, that is, you will be able to learn more than without this technique. The disadvantages of this technique is slowing down the forward step, that is, slowing down training.

You will be training larger models (for example 7B in colab), but at the expense of training speed.

config = Config(
  model_name_or_path="HuggingFaceH4/zephyr-7b-beta",
  use_gradient_checkpointing=True,
)

This speeds up training and GPU memory consumption, but it does not work with all models and GPUs. You also need to install flash-attn for this. This can be done using:

pip install "xllm[train]"

config = Config(
  model_name_or_path="meta-llama/Llama-2-7b-hf",
  use_flash_attention_2=True,
)

Recommendations

• Incredibly effective method is LoRA (apply_lora). It allows for a tremendous reduction in training costs and, moreover, helps very effectively combat catastrophic forgetting.
• Then, I advise using load_in_4bit and prepare_model_for_kbit_training together. This also significantly reduces memory consumption.
• Lastly, I would recommend apply use_gradient_checkpointing. This method also greatly reduces memory consumption, but at the expense of slowing down training.
• If your project allows, I suggest enabling push_to_hub and hub_private_repo, also specifying the model name in hub_model_id and save_steps. Example: "BobaZooba/SupaDupaLlama-7B-LoRA". During training, every checkpoint of your model will be saved in the HuggingFace Hub. If you specified apply_lora, then only the LoRA weights will be saved, which you can later easily fuse with the main model, for example, using xllm.
• If your GPU allows it add use_flash_attention_2
• I also recommend using report_to_wandb, also specifying wandb_project (the project name in W&B) and wandb_entity (user or organization name in W&B).
• Note that for push_to_hub, you need to log in to the HuggingFace Hub beforehand or specify the token (HUGGING_FACE_HUB_TOKEN) in the .env file. Similarly, when using report_to_wandb, you will need to log in to W&B. You can either specify the token (WANDB_API_KEY) in the .env file or you will be prompted to enter the token on the command line.

Features

• QLoRA

• Gradient checkpointing

• Flash Attention 2

• Stabilize training

• Push checkpoints to HuggingFace Hub

• W&B report

  config = Config(
    model_name_or_path="meta-llama/Llama-2-7b-hf",
    tokenizer_padding_side="right",  # good for llama2
  
    warmup_steps=1000,
    max_steps=10000,
    logging_steps=1,
    save_steps=1000,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=2,
    max_length=2048,
  
    stabilize=True,
    use_flash_attention_2=True,
  
    apply_lora=True,
  
    load_in_4bit=True,
    prepare_model_for_kbit_training=True,
  
    use_gradient_checkpointing=True,
  
    push_to_hub=False,
    hub_private_repo=True,
    hub_model_id="BobaZooba/SupaDupaLlama-7B-LoRA",
  
    report_to_wandb=False,
    wandb_project="xllm-demo",
    wandb_entity="bobazooba",
  )
  

This operation is only for models with a LoRA adapter.

You can explicitly specify to fuse the model after training.

config = Config(
  model_name_or_path="HuggingFaceH4/zephyr-7b-beta",
  apply_lora=True,
  fuse_after_training=True,
)

Even when you are using QLoRa

config = Config(
  model_name_or_path="HuggingFaceH4/zephyr-7b-beta",
  apply_lora=True,
  load_in_4bit=True,
  prepare_model_for_kbit_training=True,
  fuse_after_training=True,
)

Or you can fuse the model yourself after training.

experiment.fuse_lora()

DeepSpeed is needed for training models on multiple GPUs. DeepSpeed allows you to efficiently manage the resources of several GPUs during training. For example, you can distribute the gradients and the state of the optimizer to several GPUs, rather than storing a complete set of gradients and the state of the optimizer on each GPU. Starting training using DeepSpeed can only happen from the command line.

train.py

from xllm import Config
from xllm.datasets import GeneralDataset
from xllm.cli import cli_run_train

if __name__ == '__main__':
  train_data = ["Hello!"] * 100
  train_dataset = GeneralDataset.from_list(data=train_data)
  cli_run_train(config_cls=Config, train_dataset=train_dataset)

Run train (in the num_gpus parameter, specify as many GPUs as you have)

deepspeed --num_gpus=8 train.py --deepspeed_stage 2

You also can pass other parameters

deepspeed --num_gpus=8 train.py \
  --deepspeed_stage 2 \
  --apply_lora True \
  --stabilize True \
  --use_gradient_checkpointing True

1. Downloading and preparing data and model

python3 MY_PROJECT/cli/prepare.py \
  --dataset_key MY_DATASET \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --path_to_env_file ./.env

2. Run train using DeepSpeed on multiple GPUs

deepspeed --num_gpus=8 MY_PROJECT/cli/train.py \
  --use_gradient_checkpointing True \
  --deepspeed_stage 2 \
  --stabilize True \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --use_flash_attention_2 False \
  --load_in_4bit True \
  --apply_lora True \
  --raw_lora_target_modules all \
  --per_device_train_batch_size 8 \
  --warmup_steps 1000 \
  --save_total_limit 0 \
  --push_to_hub True \
  --hub_model_id MY_HF_HUB_NAME/LORA_MODEL_NAME \
  --hub_private_repo True \
  --report_to_wandb True \
  --path_to_env_file ./.env

3. Fuse LoRA

python3 MY_PROJECT/cli/fuse.py \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --lora_hub_model_id MY_HF_HUB_NAME/LORA_MODEL_NAME \
  --hub_model_id MY_HF_HUB_NAME/MODEL_NAME \
  --hub_private_repo True \
  --force_fp16 True \
  --fused_model_local_path ./fused_model/ \
  --path_to_env_file ./.env

4. [Optional] GPTQ quantization of the trained model with fused LoRA

python3 MY_PROJECT/cli/quantize.py \
  --model_name_or_path ./fused_model/ \
  --apply_lora False \
  --stabilize False \
  --quantization_max_samples 128 \
  --quantized_model_path ./quantized_model/ \
  --prepare_model_for_kbit_training False \
  --quantized_hub_model_id MY_HF_HUB_NAME/MODEL_NAME_GPTQ \
  --quantized_hub_private_repo True \
  --path_to_env_file ./.env

from xllm import Config
from xllm.datasets import GeneralDataset
from xllm.cli import cli_run_train

if __name__ == '__main__':
  train_data = ["Hello!"] * 100
  train_dataset = GeneralDataset.from_list(data=train_data)
  cli_run_train(config_cls=Config, train_dataset=train_dataset)