pytorch/xla

xs.make_mesh

Open

#8.838 geöffnet am 14. März 2025

Auf GitHub ansehen
 (1 Kommentar) (0 Reaktionen) (1 zugewiesene Person)C++ (575 Forks)github user discovery
enhancementgood first issue

Repository-Metriken

Stars
 (2.786 Stars)
PR-Merge-Metriken
 (PR-Metriken ausstehend)

Beschreibung

🚀 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

Contributor Guide

Research-Richtung
Analysieren Sie die Implementierung von make mesh in JAX und verstehen Sie, wie PyTorch/XLA derzeit mit Device Mesh und kollektiven Operationen umgeht. Identifizieren Sie die wichtigsten Leistungsfaktoren und schlagen Sie eine optimierte Strategie zur Geräteanordnung vor.
Tech Stack
python
Domain
backendperformance
Issue Type
Funktion
SchwierigkeitGeschätzte Implementierungsschwierigkeit für neue Contributors, von 1 für sehr kleine Änderungen bis 5 für Expertenarbeit.
3
Geschätzte ZeitEin grober Zeitrahmen, um zu recherchieren, zu implementieren, zu testen und einen Pull Request vorzubereiten.
1-2 Tage
AktivitätsstatusWie verfügbar das Issue gerade wirkt: frisch, aktiv, stale, blockiert oder wartet auf Maintainer-Input.
Aktiv
KlarheitWie klar das Issue die erwartete Änderung, Akzeptanzkriterien und den nächsten Schritt erklärt.
Meist klar
Voraussetzungen
PythonC++PyTorch
EinsteigerfreundlichkeitEin Score von 1 bis 100, der schätzt, wie zugänglich dieses Issue für First-Time Contributors ist.
40