CPU

CPU

GPU를 쓸 수 없거나 비용 효율적인 구성을 원할 때 CPU 훈련이 잘 동작해요. 최신 Intel CPU는 CPU 백엔드를 위한 PyTorch의 AMP(Automatic Mixed Precision)로 bf16 혼합 정밀도 훈련을 지원해서, 메모리 사용을 줄이고 훈련을 빠르게 합니다.

출처: 문서

본문

단일 CPU가 너무 느리다면 CPU 훈련을 여러 소켓이나 노드로 확장할 수 있어요. 아래 예시는 세 가지 시나리오를 다룹니다.

  • 단일 CPU
  • 한 머신의 여러 프로세스(CPU 소켓당 프로세스 하나)
  • 여러 머신에 걸친 여러 프로세스

모든 분산 예시는 Intel oneAPI HPC Toolkit의 Intel MPI로 통신하고, Trainer와 함께 DDP 전략을 사용합니다.

Trainer는 CPU에서 bf16 혼합 정밀도 훈련을 지원합니다. CPU 훈련에서는 fp16보다 bf16을 선호하는데, 수치적으로 더 안정적이기 때문이에요. PyTorch의 CPU autocast를 켜려면 --bf16을 넘기고, CPU 훈련을 강제하려면 --use_cpu를 넘깁니다. 아래 예시는 run_qa.py 스크립트를 실행합니다.

python run_qa.py \
 --model_name_or_path google-bert/bert-base-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 12 \
 --learning_rate 3e-5 \
 --num_train_epochs 2 \
 --max_seq_length 384 \
 --doc_stride 128 \
 --output_dir /tmp/debug_squad/ \
 --bf16 \
 --use_cpu

같은 매개변수를 TrainingArguments에 직접 넘길 수도 있어요.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    use_cpu=True,
)

듀얼 소켓 CPU에서는 소켓당 프로세스 하나를 실행합니다. 각 소켓에 로컬인 메모리 접근을 유지하면 처리량이 향상돼요. 아래 예시는 소켓당 프로세스 하나로 단일 머신에서 두 프로세스를 실행합니다.

[!TIP] OMP_NUM_THREADS를 한 소켓의 물리적 코어 수에서 OS가 예약하는 코어 하나를 뺀 값으로 설정하세요. 예를 들어 24코어 소켓에서는 OMP_NUM_THREADS=23으로 설정합니다.

export MASTER_ADDR=127.0.0.1
mpirun -n 2 -genv OMP_NUM_THREADS=23 \
python3 run_qa.py \
 --model_name_or_path google-bert/bert-large-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 12 \
 --learning_rate 3e-5 \
 --num_train_epochs 2 \
 --max_seq_length 384 \
 --doc_stride 128 \
 --output_dir /tmp/debug_squad/

각 노드의 IP 주소를 나열한 hostfile을 사용해 두 Xeon 머신(node0과 node1)에 걸쳐 훈련을 네 프로세스로 확장합니다. -n 4 플래그는 총 프로세스 수를 설정합니다. -ppn 2는 노드당 두 프로세스를 설정합니다(소켓당 하나).

메인 프로세스 역할을 하는 node0에서 이 스크립트를 실행하세요.

[!TIP] OMP_NUM_THREADS를 한 소켓의 물리적 코어 수에서 OS가 예약하는 코어 하나를 뺀 값으로 설정하세요. 예를 들어 24코어 소켓에서는 OMP_NUM_THREADS=23으로 설정합니다.

각 노드의 IP 주소가 담긴 hostfile을 만듭니다.

cat hostfile
xxx.xxx.xxx.xxx #node0 ip
xxx.xxx.xxx.xxx #node1 ip

그런 다음 훈련 스크립트를 실행합니다.

export MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip
mpirun -f hostfile -n 4 -ppn 2 \
 -genv OMP_NUM_THREADS=23 \
python3 run_qa.py \
 --model_name_or_path google-bert/bert-large-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 12 \
 --learning_rate 3e-5 \
 --num_train_epochs 2 \
 --max_seq_length 384 \
 --doc_stride 128 \
 --output_dir /tmp/debug_squad/ \
 --use_cpu \
 --bf16

Kubernetes

분산 CPU 훈련은 PyTorchJob을 사용해 Kubernetes 클러스터에서도 실행할 수 있어요.

배포하기 전에 다음 설정을 완료하세요.

  1. Kubeflow가 설치된 Kubernetes 클러스터에 접근할 수 있는지 확인합니다.
  2. 클러스터와 상호작용하도록 kubectl을 설치하고 설정합니다.
  3. 데이터셋과 모델 파일을 저장할 PersistentVolumeClaim (PVC)을 설정합니다.
  4. 훈련 스크립트와 의존성을 위한 Docker 이미지를 빌드합니다.

아래 예시 Dockerfile은 멀티노드 CPU 훈련을 위한 MPI 지원을 포함한 Intel 최적화 PyTorch 베이스 이미지에서 시작합니다. 성능 라이브러리 두 개도 설치합니다.

  • google-perftools (libtcmalloc): 기본 시스템 할당자와 비교해 할당 오버헤드를 줄이는 메모리 할당자.
  • libomp-dev (libiomp5): GNU OpenMP 기본값보다 더 나은 CPU 스레드 관리를 제공하는 Intel의 OpenMP 런타임.
FROM intel/intel-optimized-pytorch:2.4.0-pip-multinode

