Ray를 사용하여 ML 워크로드 확장

이 문서에서는 TPU에서 Ray 및 JAX를 사용하여 머신러닝(ML) 워크로드를 실행하는 방법을 자세히 설명합니다. Ray에서 TPU를 사용하는 방법에는 기기 중심 모드 (PyTorch/XLA)와 호스트 중심 모드(JAX)라는 두 가지 모드가 있습니다.

이 문서에서는 이미 TPU 환경이 설정되어 있다고 가정합니다. 자세한 내용은 다음 리소스를 참조하세요.

Cloud TPU: Cloud TPU 환경 설정 및 TPU 리소스 관리
Google Kubernetes Engine(GKE): GKE Autopilot에 TPU 워크로드 배포 또는 GKE Standard에 TPU 워크로드 배포

기기 중심 모드(PyTorch/XLA)

기기 중심 모드는 기존 PyTorch의 프로그래매틱 스타일을 대부분 유지합니다. 이 모드에서 새로운 XLA 기기 유형을 추가하면 다른 PyTorch 기기처럼 작동합니다. 각 개별 프로세스는 하나의 XLA 기기와 상호작용합니다.

이 모드는 GPU를 사용하는 PyTorch에 이미 익숙하고 유사한 코딩 추상화를 사용하려는 경우에 적합합니다.

다음 섹션에서는 Ray를 사용하지 않고 하나 이상의 기기에서 PyTorch/XLA 워크로드를 실행하는 방법과 Ray를 사용하여 여러 호스트에서 동일한 워크로드를 실행하는 방법을 설명합니다.

TPU 만들기

TPU 생성 파라미터의 환경 변수를 만듭니다.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-8
export RUNTIME_VERSION=v2-alpha-tpuv5

환경 변수 설명

변수	설명
`PROJECT_ID`	Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
`TPU_NAME`	TPU의 이름입니다.
`ZONE`	TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
`ACCELERATOR_TYPE`	가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
`RUNTIME_VERSION`	Cloud TPU 소프트웨어 버전입니다.

다음 명령어를 사용하여 코어가 8개인 v5p TPU VM을 만듭니다.

gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE  \
   --version=$RUNTIME_VERSION

다음 명령어를 사용하여 TPU VM에 연결합니다.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

GKE를 사용하는 경우 설정 정보는 GKE 기반 KubeRay 가이드를 참조하세요.

설치 요구사항

TPU VM에서 다음 명령어를 실행하여 필요한 종속 항목을 설치합니다.

다음을 파일에 저장합니다. 예를 들면 requirements.txt입니다.

--find-links https://storage.googleapis.com/libtpu-releases/index.html
--find-links https://storage.googleapis.com/libtpu-wheels/index.html
torch~=2.6.0
torch_xla[tpu]~=2.6.0
ray[default]==2.40.0

필수 종속 항목을 설치하려면 다음을 실행합니다.
```
pip install -r requirements.txt
```

GKE에서 워크로드를 실행하는 경우 필요한 종속 항목을 설치하는 Dockerfile을 만드는 것이 좋습니다. 예시를 보려면 GKE 문서의 TPU 슬라이스 노드에서 워크로드 실행을 참조하세요.

단일 기기에서 PyTorch/XLA 워크로드 실행

다음 예시에서는 단일 기기, 즉 TPU 칩에서 XLA 텐서를 만드는 방법을 보여줍니다. 이는 PyTorch가 다른 기기 유형을 처리하는 방식과 유사합니다.

다음 코드 스니펫을 파일에 저장합니다. 예를 들면 workload.py입니다.
```
import torch
import torch_xla
import torch_xla.core.xla_model as xm

t = torch.randn(2, 2, device=xm.xla_device())
print(t.device)
print(t)
```
import torch_xla 가져오기 문은 PyTorch/XLA를 초기화하고 xm.xla_device() 함수는 현재 XLA 기기인 TPU 칩을 반환합니다.
PJRT_DEVICE 환경 변수를 TPU로 설정합니다.
```
export PJRT_DEVICE=TPU
```
스크립트를 실행합니다.
```
python workload.py
```
결과는 다음과 유사합니다. 출력에서 XLA 기기를 찾았다는 메시지가 표시되는지 확인합니다.
```
xla:0
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
```

여러 기기에서 PyTorch/XLA 실행

여러 기기에서 실행되도록 이전 섹션에서 코드 스니펫을 업데이트합니다.

import torch
import torch_xla
import torch_xla.core.xla_model as xm

def _mp_fn(index):
    t = torch.randn(2, 2, device=xm.xla_device())
    print(t.device)
    print(t)

