Революція в розподілених навантаженнях з JobSet на Kubernetes

Головна — News — Революція в розподілених навантаженнях з JobSet на Kubernetes

Автори: Даніель Вега-Майре (Google), Абдулла Гарайбе (Google), Кевін Хеннон (Red Hat)

Представляємо JobSet

JobSet — це інноваційний API з відкритим вихідним кодом, який спрощує представлення та оркестрацію розподілених завдань у Kubernetes. Цей інструмент розроблений для обслуговування як розподіленого навчання машин (ML), так і навантажень високопродуктивних обчислень (HPC). JobSet надає єдину модель API, що дозволяє здійснювати складні налаштування та оптимізації для масштабних контейнеризованих навантажень. У сучасну епоху прискорених досліджень у сфері штучного інтелекту та зростаючих вимог HPC, цей API пропонує вдосконалений підхід до управління розподіленими ресурсами та планування завдань на великих кластерах.

Чому JobSet?

Нещодавно спільнота Kubernetes покращила свою рідну екосистему пакетної обробки, привернувши увагу інженерів ML завдяки потужним можливостям планування та управління ресурсами. Зростання колосальних моделей ML — особливо великих мовних моделей (LLM) — підкреслило необхідність розподілу навантажень на десятки тисяч прискорювачів. Оскільки моделі часто охоплюють тисячі хостів, традиційні парадигми планування на одному хості більше не є достатніми.

Контейнеризований код навчання моделей, розгорнутий у Kubernetes, використовує примітиви розподіленої комунікації, такі як all-gather та all-reduce. Шаруючи параметри моделей і набори даних між вузлами, ці навантаження можуть ефективно синхронізувати градієнти та ділитися метаданими. Kubernetes відмінно справляється з плануванням і управлінням такими контейнеризованими додатками, забезпечуючи безперебійне управління життєвим циклом та підтримуючи визначені користувачем API, об’єкти та контролери для створення індивідуальних рішень для розподіленого навчання.

Незважаючи на досягнення у рідних Job API Kubernetes, існуючі примітиви та фрагментовані рішення для оркестрації навчання (наприклад, численні специфічні для фреймворків ресурси Kubeflow) не змогли встигнути за еволюцією розподілених навчальних парадигм. JobSet заповнює ці прогалини, спираючись на стабільний Job API, щоб надати додаткові абстракції, необхідні для сучасних випадків використання ML і HPC.

Як працює JobSet

JobSet моделює розподілене навантаження як групу завдань Kubernetes, дозволяючи користувачам вказувати кілька шаблонів подів для різних ролей (таких як робітник-драйвер, сервери параметрів тощо). Механізм базується на концепції ReplicatedJob. По суті, ReplicatedJob діє як шаблон завдання, доповнений бажаною кількістю реплік, що дозволяє декларативно створювати ідентичні дочірні завдання на різних островах прискорювачів. Це усуває необхідність у громіздкому скриптуванні або складних налаштуваннях Helm Chart під час генерації кількох варіантів одного й того ж завдання.

Додаткові можливості JobSet включають:

Репліковані Завдання: У великих дата-центрах спеціалізовані апаратні прискорювачі, такі як GPU та TPU, часто згруповані в специфічних топологічних доменах (або островах), з’єднаних з високошвидкісними мережами (наприклад, NVLink або ICI). JobSet дозволяє розподілити навантаження на кілька ідентичних дочірніх завдань, кожне з яких націлене на певний острів прискорювача, щоб комунікація між вузлами залишалася локальною та ефективною.

Автоматичне управління безголовими сервісами: Полегшуючи пряме спілкування між подами, JobSet автоматично створює та керує безголовими сервісами. Це покращує з’єднуваність, забезпечуючи, щоб поди могли спілкуватися, використовуючи стабільні імена хостів протягом усього життєвого циклу навантаження.

Конфігуровані політики успіху та невдачі: Адміністратори можуть визначати точні умови, за яких JobSet позначається як завершений — наприклад, лише після успішного завершення всіх подів в конкретному ReplicatedJob — або налаштувати максимальні обмеження на перезавантаження. Ці можливості забезпечують гнучкість для захисту довготривалих розподілених навчальних завдань, які поновлюються з останньої контрольної точки, якщо виникає помилкова ситуація.

Ексклюзивне розміщення за топологічним доменом: Щоб ще більше оптимізувати продуктивність, функції JobSet дозволяють вказувати 1:1 відповідність між дочірнім завданням і топологічним доменом, гарантуячи, що поди розміщені на одному фізичному або мережевому домені. Це особливо корисно для розподіленого навчання з паралельними даними, зменшуючи накладні витрати на міждоменну комунікацію, використовуючи при цьому високу пропускну здатність внутрішньо-острівкової з’єднаності.

Інтеграція з Kueue: JobSet безшовно інтегрується з Kueue, системою черг завдань, яка дозволяє перевантаження кластерів, чергування навантажень, запобігання блокування та багатокористувацькість. Це гарантує ефективне планування ресурсів та виконання навантажень у спільних середовищах.

Технічне поглиблення: Розподілена оркестрація та комунікація

З технічної точки зору одним із найзначніших внесків JobSet є його абстракція розподіленого навчального навантаження у кілька завдань Kubernetes. Кожен ReplicatedJob функціонує як повністю визначений шаблон завдання, що дозволяє тонко налаштовувати такі параметри, як:

Поди з кількома шаблонами: Різні поди в одному навантаженні можуть вимагати різних контейнерів, ресурсів або політик невдачі. Модель драйвер-робітник є яскравим прикладом, де драйвер може вимагати іншого життєвого циклу в порівнянні з робітниками.

