NamedSharding doesn't work in multihost

fattorib · Jun 6, 2023 · 69b38a9 · 69b38a9
1 parent b6e35d0
commit 69b38a9
Showing 1 changed file with 0 additions and 7 deletions.
diff --git a/main.py b/main.py
@@ -293,13 +293,6 @@ def compute_tokens_seen(absolute_step, max_context):
             import gc
             gc.collect()
 
-            named_sharding_params = jax.tree_map(lambda x: jax.sharding.NamedSharding(mesh, x), param_spec)
-            named_sharding_opt = jax.tree_map(lambda x: jax.sharding.NamedSharding(mesh, x), opt_state_spec)            
-
-            params = jax.device_put(params, named_sharding_params)
-            opt_state = jax.device_put(opt_state, named_sharding_opt)
-
-
         else:
             raise NotImplementedError(
                 "Checkpointing not currently implemented for GPU."