Présentation des parcours sur Cloud

Pathways est un système conçu pour permettre la création de systèmes de machine learning à grande échelle, multitâches et à activation partielle. Il permet d'utiliser des milliers ou des dizaines de milliers d'accélérateurs, avec la possibilité d'allouer dynamiquement des quantités variables de puissance de calcul pour différentes tâches en fonction de leurs besoins de traitement.

Pathways simplifie les calculs de machine learning à grande échelle en permettant à un seul client JAX d'orchestrer des charges de travail sur plusieurs grandes tranches de TPU, pouvant potentiellement s'étendre sur des milliers de puces TPU.

Pathways est utilisé en interne chez Google pour entraîner de grands modèles comme Gemini. Pathways on Cloud offre les mêmes avantages aux clients Google Cloud .

Avant de commencer

Vérifiez que vous disposez bien des éléments suivants :

Ce document explique comment utiliser les TPU gérés Pathways sur Google Kubernetes Engine (GKE) pour les charges de travail par lot, en temps réel et interactives. Il part du principe que vous savez déjà comment utiliser les TPU avec GKE, y compris les TPU multislices sur Google Kubernetes Engine, et que vous avez une expérience générale avec les TPU multislices.

Contrôleur unique et multicontrôleur

Il existe principalement deux façons de gérer et d'orchestrer les calculs sur plusieurs appareils :

Fonctionnalité	Manette unique (Parcours)	Plusieurs contrôleurs (par défaut dans JAX)
Contrôle	Point de contrôle unique : un seul programme client sert de contrôleur central.	Contrôle distribué : plusieurs processus participent, chacun avec sa propre instance d'interpréteur Python.
Afficher	Vue unifiée : le client voit tous les appareils comme un seul système unifié.	Vue localisée : chaque processus Python ne voit que les appareils qui lui sont connectés.
Programmation	Programmation simplifiée : les utilisateurs interagissent avec un seul client, ce qui fait apparaître le système comme une seule grande machine avec de nombreux accélérateurs locaux.	SPMD : utilise principalement le paradigme SPMD, qui exige que tous les appareils exécutent le même programme.
Flexibilité	Il est compatible avec des schémas de calcul plus complexes que SPMD, y compris le parallélisme de pipeline asymétrique et la parcimonie de calcul.	Peut être moins flexible dans la gestion des ressources, en particulier sur différentes tranches de TPU.

Composants des parcours

La section suivante décrit les principaux composants de l'architecture Pathways.

Gestionnaire de ressources Pathways

Il s'agit du plan de contrôle central du système Pathways. Il gère toutes les ressources d'accélérateur et est responsable de la coordination de l'allocation des accélérateurs pour les jobs utilisateur. Il surveille l'état des nœuds de calcul et gère la planification, la mise en veille et la reprise des tâches. Il sert de point de contact unique pour les erreurs et l'état du système. Ce composant ne nécessite que des ressources de processeur.

Client Pathways

Il s'agit d'une implémentation du runtime du framework intermédiaire (IFRT, Interim Framework Runtime) qui sert de point d'entrée dans le système Pathways. Il reçoit les opérations de haut niveau (HLO) de votre programme. Le client Pathways est chargé de la coordination avec le gestionnaire de ressources Pathways pour déterminer où placer les programmes compilés pour l'exécution en fonction du code utilisateur. Il présente une vue unifiée du système à un client JAX donné. Ce composant ne nécessite que des ressources de processeur.

Employé Pathways

Il s'agit des processus qui s'exécutent sur les machines d'accélérateur (VM TPU). Ils reçoivent les exécutables compilés de votre programme à partir du serveur proxy IFRT et effectuent les calculs sur les TPU. Les employés Pathways renvoient les données à votre programme via le serveur proxy IFRT. Ce composant nécessite des ressources d'accélérateur.

Client proxy IFRT

Il s'agit d'une implémentation OSS de l'API Interim Framework Runtime (IFRT) qui découple le code utilisateur du runtime sous-jacent et améliore la portabilité et la transparence du code. JAX utilise cette implémentation comme alternative à son environnement d'exécution multi-contrôleur par défaut. Le client proxy IFRT sert de pont de communication entre votre programme et les composants Pathways. Il envoie des requêtes au serveur proxy IFRT et reçoit les résultats. Il s'agit d'une implémentation OSS de l'API IFRT. Ce composant ne nécessite que des ressources de processeur.

Serveur proxy IFRT