RUN apt-get update -y && \
    apt-get install -y --no-install-recommends --fix-missing \
    google-perftools \
    libomp-dev

WORKDIR /workspace

# Download and extract the transformers code
ARG HF_TRANSFORMERS_VER="4.46.0"
RUN pip install --no-cache-dir \
    transformers==${HF_TRANSFORMERS_VER} && \
    mkdir transformers && \
    curl -sSL --retry 5 https://github.com/huggingface/transformers/archive/refs/tags/v${HF_TRANSFORMERS_VER}.tar.gz | tar -C transformers --strip-components=1 -xzf -

배포하기 전에 이미지를 빌드하고 클러스터 노드에서 접근 가능한 레지스트리에 푸시하세요.

PyTorchJob

PyTorchJob은 PyTorch 분산 훈련 작업을 관리하는 Kubernetes 커스텀 리소스입니다. 워커 파드 생애주기, 재시작 정책, 프로세스 조정을 처리합니다. 훈련 스크립트는 모델과 데이터만 처리하면 돼요.

아래 예시 yaml 파일은 bf16을 켠 채 run_qa.py 스크립트를 실행하는 워커 네 개를 설정합니다.

CPU 리소스 limits와 requests를 설정할 때는 한 단위가 물리적 코어 또는 가상 코어 하나와 같은 CPU 단위를 사용하세요. Guaranteed quality of service를 위해 limits와 requests를 같은 값으로 설정합니다. kubelet과 시스템 프로세스를 위해 일부 코어는 할당하지 않은 채로 두세요. OMP_NUM_THREADS를 할당된 CPU 단위 수와 일치하도록 설정하면 PyTorch가 사용 가능한 모든 코어를 사용합니다.

훈련 스크립트와 클러스터의 노드 수에 따라 yaml 파일을 조정하세요.

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: transformers-pytorchjob
spec:
  elasticPolicy:
    rdzvBackend: c10d
    minReplicas: 1
    maxReplicas: 4
    maxRestarts: 10
  pytorchReplicaSpecs:
    Worker:
      replicas: 4  # The number of worker pods
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: <image name>:<tag>  # Specify the docker image to use for the worker pods
              imagePullPolicy: IfNotPresent
              command: ["/bin/bash", "-c"]
              args:
                - >-
                  cd /workspace/transformers;
                  pip install -r /workspace/transformers/examples/pytorch/question-answering/requirements.txt;
                  torchrun /workspace/transformers/examples/pytorch/question-answering/run_qa.py \
                    --model_name_or_path distilbert/distilbert-base-uncased \
                    --dataset_name squad \
                    --do_train \
                    --do_eval \
                    --per_device_train_batch_size 12 \
                    --learning_rate 3e-5 \
                    --num_train_epochs 2 \
                    --max_seq_length 384 \
                    --doc_stride 128 \
                    --output_dir /tmp/pvc-mount/output_$(date +%Y%m%d_%H%M%S) \
                    --bf16;
              env:
              - name: LD_PRELOAD
                value: "/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4.5.9:/usr/local/lib/libiomp5.so"
              - name: HF_HUB_CACHE
                value: "/tmp/pvc-mount/hub_cache"
              - name: HF_DATASETS_CACHE
                value: "/tmp/pvc-mount/hf_datasets_cache"
              - name: LOGLEVEL
                value: "INFO"
              - name: OMP_NUM_THREADS  # Set to match the number of allocated CPU units
                value: "240"
              resources:
                limits:
                  cpu: 240  # Update the CPU and memory limit values based on your nodes
                  memory: 128Gi
                requests:
                  cpu: 240  # Update the CPU and memory request values based on your nodes
                  memory: 128Gi
              volumeMounts:
              - name: pvc-volume
                mountPath: /tmp/pvc-mount
              - mountPath: /dev/shm
                name: dshm
          restartPolicy: Never
          nodeSelector:  # Optionally use nodeSelector to match a certain node label for the worker pods
            node-type: gnr
          volumes:
          - name: pvc-volume
            persistentVolumeClaim:
              claimName: transformers-pvc
          - name: dshm
            emptyDir:
              medium: Memory

배포 (Deploy)

다음 명령으로 PyTorchJob을 클러스터에 배포합니다.

export NAMESPACE=<specify your namespace>

kubectl create -f pytorchjob.yaml -n ${NAMESPACE}

상태를 모니터링하려면 네임스페이스의 파드를 나열하세요. 컨테이너 이미지를 가져오는 동안 파드는 Pending으로 시작했다가 Running으로 전이합니다.

kubectl get pods -n ${NAMESPACE}

NAME                                                     READY   STATUS                  RESTARTS          AGE
...
transformers-pytorchjob-worker-0                         1/1     Running                 0                 7m37s
transformers-pytorchjob-worker-1                         1/1     Running                 0                 7m37s
transformers-pytorchjob-worker-2                         1/1     Running                 0                 7m37s
transformers-pytorchjob-worker-3                         1/1     Running                 0                 7m37s
...

워커 파드의 로그를 스트리밍해 훈련 진행 상황을 따라가세요.

kubectl logs transformers-pytorchjob-worker-0 -n ${NAMESPACE} -f

훈련이 끝나면 PVC 또는 저장 위치에서 훈련된 모델을 복사합니다. 그런 다음 PyTorchJob 리소스를 삭제해요.

kubectl delete -f pytorchjob.yaml -n ${NAMESPACE}

다음 단계 (Next steps)

더 알아보기 (Learn more)