Масштабирование обучения с подкреплением с использованием GKE и Managed Lustre

1. Введение

Если вы предпочитаете запускать скрипты из пакета напрямую, без пошагового руководства, вы можете найти их в репозитории GoogleCloudPlatform/devrel-demos .

В этом практическом занятии вы узнаете, как развернуть высокопроизводительный конвейер обучения для обучения с подкреплением (Reinforcement Learning, RL) с использованием Google Kubernetes Engine (GKE) и Managed Lustre .

В задачах обучения с подкреплением, особенно в тех, которые используют такие алгоритмы, как групповая относительная оптимизация политики (GRPO), генерируются огромные объемы данных во время «генерации опыта» и требуют частого сохранения контрольных точек. Стандартное объектное хранилище может создавать узкие места во время этих всплесков операций ввода-вывода, в результате чего дорогостоящие ускорители простаивают.

Для устранения этих узких мест и достижения более высокой пропускной способности обучения вы будете использовать Managed Lustre — параллельную файловую систему.

Что вы будете делать

  • Настройте переменные среды для кластера Ray на базе графического процессора.
  • Создайте кластер Spot GPU в GKE и управляемый экземпляр Lustre с помощью Cluster Toolkit .
  • Разверните кластер KubeRay и смонтируйте файловую систему Lustre.
  • Отправьте данные о нагрузке для обучения в NeMo-RL .
  • Обеспечьте высокую пропускную способность и низкую задержку при создании контрольных точек с помощью облачного мониторинга.

Архитектурная схема GKE, KubeRay и Managed Lustre.

Что вам понадобится

  • Веб-браузер, например Chrome .
  • Проект Google Cloud с включенной функцией выставления счетов.

Данный практический семинар предназначен для опытных технических специалистов, инженеров платформ и исследователей в области искусственного интеллекта, знакомых с GKE и концепциями хранения данных.

Ориентировочная общая продолжительность: от 45 до 60 минут плюс 2 часа тренировочного времени.

2. Прежде чем начать

Создайте проект в Google Cloud.

  1. В консоли Google Cloud выберите или создайте проект Google Cloud .
  2. Убедитесь, что для вашего облачного проекта включена функция выставления счетов.

Запустить Cloud Shell

Cloud Shell — это среда командной строки, работающая в Google Cloud и поставляемая с предустановленными необходимыми инструментами.

  1. В верхней части консоли Google Cloud нажмите кнопку «Активировать Cloud Shell» .
  2. После подключения к Cloud Shell подтвердите свою аутентификацию: 
    gcloud auth list
  3. Убедитесь, что ваш проект настроен: 
    gcloud config get project
  4. Если параметры вашего проекта заданы не так, как ожидалось, настройте их следующим образом: 
    export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Установите Cluster Toolkit.

В этом практическом задании для развертывания кластера GKE используется Cluster Toolkit ( gcluster ). Инструкции по настройке Cluster Toolkit см. в руководстве по настройке Cluster Toolkit .

Включить API

Выполните эту команду в Cloud Shell, чтобы включить все необходимые API:

gcloud services enable \
  container.googleapis.com \
  lustre.googleapis.com \
  compute.googleapis.com \
  servicenetworking.googleapis.com

3. Настройка переменных среды

Для обеспечения единообразия команд в этом практическом задании настройте несколько переменных окружения.

Создайте файл с именем env.sh и заполните его вашей конфигурацией. Вы можете использовать следующий шаблон:

# Environment Variables for the RL Demo execution
export PROJECT_ID="{{'<var>'}}PROJECT_ID{{'</var>'}}"
export ZONE="us-east1-b"
export REGION="us-east1"
export CLUSTER_NAME="ray-a4-gpu-spot"
export HF_TOKEN="{{'<var>'}}YOUR_HF_TOKEN{{'</var>'}}" # Required for downloading models
export WANDB_API_KEY="{{'<var>'}}YOUR_WANDB_API_KEY{{'</var>'}}" # Optional

# Topology defaults
export NUM_NODES="8"
export GPUS_PER_NODE="8" # Fixed for A4/B200 architecture

Замените <YOUR_PROJECT_ID> и <YOUR_HF_TOKEN> на ваши фактические значения.

Подключите файл, чтобы загрузить переменные в текущую сессию:

source env.sh

4. Разверните кластер GKE и управляемый Lustre с помощью Cluster Toolkit.

На этом этапе вы используете Cluster Toolkit ( gcluster ) для развертывания кластера GKE с графическими процессорами Spot и автоматического выделения управляемого хранилища Lustre с драйвером Lustre CSI и предварительно настроенным PersistentVolumeClaim ( lustre-pvc ).

Подготовьте план.

