---
# System prepended metadata

title: Sagemaker
tags: [llm]

---

### Architecture setup
#### VPC
- Internet gateway 생성
- subnet은 최소 2개: public / private 필요
- 다양한 sagemaker training instance를 활용하기 위해서는 availability zone A 사용 권장

#### Subnets
- public subnet:
    - Internet gateway로 열려있는 subnet으로 외부망 통신 가능
    - Routing table:
        - S3 gateway endpoint 설정
        - IPv4 -> internet gateway
    - EFS interface endpoint 설정: Sagemaker jupyerlab에서 EFS를 연결해서 training instance에 할당하기 위한 용도
        - Security group:
            - sagemaker 전용 security group 할당: sagemaker에서 EFS 통신 가능하도록
    - NAT gateway public subnet으로 생성
- private subnet:
    - Sagemaker training instance가 배정될 subnet
    - 외부망 통신을 위해 public subnet에 할당되어있는 NAT gateway를 사용
    - Routing table:
        - S3 gateway endpoint 설정
        - IPv4 -> NAT gateway
        
#### Security group
- EFS security group
    - inbound:
        - 모든 NFS 트래픽 / target: Sagemaker security group
    - outbound:
        - IPv4 / target: 0.0.0.0/0
- Sagemaker security group
    - inbound:
        - 모든 트래픽 / target: 자기 자신
        - 모든 NFS 트래픽 / target: EFS security group
    - outbound:
        - IPv4 / target: 0.0.0.0/0

#### EFS
- Pretrained model 및 cache directory로 활용
- Network setup:
    - Subnet: Sagemaker와 training instance에서 사용하는 AZ와 Subnet 추가
    - Security group: EFS security group

#### IAM user & role
- Developer user 생성: notebook에서 aws command를 사용 용도
    - 권한:
        - AmazonElasticContainerRegistryPublicFullAccess: custom docker image ECR 관리
        - AmazonS3FullAccess: S3 training data 관리
        - AmazonSageMakerFullAccess: Sagemaker session 관리
        - ECRFullAccess
- Sagemaker executor: sagemaker notebook 생성시 사용할 role
    - 권한:
        - AmazonS3FullAccess
        - AmazonSageMakerFullAccess
        - AmazonSageMakerCanvasFullAccess
        - AmazonSageMakerCanvasAIServicesAccess
        
#### Sagemaker notebook
- Disk size: 200Gb (custom docker image를 생성하기 위해 넉넉히 구성)
- IAM role: 위에서 생성한 Sagemaker executor
- VPC: 위에서 생성한 VPC
- Subnet: 위에서 생성한 public subnet
- Security group: 위에서 생성한 Sagemaker security group


### Pretraining with megatron
#### Training data preprocessing

#### Training image
- dockerfile:
    apex install 방법은 3번 줄 참고
    ```dockerfile=
    FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.0.0-gpu-py310
   
    RUN git clone https://github.com/NVIDIA/apex \
      && cd apex \
      # recent commit, torch._six fix is included https://github.com/NVIDIA/apex/commit/a78ccf0b3e3f7130b3f157732dc8e8e651389922
      && git checkout 27de66c \
      && pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ \
      && cd .. \

    RUN pip install regex

    WORKDIR /
    ```


### Pretraining with deepspeed
#### Training image
- dockerfile:
    deepspeed를 설치할때 변수(15, 16번 줄) 중요
    ```dockerfile=
    FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.0.1-gpu-py310-cu118-ubuntu20.04-ec2

    RUN pip install --no-cache-dir rouge_score \
        && pip install --no-cache-dir fire \
        && pip install --no-cache-dir sentencepiece \
        && pip install --no-cache-dir wandb \
        && pip install --no-cache-dir openai

    RUN pip install --no-cache-dir transformers==4.28.0
    RUN pip uninstall -y accelerate \
        && pip install --no-cache-dir accelerate>=0.20.3

    RUN pip uninstall -y deepspeed

    ARG DS_BUILD_OPS=0
    ARG DS_BUILD_FUSED_ADAM=1

    RUN mkdir -p /tmp && \
        cd /tmp && \
        git clone https://github.com/microsoft/DeepSpeed.git && \
        cd DeepSpeed && \
        pip install -r requirements/requirements-dev.txt && \
        pip install -r requirements/requirements.txt && \
        DS_BUILD_OPS=$DS_BUILD_OPS DS_BUILD_FUSED_ADAM=$DS_BUILD_FUSED_ADAM pip install -v . && \
        cd .. && \
        rm -rf DeepSpeed

    WORKDIR /
    ```
