Something went wrong. Try again.
lol
Something went wrong. Try again.
2.3 kB · 68 lines
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869diff --git a/brax/training/agents/apg/train.py b/brax/training/agents/apg/train.pyindex f5fcb0e..87b198f 100644--- a/brax/training/agents/apg/train.py+++ b/brax/training/agents/apg/train.py@@ -310,7 +310,7 @@ def train( specs.Array((env.observation_size,), jnp.dtype(dtype)) ), )- training_state = jax.device_put_replicated(+ training_state = pmap.device_put_replicated( training_state, jax.local_devices()[:local_devices_to_use] ) diff --git a/brax/training/agents/ppo/train.py b/brax/training/agents/ppo/train.pyindex 9aec960..6624733 100644--- a/brax/training/agents/ppo/train.py+++ b/brax/training/agents/ppo/train.py@@ -753,7 +753,7 @@ def train( {}, ) - training_state = jax.device_put_replicated(+ training_state = pmap.device_put_replicated( training_state, jax.local_devices()[:local_devices_to_use] ) diff --git a/brax/training/agents/sac/train.py b/brax/training/agents/sac/train.pyindex be716e9..8dcf3bf 100644--- a/brax/training/agents/sac/train.py+++ b/brax/training/agents/sac/train.py@@ -108,7 +108,7 @@ def _init_training_state( alpha_params=log_alpha, normalizer_params=normalizer_params, )- return jax.device_put_replicated(+ return pmap.device_put_replicated( training_state, jax.local_devices()[:local_devices_to_use] ) diff --git a/brax/training/pmap.py b/brax/training/pmap.pyindex 82760fc..af62ef8 100644--- a/brax/training/pmap.py+++ b/brax/training/pmap.py@@ -19,12 +19,23 @@ from typing import Any import jax import jax.numpy as jnp+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P+import numpy as np+++def device_put_replicated(x, devices):+ """Drop-in replacement for jax.device_put_replicated supporting pytrees."""+ mesh = Mesh(np.array(devices), ('x',))+ sharding = NamedSharding(mesh, P('x'))+ return jax.tree.map(+ lambda y: jax.device_put(jnp.stack([y] * len(devices)), sharding), x+ ) def bcast_local_devices(value, local_devices_to_use=1): """Broadcasts an object to all local devices.""" devices = jax.local_devices()[:local_devices_to_use]- return jax.device_put_replicated(value, devices)+ return device_put_replicated(value, devices) def synchronize_hosts():