Перед развертыванием ознакомьтесь с шаблоном examples/gke-a4/gke-a4.yaml (подробности см. в разделе «Создание кластера A4 »):

  1. Включить Managed Lustre : раскомментируйте разделы модуля managed-lustre и lustre-pvc в gke-a4.yaml .
  2. Включить дополнение RayOperator : установите enable_ray_operator: true в настройках модуля gke_cluster в gke-a4.yaml .

Развертывание инфраструктуры

Настройте разрешенный CIDR для доступа к Cloud Shell и выполните развертывание с помощью gcluster deploy :

export AUTHORIZED_CIDR="$(curl -s ifconfig.me)/32"

gcluster deploy examples/gke-a4/gke-a4.yaml \
  --vars project_id=${PROJECT_ID},deployment_name=${CLUSTER_NAME},region=${REGION},zone=${ZONE},static_node_count=${NUM_NODES},authorized_cidr=${AUTHORIZED_CIDR},spot=true

Дождитесь завершения развертывания. Cluster Toolkit автоматически предоставляет сеть VPC, пиринг частного доступа к сервисам (PSA), управляемую файловую систему Lustre, драйвер Lustre CSI, надстройку RayOperator и запрос на хранилище Kubernetes ( lustre-pvc ) в рамках одного скоординированного развертывания.

5. Разверните кластер Ray на GKE.

На этом этапе вы развернете кластер KubeRay на узлах GKE и смонтируете файловую систему Lustre, используя PersistentVolumeClaim ( lustre-pvc ), автоматически предоставляемый Cluster Toolkit.

Создание конфигурации RayCluster

Создайте файл с именем ray-cluster.yaml . В нем указываются головной и рабочий узлы KubeRay, используется тип ускорителя nvidia-b200 , а том Lustre монтируется в каталог /lustre .

cat << EOF > ray-cluster.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ${CLUSTER_NAME}
  namespace: default
spec:
  rayVersion: '2.54.0'
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-head
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          ports:
          - containerPort: 6379
            name: gcs-server
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "32"
              memory: "1000Gi"
            requests:
              cpu: "8"
              memory: "64Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
  workerGroupSpecs:
  - groupName: gpu-worker-group
    replicas: ${NUM_NODES}
    minReplicas: ${NUM_NODES}
    maxReplicas: ${NUM_NODES}
    rayStartParams: {}
    template:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-b200
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: ray-worker
          image: nvcr.io/nvidia/nemo-rl:v0.4.0
          resources:
            limits:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
            requests:
              nvidia.com/gpu: "8"
              cpu: "100"
              memory: "1000Gi"
          volumeMounts:
          - mountPath: /lustre
            name: lustre-storage
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: lustre-storage
          persistentVolumeClaim:
            claimName: lustre-pvc
        - name: dshm
          emptyDir:
            medium: Memory
EOF

Подключитесь к кластеру

Убедитесь, что ваша сессия Cloud Shell аутентифицирована в кластере GKE:

gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --region ${REGION} \
  --project ${PROJECT_ID}

Применить конфигурацию RayCluster

Примените конфигурацию кластера Ray: 

kubectl apply -f ray-cluster.yaml

Проверьте состояние кластера.

Отслеживайте создание подов: 

kubectl get pods -w

Дождитесь, пока головной и рабочий модули Running .

6. Предоставьте данные о нагрузке по обучению с подкреплением.

На этом этапе вы отправите задание на обучение NeMo-RL GRPO в свой кластер Ray.

Подключитесь к панели управления Ray.

Для отправки заданий и просмотра метрик необходимо подключиться к панели мониторинга Ray. Поскольку панель мониторинга находится в GKE, используйте переадресацию портов для доступа к ней из Cloud Shell: 

# Run this in a separate Cloud Shell tab or in the background
kubectl port-forward service/${CLUSTER_NAME}-head-svc 8265:8265 &

Создайте исполняемый скрипт.

Создайте файл с именем run_nemo_rl.sh . Этот скрипт будет выполняться на рабочих узлах кластера Ray. Мы используем cat << EOF для заполнения переменных окружения, которые вы установили ранее. 

cat << EOF > run_nemo_rl.sh
#!/bin/bash
set -ex

# Override job runtime conflicts (NeMo-RL passes os.environ to ray.init)
export RAY_OVERRIDE_JOB_RUNTIME_ENV=1

echo "--- Running on Ray Cluster ---"
cd /opt/nemo-rl

# Ensure directories exist on the high-speed Lustre drive
mkdir -p /lustre/huggingface_cache
mkdir -p /lustre/nemo_rl_qwen_72b_ds_cp