Ce serveur gRPC reçoit les requêtes du client proxy IFRT et les transmet au client Pathways, qui gère la distribution réelle du travail. Ce composant ne nécessite que des ressources de processeur.

Serveur side-car

Ce serveur gRPC est colocalisé avec le nœud de calcul Pathways sur la VM de l'accélérateur pour exécuter le code Python spécifié par l'utilisateur directement sur la VM de l'accélérateur. Cela permet de réduire la latence de transfert de données du contrôleur vers les accélérateurs. Le serveur side-car interagit avec le worker Pathways via un protocole personnalisé versionné sur le transport gRPC.

Affiche la relation entre les composants Pathways. — Composants des parcours

API PathwaysJob

L'API PathwaysJob est une API Kubernetes native OSS que vous utilisez pour déployer des charges de travail d'entraînement de ML et d'inférence par lot. Le contrôleur de PathwaysJob utilise l'API JobSet pour gérer le cycle de vie et la coordination de tous les composants Pathways. Cette définition de ressource personnalisée (CRD, Custom Resource Definition) vous offre une interface de haut niveau pour définir vos charges de travail Pathways, en vous évitant d'avoir à gérer directement les spécifications de chaque pod pour les scénarios courants. Pour obtenir une liste complète de tous les paramètres et de leur signification spécifique, consultez la documentation de l'API PathwaysJob sur GitHub.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

Le tableau suivant décrit les paramètres de l'API PathwaysJob :

