ml.yaml

project: lora-cutoff

config: 
  - llama-2-7b
  - checkpoint-3
  - lora

llama-2-7b:
  pretrained_model_name_or_path: meta-llama/Llama-2-7b-hf

defaults:
  ################################################################################
  # SFT parameters
  ################################################################################

  # Maximum sequence length to use
  max_seq_length: 768

  # Pack multiple short examples in the same input sequence to increase efficiency
  packing: False

  ################################################################################
  # Trainer parameters
  ################################################################################

  # Output directory where the model predictions and checkpoints will be stored
  output_dir: ./results

  # Enable fp16/bf16 training (set bf16 to True with an A100)
  fp16: False

  bf16: True

  # Maximum gradient normal (gradient clipping)
  max_grad_norm: 0.3

  # Weight decay to apply to all layers except bias/LayerNorm weights
  weight_decay: 0.001

  # Learning rate schedule
  lr_scheduler_type: cosine

  # Number of training steps (overrides num_train_epochs)
  max_steps: -1

  # Ratio of steps for a linear warmup (from 0 to learning rate)
  warmup_ratio: 0.03

  # Group sequences into batches with same length (saves memory and speeds up training considerably)
  group_by_length: True

  # Where else to log metrics (tensorboard, mlflow, wandb)
  report_to: tensorboard

full-parameter:
  ################################################################################
  # Instance parameters
  ################################################################################

  # Instances type used for the training job
  instance_type: ml.p4d.24xlarge

  # How long to allow the job to run (in seconds)
  max_run: 7200

  # Whether to use spot instances to reduce costs
  use_spot_instances: False
  # max_wait: 9800

  ################################################################################
  # Distributed training parameters
  ################################################################################

  # Whether to use distributed training
  enable_distributed: True

  # Number of GPUs to use
  num_processes: 8

  ################################################################################
  # DeepSpeed parameters
  ################################################################################

  # Whether to use DeepSpeed
  enable_deepspeed: True

  zero_stage: 3

  offload: False

  ################################################################################
  # Trainer parameters
  ################################################################################

  # Number of training epochs
  num_train_epochs: 4

  # Batch size per GPU for training
  per_device_train_batch_size: 16

  # Batch size per GPU for evaluation
  per_device_eval_batch_size: 16

  # Number of update steps to accumulate the gradients for
  gradient_accumulation_steps: 6

  # Enable gradient checkpointing
  gradient_checkpointing: True

  # Initial learning rate (AdamW optimizer)
  learning_rate: 1e-5

  # When to evaluate the model (steps or epoch)
  evaluation_strategy: epoch

  # Save checkpoint every X steps
  save_steps: 3

  # Log every X steps
  logging_steps: 1

checkpoint-3:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-3/?channel=model&input_mode=FastFile

checkpoint-6:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-6/?channel=model&input_mode=FastFile

checkpoint-9:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-9/?channel=model&input_mode=FastFile

checkpoint-12:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-12/?channel=model&input_mode=FastFile

checkpoint-15:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-15/?channel=model&input_mode=FastFile

checkpoint-18:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-18/?channel=model&input_mode=FastFile

checkpoint-21:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-21/?channel=model&input_mode=FastFile

checkpoint-24:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-24/?channel=model&input_mode=FastFile

checkpoint-27:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-27/?channel=model&input_mode=FastFile

checkpoint-30:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-30/?channel=model&input_mode=FastFile

checkpoint-33:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-33/?channel=model&input_mode=FastFile

checkpoint-36:
  pretrained_model_name_or_path: /opt/ml/input/data/model
  input:
    - s3://sagemaker-us-east-1-283314153768/lora-cutoff-2023-09-25-01-32-10-781/output/model/checkpoint-36/?channel=model&input_mode=FastFile

lora:
  ################################################################################
  # Instance parameters
  ################################################################################

  # Instances type used for the training job
  instance_type: ml.g5.2xlarge

  # Whether to use spot instances to reduce costs
  use_spot_instances: True

  max_run: 7200

  # How long to wait for a spot instance to become available (in seconds)
  max_wait: 9800

  ################################################################################
  # Trainer parameters
  ################################################################################

  # Number of training epochs
  num_train_epochs: 1

  # Batch size per GPU for training
  per_device_train_batch_size: 8

  # Batch size per GPU for evaluation
  per_device_eval_batch_size: 8

  # Number of update steps to accumulate the gradients for
  gradient_accumulation_steps: 4

  # Enable gradient checkpointing
  gradient_checkpointing: True

  # Adam optimizer to use
  optim: paged_adamw_32bit

  # Initial learning rate (AdamW optimizer)
  learning_rate: 3e-5

  # When to evaluate the model (steps or epoch)
  evaluation_strategy: steps

  eval_steps: 0.25
  
  logging_steps: 0.025

  ################################################################################
  # LoRA parameters
  ################################################################################

  # Whether to use LoRA
  enable_lora: True

  # LoRA attention dimension
  lora_r: 8

  # Alpha parameter for LoRA scaling
  lora_alpha: 16

  # Dropout probability for LoRA layers
  lora_dropout: 0.05

  target_modules: ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head"]

  ################################################################################
  # Bits and Bytes
  ################################################################################

  # Activate 4-bit precision base model loading
  load_in_4bit: True

  # Compute dtype for 4-bit base models
  bnb_4bit_compute_dtype: bfloat16

  # Quantization type (fp4 or nf4)
  bnb_4bit_quant_type: nf4

  # Activate nested quantization for 4-bit base models (double quantization)
  bnb_4bit_use_double_quant: True