if __name__ == '__main__':
    torch_xla.launch(_mp_fn, args=())

스크립트를 실행합니다.

python workload.py

TPU v5p-8에서 코드 스니펫을 실행하면 출력은 다음과 유사하게 표시됩니다.

xla:0
xla:0
xla:0
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')
xla:0
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')

torch_xla.launch()는 함수와 매개변수 목록이라는 두 가지 인수를 사용합니다. 사용 가능한 각 XLA 기기의 프로세스를 만들고 인수에 지정된 함수를 호출합니다. 이 예에서는 4개의 TPU 디바이스를 사용할 수 있으므로 torch_xla.launch()는 4개의 프로세스를 만들고 각 기기에서 _mp_fn()을 호출합니다. 각 프로세스는 하나의 기기에만 액세스할 수 있으므로 각 기기의 색인은 0이고 모든 프로세스에 xla:0이 출력됩니다.

Ray를 사용하여 여러 호스트에서 PyTorch/XLA 실행

다음 섹션에서는 더 큰 멀티 호스트 TPU 슬라이스에서 동일한 코드 스니펫을 실행하는 방법을 보여줍니다. 멀티 호스트 TPU 아키텍처에 대한 자세한 내용은 시스템 아키텍처를 참조하세요.

이 예시에서는 Ray를 수동으로 설정합니다. Ray 설정에 이미 익숙하다면 마지막 섹션인 Ray 워크로드 실행으로 건너뛰어도 됩니다. 프로덕션 환경에 Ray를 설정하는 방법에 관한 자세한 내용은 다음 리소스를 참조하세요.

멀티 호스트 TPU VM 만들기

TPU 생성 파라미터의 환경 변수를 만듭니다.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-16
export RUNTIME_VERSION=v2-alpha-tpuv5

환경 변수 설명

변수	설명
`PROJECT_ID`	Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
`TPU_NAME`	TPU의 이름입니다.
`ZONE`	TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
`ACCELERATOR_TYPE`	가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
`RUNTIME_VERSION`	Cloud TPU 소프트웨어 버전입니다.

다음 명령어를 사용하여 2개의 호스트(호스트당 4개의 TPU 칩이 있는 v5p-16)가 있는 멀티 호스트 TPU v5p를 만듭니다.

gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION

Ray 설정

TPU v5p-16에는 2개의 TPU 호스트가 있으며, 각 호스트에는 4개의 TPU 칩이 있습니다. 이 예시에서는 한 호스트에서 Ray 헤드 노드를 시작하고 두 번째 호스트를 Ray 클러스터에 워커 노드로 추가합니다.

SSH를 사용하여 첫 번째 호스트에 연결합니다.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0

설치 요구사항 섹션과 동일한 요구사항 파일로 종속 항목을 설치합니다.
```
pip install -r requirements.txt
```

Ray 프로세스를 시작합니다.

ray start --head --port=6379

결과는 다음과 유사합니다.

Enable usage stats collection? This prompt will auto-proceed in 10 seconds to avoid blocking cluster startup. Confirm [Y/n]: y
Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details.

Local node IP: 10.130.0.76

--------------------
Ray runtime started.
--------------------

Next steps
To add another node to this Ray cluster, run
    ray start --address='10.130.0.76:6379'

To connect to this Ray cluster:
    import ray
    ray.init()

To terminate the Ray runtime, run
    ray stop

To view the status of the cluster, use
    ray status

이제 이 TPU 호스트가 Ray 헤드 노드입니다. 다음과 같이 Ray 클러스터에 노드를 추가하는 방법을 보여주는 줄을 메모합니다.

To add another node to this Ray cluster, run
    ray start --address='10.130.0.76:6379'

나중에 이 명령어를 사용합니다.

Ray 클러스터 상태를 확인합니다.

ray status

결과는 다음과 유사합니다.

======== Autoscaler status: 2025-01-14 22:03:39.385610 ========
Node status
---------------------------------------------------------------
Active:
1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
Pending:
(no pending nodes)
Recent failures:
(no failures)

Resources
---------------------------------------------------------------
Usage:
0.0/208.0 CPU
0.0/4.0 TPU
0.0/1.0 TPU-v5p-16-head
0B/268.44GiB memory
0B/119.04GiB object_store_memory
0.0/1.0 your-tpu-name

Demands:
(no resource demands)

지금까지 헤드 노드만 추가했기 때문에 클러스터에는 4개의 TPU(0.0/4.0 TPU)만 포함되어 있습니다.

이제 헤드 노드가 실행 중이므로 클러스터에 두 번째 호스트를 추가할 수 있습니다.

