This project simplifies personalized Gen-AI SaaS apps. We fine-tune pre-trained models for users, use single GPUs, and ensure real-time responsiveness. A base txt2img Stable Diffusion model from SageMaker JumpStart is used. Challenges include traffic spikes, low-latency, and cost-efficiency. We aim for efficient, user-centric AI solutions.
#!/usr/bin/env python3
import os

import aws_cdk as cdk
# Set the Kaggle API credentials as environment variables
# os.environ["KAGGLE_USERNAME"] = "your_kaggle_username"
# os.environ["KAGGLE_KEY"] = "your_kaggle_api_key"
os.environ["KAGGLE_USERNAME"] = "voantonino"
os.environ["KAGGLE_KEY"] = "0da1e9d995fd7fcf18c3d1b83fa70554"
from genai_personalized.genai_personalized_stack import GenAIPersonalizedStack


app = cdk.App()
GenAIPersonalizedStack(app, "GenAIPersonalizedStackStack",
    # If you don't specify 'env', this stack will be environment-agnostic.
    # Account/Region-dependent features and context lookups will not work,
    # but a single synthesized template can be deployed anywhere.

    # Uncomment the next line to specialize this stack for the AWS Account
    # and Region that are implied by the current CLI configuration.

    #env=cdk.Environment(account=os.getenv('CDK_DEFAULT_ACCOUNT'), region=os.getenv('CDK_DEFAULT_REGION')),

    # For more information, see https://docs.aws.amazon.com/cdk/latest/guide/environments.html
    )

app.synth()

app.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import sys
import os
import subprocess
import json
import sys
import logging
from argparse import ArgumentParser

logger = logging.getLogger(__name__)

def parse_args():
    parser = ArgumentParser(
        description=("SageMaker DeepSpeed Launch helper utility that will spawn deepspeed training scripts")
    )
    parser.add_argument(
        "--training_script",
        type=str,
        help="Path to the training program/script to be run in parallel, can be either absolute or relative",
    )

    # rest from the training program
    parsed, nargs = parser.parse_known_args()

    return parsed.training_script, nargs


def main():
    num_gpus     = int(os.environ.get("SM_NUM_GPUS", 0))
    hosts        = json.loads(os.environ.get("SM_HOSTS", "{}"))
    num_nodes    = len(hosts)
    current_host = os.environ.get("SM_CURRENT_HOST", 0)
    rank         = hosts.index(current_host)
    
    print(f"num_gpus = {num_gpus}, num_nodes = {num_nodes}, current_host = {current_host}, rank = {rank}")
    
    s5cmd_command = "chmod +x ./scripts/s5cmd"
    print(f"s5cmd_command = {s5cmd_command}")
    launch_func(s5cmd_command)
    
    train_script, args = parse_args()
    deepspeed_command = f"deepspeed --num_gpus={num_gpus} {train_script} {' '.join(args)}"
    print(f"deepspeed_command = {deepspeed_command}")
    launch_func(deepspeed_command)

def launch_func(command):
    try:
        subprocess.run(command, shell=False)
    except Exception as e:
        logger.info(e)

if __name__ == "__main__":
    main()
finetuning/deepspeed/launcher.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os
import argparse
import numpy as np
from transformers import (
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    AutoTokenizer,
    set_seed,
)
from datasets import load_from_disk
import torch
import evaluate
import nltk
from nltk.tokenize import sent_tokenize

from huggingface_hub import HfFolder
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

# nltk.download("punkt", quiet=True)

# Metric
metric = evaluate.load("rouge")

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [label.strip() for label in labels]

    # rougeLSum expects newline after each sentence
    preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
    labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]

    return preds, labels


def parse_arge():
    """Parse the arguments."""
    parser = argparse.ArgumentParser()
    
    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--model_id", type=str, default="google/flan-t5-xl", help="Model id to use for training.")
    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs to train for.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--per_device_eval_batch_size", type=int, default=8, help="Batch size to use for testing.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient Accumulation Steps")
    parser.add_argument("--generation_max_length", type=int, default=140, help="Maximum length to use for generation")
    parser.add_argument("--generation_num_beams", type=int, default=4, help="Number of beams to use for generation.")
    parser.add_argument("--learning_rate", type=float, default=3e-3, help="Learning rate to use for training.")
    parser.add_argument("--seed", type=int, default=42, help="Seed to use for training.")
    parser.add_argument("--deepspeed", type=str, default=None, help="Path to deepspeed config file.")
    parser.add_argument("--gradient_checkpointing", type=bool, default=True, help="Whether to use gradient checkpointing.")
    parser.add_argument(
        "--bf16",
        type=bool,
        default=True if torch.cuda.get_device_capability()[0] == 8 else False,
        help="Whether to use bf16.",
    )
    parser.add_argument("--save_model_s3_path", type=str, default=None, help="s3 path to save model")
    # Data, model, and output directories
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
    
    args = parser.parse_known_args()
    return args


def training_function(args):
    # set seed
    set_seed(args.seed)

    # load dataset from disk and tokenizer
    train_dataset = load_from_disk(args.training_dir)
    eval_dataset = load_from_disk(args.test_dir)
    
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)
    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(
        args.model_id,
        use_cache=False if args.gradient_checkpointing else True,  # this is needed for gradient checkpointing
        cache_dir = "/opt/ml/input/" # changed for SM to have enough storage space
    )

    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = -100
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer, model=model, label_pad_token_id=label_pad_token_id, pad_to_multiple_of=8
    )

    # Define compute metrics function
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        # Replace -100 in the labels as we can't decode them.
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
        result = {k: round(v * 100, 4) for k, v in result.items()}
        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
        result["gen_len"] = np.mean(prediction_lens)
        return result

    output_dir = '/tmp/output/checkpoints'

    training_args = Seq2SeqTrainingArguments(
        output_dir                  = output_dir,
        per_device_train_batch_size = args.per_device_train_batch_size,
        per_device_eval_batch_size  = args.per_device_eval_batch_size,
        gradient_accumulation_steps = args.gradient_accumulation_steps,
        predict_with_generate       = True,
        generation_max_length       = args.generation_max_length,
        generation_num_beams        = args.generation_num_beams,
        fp16                        = False, # Overflows with fp16
        bf16                        = args.bf16,
        learning_rate               = args.learning_rate,
        num_train_epochs            = args.epochs,
        deepspeed                   = args.deepspeed,
        gradient_checkpointing      = args.gradient_checkpointing,
        
        # logging & evaluation strategies
        logging_dir                 = f"{args.output_data_dir}/logs",
        logging_strategy            = "steps",
        logging_steps               = 500,
        evaluation_strategy         = "epoch", # Default - steps
        save_strategy               = "epoch", # Default - steps
        save_total_limit            = 2,
        load_best_model_at_end      = True,
        # metric_for_best_model      = "overall_f1", # default is loss (eval)
    )
    
    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics, #TODO: Commenting to check if this why evaluation is slow
    )

    # Start training
    trainer.train()
    
    best_ckpt_path = trainer.state.best_model_checkpoint
    print('Best Model Checkpoint:', best_ckpt_path)
    
    print("Starting: Save best model artifacts")
    save_model_dir = '/tmp/output/asset'
    tokenizer.save_pretrained(save_model_dir)
    trainer.save_model(save_model_dir)
    print("Ending: Save best model artifacts")
    
    WORLD_RANK = int(os.environ['RANK'])
    if WORLD_RANK == 0:
        output_dir = output_dir + '/'
        checkpoints_folder_path = args.save_model_s3_path + 'checkpoints/'
        os.system("./scripts/s5cmd sync --delete {0} {1}".format(output_dir, checkpoints_folder_path))
        save_model_dir = save_model_dir + '/'
        final_asset_folder_path = args.save_model_s3_path + 'asset/'
        os.system("./scripts/s5cmd sync --delete {0} {1}".format(save_model_dir, final_asset_folder_path))


if __name__ == "__main__":
    # main()
    args, _ = parse_arge()
    if int(os.environ['LOCAL_RANK']) == 0:
        nltk.download("punkt", quiet=True)
    training_function(args)

finetuning/deepspeed/scripts/qa-deepspeed-s5cmd.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os
import argparse
import nltk
from nltk.tokenize import sent_tokenize
nltk.download("punkt", quiet=True)
import numpy as np
import torch
import evaluate
from transformers import (
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    AutoTokenizer,
    set_seed,
)
from datasets import load_from_disk
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

# Metric
metric = evaluate.load("rouge")

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [label.strip() for label in labels]

    # rougeLSum expects newline after each sentence
    preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
    labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]

    return preds, labels


def parse_arge():
    """Parse the arguments."""
    parser = argparse.ArgumentParser()
    
    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--model_id", type=str, default="google/flan-t5-xl", help="Model id to use for training.")
    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs to train for.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--per_device_eval_batch_size", type=int, default=8, help="Batch size to use for testing.")
    parser.add_argument("--generation_max_length", type=int, default=140, help="Maximum length to use for generation")
    parser.add_argument("--generation_num_beams", type=int, default=4, help="Number of beams to use for generation.")
    parser.add_argument("--learning_rate", type=float, default=3e-3, help="Learning rate to use for training.")
    parser.add_argument("--seed", type=int, default=42, help="Seed to use for training.")
    parser.add_argument("--deepspeed", type=str, default=None, help="Path to deepspeed config file.")
    parser.add_argument("--gradient_checkpointing", type=bool, default=True, help="Whether to use gradient checkpointing.")
    parser.add_argument(
        "--bf16",
        type=bool,
        default=True if torch.cuda.get_device_capability()[0] == 8 else False,
        help="Whether to use bf16.",
    )
    # Data, model, and output directories
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
    
    args = parser.parse_known_args()
    return args