#### Config files
- model: facebook/opt-6.7b
- hyperparameter:
    ```python=
    'model_name_or_path' : 'facebook/opt-6.7b' ,
    'data_path' : '/opt/ml/input/data/training/train_data.json',
    'bf16' : False,  ## p4d 이상 일때만 가능
    'output_dir' : '/opt/ml/model',
    'num_train_epochs' : 1,
    'per_device_train_batch_size' : 4,
    'per_device_eval_batch_size' : 4,
    'gradient_accumulation_steps' : 8,
    'evaluation_strategy' : 'no',
    'save_strategy' : 'steps',
    'save_steps' : 2000,
    'save_total_limit' : 1,
    'learning_rate' : 2e-5,
    'weight_decay' : 0.,
    'warmup_ratio' : 0.03,
    'deepspeed' : "/opt/ml/code/configs/default_offload_opt_param-original.json",
    'tf32' : False,  ## p4d 이상 일때만 가능
    'cache_dir' : '/opt/ml/input/data/cache_dir',
    'report_to' : 'none',
    ```
  deepspeed config:
    ```json=
    {
      "bf16": {
        "enabled": "auto"
      },
      "optimizer": {
        "type": "AdamW",
        "params": {
          "lr": "auto",
          "betas": "auto",
          "eps": "auto",
          "weight_decay": "auto"
        }
      },
      "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
          "total_num_steps": "auto",
          "warmup_min_lr": "auto",
          "warmup_max_lr": "auto",
          "warmup_num_steps": "auto"
        }
      },
      "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
          "device": "cpu",
          "pin_memory": true
        },
        "offload_param": {
          "device": "cpu",
          "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": false
      },
      "gradient_accumulation_steps": "auto",
      "gradient_clipping": "auto",
      "steps_per_print": 5,
      "train_batch_size": "auto",
      "train_micro_batch_size_per_gpu": "auto",
      "wall_clock_breakdown": false
    }
    ```
    

### Finetuining with LMFlow
#### Data preprocessing
- Loose json format 준비: line delimited json
    ```json=
    {"text": "노","label": "confirm.no"}
    {"text": "오우 노","label": "confirm.no"}
    {"text": "Never","label": "confirm.no"}
    {"text": "네버","label": "confirm.no"}
    ```