SSH를 사용하여 두 번째 호스트에 연결합니다.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1

설치 요구사항 섹션과 동일한 요구사항 파일로 종속 항목을 설치합니다.
```
pip install -r requirements.txt
```

Ray 프로세스를 시작합니다. 이 노드를 기존 Ray 클러스터에 추가하려면 ray start 명령어의 출력에서 명령어를 사용합니다. 다음 명령어에서 IP 주소와 포트를 바꿔야 합니다.

ray start --address='10.130.0.76:6379'

결과는 다음과 유사합니다.

Local node IP: 10.130.0.80
[2025-01-14 22:30:07,397 W 75572 75572] global_state_accessor.cc:463: Retrying to get node with node ID 35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1

--------------------
Ray runtime started.
--------------------

To terminate the Ray runtime, run
ray stop

Ray 상태를 다시 확인합니다.

ray status

결과는 다음과 유사합니다.

======== Autoscaler status: 2025-01-14 22:45:21.485617 ========
Node status
---------------------------------------------------------------
Active:
1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
1 node_35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
Pending:
(no pending nodes)
Recent failures:
(no failures)

Resources
---------------------------------------------------------------
Usage:
0.0/416.0 CPU
0.0/8.0 TPU
0.0/1.0 TPU-v5p-16-head
0B/546.83GiB memory
0B/238.35GiB object_store_memory
0.0/2.0 your-tpu-name

Demands:
(no resource demands)

두 번째 TPU 호스트가 이제 클러스터의 노드가 됩니다. 사용 가능한 리소스 목록에 8개의 TPU(0.0/8.0 TPU)가 표시됩니다.

Ray 워크로드 실행

Ray 클러스터에서 실행되도록 코드 스니펫을 업데이트합니다.

import os
import torch
import torch_xla
import torch_xla.core.xla_model as xm
import ray

import torch.distributed as dist
import torch_xla.runtime as xr
from torch_xla._internal import pjrt

# Defines the local PJRT world size, the number of processes per host.
LOCAL_WORLD_SIZE = 4
# Defines the number of hosts in the Ray cluster.
NUM_OF_HOSTS = 4
GLOBAL_WORLD_SIZE = LOCAL_WORLD_SIZE * NUM_OF_HOSTS

def init_env():
    local_rank = int(os.environ['TPU_VISIBLE_CHIPS'])

    pjrt.initialize_multiprocess(local_rank, LOCAL_WORLD_SIZE)
    xr._init_world_size_ordinal()

# This decorator signals to Ray that the `print_tensor()` function should be run on a single TPU chip.
@ray.remote(resources={"TPU": 1})
def print_tensor():
    # Initializes the runtime environment on each Ray worker. Equivalent to
    # the `torch_xla.launch call` in the Run PyTorch/XLA on multiple devices section.
    init_env()

    t = torch.randn(2, 2, device=xm.xla_device())
    print(t.device)
    print(t)

ray.init()

# Uses Ray to dispatch the function call across available nodes in the cluster.
tasks = [print_tensor.remote() for _ in range(GLOBAL_WORLD_SIZE)]
ray.get(tasks)

ray.shutdown()

Ray 헤드 노드에서 스크립트를 실행합니다. ray-workload.py를 스크립트 경로로 바꿉니다.

python ray-workload.py

결과는 다음과 유사합니다.

WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
xla:0
xla:0
xla:0
xla:0
xla:0
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
xla:0
xla:0
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
xla:0
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')

출력은 멀티 호스트 TPU 슬라이스에 있는 각 XLA 기기(이 예에서는 8개 기기)에서 함수가 성공적으로 호출되었음을 나타냅니다.

호스트 중심 모드(JAX)

다음 섹션에서는 JAX를 사용한 호스트 중심 모드를 설명합니다. JAX는 함수형 프로그래밍 패러다임을 사용하며 상위 수준의 단일 프로그램 다중 데이터(SPMD) 시맨틱을 지원합니다. JAX 코드는 각 프로세스가 단일 XLA 기기와 상호작용하는 대신 단일 호스트의 여러 기기에서 동시에 작동하도록 설계되었습니다.

JAX는 고성능 컴퓨팅용으로 설계되었으며 대규모 학습 및 추론에 TPU를 효율적으로 활용할 수 있습니다. 함수형 프로그래밍 개념에 익숙하다면 이 모드는 JAX의 잠재력을 최대한 활용할 수 있는 이상적인 모드입니다.

