Enforcing an expensive function to be only called in host during training loop

[riven314]

I am one of the participant of Huggingface Community Week. I got a question about pmap and training loop, would really appreciate if someone could offer me some advice on that:

I have an memory & time expensive step that needa run per N iterations, e.g.:

for it, batch in enumerate(dataset):
  state = pmap_train(keys, model, state, batch, lr)
  if it % N == 0:
     run_expensive_step(keys, model, state, batch, clip_model, lr)

Questions

1. will this run_expensive_step create a big bottleneck on pmap_train call, or it wont be blocked (maybe becoz of async nature?)...
2. is there any way to enforce run_expensive_step to be called only on the host core? (coz my host have 300+GB and run_expensive_step is super memory hungry)
3. model has been passed through flax.jax_utils.replicate and batch has been passed through shard, but clip_model haven't been flax.jax_utils.replicate. Do I needa reduce model and batch dimension before each run_expensive_step call?

[changjonathanc]

3. You can use this in flax:

def unreplicate(tree):
  """Returns a single instance of a replicated array."""
  return jax.tree_map(lambda x: x[0], tree)

source: https://github.com/google/flax/blob/1844de4d4b5ef819c76cb21ded038cfc10624aa5/flax/jax_utils.py#L78-L80