- LMFlow의 dataset format으로 변환 [reference]( https://optimalscale.github.io/LMFlow/examples/DATASETS.html)
    - LMFlow data format
        ```json=
        {
          "type": "text_only",
          "instances": [
            {
              "text": "####본문: 보험이 뭐야\n####요약: ask.definition.insurance <|endoftext|>"
            },
            {
              "text": "####본문: 보험의 정의\n####요약: ask.definition.insurance <|endoftext|>"
            },
            ...
        ```
    - data_converter.py
        ```python=
        def aihub_summarization(src_fn: str, out_fn: str, type: str ='text_only', eos_token: str=''):
        assert out_fn.endswith(".json")
        os.makedirs(os.path.dirname(out_fn), exist_ok=True)

        instances = []
        with open(src_fn, "r") as fp:
            for line in tqdm(fp):
                data = json.loads(line)
                label = data["label"].strip()
                text = data["text"].strip()
                if type == 'text_only':
                    instances.append({
                        "text": f"####본문: {text}\n####요약: {label} {eos_token}",
                    })
                elif type == 'text2text':
                    instances.append({
                        "input": f"####본문: {text}\n####요약: ",
                        "output": f"{label} {eos_token}",     
                    })
                else:
                    raise ValueError("The type is neither 'text_only' nor 'text2text'")

        with open(out_fn, "w", encoding="utf-8") as fp:
            json.dump({
                "type": type,
                "instances": instances,
            }, fp, indent=2, ensure_ascii=False)
        ```
    
- LMFlow : https://github.com/OptimalScale/LMFlow


### Inference
#### Inference facebook 6.7B model
- 먼저 model.tar.gz의 구조를 보면 아래와 같다
    나는 endpoint에 올라갈 data양을 줄여주기위해 facebook 모델에서 필요없는 데이터는 전부 삭제했고 pytorch model만 남겨놓았다
    ![](https://hackmd.io/_uploads/SyhbqWuAn.png)

- 위 tar.gz에 inference에 필요한 code를 추가해야 한다.
    code라는 directory를 생성하고 그 안에 inference.py와 requirements.txt 파일을 생성한다.

- inference.py
    해당 파일에서는 endpoint가 inference할 때 data_preprocess, model_load, data_postprocess등 다양한 처리를 추가할 수 있다.
    ```python=
    from typing import Dict, List, Any
    from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
    import torch
    import json

    JSON_CONTENT_TYPE = 'application/json'

    def model_fn(model_dir):
        # load model and processor from model_dir
        model = AutoModelForCausalLM.from_pretrained(
                    model_dir,
                    torch_dtype=torch.float16,
                    low_cpu_mem_usage=True
                ).cuda()
        tokenizer = AutoTokenizer.from_pretrained(model_dir)

        return model, tokenizer


    def input_fn(serialized_input_data, content_type=JSON_CONTENT_TYPE):
       if content_type == JSON_CONTENT_TYPE:
           input_data = json.loads(serialized_input_data)
           return input_data
    
       else:
           raise Exception('Requested unsupported ContentType in Accept: ' + content_type)
           return


    def predict_fn(data, model_and_tokenizer):
        # unpack model and tokenizer
        model, tokenizer = model_and_tokenizer

        # process input
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None)

        # preprocess
        input_ids = tokenizer(inputs, return_tensors="pt").input_ids.cuda()

        set_seed(32)

        # pass inputs with all kwargs in data
        if parameters is not None:
            # outputs = model.generate(input_ids, **parameters)
            outputs = model.generate(
                input_ids=input_ids,
                do_sample=True,
                **parameters
            )
        else:
            # outputs = model.generate(input_ids)
            outputs = model.generate(
                input_ids=input_ids,
                do_sample=True
            )

        # postprocess the prediction
        prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)

        return [{"generated_text": prediction}]
    ```

- requirements.txt에는 필요한 depedency를 추가한다
    accelerate은 여러 gpu에 model 효율적으로 load되게 하는데 활용된다.
    ```python=
    accelerate==0.18.0
    transformers==4.28
    ```

- 여기까지 진행하면 model.tar.gz에 archive 해야할 directory 구조는 아래와 같다.
    ![](https://hackmd.io/_uploads/SyxVnbOC3.png)

- model.tar.gz로 archive
    여기에서 매우 중요한건 model.tar.gz 가장 상단에 위 디렉토리 구조가 위치하도록 해야한다. 만약 내부에 hf라던지 다른 폴더 하위에 위 구조가 위치하게 되면 endpoint는 model 데이터를 찾지 못 한다.
    ```tar -zcvf model.tar.gz .```

- S3 deploy
    ```aws s3 cp model.tar.gz s3://{address}```
    
- Endpoint deploy script
    ```python=
    import boto3
    import sagemaker
    import json

    boto3.setup_default_session(region_name='us-east-1')
    sagemaker.Session(boto3.session.Session())

    hub = {
        'HF_TASK':'text-generation'
    }

    huggingface_model=HuggingFaceModel(
        model_data="s3://{model.tar.gz path}",
        transformers_version="4.28",
        pytorch_version="2.0",
        py_version="py310",
        env=hub,
        model_server_workers=1,
        role={sagemaker executor role arn}
    )

    from time import strftime
    current_time = strftime("%Y%m%d-%H%M")
    print(current_time)

    huggingface_model.deploy(
        initial_instance_count=1,
        instance_type="ml.g5.24xlarge",
        endpoint_name={endpoint_name},
    )
    ```
    
- Load endpoint and inference
    ```python=
    from sagemaker.predictor import Predictor
    import json
    session, bucket, role = prepare_session()
    payload = json.dumps({
        "inputs": "Can you please let us know about your hometown",
    })
    endpoint_name = '{위에서 생성한 endpoint name}'
    predictor = Predictor(
        endpoint_name=endpoint_name,
        sagemaker_session=session,
    )

    print(predictor.predict(payload, initial_args={'ContentType': 'application/json'}))
    ```

#### Inference LLM model w/ PEFT adaptor
- Load model
    ```python=
    def load_model(model_path, adaptor_path):
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForCausalLM.from_pretrained(model_path)
        model = PeftModel.from_pretrained(
                model,
                adaptor_path,
            )
        return tokenizer, model
    ```
- Inference
    ```python=
    def inference(tokenizer, model, text):
    input_ids = tokenizer.encode(text, return_tensors='pt')
    gen_ids = model.generate(input_ids=input_ids,
                             max_length=1024,
                             repetition_penalty=1.2,
                             pad_token_id=tokenizer.pad_token_id,
                             eos_token_id=tokenizer.eos_token_id,
                             use_cache=True)
    return tokenizer.decode(gen_ids[0])
    ```


### Endpoint configuration w/ inferentia
- Inferentia를 사용하는 이유는 간단함, 비용을 1/4까지 줄일 수 있음 [Ref](https://hyperconnect.github.io/2022/12/13/infra-cost-optimization-with-aws-inferentia.html)
- Supporting model: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/model-architecture-fit.html
- Decoder만 사용한 GPT기반 모델은 us-east-1 리전을 사용
- Inferentia instance type에 대한 service quota 신청
  7B 기준으로 ml.inf2.8xlarge ~ ml.inf2.24xlarge 정도면 될듯
- Inferentia를 사용하려면 모델을 NeuronSDK로 compile해서 deploy해야 됨
  이를 위해 EC2 ml.inf2.48xlarge 인스턴스까지 필요할 수 있음
  따라서 EC2 On-Demand Inf instances quota value를 196로 증설해 놓아야함