def training_function(args):
    # set seed
    set_seed(args.seed)

    # load dataset from disk and tokenizer
    train_dataset = load_from_disk(args.training_dir)
    eval_dataset = load_from_disk(args.test_dir)
    
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)
    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(
        args.model_id,
        use_cache=False if args.gradient_checkpointing else True,  # this is needed for gradient checkpointing
        cache_dir = "/opt/ml/input/" # changed for SM to have enough storage space
    )

    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = -100
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer, model=model, label_pad_token_id=label_pad_token_id, pad_to_multiple_of=8
    )

    # Define compute metrics function 
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        # Replace -100 in the labels as we can't decode them.
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
        result = {k: round(v * 100, 4) for k, v in result.items()}
        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
        result["gen_len"] = np.mean(prediction_lens)
        return result

    
    training_args = Seq2SeqTrainingArguments(
        output_dir                  = args.output_dir,
        per_device_train_batch_size = args.per_device_train_batch_size,
        per_device_eval_batch_size  = args.per_device_eval_batch_size,
        predict_with_generate       = True,
        generation_max_length       = args.generation_max_length,
        generation_num_beams        = args.generation_num_beams,
        fp16                        = False, # Overflows with fp16
        bf16                        = args.bf16,
        learning_rate               = args.learning_rate,
        num_train_epochs            = args.epochs,
        deepspeed                   = args.deepspeed,
        gradient_checkpointing      = args.gradient_checkpointing,
        
        # logging & evaluation strategies
        logging_dir                 = f"{args.output_data_dir}/logs",
        logging_strategy            = "steps",
        logging_steps               = 500,
        evaluation_strategy         = "epoch",
        save_strategy               = "epoch",
        save_total_limit            = 2,
        load_best_model_at_end      = True,
    )

    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model           = model,
        args            = training_args,
        train_dataset   = train_dataset,
        eval_dataset    = eval_dataset,
        data_collator   = data_collator,
        compute_metrics = compute_metrics,
    )

    # Start training
    trainer.train()
    
    # Saves the model to s3 uses os.environ["SM_MODEL_DIR"] to make sure checkpointing works
    trainer.save_model(os.environ["SM_MODEL_DIR"])
    tokenizer.save_pretrained(os.environ["SM_MODEL_DIR"])

def main():
    args, _ = parse_arge()
    training_function(args)


if __name__ == "__main__":
    main()

finetuning/deepspeed/scripts/qa-deepspeed.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os
import subprocess
import sys
import json
import argparse
import logging
import numpy as np
from datasets import load_from_disk, load_metric
from transformers import ( AutoTokenizer,
                          AutoModelForSeq2SeqLM,
                          DataCollatorForSeq2Seq,
                          Seq2SeqTrainingArguments,
                          Seq2SeqTrainer)
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType
import torch
import gc

def report_gpu():
    print(torch.cuda.list_gpu_processes())
    gc.collect()
    torch.cuda.empty_cache()

def training_function(args):
    print(f'Args: {args}')
    
    # is needed for Amazon SageMaker Training Compiler
    os.environ["GPU_NUM_DEVICES"] = args.n_gpus
    
    # Set up logging
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        level=logging.getLevelName("INFO"),
        handlers=[logging.StreamHandler(sys.stdout)],
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # load datasets
    train_dataset = load_from_disk(args.training_dir)
    logger.info(f"loaded train_dataset \n{train_dataset}\n")
    
    # huggingface hub model id
    model_id = "philschmid/flan-t5-xxl-sharded-fp16"
    
    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, load_in_8bit=True, device_map="auto")
    
    # Define LoRA Config
    lora_config = LoraConfig(
         r=16,
         lora_alpha=32,
         target_modules=["q", "v"],
         lora_dropout=0.05,
         bias="none",
         task_type=TaskType.SEQ_2_SEQ_LM
    )
    # prepare int-8 model for training
    model = prepare_model_for_int8_training(model)

    # add LoRA adaptor
    model = get_peft_model(model, lora_config)
    # logger.info(f'Printing trainable parameters \n{model.print_trainable_parameters()}\n')
    
    # define tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)
    logger.info('Tokenizer defined!')
    
    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = -100
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8
    )

    # define training args
    output_dir = '/tmp/output/checkpoints'
    training_args = Seq2SeqTrainingArguments(
        # auto_find_batch_size = True,
        output_dir           = output_dir,
        per_device_train_batch_size = args.per_device_train_batch_size,
        gradient_accumulation_steps = args.gradient_accumulation_steps,
        gradient_checkpointing = True,
        learning_rate        = args.learning_rate, # higher learning rate
        num_train_epochs     = args.epochs,
        logging_dir          = f"{args.output_data_dir}/logs",
        logging_strategy     = "steps",
        logging_steps        = 500,
        save_strategy        = "no"
    )
    
    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model         = model,
        args          = training_args,
        data_collator = data_collator,
        train_dataset = train_dataset,
    )
    model.config.use_cache = False
    
    # train model
    trainer.train()
    logger.info('Model training completed!')
    
    report_gpu()
    
    save_model_dir = '/tmp/output/asset'
    tokenizer.save_pretrained(save_model_dir)
    trainer.save_model(save_model_dir)
    print('Trainer completed saving the model')
    
    save_model_dir = save_model_dir + '/'
    final_asset_folder_path = args.save_model_s3_path + 'asset/'
    os.system("./scripts/s5cmd sync --delete {0} {1}".format(save_model_dir, final_asset_folder_path))
    print('Completed saving the model with s5cmd')
    
def launch_func(command):
    try:
        subprocess.run(command, shell=False)
    except Exception as e:
        logger.info(e)
    return

def parse_arge():
    parser = argparse.ArgumentParser()

    # # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--learning_rate", type=float, default=5e-5)
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient Accumulation Steps")
    parser.add_argument("--epochs", type=int, default=2)
    parser.add_argument("--model_id", type=str)
    parser.add_argument("--save_model_s3_path", type=str, default=None, help="s3 path to save model")

    # Data, model, and output directories
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    # parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

    args = parser.parse_known_args()
    return args # NOTE: This is a tuple

def main():
    args, _ = parse_arge()
    s5cmd_command = "chmod +x ./scripts/s5cmd"
    print(f"s5cmd_command = {s5cmd_command}")
    launch_func(s5cmd_command)
    training_function(args)
    
if __name__ == "__main__":
    main()
finetuning/peft/scripts/qa-peft-s5cmd.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os

import torch
from torch.utils.data import DataLoader
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_linear_schedule_with_warmup, DataCollatorForSeq2Seq,AutoConfig

from datasets import load_from_disk
from tqdm import tqdm
from utils import parse_args
from utils import parse_args,is_local_main_process,is_main_process,wait_for_everyone,main_process_first
import math
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp
import time
import numpy as np

from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType # New

@smp.step
def train_step(model, batch):
    loss = model(**batch)["loss"]
    #loss = outputs.loss
    model.backward(loss)
    return loss

@smp.step
def test_step(model, batch):
    loss = model(**batch)["loss"]
    return loss

def compute_num_params(model):
    num_params = 0
    seen = set()
    for p in model.parameters():
        if p not in seen:
            seen.add(p)
            if hasattr(p, "ds_shape"):
                num_params += np.prod(p.ds_shape) 
            else:
                num_params += np.prod(p.size())
    
    return num_params 

