JAX-Backend verwenden

In diesem Leitfaden wird erläutert, wie Sie das JAX-Backend in Meridian verwenden.

Einführung in das JAX-Backend

Meridian verwendet JAX als Standard-Backend für numerische Kernoperationen und das probabilistische Markov-Chain-Monte-Carlo-Verfahren (MCMC) zur Stichprobenerhebung (ab Meridian 2.0). JAX fördert einen funktionalen Programmierstil und nutzt die XLA-Kompilierung (Accelerated Linear Algebra), um erweiterte Leistungsoptimierungen und Arbeitsspeichereffizienz zu ermöglichen.

Das alte TensorFlow-Backend wurde eingestellt und wird in einer zukünftigen Version entfernt.

Tutorial: Unter Getting started with JAX können Sie JAX in Aktion sehen.

Backend-Konfiguration

Standardmäßig wird Meridian auf JAX ausgeführt. Sie müssen keine Umgebungsvariablen konfigurieren, um JAX zu verwenden.

Altes TensorFlow-Backend (eingestellt)

Falls Sie vorübergehend das alte TensorFlow-Backend verwenden müssen, legen Sie die Umgebungsvariable MERIDIAN_BACKEND auf 'tensorflow' fest, bevor Sie Meridian importieren:

import os

# Select legacy TensorFlow backend (deprecated)
os.environ['MERIDIAN_BACKEND'] = 'tensorflow'

# Now it is safe to import Meridian modules
from meridian.model import model
from meridian.data import load

Konfiguration der Genauigkeit

Standardmäßig wird Meridian mit 64-Bit-Genauigkeit (float64) auf JAX ausgeführt.

Sie können stattdessen aber auch 32-Bit-Genauigkeit (float32) verwenden, zum Beispiel:

• Schnellere Trainingslaufzeit: 32-Bit-Gleitkommaoperationen können auf Hardwarebeschleunigern wie GPUs oder TPUs schneller ausgeführt werden.
• Geringere Arbeitsspeichernutzung: Die 32-Bit-Genauigkeit reduziert die Arbeitsspeicherbelegung während der MCMC-Stichprobenerhebung.

Wenn Sie 32-Bit-Genauigkeit verwenden möchten, legen Sie die Umgebungsvariable MERIDIAN_ENABLE_JAX_X64 auf 'False' (oder '0') fest, bevor Sie Meridian importieren:

import os

# Disable 64-bit precision (enable 32-bit precision)
os.environ['MERIDIAN_ENABLE_JAX_X64'] = 'False'

# Now it is safe to import Meridian modules
from meridian.model import model

Wenn die Umgebungsvariable MERIDIAN_ENABLE_JAX_X64 nicht festgelegt oder auf 'True' bzw. '1' festgelegt ist, verwendet Meridian standardmäßig die 64-Bit-Genauigkeit.

Typkonsistenz

Wichtig: JAX erzwingt eine strikte Typprüfung und führt bei vielen Vorgängen keine automatische Typaufwertung oder -umwandlung für nicht übereinstimmende Gleitkommatypen durch. Das Mischen von float32- und float64-Werten führt zu einem TypeError (zum Beispiel TypeError: Cannot mix float32 and float64).

Da Meridian standardmäßig mit 64-Bit-Genauigkeit arbeitet, müssen alle nutzerseitig bereitgestellten Werte, benutzerdefinierten Arrays und Verteilungsparameter typkonsistent sein:

• Gleitkommaliterale und -Arrays: Standardmäßige Python-Gleitkommaliterale (zum Beispiel 0.2, 0.9) sind standardmäßig 64-Bit-Gleitkommazahlen. Wenn Sie NumPy-Arrays für Prior-Verteilungen oder Eingaben erstellen, verwenden Sie np.float64 oder dtype=np.float64 passend zur Standardgenauigkeit.
• Benutzerdefinierte Prior-Verteilungen mit passender Genauigkeit erstellen: Wenn Sie benutzerdefinierte Prior-Verteilungen in PriorDistribution definieren, müssen alle Verteilungsparameter (zum Beispiel loc, scale, concentration0 und concentration1) der aktiven Genauigkeit entsprechen (standardmäßig 64-Bit-Gleitkommazahl).

API-Unterschiede bei Verwendung von JAX im Vergleich zu TensorFlow

Wenn Sie das JAX-Backend verwenden, gibt es einige wichtige API-Unterschiede, die berücksichtigt werden müssen:

Prior-Verteilungen

Meridian-Modelle verwenden TensorFlow Probability unter JAX (tensorflow_probability.substrates.jax). Wenn Sie benutzerdefinierte Prior-Verteilungen unter JAX konfigurieren, importieren Sie tensorflow_probability.substrates.jax as tfp_jax und erstellen Sie Verteilungen mit tfp_jax.distributions.

