pytorch/xla

xs.make_mesh

Open

#8838 aperta il 14 mar 2025

Vedi su GitHub
 (1 commento) (0 reazioni) (1 assegnatario)C++ (575 fork)github user discovery
enhancementgood first issue

Metriche repository

Star
 (2786 star)
Metriche merge PR
 (Metriche PR in attesa)

Descrizione

🚀 Feature

We should provide a xs.make_mesh method similar to JAX: https://docs.jax.dev/en/latest/_autosummary/jax.make_mesh.html

Motivation

The ordering of devices affects the performance of collective operations:

  • Weights are sharded differently and the layout of tensors after an all-gather is different.
  • The performance of a collective op itself depends on what network connections it uses.

Pitch

The xs.make_mesh method should attempt to create a reasonably optimized device ordering given a mesh shape referencing how JAX does it.

We could also deviate from JAX and add additional options that fit our own needs.

cc @yaochengji

Guida contributor

Direzione di ricerca
Analizzare l'implementazione di make mesh di JAX e comprendere come PyTorch/XLA gestisce attualmente la mesh dei dispositivi e le operazioni collettive. Identificare i fattori chiave delle prestazioni e proporre una strategia di ordinamento dei dispositivi ottimizzata.
Tech stack
python
Dominio
backendperformance
Tipo issue
Funzionalità
DifficoltàQuanto dovrebbe essere impegnativa per un nuovo contributor.
3
Tempo stimatoTempo stimato per completare e verificare uno scope piccolo.
1-2 giorni
Stato attivitàQuanto è probabile che la issue sia ancora attiva e reviewable.
Attiva
ChiarezzaQuanto chiaramente la issue descrive problema, scope e risultato atteso.
Abbastanza chiara
Prerequisiti
PythonC++PyTorch
Adatta ai principiantiQuanto la issue è adatta a contributor alla prima esperienza.
40