Fine-Tuning an Open Source LLM in Amazon SageMaker with W&B

# the source where the data is being stored
RAW_TRAIN_DATASET_ARTIFACT = 'capecape/wandbot/run-m6nz6yrl-wandbot_questions:v0'

# we create a run
wandb.init(project=WANDB_PROJECT, job_type="text_formatting")

dataset_artifact = wandb.use_artifact(RAW_TRAIN_DATASET_ARTIFACT) # <-- this way we get tracebility
table = dataset_artifact.get("wandbot_questions")

df = pd.DataFrame(table.data, columns=table.columns)
df = df.dropna()
df = df.assign(context_len = lambda df: df.page_content.str.len()/3.6). #<-- average 3.6 chars/token
df.head()

B_INST, E_INST = "[INST] ", " [/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
EOS = "</s>"

prompt_format = (
    B_INST
    + B_SYS
    + "You are an AI assistant designed to assist developers with everyday tasks related to Weight & Biases "
    + "and provide helpful information. As an expert in the open-source python SDK wandb answer the following "
    + "question based on the context below. Answer in formatted Markdown.\n"
    + "{page_content}"
    + E_SYS
    + "{question}"
    + E_INST
    + "\n[W&B]\n"
    + "{answer}"
    + "\n[/W&B]"
    + EOS
)

def format_text(row): return prompt_format.format_map(row)

df["text"] = df.apply(format_text, axis=1)

# and save a local copy as jsonl
df.to_json("wandb_questions_ds.jsonl", orient='records', lines=True)

table = wandb.Table(dataframe=df)
wandb.log({"wandb_questions_ds": table})

# let's also save a the dataset at this stage
at = wandb.Artifact(
    name="wandb_questions_ds", 
    type="dataset",
    description="A wandbot dataset of questions and answers about W&B for training (non tokenized)",
    metadata={"prompt_format": prompt_format,
              "length": len(df),
             }
)
at.add_file("wandb_questions_ds.jsonl")
wandb.log_artifact(at)

from transformers import AutoTokenizer
from datasets import load_dataset
MODEL_NAME = "codellama/CodeLlama-7b-Instruct-hf"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=True)
tokenizer.pad_token = tokenizer.eos_token

wandb.init(project=WANDB_PROJECT, job_type="tokenizing")
artifact = wandb.use_artifact('capecape/aws_llm_workshop/wandb_questions_ds:v0', type='dataset')
artifact_dir = artifact.download()

train_dataset = load_dataset(
    path=".", 
    data_files=f"{artifact_dir}/wandb_questions_ds.jsonl", 
    split="train")

train_dataset = train_dataset.select_columns(["text"])

lm_dataset = ... # <--- the tokenized and packed dataset

# save train_dataset to s3
training_input_path = f's3://{sess.default_bucket()}/processed/wandbot/train'

lm_dataset.save_to_disk(training_input_path)
print(f"training dataset to: {training_input_path}")

at = wandb.Artifact(
    name="wandbot_dataset_tokenized", 
    type="dataset",
    description="A wandbot dataset of questions and answers about W&B - CodeLLama tokenized",
    metadata={"model_name": MODEL_NAME, "tokenizer": MODEL_NAME},
)
at.add_reference(training_input_path) # <--- we don't copy anything, just store a reference to the bucket
wandb.log_artifact(at)

from sagemaker.huggingface import HuggingFace
from huggingface_hub import HfFolder

# define Training Job Name 
MODEL_NAME = "codellama/CodeLlama-7b-Instruct-hf"
job_name = 'wandb-qlora-codellama7'

lr = 2e-4

