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 calcul pour différentes tâches en fonction de leurs exigences de traitement.

Pathways simplifie les calculs de machine learning à grande échelle en permettant à un seul client JAX d'orchestrer les 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 Google Cloud clients.

Avant de commencer

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

Ce document explique comment utiliser les TPU gérés par Pathways sur Google Kubernetes Engine (GKE) pour les charges de travail par lot, en temps réel et interactives. Il suppose que vous savez déjà comment utiliser les TPU avec GKE y compris les TPU à tranche unique et multitrances sur Google Kubernetes Engine, ainsi que l'expérience générale avec les TPU multitrances

Contrôleur unique et multicontrôleur

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

Fonctionnalité	Contrôleur unique (Pathways)	Multicontrôleur (JAX par défaut)
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 y 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, ce qui nécessite que tous les appareils exécutent le même programme.
Flexibilité	Prend en charge des modèles de calcul plus complexes au-delà de SPMD, y compris le parallélisme de pipeline asymétrique et la rareté de calcul.	Peut être moins flexible dans la gestion des ressources, en particulier sur différentes tranches de TPU.

Composants Pathways

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 tâches utilisateur. Il surveille l'état des nœuds de calcul et gère la planification, la mise en pause 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 de l'Interim Framework Runtime (IFRT) qui sert de point d'entrée dans le système Pathways. Il reçoit des opérations de haut niveau (HLO) de votre programme. Le client Pathways est chargé de coordonner 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.

Nœud de calcul Pathways

Il s'agit des processus qui s'exécutent sur les machines d'accélérateur (VM TPU). Ils reçoivent des exécutables compilés de votre programme à partir du serveur proxy IFRT et effectuent les calculs sur les TPU. Les nœuds de calcul 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 dissocie le code utilisateur de l'environnement d'exécution sous-jacent et améliore la portabilité et la transparence du code. JAX utilise cette implémentation comme alternative à son environnement d'exécution multicontrô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 en 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 d'accélérateur pour exécuter directement le code Python spécifié par l'utilisateur sur la VM d'accélérateur afin de réduire la latence de transfert des données du contrôleur aux accélérateurs. Le serveur side-car interagit avec le nœud de calcul Pathways via un protocole personnalisé versionné sur le transport gRPC.

Affiche la relation entre les composants Pathways. — Composants Pathways

API PathwaysJob

L'API PathwaysJob est une API OSS native de Kubernetes 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 exploite 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) vous offre une interface de haut niveau pour définir vos charges de travail Pathways, en éliminant la nécessité de gérer directement les spécifications de pod individuelles pour les scénarios courants. Pour obtenir la liste complète de tous les paramètres et de leurs significations spécifiques, 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 en tant que `PathwaysJob`.
`metadata.name`	Nom de l'objet `PathwaysJob` dans Kubernetes, généralement selon le modèle pathways-.
`spec`	Définit l'état et la configuration souhaités pour `PathwaysJob`.
`spec.maxRestarts`	Nombre maximal de fois où le système peut redémarrer automatiquement `PathwaysJob` en cas d'échec.
`spec.pathwaysVersion`	(Facultatif) Spécifie la version souhaitée du framework JAX à utiliser dans l'environnement Pathways pour cette tâche (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 allouées aux nœuds de calcul (par exemple, $TOPOLOGY peut être 2x2, 4x4, 2x2x2).
`spec.workers[].numSlices`	Nombre de tranches de TPU à provisionner pour le pool de nœuds de calcul (par exemple, $WORKLOAD_NODEPOOL_COUNT peut être 2).
`spec.workers[].maxSliceRestarts`	(Facultatif) Nombre maximal de fois où un nœud de calcul individuel d'une tranche 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 la tâche principale.
`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 au démarrage.
`spec.customComponents[].customEnv`	Tableau de variables d'environnement personnalisées à définir dans le conteneur. Chaque élément comporte un nom et une valeur.
`spec.pathwaysDir`	Bucket Cloud Storage utilisé par le `PathwaysJob` pour stocker les artefacts de compilation et d'autres données temporaires. Ce bucket doit être créé avant l'exécution de votre charge de travail.
`spec.controller`	Paramètres de configuration du contrôleur Pathways, qui gère l'exécution globale des tâches.
`spec.controller.deploymentMode`	Spécifie comment 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 avec un nœud de calcul TPU.
`spec.controller.elasticSlices`	(Facultatif) Nombre maximal de tranches de TPU pouvant devenir indisponibles lors de l'exécution de la tâche avant d'être considérées comme non opérationnelles.
`spec.controller.template`	(Facultatif) Définit le modèle de pod pour la tâche utilisateur. Cette option 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 la tâche utilisateur.
`spec.controller.template.spec.containers`	Tableau définissant les conteneurs qui s'exécuteront dans la tâche utilisateur
`spec.controller.template.spec.containers[].name`	Nom du conteneur dans la tâche 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, elle 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 des images de conteneur 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/nœud de calcul de ressources 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 Pathways :

  - 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 et n'est pas nécessaire pour les autres conteneurs.
--instance_count : nombre de tranches de TPU.
--instance_type : type de TPU et topologie de la tranche. 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 que le serveur proxy utilise 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 Cloud Storage utilisé pour les fichiers temporaires.

Nœud de calcul Pathways

Vous pouvez utiliser le containerSpec suivant pour déployer des nœuds de calcul 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 que les nœuds de calcul TPU utilisent 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 nœuds de calcul Pathways peuvent tous avoir des ports différents, mais dans cet exemple, le gestionnaire de ressources Pathways et le nœud de calcul 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 Pathways

Gestionnaire de ressources Pathways

Client Pathways

Nœud de calcul Pathways

Client proxy IFRT

Serveur proxy IFRT

Serveur side-car

API PathwaysJob

Composants Pathways sur GKE

Gestionnaire de ressources Pathways

Serveur proxy IFRT

Nœud de calcul Pathways

Étape suivante

Présentation des parcours sur Cloud