echo "Launching NeMo-RL GRPO training..."
uv run python examples/run_grpo_math.py \
  --config examples/configs/grpo_math_70B_megatron.yaml \
  policy.model_name='Qwen/Qwen2.5-72B-Instruct' \
  policy.megatron_cfg.converter_type='Qwen2ForCausalLM' \
  logger.wandb_enabled=False \
  cluster.num_nodes=${NUM_NODES} \
  cluster.gpus_per_node=${GPUS_PER_NODE} \
  logger.wandb.name='nemo-rl-grpo-test1' \
  grpo.max_num_steps=20 \
  grpo.num_generations_per_prompt=8 \
  grpo.num_prompts_per_step=32 \
  policy.train_global_batch_size=256 \
  checkpointing.enabled=True \
  checkpointing.save_period=2 \
  checkpointing.keep_top_k=2 \
  checkpointing.metric_name=null \
  checkpointing.checkpoint_dir=/lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 \
  data.dataset_name='DeepScaler'
EOF
chmod +x run_nemo_rl.sh

Создать файл игнорирования луча

Создайте файл .rayignore , чтобы предотвратить загрузку Ray больших или ненужных папок: 

cat << EOF > .rayignore
cluster-toolkit/
.git/
*.sh.log
EOF

Создание конфигурации среды выполнения

Создайте JSON-файл для передачи переменных окружения в задание Ray: 

cat << EOF > ray_runtime_env_nemo.json
{
  "env_vars": {
    "HF_TOKEN": "${HF_TOKEN}",
    "WANDB_API_KEY": "${WANDB_API_KEY}",
    "HF_HOME": "/lustre/huggingface_cache",
    "GLOO_SOCKET_IFNAME": "eth0",
    "NCCL_SOCKET_IFNAME": "eth0"
  }
}
EOF

Разместить вакансию

Используйте Ray CLI для отправки задания на конечную точку панели мониторинга. Если команда ray не найдена в Cloud Shell, вы можете установить ее с помощью pip install ray : 

ray job submit \
    --address="http://localhost:8265" \
    --working-dir . \
    --runtime-env ray_runtime_env_nemo.json \
    -- bash run_nemo_rl.sh

В терминале Cloud Shell вы увидите поток логов. Задача загрузит модель, инициализирует рабочие процессы Ray и запустит цикл обучения GRPO.

7. Мониторинг результатов тренировок.

На этом этапе вы сможете понаблюдать за производительностью файловой системы Lustre во время обучения и создания контрольных точек.

Проверьте журналы обучения

По мере прохождения обучения вы будете видеть в логах сообщения о сохранении контрольных точек в /lustre/nemo_rl_qwen_72b_ds_cp/nemo-rl-grpo-test1 . Обратите внимание, что создание контрольных точек происходит асинхронно и не блокирует работу рабочих процессов Ray на длительное время.

Чтобы оценить скорость создания контрольных точек, ищите в логах строки, указывающие на количество сохраненных контрольных точек.

Просматривайте показатели Lustre в облачной консоли.

Чтобы просмотреть метрики для вашего экземпляра Lustre:

  1. В консоли Google Cloud найдите «Управляемая служба для Lustre» .
  2. Щелкните по имени вашего экземпляра (например, ${CLUSTER_NAME}-lustre или rl-demo-gpu-lustre ).
  3. Перейдите на вкладку «Мониторинг» .

Здесь вы можете наблюдать:

  • Пропускная способность (байты/сек) : отслеживайте пики во время создания контрольных точек.
  • Пропускная способность : отслеживайте, сколько места занимают контрольно-пропускные пункты.

диаграмма характеристик блеска Lustre способен записывать данные с очень высокой скоростью, создавая контрольные точки за минимальное время.

8. Ресурсы для уборки

Выполните следующую команду в Cloud Shell, чтобы за один шаг уничтожить всю выделенную инфраструктуру (кластер GKE, пулы узлов GPU, управляемый экземпляр Lustre и сеть VPC): 

gcluster destroy "${CLUSTER_NAME}"

Эта команда выполняется синхронно в фоновом режиме, удаляя всю инфраструктуру, управляемую развертыванием, и выводя журналы выполнения в ваш терминал. Дождитесь полного завершения выполнения команды, прежде чем закрывать сессию Cloud Shell.

9. Поздравляем!

Вы успешно завершили практическое занятие по масштабируемому обучению с подкреплением с использованием GKE и Managed Lustre !

Что вы узнали

  • Как использовать Cluster Toolkit для развертывания кластера GKE с графическими процессорами, включающего спотовые экземпляры и управляемое хранилище Lustre.
  • Как развернуть кластер KubeRay и смонтировать хранилище Lustre.
  • Как отправить заявку на обучение в рамках программы NeMo-RL GRPO.
  • Как отслеживать производительность хранилища во время обучения.

Следующие шаги