def main():
    args = parse_args()

    model_id = "philschmid/flan-t5-xxl-sharded-fp16" # New
    config = AutoConfig.from_pretrained(model_id) 
    
    with smp.model_creation(dtype=torch.bfloat16):
        model =  AutoModelForSeq2SeqLM.from_pretrained(model_id,
                                                       cache_dir=f"/tmp",
                                                       config=config,
                                                       torch_dtype=torch.bfloat16,
                                                       # load_in_8bit=True, device_map="auto" # New
                                                      )
        # New
        lora_config = LoraConfig(
             r=16,
             lora_alpha=32,
             target_modules=["q", "v"],
             lora_dropout=0.05,
             bias="none",
             task_type=TaskType.SEQ_2_SEQ_LM
        )
        model = prepare_model_for_int8_training(model)
        model = get_peft_model(model, lora_config)
        
        model.config.use_cache = False

    num_params = compute_num_params(model)

    tokenizer = AutoTokenizer.from_pretrained(args.model_id)

    train_dataset = load_from_disk(args.training_dir)
    eval_dataset = load_from_disk(args.test_dir)

    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

    train_sampler = torch.utils.data.DistributedSampler(
                        train_dataset,
                        shuffle=True,
                        seed=args.seed,
                        rank=smp.dp_rank(),
                        num_replicas=smp.dp_size(),
                        drop_last=True,
                    )
    
    eval_sampler = torch.utils.data.DistributedSampler(
                        eval_dataset,
                        shuffle=True,
                        seed=args.seed,
                        rank=smp.dp_rank(),
                        num_replicas=smp.dp_size(),
                        drop_last=True,
                    )
    
    train_dataloader = DataLoader(train_dataset,
                                  sampler=train_sampler,
                                  collate_fn=data_collator,
                                  batch_size=args.per_device_train_batch_size,
                                  pin_memory=True
    )
    eval_dataloader = DataLoader(eval_dataset, 
                                 sampler = eval_sampler,
                                 collate_fn=data_collator, 
                                 batch_size=args.per_device_eval_batch_size, 
                                 pin_memory=True
    )

    model = smp.DistributedModel(model, trace_device="gpu")

    # enable activation checkpointing.
    m = model.get_module()
    for c in m.encoder.block.children():
        smp.set_activation_checkpointing(c)
    for c in m.decoder.block.children():
        smp.set_activation_checkpointing(c)

    optimizer = torch.optim.AdamW(model.parameters(),
                                  lr=args.learning_rate
                                 )
    
    lr_scheduler = get_linear_schedule_with_warmup(optimizer=optimizer,
                                                   num_warmup_steps=0,
                                                   num_training_steps=(len(train_dataloader) * args.epochs),
    )
    optimizer = smp.DistributedOptimizer(optimizer)

    # Train!
    total_batch_size = args.per_device_train_batch_size * smp.size() * args.gradient_accumulation_steps

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    total_steps = num_update_steps_per_epoch * args.epochs
    
    if is_main_process(smp.rank()):
        print("***** Running training *****")
        print(f"  Num examples = {len(train_dataset)}")
        print(f"Number of eval examples ={len(eval_dataset)}")
        print(f"  Num Epochs = {args.epochs}")
        print(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
        print(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
        print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        print(f"  max_train_steps = {args.max_train_steps}")
        print(f"  total_steps = {total_steps}")

    progress_bar = tqdm(range(total_steps), disable=not is_local_main_process(smp.local_rank()))
    completed_steps = 0
    total_loss = 0

    for epoch in range(args.epochs):
        model.train()
        total_loss = 0
        for step, batch in enumerate(train_dataloader):
            device = torch.device("cuda")
            batch = {k: v.to(device) for k, v, in batch.items()}
            loss = train_step(model,batch)
            
            loss = loss.reduce_mean()
            total_loss += loss.detach().float()

            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

            progress_bar.update(1)
            completed_steps += 1

            # if completed_steps >= args.max_train_steps:
            #     break

        model.eval()
        eval_loss = 0
        losses = []

        for step, batch in enumerate(tqdm(eval_dataloader)):
            with torch.no_grad():
                device = torch.device("cuda")
                batch = {k: v.to(device) for k, v, in batch.items()}
                loss = test_step(model, batch).reduce_mean()
            losses.append(loss)
            
            if step > 50: # TODO: Might wanna comment this part
                break
        
        try:
            eval_loss = torch.mean(torch.tensor(losses,dtype=float))
            perplexity = math.exp(eval_loss)
        except OverflowError:
            perplexity = float("inf")

        print(f"epoch {epoch}: Eval perplexity: {perplexity} Eval loss: {eval_loss}")

        
    smp.save_checkpoint(args.checkpoint_dir,
                tag=f"flan_t5_weights.pt",
                partial=False,
                model=model,
                optimizer=optimizer)
        
    print("saving the final model")

    wait_for_everyone()

    if is_main_process(smp.rank()):
        tokenizer.save_pretrained(args.checkpoint_dir)


if __name__ == "__main__":
    main()
finetuning/peft/scripts/qa-peft-smp.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os
import subprocess
import sys
import json
import argparse
import logging
import numpy as np
from datasets import load_from_disk, load_metric
from transformers import ( AutoTokenizer,
                          AutoModelForSeq2SeqLM,
                          DataCollatorForSeq2Seq,
                          Seq2SeqTrainingArguments,
                          Seq2SeqTrainer)
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType
import torch
import gc

def report_gpu():
    print(torch.cuda.list_gpu_processes())
    gc.collect()
    torch.cuda.empty_cache()

def training_function(args):
    print(f'Args: {args}')
    
    # is needed for Amazon SageMaker Training Compiler
    os.environ["GPU_NUM_DEVICES"] = args.n_gpus
    
    # Set up logging
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        level=logging.getLevelName("INFO"),
        handlers=[logging.StreamHandler(sys.stdout)],
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # load datasets
    train_dataset = load_from_disk(args.training_dir)
    logger.info(f"loaded train_dataset \n{train_dataset}\n")
    
    # huggingface hub model id
    model_id = "philschmid/flan-t5-xxl-sharded-fp16"
    
    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, load_in_8bit=True, device_map="auto")
    
    # Define LoRA Config
    lora_config = LoraConfig(
         r=16,
         lora_alpha=32,
         target_modules=["q", "v"],
         lora_dropout=0.05,
         bias="none",
         task_type=TaskType.SEQ_2_SEQ_LM
    )
    # prepare int-8 model for training
    model = prepare_model_for_int8_training(model)

    # add LoRA adaptor
    model = get_peft_model(model, lora_config)
    logger.info(f'Printing trainable parameters \n{model.print_trainable_parameters()}\n')

    # define tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)
    logger.info('Tokenizer defined!')
    
    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = -100
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8
    )

    # define training args
    training_args = Seq2SeqTrainingArguments(
        # auto_find_batch_size = True,
        output_dir           = args.output_dir,
        per_device_train_batch_size = args.per_device_train_batch_size,
        gradient_accumulation_steps = args.gradient_accumulation_steps,
        gradient_checkpointing = True,
        learning_rate        = args.learning_rate, # higher learning rate
        num_train_epochs     = args.epochs,
        logging_dir          = f"{args.output_data_dir}/logs",
        logging_strategy     = "steps",
        logging_steps        = 500,
        save_strategy        = "no",
    )

    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model         = model,
        args          = training_args,
        data_collator = data_collator,
        train_dataset = train_dataset,
    )
    
    model.config.use_cache = False
    
    # train model
    trainer.train()
    logger.info('Model training completed!')

    report_gpu()
        
    # Saves the tokenizer, model to s3 uses os.environ["SM_MODEL_DIR"] to make sure checkpointing works
    trainer.save_model(os.environ["SM_MODEL_DIR"])
    tokenizer.save_pretrained(os.environ["SM_MODEL_DIR"])
    print('Completed saving the model')
    
def launch_func(command):
    try:
        subprocess.run(command, shell=False)
    except Exception as e:
        logger.info(e)
    return

def parse_arge():
    parser = argparse.ArgumentParser()

    # # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--learning_rate", type=float, default=5e-5)
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient Accumulation Steps")
    parser.add_argument("--epochs", type=int, default=2)
    parser.add_argument("--model_id", type=str)
    parser.add_argument("--save_model_s3_path", type=str, default=None, help="s3 path to save model")

    # Data, model, and output directories
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    # parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

    args = parser.parse_known_args()
    return args # NOTE: This is a tuple

def main():
    args, _ = parse_arge()
    training_function(args)
    
if __name__ == "__main__":
    main()
finetuning/peft/scripts/qa-peft.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import argparse
from contextlib import contextmanager
import torch
import os
import json
from transformers import (
    MODEL_MAPPING,
    SchedulerType
)

# Distributed training helper methods.
def wait_for_everyone():
    #deepspeed.comm.barrier
    torch.distributed.barrier()


def is_main_process(rank):
    if rank == 0:
        return True
    else:
        return False

def _goes_first(is_main):
    if not is_main:
        wait_for_everyone()
    yield
    if is_main:
        wait_for_everyone()

@contextmanager
def main_process_first(rank):
    """
    Lets the main process go first inside a with block.
    The other processes will enter the with block after the main process exits.
    """
    yield from _goes_first(is_main_process(rank))

def is_local_main_process(local_rank):
    return local_rank == 0

MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