Attribut	Description
`apiVersion`	Spécifie la version de l'API pour la définition de ressource personnalisée (CRD) `PathwaysJob` : pathways-job.pathways.domain/v1.
`kind`	Identifie l'objet Kubernetes comme un `PathwaysJob`.
`metadata.name`	Nom de l'objet `PathwaysJob` dans Kubernetes, qui suit généralement le modèle pathways-.
`spec`	Définit l'état et la configuration souhaités pour `PathwaysJob`.
`spec.maxRestarts`	Nombre maximal de fois où `PathwaysJob` peut être redémarré automatiquement par le système en cas d'échec.
`spec.pathwaysVersion`	(Facultatif) Spécifie la version souhaitée du framework JAX à utiliser dans l'environnement Pathways pour ce job (par exemple, jax-0.5.3).
`spec.workers`	Tableau définissant la configuration du pool de nœuds de calcul de `PathwaysJob`, qui utilise généralement des ressources TPU.
`spec.workers[].type`	Type de machine TPU à utiliser pour les nœuds de calcul (par exemple, $TPU_MACHINE_TYPE peut être ct6e-standard-4t)
`spec.workers[].topology`	Topologie des tranches de TPU attribuées aux nœuds de calcul (par exemple, $TOPOLOGY peut être 2x2, 4x4 ou 2x2x2).
`spec.workers[].numSlices`	Nombre de tranches TPU à provisionner pour le pool de nœuds de calcul (par exemple, $WORKLOAD_NODEPOOL_COUNT peut être défini sur 2).
`spec.workers[].maxSliceRestarts`	(Facultatif) Nombre maximal de fois qu'un nœud de calcul individuel d'un segment peut être redémarré en cas d'échec.
`spec.customComponents`	(Facultatif) Tableau qui vous permet de définir et de déployer des composants personnalisés (tels que des serveurs proxy, des serveurs Pathways ou des nœuds de calcul supplémentaires) en même temps que le job principal.
`spec.customComponents[].componentType`	Spécifie le type de composant personnalisé défini (par exemple, proxy_server, pathways_server, worker, colocated_python_sidecar).
`spec.customComponents[].image`	Image Docker à utiliser pour le conteneur de ce composant personnalisé.
`spec.customComponents[].customFlags`	Tableau d'indicateurs de ligne de commande personnalisés qui seront transmis au conteneur à son démarrage.
`spec.customComponents[].customEnv`	Tableau de variables d'environnement personnalisées à définir dans le conteneur. Chaque élément possède un nom et une valeur.
`spec.pathwaysDir`	Bucket Cloud Storage utilisé par `PathwaysJob` pour stocker les artefacts de compilation et d'autres données temporaires. Vous devez créer ce bucket avant d'exécuter votre charge de travail.
`spec.controller`	Paramètres de configuration du contrôleur Pathways, qui gère l'exécution globale du job.
`spec.controller.deploymentMode`	Spécifie la façon dont les ressources de processeur du contrôleur Pathways (gestionnaire de ressources Pathways et serveur proxy) sont déployées. Le mode par défaut les déploie sur un nœud de processeur dédié, tandis que colocate_head_with_workers les déploie à côté d'un nœud de calcul TPU.
`spec.controller.elasticSlices`	(Facultatif) Nombre maximal de tranches de TPU pouvant devenir indisponibles pendant l'exécution du job avant d'être considérées comme non opérationnelles.
`spec.controller.template`	(Facultatif) Définit le modèle de pod pour le job utilisateur. Cette opération est obligatoire pour les charges de travail par lot, mais pas pour les charges de travail interactives.
`spec.controller.template.spec`	Spécification du pod pour le job utilisateur.
`spec.controller.template.spec.containers`	Tableau définissant les conteneurs qui s'exécuteront dans le job utilisateur
`spec.controller.template.spec.containers[].name`	Nom du conteneur dans le job utilisateur (dans cet exemple, il s'agit de "main").
`spec.controller.template.spec.containers[].image`	Image Docker à utiliser pour le conteneur dans le conteneur principal (dans cet exemple, il s'agit de python:3.11).
`spec.controller.template.spec.containers[].command`	Commande à exécuter au démarrage du conteneur principal. Dans cet exemple, il installe `pathwaysutils`, initialise Pathways et affiche les appareils JAX.

Composants Pathways sur GKE

Cette section mappe les composants Pathways aux composants Google Kubernetes Engine tels que les conteneurs et les pods.

Vous trouverez les images de conteneurs Pathways aux emplacements suivants.

Type de conteneur	Emplacement
Serveur proxy IFRT	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>`
Gestionnaire/Employé Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>`

Gestionnaire de ressources Pathways

Après avoir créé un cluster GKE, vous pouvez utiliser le containerSpec suivant pour déployer le gestionnaire de ressources de parcours :

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

Descriptions des arguments :

--server_port : le gestionnaire de ressources Pathways utilise ce port pour communiquer avec d'autres composants Pathways.
--node_type : type de nœud. Cette valeur doit être définie sur "resource_manager" pour le gestionnaire de ressources Pathways. Elle n'est pas nécessaire pour les autres conteneurs.
--instance_count : nombre de tranches de TPU.
--instance_type : type et topologie de la tranche de TPU. Au format tpu{TPU type}:{TPU topology}, par exemple tpuv5e:4x4.
--gcs_scratch_location : bucket Cloud Storage utilisé pour les fichiers temporaires.

Serveur proxy IFRT

Vous pouvez utiliser le containerSpec suivant pour déployer un serveur proxy IFRT :

 - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

Descriptions des arguments :

--resource_manager_address : nom d'hôte et port utilisés par le serveur proxy pour communiquer avec le gestionnaire de ressources Pathways. Le port doit être identique à la valeur --server_port utilisée pour le conteneur du gestionnaire de ressources Pathways.
--server_port : le serveur proxy IFRT utilise ce port pour communiquer avec le client proxy IFRT.
--gcs_scratch_location : bucket GCS utilisé pour les fichiers temporaires.

Employé Pathways

Vous pouvez utiliser les containerSpec suivants pour déployer des workers Pathways :

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

Descriptions des arguments :

--resource_manager_address : nom d'hôte et port utilisés par les nœuds de calcul TPU pour communiquer avec le gestionnaire de ressources Pathways. Le port doit être identique à la valeur --server_port utilisée pour le conteneur du gestionnaire de ressources Pathways.
--server_port : les nœuds de calcul utilisent ce port pour communiquer avec le serveur proxy et le gestionnaire de ressources Pathways.
--gcs_scratch_location : bucket Cloud Storage utilisé pour les fichiers temporaires.

Le gestionnaire de ressources Pathways, le serveur proxy IFRT et les workers Pathways peuvent tous avoir des ports différents, mais dans cet exemple, le gestionnaire de ressources Pathways et le worker Pathways partagent le même port.

Présentation des parcours sur Cloud Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.

Avant de commencer

Contrôleur unique et multicontrôleur

Composants des parcours

Gestionnaire de ressources Pathways

Client Pathways

Employé Pathways

Client proxy IFRT

Serveur proxy IFRT

Serveur side-car

API PathwaysJob

Composants Pathways sur GKE

Gestionnaire de ressources Pathways

Serveur proxy IFRT

Employé Pathways

Étapes suivantes

Présentation des parcours sur Cloud