이 안내에서는 JAX 및 기타 관련 패키지가 포함된 소프트웨어 환경을 비롯하여 Ray 및 TPU 환경이 이미 설정되어 있다고 가정합니다. Ray TPU 클러스터를 만들려면 KubeRay용 TPU가 있는 Google Cloud GKE 클러스터 시작의 안내를 따르세요. KubeRay에서 TPU를 사용하는 방법에 관한 자세한 내용은 KubeRay에서 TPU 사용을 참조하세요.

단일 호스트 TPU에서 JAX 워크로드 실행

다음 스크립트 예에서는 v6e-4와 같은 단일 호스트 TPU가 있는 Ray 클러스터에서 JAX 함수를 실행하는 방법을 보여줍니다. 멀티 호스트 TPU가 있는 경우 이 스크립트는 JAX의 멀티 컨트롤러 실행 모델로 인해 응답을 중지합니다. 멀티 호스트 TPU에서 Ray를 실행하는 방법에 관한 자세한 내용은 멀티 호스트 TPU에서 JAX 워크로드 실행을 참조하세요.

TPU 생성 파라미터의 환경 변수를 만듭니다.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-a
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e

환경 변수 설명

변수	설명
`PROJECT_ID`	Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
`TPU_NAME`	TPU의 이름입니다.
`ZONE`	TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
`ACCELERATOR_TYPE`	가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
`RUNTIME_VERSION`	Cloud TPU 소프트웨어 버전입니다.

다음 명령어를 사용하여 4개의 코어가 있는 v6e TPU VM을 만듭니다.

gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE  \
   --version=$RUNTIME_VERSION

다음 명령어를 사용하여 TPU VM에 연결합니다.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

TPU에 JAX 및 Ray를 설치합니다.

pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

다음 코드를 파일에 저장합니다. 예를 들면 ray-jax-single-host.py입니다.
```
import ray
import jax

@ray.remote(resources={"TPU": 4})
def my_function() -> int:
    return jax.device_count()

h = my_function.remote()
print(ray.get(h)) # => 4
```
GPU로 Ray를 실행하는 데 익숙하다면 TPU를 사용할 때 몇 가지 중요한 차이점이 있습니다.
- num_gpus를 설정하는 대신 TPU를 커스텀 리소스로 지정하고 TPU 칩 수를 설정합니다.
- Ray 워커 노드당 칩 수를 사용하여 TPU를 지정합니다. 예를 들어 v6e-4를 사용하는 경우 TPU를 4로 설정하여 원격 함수를 실행하면 전체 TPU 호스트가 사용됩니다.
- 이는 호스트당 하나의 프로세스로 GPU가 일반적으로 실행되는 방식과 다릅니다. TPU를 4가 아닌 숫자로 설정하는 것은 권장되지 않습니다.
  - 예외: 단일 호스트 v6e-8 또는 v5litepod-8인 경우 이 값을 8로 설정해야 합니다.
스크립트를 실행합니다.
```
python ray-jax-single-host.py
```

멀티 호스트 TPU에서 JAX 워크로드 실행

다음 예시 스크립트에서는 멀티 호스트 TPU가 있는 Ray 클러스터에서 JAX 함수를 실행하는 방법을 보여줍니다. 예시 스크립트는 v6e-16을 사용합니다.

TPU 생성 파라미터의 환경 변수를 만듭니다.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-a
export ACCELERATOR_TYPE=v6e-16
export RUNTIME_VERSION=v2-alpha-tpuv6e

환경 변수 설명

변수	설명
`PROJECT_ID`	Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
`TPU_NAME`	TPU의 이름입니다.
`ZONE`	TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
`ACCELERATOR_TYPE`	가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
`RUNTIME_VERSION`	Cloud TPU 소프트웨어 버전입니다.

다음 명령어를 사용하여 코어가 16개인 v6e TPU VM을 만듭니다.

gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE  \
   --version=$RUNTIME_VERSION

모든 TPU 워커에 JAX 및 Ray를 설치합니다.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --worker=all \
   --command="pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"

다음 코드를 파일에 저장합니다. 예를 들면 ray-jax-multi-host.py입니다.
```
import ray
import jax

@ray.remote(resources={"TPU": 4})
def my_function() -> int:
    return jax.device_count()

ray.init()
num_tpus = ray.available_resources()["TPU"]
num_hosts = int(num_tpus) # 4
h = [my_function.remote() for _ in range(num_hosts)]
print(ray.get(h)) # [16, 16, 16, 16]
```
GPU로 Ray를 실행하는 데 익숙하다면 TPU를 사용할 때 몇 가지 중요한 차이점이 있습니다.
- GPU의 PyTorch 워크로드와 유사합니다.
  - TPU의 JAX 워크로드는 멀티 컨트롤러, 단일 프로그램, 다중 데이터(SPMD) 방식으로 실행됩니다.
  - 기기 간의 집합은 머신러닝 프레임워크에서 처리합니다.