# args parsing
def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a FLAN T5 model on a Seq2Seq task")

    # add model id and dataset path argument
    parser.add_argument("--model_id", type=str, default="google/flan-t5-xl", help="Model id to use for training.")
    
     # add training hyperparameters for epochs, batch size, learning rate, and seed
    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs to train for.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--per_device_eval_batch_size", type=int, default=8, help="Batch size to use for testing.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient Accumulation Steps")
    parser.add_argument("--generation_max_length", type=int, default=140, help="Maximum length to use for generation")
    parser.add_argument("--generation_num_beams", type=int, default=4, help="Number of beams to use for generation.")
    parser.add_argument("--learning_rate", type=float, default=3e-3, help="Learning rate to use for training.")
    parser.add_argument("--seed", type=int, default=42, help="Seed to use for reproducible training.")
    parser.add_argument("--deepspeed", type=str, default=None, help="Path to deepspeed config file.")
    parser.add_argument("--gradient_checkpointing", type=bool, default=True, help="Whether to use gradient checkpointing.")
    parser.add_argument(
        "--bf16",
        type=bool,
        default=True if torch.cuda.get_device_capability()[0] == 8 else False,
        help="Whether to use bf16.",
    )
    parser.add_argument("--save_model_s3_path", type=str, default=None, help="s3 path to save model")
    parser.add_argument("--optim", type=str, default='adamw_hf', help="The optimizer to use: adamw_hf, adamw_torch, adamw_torch_fused, adamw_apex_fused, adamw_anyprecision or adafactor")
    parser.add_argument(
        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument(
        "--preprocessing_num_workers",
        type=int,
        default=None,
        help="The number of processes to use for the preprocessing.",
    )
    parser.add_argument("--group_texts",default=False,help="Whether to group texts together when tokenizing")
    parser.add_argument(
        "--block_size",
        type=int,
        default=None,
        help=(
            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
            " this size for training. Default to the model max input length for single sentence inputs (take into"
            " account special tokens)."
        ),
    )
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=500,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    
    parser.add_argument("checkpoint_dir",type=str,default="/opt/ml/checkpoints")
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

    args,_ = parser.parse_known_args()

    return args
finetuning/peft/scripts/utils.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os

import torch
from torch.utils.data import DataLoader
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_linear_schedule_with_warmup, DataCollatorForSeq2Seq,AutoConfig

from datasets import load_from_disk
from tqdm import tqdm
from utils import parse_args
from utils import parse_args,is_local_main_process,is_main_process,wait_for_everyone,main_process_first
import math
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp
import time
import numpy as np

@smp.step
def train_step(model, batch):
    loss = model(**batch)["loss"]
    #loss = outputs.loss
    model.backward(loss)
    return loss

@smp.step
def test_step(model, batch):
    loss = model(**batch)["loss"]
    return loss

def compute_num_params(model):
    num_params = 0
    seen = set()
    for p in model.parameters():
        if p not in seen:
            seen.add(p)
            if hasattr(p, "ds_shape"):
                num_params += np.prod(p.ds_shape) 
            else:
                num_params += np.prod(p.size())
    
    return num_params 

def main():
    args = parse_args()

    config = AutoConfig.from_pretrained(args.model_id)
    
    with smp.model_creation(
        tensor_parallelism=smp.tp_size() > 1, # Note: To enable tensor parallelism
        dtype=torch.bfloat16
        ):
        
        # smp.delay_param_initialization(enabled=True) # TODO # Note: To enable tensor parallelism
        
        model =  AutoModelForSeq2SeqLM.from_pretrained(args.model_id,
                                                       cache_dir=f"/tmp",
                                                       config=config,
                                                       torch_dtype=torch.bfloat16)
        model.config.use_cache = False

    num_params = compute_num_params(model)

    tokenizer = AutoTokenizer.from_pretrained(args.model_id)

    train_dataset = load_from_disk(args.training_dir)
    eval_dataset = load_from_disk(args.test_dir)

    data_collator = DataCollatorForSeq2Seq(tokenizer,
                                           model=model
                                          )

    train_sampler = torch.utils.data.DistributedSampler(
                train_dataset,
                shuffle=True,
                seed=args.seed,
                rank=smp.rdp_rank(), # Note: To enable tensor parallelism
                num_replicas=smp.rdp_size(), # Note: To enable tensor parallelism
                drop_last=True,
            )
    
    eval_sampler = torch.utils.data.DistributedSampler(
                eval_dataset,
                shuffle=True,
                seed=args.seed,
                rank=smp.rdp_rank(), # Note: To enable tensor parallelism
                num_replicas=smp.rdp_size(), # Note: To enable tensor parallelism
                drop_last=True,
            )
    
    train_dataloader = DataLoader(
        train_dataset, sampler=train_sampler, collate_fn=data_collator, batch_size=args.per_device_train_batch_size, pin_memory=True
    )
    eval_dataloader = DataLoader(
        eval_dataset, sampler = eval_sampler,collate_fn=data_collator, batch_size=args.per_device_eval_batch_size, pin_memory=True
    )
    model = smp.DistributedModel(model, 
                                 trace_device="gpu",
                                 backward_passes_per_step=args.gradient_accumulation_steps # Note: To enable tensor parallelism
                                )

    # enable activation checkpointing.
    m = model.get_module()
    for c in m.encoder.block.children():
        smp.set_activation_checkpointing(c)
    for c in m.decoder.block.children():
        smp.set_activation_checkpointing(c)

    # Start: This piece of code happend to be before model = smp.DistributedModel...
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
    lr_scheduler = get_linear_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=0,
        num_training_steps=(len(train_dataloader) * args.epochs),
    )
    # End: This piece of code happend to be before model = smp.DistributedModel...
    optimizer = smp.DistributedOptimizer(optimizer)

    # Train!
    total_batch_size = args.per_device_train_batch_size *smp.size() * args.gradient_accumulation_steps
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    total_steps = num_update_steps_per_epoch * args.epochs

    if is_main_process(smp.rank()):
        print("***** Running training *****")
        print(f"  Num examples = {len(train_dataset)}")
        print(f"Number of eval examples ={len(eval_dataset)}")
        print(f"  Num Epochs = {args.epochs}")
        print(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
        print(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
        print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        print(f"  max_train_steps = {args.max_train_steps}")
        print(f"  total_steps = {total_steps}")

    progress_bar = tqdm(range(total_steps), disable=not is_local_main_process(smp.local_rank()))
    completed_steps = 0
    total_loss = 0

    for epoch in range(args.num_train_epochs):
        model.train()
        total_loss = 0
        for step, batch in enumerate(train_dataloader):
            device = torch.device("cuda")
            batch = {k: v.to(device) for k, v, in batch.items()}
            loss = train_step(model,batch)
            
            loss = loss.reduce_mean()
            total_loss += loss.detach().float()

            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            
            progress_bar.update(1)
            completed_steps += 1
            if completed_steps >= args.max_train_steps:
                break
        model.eval()
        eval_loss = 0
        losses = []
        for step, batch in enumerate(tqdm(eval_dataloader)):
            with torch.no_grad():
                device = torch.device("cuda")
                batch = {k: v.to(device) for k, v, in batch.items()}
                loss = test_step(model, batch).reduce_mean()
            losses.append(loss)
            if step > 50:
                break
        try:
            eval_loss = torch.mean(torch.tensor(losses,dtype=float))
            perplexity = math.exp(eval_loss)
        except OverflowError:
            perplexity = float("inf")

        print(f"epoch {epoch}: Eval perplexity: {perplexity} Eval loss: {eval_loss}")
    smp.save_checkpoint(args.checkpoint_dir,
                tag=f"flan_t5_weights.pt",
                partial=False,
                model=model,
                optimizer=optimizer)
    print("saving the final model")
    wait_for_everyone()
    if is_main_process(smp.rank()):
        tokenizer.save_pretrained(args.checkpoint_dir)


if __name__ == "__main__":
    main()
finetuning/smp/scripts/qa-smp-tp.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os

import torch
from torch.utils.data import DataLoader
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, get_linear_schedule_with_warmup, DataCollatorForSeq2Seq,AutoConfig

from datasets import load_from_disk
from tqdm import tqdm
from utils import parse_args
from utils import parse_args,is_local_main_process,is_main_process,wait_for_everyone,main_process_first
import math
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp
import time
import numpy as np

@smp.step
def train_step(model, batch):
    loss = model(**batch)["loss"]
    #loss = outputs.loss
    model.backward(loss)
    return loss

@smp.step
def test_step(model, batch):
    loss = model(**batch)["loss"]
    return loss

def compute_num_params(model):
    num_params = 0
    seen = set()
    for p in model.parameters():
        if p not in seen:
            seen.add(p)
            if hasattr(p, "ds_shape"):
                num_params += np.prod(p.ds_shape) 
            else:
                num_params += np.prod(p.size())
    
    return num_params 

def main():
    args = parse_args()

    config = AutoConfig.from_pretrained(args.model_id)
    
    with smp.model_creation(dtype=torch.bfloat16):
        model =  AutoModelForSeq2SeqLM.from_pretrained(args.model_id,
                                                       cache_dir=f"/tmp",
                                                       config=config,
                                                       torch_dtype=torch.bfloat16
                                                      )
        model.config.use_cache = False

    num_params = compute_num_params(model)

    tokenizer = AutoTokenizer.from_pretrained(args.model_id)

    train_dataset = load_from_disk(args.training_dir)
    eval_dataset = load_from_disk(args.test_dir)

    data_collator = DataCollatorForSeq2Seq(tokenizer,
                                           model=model
                                          )

    train_sampler = torch.utils.data.DistributedSampler(
                        train_dataset,
                        shuffle=True,
                        seed=args.seed,
                        rank=smp.dp_rank(),
                        num_replicas=smp.dp_size(),
                        drop_last=True,
                    )
    
    eval_sampler = torch.utils.data.DistributedSampler(
                        eval_dataset,
                        shuffle=True,
                        seed=args.seed,
                        rank=smp.dp_rank(),
                        num_replicas=smp.dp_size(),
                        drop_last=True,
                    )
    
    train_dataloader = DataLoader(train_dataset,
                                  sampler=train_sampler,
                                  collate_fn=data_collator,
                                  batch_size=args.per_device_train_batch_size,
                                  pin_memory=True
    )
    eval_dataloader = DataLoader(eval_dataset, 
                                 sampler = eval_sampler,
                                 collate_fn=data_collator, 
                                 batch_size=args.per_device_eval_batch_size, 
                                 pin_memory=True
    )

    optimizer = torch.optim.AdamW(model.parameters(),
                                  lr=args.learning_rate
                                 )
    
    lr_scheduler = get_linear_schedule_with_warmup(optimizer=optimizer,
                                                   num_warmup_steps=0,
                                                   num_training_steps=(len(train_dataloader) * args.epochs),
    )

    model = smp.DistributedModel(model, trace_device="gpu")

    # enable activation checkpointing.
    m = model.get_module()
    for c in m.encoder.block.children():
        smp.set_activation_checkpointing(c)
    for c in m.decoder.block.children():
        smp.set_activation_checkpointing(c)

    optimizer = smp.DistributedOptimizer(optimizer)

    # Train!
    total_batch_size = args.per_device_train_batch_size * smp.size() * args.gradient_accumulation_steps

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    total_steps = num_update_steps_per_epoch * args.epochs
    
    if is_main_process(smp.rank()):
        print("***** Running training *****")
        print(f"  Num examples = {len(train_dataset)}")
        print(f"Number of eval examples ={len(eval_dataset)}")
        print(f"  Num Epochs = {args.epochs}")
        print(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
        print(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
        print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        print(f"  max_train_steps = {args.max_train_steps}")
        print(f"  total_steps = {total_steps}")

    progress_bar = tqdm(range(total_steps), disable=not is_local_main_process(smp.local_rank()))
    completed_steps = 0
    total_loss = 0

    for epoch in range(args.epochs):
        model.train()
        total_loss = 0
        for step, batch in enumerate(train_dataloader):
            device = torch.device("cuda")
            batch = {k: v.to(device) for k, v, in batch.items()}
            loss = train_step(model,batch)
            
            loss = loss.reduce_mean()
            total_loss += loss.detach().float()

            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

            progress_bar.update(1)
            completed_steps += 1

            # if completed_steps >= args.max_train_steps:
            #     break

        model.eval()
        eval_loss = 0
        losses = []

        for step, batch in enumerate(tqdm(eval_dataloader)):
            with torch.no_grad():
                device = torch.device("cuda")
                batch = {k: v.to(device) for k, v, in batch.items()}
                loss = test_step(model, batch).reduce_mean()
            losses.append(loss)
            if step > 50:
                break
        
        try:
            eval_loss = torch.mean(torch.tensor(losses,dtype=float))
            perplexity = math.exp(eval_loss)
        except OverflowError:
            perplexity = float("inf")

        print(f"epoch {epoch}: Eval perplexity: {perplexity} Eval loss: {eval_loss}")

        
    smp.save_checkpoint(args.checkpoint_dir,
                tag=f"flan_t5_weights.pt",
                partial=False,
                model=model,
                optimizer=optimizer)
        
    print("saving the final model")

    wait_for_everyone()

    if is_main_process(smp.rank()):
        tokenizer.save_pretrained(args.checkpoint_dir)


if __name__ == "__main__":
    main()
finetuning/smp/scripts/qa-smp.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import argparse
from contextlib import contextmanager
import torch
import os
import json
from transformers import (
    MODEL_MAPPING,
    SchedulerType
)

# Distributed training helper methods.
def wait_for_everyone():
    #deepspeed.comm.barrier
    torch.distributed.barrier()


def is_main_process(rank):
    if rank == 0:
        return True
    else:
        return False

def _goes_first(is_main):
    if not is_main:
        wait_for_everyone()
    yield
    if is_main:
        wait_for_everyone()

@contextmanager
def main_process_first(rank):
    """
    Lets the main process go first inside a with block.
    The other processes will enter the with block after the main process exits.
    """
    yield from _goes_first(is_main_process(rank))

def is_local_main_process(local_rank):
    return local_rank == 0

MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)


# args parsing
def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a FLAN T5 model on a Seq2Seq task")

    # add model id and dataset path argument
    parser.add_argument("--model_id", type=str, default="google/flan-t5-xl", help="Model id to use for training.")
    
     # add training hyperparameters for epochs, batch size, learning rate, and seed
    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs to train for.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--per_device_eval_batch_size", type=int, default=8, help="Batch size to use for testing.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient Accumulation Steps")
    parser.add_argument("--generation_max_length", type=int, default=140, help="Maximum length to use for generation")
    parser.add_argument("--generation_num_beams", type=int, default=4, help="Number of beams to use for generation.")
    parser.add_argument("--learning_rate", type=float, default=3e-3, help="Learning rate to use for training.")
    parser.add_argument("--seed", type=int, default=42, help="Seed to use for reproducible training.")
    parser.add_argument("--deepspeed", type=str, default=None, help="Path to deepspeed config file.")
    parser.add_argument("--gradient_checkpointing", type=bool, default=True, help="Whether to use gradient checkpointing.")
    parser.add_argument(
        "--bf16",
        type=bool,
        default=True if torch.cuda.get_device_capability()[0] == 8 else False,
        help="Whether to use bf16.",
    )
    parser.add_argument("--save_model_s3_path", type=str, default=None, help="s3 path to save model")
    parser.add_argument("--optim", type=str, default='adamw_hf', help="The optimizer to use: adamw_hf, adamw_torch, adamw_torch_fused, adamw_apex_fused, adamw_anyprecision or adafactor")
    parser.add_argument(
        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument(
        "--preprocessing_num_workers",
        type=int,
        default=None,
        help="The number of processes to use for the preprocessing.",
    )
    parser.add_argument("--group_texts",default=False,help="Whether to group texts together when tokenizing")
    parser.add_argument(
        "--block_size",
        type=int,
        default=None,
        help=(
            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
            " this size for training. Default to the model max input length for single sentence inputs (take into"
            " account special tokens)."
        ),
    )
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=500,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    
    parser.add_argument("checkpoint_dir",type=str,default="/opt/ml/checkpoints")
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--output_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

    args,_ = parser.parse_known_args()

    return args
finetuning/smp/scripts/utils.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import os
from aws_cdk import (
    CfnOutput,
    Duration,
    Stack,
    RemovalPolicy,
    aws_iam as iam,
    aws_lambda as lambda_,
    aws_stepfunctions as sfn,
    aws_stepfunctions_tasks as sfn_tasks,
    aws_s3 as s3,
    aws_s3_deployment as s3deploy,
    aws_dynamodb as dynamodb,
    aws_apigateway as apigateway,
    cloudformation_include as cfn_inc,
)
from constructs import Construct
from kaggle.api.kaggle_api_extended import KaggleApi
import tempfile


class GenAIPersonalizedStack(Stack):
    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)
        genai_personalized_bucket = s3.Bucket(self, "genai_personalized")

        trust_policy = iam.PolicyDocument(
            statements=[
                iam.PolicyStatement(
                    effect=iam.Effect.ALLOW, actions=["sts:AssumeRole"], resources=["*"]
                )
            ]
        )

        execution_role = iam.Role(
            self,
            "SageMakerExecutionRole",
            assumed_by=iam.ServicePrincipal("sagemaker.amazonaws.com"),
            inline_policies={"TrustPolicy": trust_policy},
        )

        execution_role.add_managed_policy(
            iam.ManagedPolicy.from_aws_managed_policy_name("AmazonSageMakerFullAccess")
        )
        execution_role.add_managed_policy(
            iam.ManagedPolicy.from_aws_managed_policy_name("AmazonS3FullAccess")
        )

        dataset_name = "vishesh1412/celebrity-face-image-dataset"
        with tempfile.TemporaryDirectory() as temp_dir:
            destination_path = temp_dir + "celebrity-face-image-dataset"
            if not os.path.exists(destination_path):
                os.makedirs(destination_path)

            self.download_kaggle_dataset(dataset_name, destination_path)

            s3deploy.BucketDeployment(
                self,
                "DeployDataset",
                sources=[s3deploy.Source.asset(destination_path)],
                destination_bucket=genai_personalized_bucket,
                destination_key_prefix="training_images",
            )

        model_metadata_table = dynamodb.Table(
            self,
            "ModelMetadata",
            table_name="model-metadata",
            partition_key=dynamodb.Attribute(
                name="UserID", type=dynamodb.AttributeType.STRING
            ),
            sort_key=dynamodb.Attribute(
                name="ModelID", type=dynamodb.AttributeType.STRING
            ),
            billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST,
            removal_policy=RemovalPolicy.DESTROY,
        )

        launch_fine_tuning_job_lambda = lambda_.Function(
            self,
            "LaunchFineTuningJobLambda",
            runtime=lambda_.Runtime.PYTHON_3_8,
            handler="index.handler",
            code=lambda_.Code.from_asset("lambdas/launch_fine_tune_job"),
            environment={
                "EXECUTION_ROLE_ARN": execution_role.role_arn,
                "S3_MODEL_STORE": genai_personalized_bucket.s3_url_for_object(
                    key="model-store"
                ),
                "MODEL_METADATA_TABLE": model_metadata_table.table_name,
                "JUMPSTART_SUBMIT_DIRECTORY": "s3://jumpstart-cache-prod-eu-west-1/source-directory-tarballs/stabilityai/transfer_learning/txt2img/prepack/v1.0.3/sourcedir.tar.gz",
                "JUMPSTART_S3URI": "s3://jumpstart-cache-prod-eu-west-1/stabilityai-training/train-model-txt2img-stabilityai-stable-diffusion-v2-1-base.tar.gz",
            },
        )

        policy_statement = iam.PolicyStatement(
            actions=[
                "sagemaker:CreateTrainingJob",
                "sagemaker:DescribeTrainingJob",
                "sagemaker:StopTrainingJob",
                "sagemaker:InvokeEndpoint",
                "iam:PassRole",
                "dynamodb:*",
                "s3:*",
                "states:StartExecution",
            ],
            resources=[
                "*"
            ],
        )

        launch_fine_tuning_job_lambda.role.add_to_policy(policy_statement)
        launch_fine_tuning_job_step = sfn_tasks.LambdaInvoke(
            self, "LaunchFineTuningJob", lambda_function=launch_fine_tuning_job_lambda
        )

        wait_state = sfn.Wait(
            self, "WaitState", time=sfn.WaitTime.duration(Duration.minutes(15))
        )
        untar_model_artifact_lambda = lambda_.Function(
            self,
            "UntarModelArtifactLambda",
            runtime=lambda_.Runtime.PYTHON_3_8,
            handler="untar_file.handler",
            code=lambda_.Code.from_asset("lambdas/model_artifact_preparation"),
            memory_size=10240,
            timeout=Duration.minutes(15),
            environment={
                "S3_MODEL_STORE_BUCKET_NAME": genai_personalized_bucket.bucket_name,
                "MODEL_METADATA_TABLE": model_metadata_table.table_name,
            },
        )

        untar_model_artifact_lambda.role.add_to_policy(policy_statement)
        untar_model_artifact_step = sfn_tasks.LambdaInvoke(
            self, "UntarModelArtifact", lambda_function=untar_model_artifact_lambda
        )

        upload_model_artifact_lambda = lambda_.Function(
            self,
            "UploadModelArtifactLambda",
            runtime=lambda_.Runtime.PYTHON_3_8,
            handler="upload_model_code.handler",
            code=lambda_.Code.from_asset("lambdas/model_artifact_preparation"),
            environment={
                "S3_MODEL_STORE_BUCKET_NAME": genai_personalized_bucket.bucket_name,
                "MODEL_METADATA_TABLE": model_metadata_table.table_name,
            },
        )
        upload_model_artifact_lambda.role.add_to_policy(policy_statement)
        upload_model_artifact_step = sfn_tasks.LambdaInvoke(
            self, "UploadModelArtifact", lambda_function=upload_model_artifact_lambda
        )

        definition = (
            sfn.Chain.start(launch_fine_tuning_job_step)
            .next(wait_state)
            .next(untar_model_artifact_step)
            .next(upload_model_artifact_step)
            .next(sfn.Pass(self, "FinalState"))
        )

        state_machine = sfn.StateMachine(
            self,
            "GenAIStateMachine",
            definition=definition,
            timeout=Duration.minutes(30),
        )

        genai_personalized_bucket.grant_read(launch_fine_tuning_job_lambda)
        genai_personalized_bucket.grant_write(untar_model_artifact_lambda)

        start_sfn_workflow_lambda = lambda_.Function(
            self,
            "StartSFNWorkflow",
            runtime=lambda_.Runtime.PYTHON_3_8,
            handler="start_sfn_workflow.handler",
            code=lambda_.Code.from_asset("lambdas/http_request_lambdas"),
            environment={
                "S3_TRAINING_DATA": genai_personalized_bucket.s3_url_for_object(
                    key="training_images/Celebrity Faces Dataset/Angelina Jolie/"
                ),
                "STEP_FUNCTION_ARN": state_machine.state_machine_arn,
                "MODEL_METADATA_TABLE": model_metadata_table.table_name,
            },
        )
        start_sfn_workflow_lambda.role.add_to_policy(policy_statement)

        check_model_state_lambda = lambda_.Function(
            self,
            "CheckModelState",
            runtime=lambda_.Runtime.PYTHON_3_8,
            handler="check_model_state.handler",
            code=lambda_.Code.from_asset("lambdas/http_request_lambdas"),
            environment={"MODEL_METADATA_TABLE": model_metadata_table.table_name},
        )
        check_model_state_lambda.role.add_to_policy(policy_statement)

        # ModelExecutionRoleArn and S3PathToModels are specific to account
        import_params = {
            "SageMakerProjectName": "sd-fine-tuned-mme",
            "ModelExecutionRoleArn": execution_role.role_arn,
            "StageName": "dev",
            "EndpointInstanceCount": "1",
            "EndpointInstanceType": "ml.g5.xlarge",
            "S3PathToModels": genai_personalized_bucket.s3_url_for_object(
                key="model-code/fine-tuned-sd/code/"
            ),
            "ContainerImageInference": "763104351884.dkr.ecr.eu-west-1.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.3-cu117",
        }

        cfn_inc.CfnInclude(
            self,
            "Template",
            template_file="genai_personalized/mme-template.yaml",
            parameters=import_params,
        )

        invoke_endpoint_for_inference_lambda = lambda_.Function(
            self,
            "InvokeEndpointForInference",
            runtime=lambda_.Runtime.PYTHON_3_8,
            handler="invoke_endpoint_for_inference.handler",
            code=lambda_.Code.from_asset("lambdas/http_request_lambdas"),
            environment={
                "MULTI_MODEL_ENDPOINT_NAME": "sd-fine-tuned-mme-dev",
            },
        )
        invoke_endpoint_for_inference_lambda.role.add_to_policy(policy_statement)

        api = apigateway.RestApi(
            self,
            "GenAIPersonalizedAPI",
            rest_api_name="genai-personalized-api",
        )

        catch_all_resource = api.root.add_resource("resource")
        catch_all_resource.add_proxy()

        finetune = api.root.add_resource("finetune")
        model = api.root.add_resource("model")
        inference = api.root.add_resource("inference")
        # ...

        lambda_integration1 = apigateway.LambdaIntegration(start_sfn_workflow_lambda)
        lambda_integration2 = apigateway.LambdaIntegration(check_model_state_lambda)
        lambda_integration3 = apigateway.LambdaIntegration(invoke_endpoint_for_inference_lambda)
        # ...

        finetune.add_method("POST", lambda_integration1)
        model.add_method("GET", lambda_integration2)
        inference.add_method("POST", lambda_integration3)
        # ...

        CfnOutput(self, "StateMachineArn", value=state_machine.state_machine_arn)

    def download_kaggle_dataset(self, dataset_name, destination_path):
        api = KaggleApi()
        api.authenticate()
        api.dataset_download_files(dataset_name, path=destination_path, unzip=True)

genai_personalized/genai_personalized_stack.py
aws-samples/amazon-sagemaker-personalized-generative-ai
# API Gateway Connector for getting model state from DynamoDB, launching a new training job.

import boto3
import json
from custom_encoder import CustomEncoder
import logging
import os
logger = logging.getLogger()
logger.setLevel(logging.INFO)

logger = logging.getLogger()
logger.setLevel(logging.INFO)

dynamodbTableName = os.environ['MODEL_METADATA_TABLE']
dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table(dynamodbTableName)

stepfunction_client = boto3.client('stepfunctions')

getMethod = 'GET'
postMethod = 'POST'
patchMethod = 'PATCH'
deleteMethod = 'DELETE'
healthPath = '/health'
modelPath = '/model'
model = '/models'


def handler(event, context):
    logger.info(event)
    httpMethod = event['httpMethod']
    path = event['path']
    
    if httpMethod == getMethod and path == healthPath:
        response = buildResponse(200)
    elif httpMethod == getMethod and path == modelPath:
        response = getModel(event['queryStringParameters']
                            ['UserID'], event['queryStringParameters']['ModelID'])
    elif httpMethod == postMethod and path == modelPath:
        response = saveModel(json.loads(event['body']))
    # To modify a model, not implemented
    elif httpMethod == patchMethod and path == modelPath:
        requestBody = json.loads(event['body'])
        response = modifyModel(
            requestBody['UserID'], requestBody['updateKey'], requestBody['updateValue'])
    else:
        response = buildResponse(404, 'Not Found')
    return response

# Checks if the UserID and ModelID combination already exists, and its current state
def getModel(UserID, ModelID):
    try:
        response = table.get_item(
            Key={
                'UserID': UserID,
                'ModelID': ModelID
            }
        )
        if 'Item' in response:
            return buildResponse(200, response['Item'])
        else:
            return buildResponse(404, {'Message': 'UserID or ModelID not found'})
    except:
        logger.exception('Error Handling')

def saveModel(requestBody):
    
    model_params = {
        "model_status" : "TRAINING_REQUESTED",
        "target_model_name_mme" : "",
        "training_job_name" : ""
    }
    requestBody.update(model_params)
    
    try:
        table.put_item(Item=requestBody)
        body = {
            'Operation': 'SAVE',
            'Message': 'SUCCESS',
            'Item': requestBody
        }
        
        input_dict = {
            'user_id': requestBody['UserID'],
            'model_id' : requestBody['ModelID'],
            's3_fine_tuning_images_path': requestBody['s3_fine_tuning_images_path']
        }
        answ = stepfunction_client.start_execution(
            stateMachineArn='arn:aws:states:us-east-1:055107841600:stateMachine:ab-ft-sd-workflow',
            input=json.dumps(input_dict))
        
        return buildResponse(200, body)
    
    except:
        logger.exception('Error Handling')

def modifyModel(UserID, updateKey, updateValue):
    try:
        response = table.update_item(
            Key={
                'UserID': UserID
            },
            UpdateExpression='set %s = :value' (updateKey),
            ExpressionAttributeValues={
                ':value': updateValue
            },
            ReturnValues='UPDATED_NEW'
        )
        body = {
            'Operation': 'UPDATE',
            'Message': 'SUCCESS',
            'UpdatedAttributes': response
        }
        return buildResponse(200, body)
    except:
        logger.exception('Error Handling')

def buildResponse(statusCode, body=None):
    response = {
        'statusCode': statusCode,
        'headers': {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*'
        }
    }
    if body is not None:
        response['body'] = json.dumps(body, cls=CustomEncoder)
    return response

lambdas/http_request_lambdas/check_model_state.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import json

class CustomEncoder(json.JSONEncoder):
    
    def default(self,obj):
    
        if isinstance(obj,Decimal):
            return float(obj)
        
        return json.JSONEncoder.default(self, obj)
lambdas/http_request_lambdas/custom_encoder.py
aws-samples/amazon-sagemaker-personalized-generative-ai
# Function that connects to API gateway to get an image from the MME
# The Endpoint's name is in an environment variable (multi_model_endpoint_name) 
# The Model's name is passed as an argument in the API call
# Lastly, some processing is done to return the decoded image

import os
import io
import boto3
import json
from io import BytesIO
import base64
from base64 import b64encode

ENDPOINT_NAME = os.environ['MULTI_MODEL_ENDPOINT_NAME']
runtime = boto3.client('sagemaker-runtime')
  
def handler(event, context):
    
    TARGET_MODEL = event['queryStringParameters']['model_name']

    payload = '{"prompt" : "a photo of a zwx person"}'
    
    response = runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME,
                                       Body=payload,
                                       ContentType='application/json',
                                       TargetModel=TARGET_MODEL)

    
    result = response['Body'].read()
    stream = BytesIO(result)
    image = stream.read()
    
    return{
        #"isBase64Encoded": True,
        "statusCode": 200,
        'headers': {"Content-Type": "text/plain",
                    'Access-Control-Allow-Origin': '*'},
        'body': base64.b64encode(image).decode('utf-8')
    }
lambdas/http_request_lambdas/invoke_endpoint_for_inference.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import boto3
import json
import uuid
import os
from custom_encoder import CustomEncoder
import logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

logger = logging.getLogger()
logger.setLevel(logging.INFO)

dynamodbTableName = os.environ['MODEL_METADATA_TABLE']
dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table(dynamodbTableName)

stepfunction_client = boto3.client('stepfunctions')

postMethod = 'POST'
finetunePath = '/finetune'

def handler(event, context):
    
    logger.info(event)
    s3_training_data = os.environ['S3_TRAINING_DATA']

    body = {
        "s3_training_data": s3_training_data,
        "UserID": str(uuid.uuid4()),
        "ModelID": str(uuid.uuid4())
    }

    httpMethod = 'POST'
    path = '/finetune'
    
    if httpMethod == postMethod and path == finetunePath:
        response = saveModel(body)
    else:
        response = buildResponse(404, 'Not Found')
    return response

def saveModel(requestBody):
    step_function_arn = os.environ['STEP_FUNCTION_ARN']
    model_params = {
        "model_status" : "TRAINING_REQUESTED",
        "target_model_name_mme" : "",
        "training_job_name" : ""
    }
    requestBody.update(model_params)
    
    try:
        table.put_item(Item=requestBody)
        body = {
            'Operation': 'SAVE',
            'Message': 'SUCCESS',
            'Item': requestBody
        }
        
        input_dict = {
            'user_id': requestBody['UserID'],
            'model_id' : requestBody['ModelID'],
            's3_training_data': requestBody['s3_training_data']
        }
        answ = stepfunction_client.start_execution(
            stateMachineArn=step_function_arn,
            input=json.dumps(input_dict))
        
        return buildResponse(200, body)
    
    except:
        logger.exception('Error Handling')


def buildResponse(statusCode, body=None):
    response = {
        'statusCode': statusCode,
        'headers': {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*'
        }
    }
    if body is not None:
        response['body'] = json.dumps(body, cls=CustomEncoder)
    return response

lambdas/http_request_lambdas/start_sfn_workflow.py
aws-samples/amazon-sagemaker-personalized-generative-ai
# This function is the first in the Step Function workflow and copies an existing 
# training job (training_job_name_to_copy) and just modifies the S3Uri of the training data.
# It also updates the DynamoDB table to set status to 'MODEL_TRAINING' and adds the training job name
import boto3, os, datetime, json


jumpstart_submit_directory = os.environ['JUMPSTART_SUBMIT_DIRECTORY']
jumpstart_s3uri = os.environ['JUMPSTART_S3URI']
execution_role_arn = os.environ['EXECUTION_ROLE_ARN']
s3_model_store = os.environ['S3_MODEL_STORE']
dynamodbTableName = os.environ['MODEL_METADATA_TABLE']
dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table(dynamodbTableName)
s3_client = boto3.client('s3')

def handler(event, context):

    user_id = event['user_id']
    model_id = event['model_id']
    s3_training_data = event['s3_training_data']
    
    sm = boto3.client('sagemaker')
    job = {
        'TrainingJobName': 'None',
        'TrainingJobArn': 'None',
        'HyperParameters': {
            'adam_beta1': '"0.9"',
            'adam_beta2': '"0.999"',
            'adam_epsilon': '"1e-08"',
            'adam_weight_decay': '"0.01"',
            'batch_size': '"1"',
            'center_crop': '"False"',
            'compute_fid': '"False"',
            'epochs': '"10"',
            'gradient_accumulation_steps': '"1"',
            'learning_rate': '"2e-06"',
            'lr_scheduler': '"constant"',
            'max_grad_norm': '"1.0"',
            'max_steps': '"200"',
            'num_class_images': '"100"',
            'prior_loss_weight': '"1.0"',
            'sagemaker_container_log_level': '20',
            'sagemaker_job_name': '"sd-fine-tune-training-job-model-txt2img-2023-04-20-12-44-39-910"',
            'sagemaker_program': '"transfer_learning.py"',
            'sagemaker_region': '"eu-west-1"',
            'sagemaker_submit_directory': f'"{jumpstart_submit_directory}"',
            'seed': '"0"',
            'with_prior_preservation': '"False"'
        },
        'AlgorithmSpecification': {
            'TrainingImage': '763104351884.dkr.ecr.eu-west-1.amazonaws.com/huggingface-pytorch-training:1.10.2-transformers4.17.0-gpu-py38-cu113-ubuntu20.04',
            'TrainingInputMode': 'File',
            'MetricDefinitions': [
                {'Name': 'stabilityai-txt2img:train-loss', 'Regex': 'train_avg_loss=([0-9\\.]+)'},
                {'Name': 'fid_score', 'Regex': 'fid_score=([-+]?\\d\\.?\\d*)'}
            ],
            'EnableSageMakerMetricsTimeSeries': False
        },
        'RoleArn': execution_role_arn,
        'InputDataConfig': [
            {
                'ChannelName': 'training',
                'DataSource': {
                    'S3DataSource': {
                        'S3DataType': 'S3Prefix',
                        'S3Uri': 'None',
                        'S3DataDistributionType': 'FullyReplicated'
                    }
                },
                'CompressionType': 'None',
                'RecordWrapperType': 'None'
            },
            {
                'ChannelName': 'model',
                'DataSource': {
                    'S3DataSource': {
                        'S3DataType': 'S3Prefix',
                        'S3Uri': jumpstart_s3uri,
                        'S3DataDistributionType': 'FullyReplicated'
                    }
                },
                'ContentType': 'application/x-sagemaker-model',
                'CompressionType': 'None',
                'RecordWrapperType': 'None',
                'InputMode': 'File'
            }
        ],
        'OutputDataConfig': {
            'KmsKeyId': '',
            'S3OutputPath': s3_model_store
        },
        'ResourceConfig': {
            'InstanceType': 'ml.g5.2xlarge',
            'InstanceCount': 1,
            'VolumeSizeInGB': 30
        },
        'StoppingCondition': {
            'MaxRuntimeInSeconds': 360000
        },
        'EnableNetworkIsolation': False,
        'EnableInterContainerTrafficEncryption': False,
        'EnableManagedSpotTraining': False,
        'DebugHookConfig': {
            'S3OutputPath': s3_model_store,
            'CollectionConfigurations': []
        },
        'ProfilerConfig': {
            'S3OutputPath': s3_model_store,
            'ProfilingIntervalInMilliseconds': 500,
            'DisableProfiler': False
        }
    }

    training_job_prefix = 'sd-fine-tune-training-job-model-txt2img-'
    training_job_name = training_job_prefix+str(datetime.datetime.today()).replace(' ', '-').replace(':', '-').rsplit('.')[0]
    
    job['InputDataConfig'][0]['DataSource']['S3DataSource']['S3Uri'] = s3_training_data
    job['ResourceConfig']['InstanceType'] = 'ml.g5.2xlarge' 

    json_data = {
        "instance_prompt": "a photo of a zwx person",
        "class_prompt": "a photo of a person"
    }

    json_string = json.dumps(json_data)
    parts = s3_training_data.split('/')
    bucket_name = parts[2]
    object_key = '/'.join(parts[3:]) + "dataset_info.json"

    s3_client.put_object(
        Bucket=bucket_name,
        Key=object_key,
        Body=json_string
    )

    print(f"Uploaded {object_key} to {s3_training_data}")
    print("Starting training job %s" % training_job_name)

    if 'VpcConfig' in job:
        resp = sm.create_training_job(
            TrainingJobName=training_job_name, AlgorithmSpecification=job['AlgorithmSpecification'], RoleArn=job['RoleArn'],
            InputDataConfig=job['InputDataConfig'], OutputDataConfig=job['OutputDataConfig'],
            ResourceConfig=job['ResourceConfig'], StoppingCondition=job['StoppingCondition'],
            HyperParameters=job['HyperParameters'] if 'HyperParameters' in job else {},
            VpcConfig=job['VpcConfig'],
            Tags=job['Tags'] if 'Tags' in job else [])
    else:
        resp = sm.create_training_job(
            TrainingJobName=training_job_name, AlgorithmSpecification=job['AlgorithmSpecification'], RoleArn=job['RoleArn'],
            InputDataConfig=job['InputDataConfig'], OutputDataConfig=job['OutputDataConfig'],
            ResourceConfig=job['ResourceConfig'], StoppingCondition=job['StoppingCondition'],
            HyperParameters=job['HyperParameters'] if 'HyperParameters' in job else {},
            Tags=job['Tags'] if 'Tags' in job else [])
    

    print(resp)
    
    response = table.update_item(
        Key={
            'UserID': user_id,
            'ModelID': model_id
        },
        UpdateExpression='SET model_status = :input1, training_job_name = :input2',
        ExpressionAttributeValues={
            ':input1': 'MODEL_TRAINING',
            ':input2': training_job_name
        },
        ReturnValues='UPDATED_NEW'
    )
    print(response)
    
    return {
        'training_job_name': training_job_name,
        'user_id':user_id,
        'model_id':model_id
    }
lambdas/launch_fine_tune_job/index.py
aws-samples/amazon-sagemaker-personalized-generative-ai
# To use accelerate, the model artifacts need to be untarred.
# This function untars the model artifacts, and copies them to the 'ab-model-store-32987' bucket
# It also updates the status in the DynamoDB table to 'PROCESSING_MODEL_ARTIFACTS'
import boto3
import botocore
import tarfile
import os

from io import BytesIO
s3_client = boto3.client('s3')

s3_model_store_bucket_name = os.environ['S3_MODEL_STORE_BUCKET_NAME']
dynamodbTableName = os.environ['MODEL_METADATA_TABLE']
dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table(dynamodbTableName)

def handler(event, context):

    training_job_name = event['Payload']['training_job_name']
    user_id = event['Payload']['user_id']
    model_id = event['Payload']['model_id']
    key = 'model-store/' + training_job_name + '/output/model.tar.gz'
    
    table.update_item(
        Key={
            'UserID': user_id,
            'ModelID': model_id
        },
        UpdateExpression='SET model_status = :input',
        ExpressionAttributeValues={
            ':input': 'PROCESSING_MODEL_ARTIFACTS'
        },
        ReturnValues='UPDATED_NEW'
    )

    input_tar_file = s3_client.get_object(Bucket = s3_model_store_bucket_name, Key = key)
    input_tar_content = input_tar_file['Body'].read()
    
    unzipped_artifacts_s3_path = ('/model-artifact/' + training_job_name + '/')

    with tarfile.open(fileobj = BytesIO(input_tar_content)) as tar:
        for tar_resource in tar:
            if (tar_resource.isfile()):
                inner_file_bytes = tar.extractfile(tar_resource).read()
                s3_client.upload_fileobj(BytesIO(inner_file_bytes), Bucket = s3_model_store_bucket_name, Key = (unzipped_artifacts_s3_path+tar_resource.name))
                
    return{
        'unzipped_artifacts_s3_path': unzipped_artifacts_s3_path,
        'user_id':user_id,
        'model_id':model_id
    }
lambdas/model_artifact_preparation/untar_file.py
aws-samples/amazon-sagemaker-personalized-generative-ai
# This function creates the necessary files for Accelerate:
# 1. The serving.properties file 
# 2. The Python file with the inference code
#
# It packages them into a tar.gz, and uploads it to the 'ab-model-code-90823' bucket, which is the S3 prefix 
# where all the models artifacts (.tar.gz) in the Multi-Model endpoint are located
import tarfile
import os.path
import boto3

model_code = """import logging
import os
from time import time
import torch
from diffusers import DiffusionPipeline
import diffusers
from djl_python.inputs import Input
from djl_python.outputs import Output
from typing import Optional
from io import BytesIO
from PIL import Image


def get_torch_dtype_from_str(dtype: str):
    if dtype == "fp16":
        return torch.float16
    raise ValueError(
        f"Invalid data type: {dtype}. DeepSpeed currently only supports fp16 for stable diffusion"
    )

class StableDiffusionService(object):

    def __init__(self):
        self.pipeline = None
        self.initialized = False
        self.logger = logging.getLogger()
        self.model_id_or_path = None
        self.data_type = None
        self.device = None
        self.world_size = None
        self.max_tokens = None
        self.save_image_dir = None

    def initialize(self, properties: dict):
        # model_id can point to huggingface model_id or local directory.
        # If option.model_id points to a s3 bucket, we download it and set model_id to the download directory.
        # Otherwise we assume model artifacts are in the model_dir
        self.model_id_or_path = properties.get("model_id") or properties.get(
            "model_dir")
        self.data_type = get_torch_dtype_from_str(properties.get("dtype"))
        self.max_tokens = int(properties.get("max_tokens", "1024"))
        self.device = int(os.getenv("LOCAL_RANK", "0"))

        if os.path.exists(self.model_id_or_path):
            config_file = os.path.join(self.model_id_or_path,
                                       "model_index.json")
            if not os.path.exists(config_file):
                raise ValueError(
                    f"{self.model_id_or_path} does not contain a model_index.json."
                    f"This is required for loading stable diffusion models from local storage"
                )

        kwargs = {"torch_dtype": torch.float16, "revision": "fp16"}
        
        start = time()
        pipeline = DiffusionPipeline.from_pretrained(self.model_id_or_path,
                                                     device_map='auto',
                                                     low_cpu_mem_usage=True,
                                                     **kwargs
                                                    )
    
        duration = time()-start 
        self.logger.info(f'Loaded model in {duration} seconds')

        self.pipeline = pipeline
        self.initialized = True


    def inference(self, inputs: Input):
        try:
            content_type = inputs.get_property("Content-Type")
            if content_type == "application/json":
                request = inputs.get_as_json()
                prompt = request.pop("prompt")
                params = request.pop("parameters", {})
                start = time()
                result = self.pipeline(prompt, **params)
                duration = time() - start
                self.logger.info(f'Inference took {duration} seconds')
            elif content_type and content_type.startswith("text/"):
                prompt = inputs.get_as_string()
                result = self.pipeline(prompt)
            else:
                # in case an image and a prompt is sent in the input
                init_image = Image.open(BytesIO(
                    inputs.get_as_bytes())).convert("RGB")
                request = inputs.get_as_json("json")
                prompt = request.pop("prompt")
                params = request.pop("parameters", {})
                result = self.pipeline(prompt, image=init_image, **params)

            img = result.images[0]
            buf = BytesIO()
            img.save(buf, format="PNG")
            byte_img = buf.getvalue()
            outputs = Output().add(byte_img).add_property(
                "content-type", "image/png")

        except Exception as e:
            logging.exception("Inference failed")
            outputs = Output().error(str(e))
        return outputs


_service = StableDiffusionService()


def handle(inputs: Input) -> Optional[Output]:
    if not _service.initialized:
        _service.initialize(inputs.get_properties())

    if inputs.is_empty():
        return None

    return _service.inference(inputs)"""

s3_client = boto3.client('s3')
s3_model_store_bucket_name = os.environ['S3_MODEL_STORE_BUCKET_NAME']
dynamodbTableName = os.environ['MODEL_METADATA_TABLE']
dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table(dynamodbTableName)
        
def handler(event,context):
    
    model_artifact_path = event['Payload']['unzipped_artifacts_s3_path']
    
    serving_properties = """engine=Python
option.s3url={path}
option.dtype=fp16""".format(path=model_artifact_path)
    
    serving_template = """engine=Python
option.s3url={{ s3url }}
option.dtype=fp16"""
    
    model_file = open("/tmp/model.py", "w+")
    model_file.write(str(model_code))
    model_file.close()
    
    serving_properties_file = open("/tmp/serving.properties", "w+")
    serving_properties_file.write(str(serving_properties))
    serving_properties_file.close()
    
    serving_template_file = open("/tmp/serving.template", "w+")
    serving_template_file.write(str(serving_template))
    serving_template_file.close()
    
    model_file_name = '/tmp/model.py'
    serving_properties_file_name = '/tmp/serving.properties'
    serving_template_file_name = '/tmp/serving.template'
    
    with tarfile.open("/tmp/model.tar.gz", "w|gz") as tf:
        tf.add(f"{model_file_name}")
        tf.add(f"{serving_properties_file_name}")
        tf.add(f"{serving_template_file_name}")
        
    user_id = event['Payload']['user_id']
    model_id = event['Payload']['model_id']
    
    model_name = 'ft-sd-'+user_id+'-'+model_id+'.tar.gz'
    s3_client.upload_file('/tmp/model.tar.gz', s3_model_store_bucket_name, ('model-code/fine-tuned-sd/code/'+model_name))
    
    response = table.update_item(
        Key={
            'UserID': user_id,
            'ModelID': model_id
        },
        UpdateExpression='SET model_status = :input1, target_model_name_mme = :input2',
        ExpressionAttributeValues={
            ':input1': 'AVAILABLE',
            ':input2': model_name
        },
        ReturnValues='UPDATED_NEW'
    )
    
    return {
        'model_name': model_name,
        'user_id' : user_id,
        "model_id" : model_id
    }

lambdas/model_artifact_preparation/upload_model_code.py
aws-samples/amazon-sagemaker-personalized-generative-ai
import aws_cdk as core
import aws_cdk.assertions as assertions

from genai_personalized.genai_personalized_stack import GenaiPersonalizedStack

# example tests. To run these tests, uncomment this file along with the example
# resource in genai_personalized/genai_personalized_stack.py
def test_sqs_queue_created():
    app = core.App()
    stack = GenaiPersonalizedStack(app, "genai-personalized")
    template = assertions.Template.from_stack(stack)

#     template.has_resource_properties("AWS::SQS::Queue", {
#         "VisibilityTimeout": 300
#     })

tests/unit/test_genai_personalized_stack.py
aws-samples/amazon-sagemaker-personalized-generative-ai