torchrun 명령어를 활용한 PytorchJob 분산 데이터 병렬 처리 (with minGPT)

minGPT 소개

A PyTorch re-implementation of GPT, both training and inference

출처: https://github.com/karpathy/minGPT

 

PytorchJob

PytorchJob은 Kubeflow 프로젝트의 Training-operator 컴포넌트에서 Pytorch 병렬처리를 위한 CRD입니다.

(참고: https://www.kubeflow.org/docs/components/training/pytorch/)

저는 https://github.com/kubeflow/training-operator/tree/master/examples/pytorch 의 예제코드를 가지고 기초적인 실습을 진행했습니다.

하지만, LLM을 위한 Transformer, Huggingface 등의 라이브러리를 사용하면서 복잡성과 디테일이 필요해지면서 PytorchJob이 어떻게 동작하는지 살펴보았습니다.

 

PytorchJob Manifest

apiVersion: "kubeflow.org/v1"
kind: "PyTorchJob"
metadata:
  name: "pytorch-mingpt-nccl"
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          containers:
            - name: pytorch
              image: pytorch_mingpt:latest
              imagePullPolicy: Always
              command: ["torchrun"]
              args: ["--nnodes", $(WORLD_SIZE), "--nproc_per_node", "1", "--rdzv-backend", "c10d", "--rdzv-endpoint", $(MASTER_ADDR):$(MASTER_PORT), "/temp/minGPT-ddp/main.py"]
              resources:
                limits:
                  nvidia.com/gpu: 1
              volumeMounts:
              - name: share-volume
                mountPath: /temp
          volumes:
          - name: share-volume
            nfs:
              path: /my/nas/path
              server: my.nas.com
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          containers:
            - name: pytorch
              image: pytorch_mingpt:latest
              imagePullPolicy: Always
              command: ["torchrun"]
              args: ["--nnodes", $(WORLD_SIZE), "--nproc_per_node", "1", "--rdzv-backend", "c10d", "--rdzv-endpoint", $(MASTER_ADDR):$(MASTER_PORT), "/temp/minGPT-ddp/main.py"]
              resources:
                limits:
                  nvidia.com/gpu: 1
              volumeMounts:
              - name: share-volume
                mountPath: /temp
          volumes:
          - name: share-volume
            nfs:
              path: /my/nas/path
              server: my.nas.com

위와 같이 Master node 1개, Worker node 2개로 replica 설정 후 apply하였습니다.

 

Master node Pod

apiVersion: v1
kind: Pod
metadata:
  ...
spec:
  containers:
    ...
    env:
    - name: PYTHONUNBUFFERED
      value: "0"
    - name: MASTER_PORT
      value: "23456"
    - name: MASTER_ADDR
      value: pytorch-mingpt-nccl-master-0
    - name: WORLD_SIZE
      value: "3"
    - name: RANK
      value: "0"
...

Worker nodes Pod

apiVersion: v1
kind: Pod
metadata:
  ...
spec:
  containers:
    ...
    env:
    - name: PYTHONUNBUFFERED
      value: "0"
    - name: MASTER_PORT
      value: "23456"
    - name: MASTER_ADDR
      value: pytorch-mingpt-nccl-master-0
    - name: WORLD_SIZE
      value: "3"
    - name: RANK
      value: "0"
...

apiVersion: v1
kind: Pod
metadata:
  ...
spec:
  containers:
    ...
    env:
    - name: PYTHONUNBUFFERED
      value: "0"
    - name: MASTER_PORT
      value: "23456"
    - name: MASTER_ADDR
      value: pytorch-mingpt-nccl-master-0
    - name: WORLD_SIZE
      value: "3"
    - name: RANK
      value: "0"
...

각 Master, Worker node Pod의 spec.containers.env 필드 부분만 추출했습니다.

PytorchJob은 각 Pod에 위와 같이 환경변수를 자동으로 주입해주는 것을 확인할 수 있었습니다.

(추가적으로, Worker Pod에서는 Master Pod의 Service Domain으로 alive check 요청을 보내는 initContainer도 띄웁니다.)

따라서, 기존 분산학습돌릴때 직접 정해주던 rank, world_size 등을 직접 선언할 필요가 없어졌습니다.

 

torchrun

"How to install Torchrun?"

torchrun은 pytorch 1.10부터 사용 가능하며 이전버전에서는 python -m torch.distributed.run 명령어로 실행 가능합니다.

출처: https://discuss.pytorch.org/t/how-to-install-torchrun/135138

 

다음으로, 위 PytorchJob Manifest 에서의 Command와 args 부분에 대한 설명입니다.

 

command로는 torchrun이 들어가고,

args 부분에서 선언되는

• nnodes는 노드의 수를 설정하는 인수이며 env로 선언되는 WORLD_SIZE로부터 변수값으로 할당하기 위해  $(WORLD_SIZE)로 표기합니다.
• nproc_per_node는 각 노드에서 실행될 프로세스 수를 설정하는 인수입니다. 각 노드에서 GPU 1장을 사용하기 위해 1로 표기합니다.
• rdzv-backend는 분산 학습을 위한 백엔드를 설정하는 옵션입니다. 분산 학습에서 백엔드는 각 학습 프로세스 간에 상호 통신 및 조정을 담당합니다. 주로 사용되는 백엔드 중 하나는 "c10d"입니다. "c10d"는 PyTorch의 분산 백엔드로, 다수의 GPU 노드 간 효율적인 통신을 지원합니다. GPU 노드간 통신을 위해 c10d로 표기합니다.
  
• rdzv-endpoint는 분산 학습을 위한 백엔드와 상호 작용할 때 사용할 엔드포인트(endpoint)를 설정하는 옵션입니다. 이 엔드포인트는 백엔드에서 제공되며, 학습 프로세스 간에 서로 찾아갈 수 있게 해줍니다. 주로 사용되는 형식은 hostname:port이며, 이를 통해 분산 학습 프로세스가 서로 연결됩니다. env로부터 불러온 값을 통해 $(MASTER_ADDR):$(MASTER_PORT)로 표기합니다.
• 참고로 master_addr은 rdzv-backend가 static 이 아니거나 rdzv-endpoint가 선언되지 않았을때에만 사용됩니다. (master_addr is only used for static rdzv_backend and when rdzv_endpoint is not specified.)
• 여기서, rdzv(rendezvous)란 분산된 시스템 내의 모든 워커가 동기화되고 통신할 수 있게 하는 것입니다. (출처: https://pytorch.org/docs/stable/elastic/rendezvous.html)

torchrun reference: https://pytorch.org/docs/stable/elastic/run.html

 

학습에 사용되는 minGPT ddp 프로젝트는 https://github.com/pytorch/examples/tree/main/distributed/minGPT-ddp 로부터 불러왔고

os.environ["NCCL_DEBUG"] = "INFO"를 추가하여 디버깅 및 더많은 로깅 정보를 가져오게하여 학습결과를 아래와 같이 출력해본 결과

마스터 노드에서는 아래와 같은 출력이,

root@BD1DEV-L-KUBESPAWNER-MASTER-001:/home/hsd3030# k logs -n kubeflow pytorch-mingpt-nccl-master-0 --all-containers -f
master_addr is only used for static rdzv_backend and when rdzv_endpoint is not specified.
0
[2023-10-29 10:13:19,772][torch.distributed.distributed_c10d][INFO] - Added key: store_based_barrier_key:1 to store for rank: 0
[2023-10-29 10:13:19,773][torch.distributed.distributed_c10d][INFO] - Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 3 nodes.
init model
Data has 55769 characters, 59 unique.
number of parameters: 27.32M
Snapshot not found. Training model from scratch
pytorch-mingpt-nccl-master-0:10:10 [0] NCCL INFO Bootstrap : Using eth0:10.244.7.217<0>
pytorch-mingpt-nccl-master-0:10:10 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
pytorch-mingpt-nccl-master-0:10:10 [0] NCCL INFO cudaDriverVersion 12000
NCCL version 2.14.3+cuda11.7
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Failed to open libibverbs.so[.1]
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO NET/Socket : Using [0]eth0:10.244.7.217<0>
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Using network Socket
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 00/02 :    0   1   2
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 01/02 :    0   1   2
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Trees [0] 2/-1/-1->0->-1 [1] 2/-1/-1->0->1
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 00/0 : 2[3000] -> 0[3000] [receive] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 01/0 : 2[3000] -> 0[3000] [receive] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 00/0 : 0[3000] -> 1[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 01/0 : 0[3000] -> 1[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Connected all rings
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 01/0 : 1[3000] -> 0[3000] [receive] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 00/0 : 0[3000] -> 2[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Channel 01/0 : 0[3000] -> 2[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO Connected all trees
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO threadThresholds 8/8/64 | 24/8/64 | 512 | 512
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
pytorch-mingpt-nccl-master-0:10:31 [0] NCCL INFO comm 0x6fdef80 rank 0 nranks 3 cudaDev 0 busId 3000 - Init COMPLETE
train
[GPU0] Epoch 1 | Iter 0 | Train Loss 4.18325
[2023-10-29 10:13:23,283][torch.nn.parallel.distributed][INFO] - Reducer buckets have been rebuilt in this iteration.
[GPU0] Epoch 1 | Iter 0 | Eval Loss 2.32799
pytorch-mingpt-nccl-master-0:10:32 [0] NCCL INFO [Service thread] Connection closed by localRank 0
pytorch-mingpt-nccl-master-0:10:10 [0] NCCL INFO comm 0x6fdef80 rank 0 nranks 3 cudaDev 0 busId 3000 - Abort COMPLETE

 

워커노드에서는 아래와 같은 출력이 나오는 것을 확인할 수 있었습니다.

root@BD1DEV-L-KUBESPAWNER-MASTER-001:/home/hsd3030# k logs -n kubeflow pytorch-mingpt-nccl-worker-0 --all-containers -f
nslookup: can't resolve '(null)': Name does not resolve
nslookup: can't resolve 'pytorch-mingpt-nccl-master-0': Name does not resolve
waiting for master
nslookup: can't resolve '(null)': Name does not resolve
 
Name:      pytorch-mingpt-nccl-master-0
Address 1: 10.244.7.217 10-244-7-217.pytorch-mingpt-nccl-master-0.kubeflow.svc.cluster.local
master_addr is only used for static rdzv_backend and when rdzv_endpoint is not specified.
0
[2023-10-29 10:13:16,704][torch.distributed.distributed_c10d][INFO] - Added key: store_based_barrier_key:1 to store for rank: 1
[2023-10-29 10:13:17,697][torch.distributed.distributed_c10d][INFO] - Rank 1: Completed store-based barrier for key:store_based_barrier_key:1 with 3 nodes.
init model
Data has 55769 characters, 59 unique.
number of parameters: 27.32M
Snapshot not found. Training model from scratch
pytorch-mingpt-nccl-worker-0:9:9 [0] NCCL INFO cudaDriverVersion 11040
pytorch-mingpt-nccl-worker-0:9:9 [0] NCCL INFO Bootstrap : Using eth0:10.244.3.32<0>
pytorch-mingpt-nccl-worker-0:9:9 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Failed to open libibverbs.so[.1]
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO NET/Socket : Using [0]eth0:10.244.3.32<0>
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Using network Socket
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Trees [0] -1/-1/-1->1->2 [1] 0/-1/-1->1->-1
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Channel 00/0 : 0[3000] -> 1[3000] [receive] via NET/Socket/0
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Channel 01/0 : 0[3000] -> 1[3000] [receive] via NET/Socket/0
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Channel 00/0 : 1[3000] -> 2[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Channel 01/0 : 1[3000] -> 2[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Connected all rings
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Channel 00/0 : 2[3000] -> 1[3000] [receive] via NET/Socket/0
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Channel 01/0 : 1[3000] -> 0[3000] [send] via NET/Socket/0
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO Connected all trees
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO threadThresholds 8/8/64 | 24/8/64 | 512 | 512
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
pytorch-mingpt-nccl-worker-0:9:28 [0] NCCL INFO comm 0x64942c0 rank 1 nranks 3 cudaDev 0 busId 3000 - Init COMPLETE
train
[GPU1] Epoch 1 | Iter 0 | Train Loss 4.18533
[2023-10-29 10:13:21,202][torch.nn.parallel.distributed][INFO] - Reducer buckets have been rebuilt in this iteration.
[GPU1] Epoch 1 | Iter 0 | Eval Loss 2.31698
pytorch-mingpt-nccl-worker-0:9:29 [0] NCCL INFO [Service thread] Connection closed by localRank 0
pytorch-mingpt-nccl-worker-0:9:9 [0] NCCL INFO comm 0x64942c0 rank 1 nranks 3 cudaDev 0 busId 3000 - Abort COMPLETE

(Epoch은 1로 수정하였습니다.)

 

 

마스터 노드의 로그에서 주요하게 살펴봐야할 점은 아래와 같습니다.

• "Added key: store_based_barrier_key:1 to store for rank: 0" 및 "Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 3 nodes." 메시지는 분산 학습 프로세스 간의 동기화 및 조정을 보여줍니다.
• "Failed to open libibverbs.so[.1]"은 InfiniBand 네트워크 라이브러리 또는 디바이스를 사용하는 프로그램 또는 라이브러리에서 InfiniBand 인터넷 라이브러리 libibverbs를 찾지 못한 경우 발생합니다. InfiniBand 디바이스를 사용하지만, 일단 무시하도록 처리했습니다.
• "Connected all rings" 및 "Connected all trees"는 GPU 노드 간의 통신이 성공적으로 연결되었음을 나타냅니다.

 

결론

이로써, 위 모든 과정을 통해 Kubeflow Training-Operator를 통해 PytorchJob이 어떻게 동작하고 어떤 설정을 해주며 학습 코드와 실행 커맨드는 어떻게 지정해줘야 하는가에 대해 자세히 살펴볼 수 있었고, LLM 학습을 위해 여러 실험환경을 제공할 수 있게 되었습니다.