利用 Mountpoint for S3 在 Kubernetes 上加速 LLM 的训练

{
   "Version": "2012-10-17",
   "Statement": [
        {
            "Sid": "MountpointFullBucketAccess",
            "Effect": "Allow",
            "Action": [
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::DOC-EXAMPLE-BUCKET1"
            ]
        },
        {
            "Sid": "MountpointFullObjectAccess",
            "Effect": "Allow",
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:AbortMultipartUpload",
                "s3:DeleteObject"
            ],
            "Resource": [
                "arn:aws:s3:::DOC-EXAMPLE-BUCKET1/*"
            ]
        }
   ]
}

CLUSTER_NAME=my-cluster
REGION=region-code
ROLE_NAME=AmazonEKS_S3_CSI_DriverRole
ROLE_ARN=AmazonS3CSIDriverPolicy
eksctl create iamserviceaccount \
    --name s3-csi-driver-sa \
    --namespace kube-system \
    --cluster $CLUSTER_NAME \
    --attach-policy-arn $ROLE_ARN \
    --approve
    --role-name $ROLE_NAME
    --region $REGION

kubectl apply -k "github.com/awslabs/mountpoint-s3-csi-driver/deploy/kubernetes/overlays/stable/"

kubectl get pods -n kube-system -l app.kubernetes.io/name=aws-mountpoint-s3-csi-driver

apiVersion: v1
kind: PersistentVolume
metadata:
  name: s3-pv
spec:
  capacity:
    storage: 1200Gi # ignored, required
  accessModes:
    - ReadWriteMany # supported options: ReadWriteMany / ReadOnlyMany
  mountOptions:
    - allow-delete
    - region us-west-2
  csi:
    driver: s3.csi.aws.com # required
    volumeHandle: s3-csi-driver-volume
    volumeAttributes:
      bucketName: YOU-BUCKET-NAME

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: s3-claim
spec:
  accessModes:
    - ReadWriteMany # supported options: ReadWriteMany / ReadOnlyMany
  storageClassName: "" # required for static provisioning
  resources:
    requests:
      storage: 1200Gi # ignored, required
  volumeName: s3-pv

aws s3 cp examples/scripts/datasets/openassistant-guanaco-train.arrow s3://YOUR-BUCKET_NAME

FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.0.1-gpu-py310-cu118-ubuntu20.04-ec2 as base
WORKDIR /llama2
RUN pip install trl
RUN git clone https://github.com/lvwerra/trl
WORKDIR /llama2/trl/
RUN pip install -r requirements.txt 
RUN sed -i 's|dataset = load_dataset(script_args.dataset_name, split="train")|dataset = load_dataset("arrow", data_files="/mount-s3/openassistant-guanaco-train.arrow",split="train")|' examples/scripts/sft.py

aws ecr get-login-password --region us-east-1 |sudo docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com

sudo docker build -t public.ecr.aws/h6r2a7o6/emr:llama2_finetune-v4 .

aws ecr-public get-login-password --region us-east-1 | sudo docker login --username AWS --password-stdin public.ecr.aws/h6r2a7o6

sudo docker push  public.ecr.aws/h6r2a7o6/emr:llama2_finetune-v4

apiVersion: v1
kind: Pod
metadata:
  name: llama2-pod
spec:
  containers:
  - name: app
    image: public.ecr.aws/h6r2a7o6/emr:llama2_finetune-v4
    command: ["/bin/sh"]
    args: ["-c",  "python3 examples/scripts/sft.py --model_name meta-llama/Llama-2-7b-hf --dataset_name timdettmers/openassistant-guanaco --load_in_4bit --use_peft --batch_size 1 --gradient_accumulation_steps 2"]
    volumeMounts:
    - name: persistent-storage
      mountPath: /mount-s3
  volumes:
  - name: persistent-storage
    persistentVolumeClaim:
      claimName: s3-claim

k apply -f llama2_finetuning.yaml；k get pods

k logs -f llama2-pod
The cache for model files in Transformers v4.22.0 has been updated. Migrating your old cache. This is a one-time only operation. You can interrupt this and resume the migration later on by calling `transformers.utils.move_cache()`.
0it [00:00, ?it/s]
/opt/conda/lib/python3.10/site-packages/trl/trainer/ppo_config.py:141: UserWarning: The `optimize_cuda_cache` arguement will be deprecated soon, please use `optimize_device_cache` instead.
  warnings.warn(
Detected kernel version 5.4.253, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.
/opt/conda/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:472: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.
  warnings.warn(
Loading checkpoint shards: 100%|██████████| 2/2 [00:07<00:00,  3.94s/it]
/opt/conda/lib/python3.10/site-packages/transformers/utils/hub.py:374: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.
  warnings.warn(
Downloading data files: 100%|██████████| 1/1 [00:00<00:00, 15650.39it/s]
Extracting data files: 100%|██████████| 1/1 [00:01<00:00,  1.06s/it]
Generating train split: 9846 examples [00:00, 19719.37 examples/s]
Map: 100%|██████████| 9846/9846 [00:02<00:00, 4094.10 examples/s]
Detected kernel version 5.4.253, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.
  0%|          | 0/14769 [00:00<?, ?it/s]You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
{'loss': 2.1358, 'learning_rate': 1.4099045297582776e-05, 'epoch': 0.0}
{'loss': 1.9417, 'learning_rate': 1.409809059516555e-05, 'epoch': 0.0}
{'loss': 1.2295, 'learning_rate': 1.4097135892748325e-05, 'epoch': 0.0}
{'loss': 2.2918, 'learning_rate': 1.4096181190331099e-05, 'epoch': 0.0}