# hyperparameters, which are passed into the training job
hyperparameters = {
    'model_id': MODEL_NAME,                           # pre-trained model
    # 'dataset_artifact': AT_ADDRESS,                   # Artifact containing the dataset at W&B
    # 'dataset_path': '/opt/ml/input/data/training',    # path where sagemaker will save training dataset
    'dataset_path': AT_ADDRESS,
    'epochs': 1,                                      # number of training epochs
    'per_device_train_batch_size': 2,                 # batch size for training
    'lr': lr,                                         # learning rate used during training
    'hf_token': HfFolder.get_token(),                 # huggingface token to access llama 2
    'merge_weights': True,                            # wether to merge LoRA into the model (needs more memory)
    'report_to': "wandb",                              # report to wandb
    'wandb_project': WANDB_PROJECT,
    "run_name":  f"{MODEL_NAME}__qlora",
}

# create the Estimator
huggingface_estimator = HuggingFace(
    entry_point          = 'run_clm.py',      # train script
    source_dir           = 'scripts',         # directory which includes all the files needed for training
    instance_type        = 'ml.g5.4xlarge',   # instances type used for the training job
    instance_count       = 1,                 # the number of instances used for training
    base_job_name        = job_name,          # the name of the training job
    role                 = role,              # Iam role used in training job to access AWS ressources, e.g. S3
    volume_size          = 300,               # the size of the EBS volume in GB
    transformers_version = '4.28',            # the transformers version used in the training job
    pytorch_version      = '2.0',             # the pytorch_version version used in the training job
    py_version           = 'py310',           # the python version used in the training job
    hyperparameters      =  hyperparameters,  # the hyperparameters passed to the training job
    environment          = { "HUGGINGFACE_HUB_CACHE": "/tmp/.cache"}, # set env variable to cache models in /tmp
)

from sagemaker.huggingface import get_huggingface_llm_image_uri

# retrieve the llm image uri
llm_image = get_huggingface_llm_image_uri(
  "huggingface",
  version="1.1.0",
  session=sess,
)

model_s3_path = "s3://sagemaker-us-east-1-372108735839/wandb-qlora-codellama7-2023-10-26-00-01-56-374/output/model/"

import json
from sagemaker.huggingface import HuggingFaceModel

instance_type = "ml.g5.2xlarge". # the same instance type we used for training
number_of_gpu = 1
health_check_timeout = 300

# Define Model and Endpoint configuration parameter
config = {
  'HF_MODEL_ID': "/opt/ml/model", # path to where sagemaker stores the model
  'SM_NUM_GPUS': json.dumps(number_of_gpu), # Number of GPU used per replica
  'MAX_INPUT_LENGTH': json.dumps(2048), # Max length of input text
  'MAX_TOTAL_TOKENS': json.dumps(3000), # Max length of the generation (including input text)
}

# create HuggingFaceModel with the image uri
llm_model = HuggingFaceModel(
  role=role,
  image_uri=llm_image,
  model_data={'S3DataSource':{'S3Uri': model_s3_path,'S3DataType': 'S3Prefix','CompressionType': 'None'}},
  env=config
)

llm = llm_model.deploy(
  initial_instance_count=1,
  instance_type=instance_type,
  container_startup_health_check_timeout=health_check_timeout, # 10 minutes to be able to load the model
)

from time import perf_counter

# hyperparameters for llm
parameters = {
    "do_sample": True,
    "top_p": 0.9,
    "temperature": 0.8,
    "max_new_tokens": 952,
    "repetition_penalty": 1.03,
    "stop": ["[/W&B]", "</s>"],
}

def call_endpoint(formatted_input, parameters):
    "Call the SM endpoint and parse the output in string format"
    t0 = perf_counter()
    payload = {
      "inputs": formatted_input,
      "parameters": parameters,
    }
    response = llm.predict(payload)
    total_time = perf_counter() - t0

    return response[0]["generated_text"], total_time

params_cols = list(parameters.keys())
table = wandb.Table(columns=["question", "original_answer", "generated_answer", "time(s)"] + params_cols) 

for s in tqdm(eval_ds):
    generated_answer, req_time = call_endpoint(s["text"], parameters)
    table.add_data(s["question"], s["answer"], generated_answer, req_time, *list(parameters.values()))

wandb.log({"evaluation_answers": table})
wandb.finish()

llm.delete_model()
llm.delete_endpoint()