Achten Sie darauf, dass alle benutzerdefinierten Verteilungsparameter 64-Bit-Genauigkeit verwenden (zum Beispiel np.float64 oder 64-Bit-Gleitkomma-Arrays), um die Typkonsistenz mit den Standardgenauigkeitseinstellungen von Meridian zu wahren.

Achtung: Das Standardmodul tfp.distributions von TensorFlow Probability ist nicht mit dem JAX-Backend kompatibel und führt zu Laufzeitfehlern. Verwenden Sie immer tfp_jax.distributions.

JAX

import numpy as np
import tensorflow_probability.substrates.jax as tfp_jax
from meridian.model import constants
from meridian.model import prior_distribution

# Parameters use 64-bit precision
roi_mu = np.float64(0.2)
roi_sigma = np.float64(0.9)
prior = prior_distribution.PriorDistribution(
    roi_m=tfp_jax.distributions.LogNormal(
        roi_mu, roi_sigma, name=constants.ROI_M
    )
)

TensorFlow (eingestellt)

import tensorflow_probability as tfp
from meridian.model import constants
from meridian.model import prior_distribution

roi_mu = 0.2
roi_sigma = 0.9
prior = prior_distribution.PriorDistribution(
    roi_m=tfp.distributions.LogNormal(
        roi_mu, roi_sigma, name=constants.ROI_M
    )
)

Explizite Seed-Anforderung

Bei Verwendung des JAX-Backends ist für stochastische Funktionen (z. B. in sample_posterior()) ein expliziter Seed erforderlich. Während TensorFlow einen globalen Zufallszahlengenerator verwendet, der automatisch einen zufälligen Seed auswählt, ist dieser Seed in JAX explizit. Wir haben keine statistisch signifikanten Unterschiede bei den ROI-Schätzungen oder Budgetverschiebungen zwischen den verschiedenen Seeds festgestellt.

# Explicitly set a seed for MCMC sampling when using the JAX backend
mmm.sample_posterior(
    n_chains=2,
    n_adapt=1000,
    n_burnin=500,
    n_keep=1000,
    seed=0,
)

Weitere Informationen zu Zufallszahlen und Seeds in JAX finden Sie in der entsprechenden Dokumentation Pseudorandom numbers.

Numerische Unterschiede und Reproduzierbarkeit

Da TensorFlow und JAX ihre Berechnungsgraphen unterschiedlich kompilieren, können bei der Umstellung auf JAX mit denselben Daten und Zufallswerten geringfügige numerische Unterschiede bei den Posterior-Schätzungen auftreten.

Auch wenn die Posterior-Verteilungen zwischen den Backends vielleicht nicht identisch sind, fallen Unterschiede in der Regel gering aus und sind für geschäftliche Messwerte wie den ROI und die Budgetzuweisung statistisch nicht signifikant. So bleibt die Integrität der Statistiken Ihres Modells beim Wechsel zum JAX-Backend erhalten.

Hinweise zur Leistung

Interne Tests haben ergeben, dass JAX initiale Modellläufe massiv beschleunigt: Bei der Verwendung von GPUs wurde die durchschnittliche Laufzeit im Vergleich zu TensorFlow um etwa 40 % und die Arbeitsspeichernutzung um etwa 70 % reduziert JAX hat außerdem die Modelliterationen optimiert und ermöglicht dadurch doppelt so schnelle Laufzeiten, eine um den Faktor 4 geringere Arbeitsspeichernutzung sowie unterbrechungsfreie Workflows, da Kernel-Neustarts entfallen.

Hinweis: Diese Ergebnisse stammen aus internen Analysen, die mit dem Colab für den Einstieg auf einer GPU durchgeführt wurden. Dabei wurde der Zustand der Laufzeit vor und nach dem Ausführen von „sample_posterior“ gemessen. Je nach Konfiguration können die Ergebnisse variieren. Tests zeigen jedoch immer wieder Effizienzsteigerungen mit JAX im Vergleich zu TensorFlow.

Dank der erhöhten Arbeitsspeichereffizienz haben Sie mehr Spielraum, um rechenintensive Parameter anzupassen. In Meridian.sample_posterior() können Sie beispielsweise das Argument unrolled_leapfrog_steps erhöhen (etwa von 1 auf 5). Dies kann die Konvergenz beschleunigen, indem die Trajektorienlänge des No-U-Turn-Samplers (NUTS) erhöht wird, ohne die Hardware-Arbeitsspeicherlimits zu überschreiten. Sie können auch den Parameter n_adapt erhöhen, um die Konvergenz während der Anpassungsphase zu unterstützen.