Групи Завдань та обізнаність про топологію: У великих розгортаннях мережна топологія стає критичним фактором. Групи завдань дозволяють адміністраторам контролювати локальність подів, забезпечуючи, щоб завдання, що потребують інтенсивної комунікації, виконувалися через шляхи з низькою затримкою та високою пропускною здатністю, зменшуючи трафік між стійками або між дата-центрами.

Управління міжподовою комунікацією: Завдяки автоматичному розгортанню безголових сервісів JobSet забезпечує надійне встановлення з’єднань між подами. Це необхідно для колективних операцій, які лежать в основі таких комунікаційних фреймворків, як MPI, Ray або Spark, де синхронізація між подами є обов’язковою.

Оптимізація продуктивності та топологічні міркування

У сучасних умовах дата-центрів балансування затримки мережі та пропускної здатності є ключовим для ефективних розподілених обчислень. Дизайн JobSet дозволяє:

Оптимізації локальності: Завдяки відображенню дочірніх завдань на окремі острови прискорювачів (такі як конкретні стійки або частини TPU), фреймворк забезпечує, щоб більшість частих і малозатратних операцій виконувалися локально, з мінімальними залежностями від повільної мережі дата-центру.

Користувацькі обмеження ресурсів: Контролюючи вибір вузлів за допомогою параметрів nodeSelector та правил афінності, JobSet відповідає специфічним для апаратного забезпечення конфігураціям — будь то групування вузлів з NVLink-підключеними GPU або TPU подів, з’єднаними через мережі ICI.

Конфігурований рівень паралелізму та завершення: Тонке налаштування параметрів, таких як parallelism, completions та backoffLimit, дозволяє інженерам точно визначати стратегії виконання навантаження. Ці параметри є критично важливими в сценаріях, де мінімізація часу простою та прискорення відновлення після збоїв є пріоритетними.

Приклад використання: Розподілене навчання ML на частинах TPU з Jax

Розподілене навчання ML на кількох частинах TPU з Jax

Наступний приклад YAML демонструє специфікацію JobSet, призначену для запуску роботи з мультичастинами TPU на 4 частинах TPU v5e. Ця конфігурація використовує Jax через OpenXLA для ефективного, JIT-компільованого розподіленого навчання ML. Хоча приклад конкретно націлений на Jax, подібні конфігурації можуть бути досягнуті з використанням PyTorch/XLA або інших фреймворків, які підтримують обчислення на TPU.

# Запуск простого навантаження Jax на 
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice
  annotations:
    # Призначити кожному дочірньому завданню ексклюзивне використання частини TPU 
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: workers
    replicas: 4  # Одна репліка на частину TPU
    template:
      spec:
        parallelism: 2  # Кількість ВМ на частину TPU
        completions: 2  # Кількість завершень на основне навантаження
        backoffLimit: 0
        template:
          spec:
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
              cloud.google.com/gke-tpu-topology: 2x4
            containers:
            - name: jax-tpu
              image: python:3.8
              ports:
              - containerPort: 8471
              - containerPort: 8080
              securityContext:
                privileged: true
              command:
              - bash
              - -c
              - |
                pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                python -c 'import jax; print("Кількість глобальних пристроїв:", jax.device_count())'
                sleep 60
              resources:
                limits:
                  google.com/tpu: 4

Майбутні роботи, вплив на галузь та як долучитися

JobSet має активний план розвитку з численними запланованими функціями, спрямованими на подальше вдосконалення розподілених ML і HPC робочих процесів. Майбутні покращення включають вдосконалене управління гетерогенними навантаженнями, покращені функції моніторингу та журналювання, адаптовані для розподіленого навчання, а також глибшу інтеграцію з новими інструментами оркестрації контейнерів.

Триваюча співпраця між інженерами Google, Red Hat та внески з відкритої спільноти підкреслюють сильну галузеву підтримку JobSet. Його еволюція функцій очікується, що стимулюватиме інновації в таких областях, як стійкість до відмов, динамічне масштабування ресурсів та навіть автоматизовані системи відновлення для критичних виробничих середовищ.

Запрошуємо розробників, дослідників та експертів галузі долучитися до обговорення ідей, повідомляти про проблеми та допомагати розширювати документацію. Приєднуйтесь до обговорення на репозиторії JobSet, беріть участь у поштовій розсилці або спілкуйтесь на Slack.

Технічні думки експертів та перспективи екосистеми

Недавні обговорення серед експертів Kubernetes та хмарних обчислень підкреслюють JobSet як ключовий інструмент для подолання розриву між теоретичними моделями оркестрації та практичними реаліями розгортання. Архітектори наголошують на важливості налаштовуваного виділення ресурсів та планування з урахуванням топології для задоволення все більш чутливих до мережі та обчислювально інтенсивних навантажень. Зростаючі вимоги з боку досліджень у сфері ШІ та HPC роблять модульну та розширювану архітектуру JobSet добре підходящою для впливу на майбутні стратегії оркестрації та найкращі практики.

Крім того, кілька лідерів думок у цій галузі зазначили, що інтеграція JobSet з інструментами моніторингу та системами розподіленого трасування може забезпечити глибше розуміння продуктивності кластерів, потенційно сприяючи подальшим оптимізаціям як у програмному, так і в апаратному дизайні.

Висновок

JobSet є значним кроком уперед в оркестрації розподіленого навчання та HPC навантажень у Kubernetes. Завдяки вирішенню існуючих неефективностей планування та впровадженню вдосконалених параметрів конфігурації, він пропонує надійну платформу як для досліджень на ранніх стадіях, так і для великих виробничих розгортань. Ми закликаємо спільноту долучитися до його подальшої еволюції та допомогти сформувати майбутнє розподілених обчислень.

0 коментарів

Залишити відповідь