- GPU의 PyTorch 워크로드와 달리 JAX는 클러스터에서 사용 가능한 기기를 전체적으로 확인할 수 있습니다.

스크립트를 모든 TPU 워커에 복사합니다.

gcloud compute tpus tpu-vm scp ray-jax-multi-host.py $TPU_NAME: --zone=$ZONE --worker=all

스크립트를 실행합니다.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --worker=all \
   --command="python ray-jax-multi-host.py"

멀티슬라이스 JAX 워크로드 실행

멀티슬라이스를 사용하면 단일 TPU Pod 내에서 또는 데이터 센터 네트워크를 통해 여러 포드에서 여러 TPU 슬라이스에 걸쳐 있는 워크로드를 실행할 수 있습니다.

ray-tpu 패키지를 사용하여 Ray와 TPU 슬라이스의 상호작용을 간소화할 수 있습니다.

pip을 사용하여 ray-tpu을 설치합니다.

pip install ray-tpu

ray-tpu 패키지 사용에 관한 자세한 내용은 GitHub 저장소의 시작하기를 참조하세요. 멀티슬라이스를 사용하는 예는 멀티슬라이스에서 실행을 참조하세요.

Ray 및 MaxText를 사용하여 워크로드 조정

MaxText에서 Ray를 사용하는 방법에 관한 자세한 내용은 MaxText로 학습 작업 실행을 참조하세요.

TPU 및 Ray 리소스

Ray는 사용량의 차이를 수용하기 위해 TPU를 GPU와 다르게 취급합니다. 다음 예시에는 총 9개의 Ray 노드가 있습니다.

Ray 헤드 노드가 n1-standard-16 VM에서 실행 중입니다.
Ray 워커 노드는 2개의 v6e-16 TPU에서 실행됩니다. 각 TPU는 4명의 작업자로 구성됩니다.

$ ray status
======== Autoscaler status: 2024-10-17 09:30:00.854415 ========
Node status
---------------------------------------------------------------
Active:
 1 node_e54a65b81456cee40fcab16ce7b96f85406637eeb314517d9572dab2
 1 node_9a8931136f8d2ab905b07d23375768f41f27cc42f348e9f228dcb1a2
 1 node_c865cf8c0f7d03d4d6cae12781c68a840e113c6c9b8e26daeac23d63
 1 node_435b1f8f1fbcd6a4649c09690915b692a5bac468598e9049a2fac9f1
 1 node_3ed19176e9ecc2ac240c818eeb3bd4888fbc0812afebabd2d32f0a91
 1 node_6a88fe1b74f252a332b08da229781c3c62d8bf00a5ec2b90c0d9b867
 1 node_5ead13d0d60befd3a7081ef8b03ca0920834e5c25c376822b6307393
 1 node_b93cb79c06943c1beb155d421bbd895e161ba13bccf32128a9be901a
 1 node_9072795b8604ead901c5268ffcc8cc8602c662116ac0a0272a7c4e04
Pending:
 (no pending nodes)
Recent failures:
 (no failures)

Resources
---------------------------------------------------------------
Usage:
 0.0/727.0 CPU
 0.0/32.0 TPU
 0.0/2.0 TPU-v6e-16-head
 0B/5.13TiB memory
 0B/1.47TiB object_store_memory
 0.0/4.0 tpu-group-0
 0.0/4.0 tpu-group-1

Demands:
 (no resource demands)

리소스 사용량 필드 설명:

CPU: 클러스터에서 사용할 수 있는 총 CPU 수입니다.
TPU: 클러스터의 TPU 칩 수입니다.
TPU-v6e-16-head: TPU 슬라이스의 워커 0에 해당하는 리소스의 특별 식별자입니다. 이는 개별 TPU 슬라이스에 액세스하는 데 중요합니다.
memory: 애플리케이션에서 사용하는 작업자 힙 메모리입니다.
object_store_memory: 애플리케이션이 ray.put을 사용하여 객체 스토어에 객체를 만들 때와 원격 함수에서 값을 반환할 때 사용되는 메모리입니다.
tpu-group-0 및 tpu-group-1: 개별 TPU 슬라이스에 대한 고유 식별자입니다. 이는 슬라이스에서 작업을 실행하는 데 중요합니다. 이 필드는 v6e-16에서 TPU 슬라이스당 호스트가 4개 있으므로 4로 설정됩니다.