<p dir="auto">We've built jaxlib 0.3.7 on conda-forge with cuda: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1234482537" data-permission-text="Title is private" data-url="https://github.com/conda-forge/jaxlib-feedstock/issues/97" data-hovercard-type="pull_request" data-hovercard-url="/conda-forge/jaxlib-feedstock/pull/97/hovercard" href="https://github.com/conda-forge/jaxlib-feedstock/pull/97">conda-forge/jaxlib-feedstock#97</a></p>
<p dir="auto">I am opening this issue for viz but also for testing/confirmation all is good. Could someone from the team here have a look with us? Easy things to check, we get something like ~11000 actions/cycles of bazel compilation (compared to ~7500 for non-cuda) and the size of our artifact is similar to that of the official jaxlib ~200 mb. I tested it locally and it seems to put the arrays on the GPU by default. So it is working, but more confirmation is always good!</p>
<p dir="auto"><em>Note</em> that we build for 10.2, 11.0, 11.1, and 11.2 on conda-forge by default, though we we can adjust this if requested, and do <em>note</em> the forward compatibility of 11.2, it is really "11.2+" meaning that it compiles for 11.2 and above (if someone installs that the 11.2 compiled version, it will in fact install cudatoolkit 11.6 with it and the latest cudnn)</p>
<p dir="auto">People interested in checking this out, should head to the:</p>
<ul dir="auto">
<li>PR: link above</li>
<li>logs: <a href="https://dev.azure.com/conda-forge/feedstock-builds/_build/results?buildId=506049" rel="nofollow">https://dev.azure.com/conda-forge/feedstock-builds/_build/results?buildId=506049</a></li>
<li>artifacts: <a href="https://dev.azure.com/conda-forge/feedstock-builds/_build/results?buildId=506049&amp;view=artifacts&amp;pathAsName=false&amp;type=publishedArtifacts" rel="nofollow">https://dev.azure.com/conda-forge/feedstock-builds/_build/results?buildId=506049&amp;view=artifacts&amp;pathAsName=false&amp;type=publishedArtifacts</a></li>
</ul>
<p dir="auto">You will see we still have a few problems (lost connections + timeouts mainly) due to the strict 6-hour time limit (it is tricky because we are very close to the limit, but hopefully we will manage to get all of them by rerunning the CI a few times)</p>
<p dir="auto">Tagging <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/hawkinsp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/hawkinsp">@hawkinsp</a> again who's been involved in issues related to the builds before + <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jakevdp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jakevdp">@jakevdp</a> who's been very helpful and engaging.</p>
<hr>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<ul dir="auto">
<li>
<p dir="auto">Duplicate issues: <a href="https://github.com/google/jax/issues/788" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/788/hovercard">#788</a> and <a href="https://github.com/google/jax/issues/4528" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/4528/hovercard">#4528</a> are not suitable for this case.</p>
</li>
<li>
<p dir="auto">How to reproduce the bug:<br>
[1] Just run <a href="https://github.com/google-research/google-research/blob/master/muNet/mu2Net.ipynb">mu2Net</a> on 8gpus A100, use BENCHMARK = 'ViT large / Chars benchmark'<br>
[2] OOM error will occur when train_step function compiled by jax.jit is executed.<br>
[3] The A100 have sufficent 80GiB memory.</p>
</li>
<li>
<p dir="auto">Core code:</p>
</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="@partial(jax.jit, static_argnames=['model', 'optimizer'], donate_argnums=[0, 2])
def train_step(params, fixed_params, opt_state, images, labels, model, optimizer):
  def loss_fn(params, fixed_params, images, labels):
    logits = model.apply({'params': format_params(params, fixed_params)},
                         images, train=USE_DROPOUT)
    labels = jax.nn.one_hot(labels, logits.shape[-1])
    return -jnp.mean(jnp.sum(labels * nn.log_softmax(logits), axis=-1))
  grads = jax.grad(loss_fn)(params, fixed_params, images, labels)
  updates, opt_state = optimizer.update(grads, opt_state, params=params)
  params = optax.apply_updates(params, updates)
  return params, opt_state

def train_loop(paths, ds_train, ds_validation, devices, exp_config):
  global LOOP_START
  global TRAIN_PARAMS_TRANSFER
  global TRAIN
  timing = {'start_time': time.time(),
            'start_time_loop': LOOP_START}
  task = paths[0].task
  # The following values should be shared by all paths in this generation batch.
  for path in paths:
    assert task == path.task
    assert paths[0].hparams['ds_image_size'] == path.hparams['ds_image_size']

  gc.collect()

  # Compile.
  compile_train_batches_arr = jax.device_put_replicated(
      get_sample_batch(
        paths[0].hparams['ds_image_size'],
        task.train_batch_size),
      devices)
  compile_eval_batches_arr = jax.device_put_replicated(
      get_sample_batch(
          paths[0].hparams['ds_image_size'],
          task.validation_batch_size),
      devices)

  for p_id, path in enumerate(paths):
    if VERBOSE:
      print('Parent')
      print(prp(path.parent))
      print(prp(path))
    path.device_id = p_id % len(devices)
    path.device = devices[path.device_id]
    print(&quot;path:&quot;, p_id, &quot;device:&quot;, path.device)
    path.optimizer = path.get_optimizer()
    path.optimizer_init_fn = jax.jit(path.optimizer.init, device=path.device)
    path.best_params_local = None
    path.best_opt_state_local = None
    path.best_quality = None
    path.best_score = path.parent.score() if path.task is path.parent.task else -np.inf
    path.evals = []

    # Launch parallel compilation of eval and train step functions.
    params_local = path.get_trainable_params()
    check_is_local(params_local)
    path.compile_params_device = jax.device_put(params_local, path.device)
    path.compile_fixed_params_device = jax.device_put(
        path.get_fixed_params(),
        path.device)
    path.compile_train = Thread(
        target=train_step,
        args=(path.compile_params_device,
              path.compile_fixed_params_device,
              path.optimizer_init_fn(params_local),
              compile_train_batches_arr['image'][path.device_id],
              compile_train_batches_arr['label'][path.device_id],
              path.model,
              path.optimizer))
    path.compile_eval = Thread(
        target=eval_step,
        args=(format_params(
                  path.compile_params_device,
                  path.compile_fixed_params_device),
              compile_eval_batches_arr['image'][path.device_id],
              compile_eval_batches_arr['label'][path.device_id],
              path.model))
    path.compile_eval.start()

  for path in paths:
    path.compile_eval.join()
    del path.compile_eval
    timing['end_compile_eval'] = time.time()
    path.compile_train.start()
  del compile_eval_batches_arr

  for path in paths:
    path.compile_train.join()
    del path.compile_train
    del path.compile_params_device
    del path.compile_fixed_params_device
    timing['end_compile'] = time.time()
  del compile_train_batches_arr

  gc.collect()

  # Parameter transfer.
  transfer_start = time.time()
  for path in paths:
    path.params_device = jax.device_put(
        path.get_trainable_params(),
        path.device)
    path.fixed_params_device = jax.device_put(
        path.get_fixed_params(),
        path.device)
    path.opt_state_device = path.optimizer_init_fn(path.params_device)
    # Set opt state.
    for c in path.components:
      if c.is_trainable():
        assert c.name in path.opt_state_device[1][0].trace.keys()
        if c.opt_state is not None:
          path.opt_state_device = (
              path.opt_state_device[0],
              (optax.TraceState(
                  trace=path.opt_state_device[1][0].trace.copy(
                      {c.name: jax.device_put(c.opt_state,
                                              path.device)})),
               path.opt_state_device[1][1]
               )
          )
    check_is_on_device(path.opt_state_device, path.device)

  TRAIN_PARAMS_TRANSFER += time.time() - transfer_start

  iter_ds_validation = iter(ds_validation)
  # TRAIN
  for t_step, train_batch in zip(
      range(exp_config.num_validations_per_path_training
            * task.num_train_batches_between_validations),
      ds_train,
  ):
    replicated_start = time.time()
    train_batch_arr = jax.device_put_replicated(train_batch, devices)
    TRAIN_PARAMS_TRANSFER += time.time() - replicated_start
    for p_id, path in enumerate(paths):
      if t_step == 0:
        timing['end_prep'] = time.time()
        t_step_0_time = time.time()
      train_step_start = time.time()
      path.params_device, path.opt_state_device = train_step(
          path.params_device,
          path.fixed_params_device,
          path.opt_state_device,
          train_batch_arr['image'][path.device_id],
          train_batch_arr['label'][path.device_id],
          path.model,
          path.optimizer)
      TRAIN += time.time() - train_step_start
      if t_step == 0 and time.time() - t_step_0_time &gt; 1:
        print(f'WARNING: First train step took: {time.time()-t_step_0_time:.2f} s')
    del train_batch, train_batch_arr

    # EVAL
    # ..."><pre class="notranslate"><span class="pl-en">@<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>, <span class="pl-s1">static_argnames</span><span class="pl-c1">=</span>[<span class="pl-s">'model'</span>, <span class="pl-s">'optimizer'</span>], <span class="pl-s1">donate_argnums</span><span class="pl-c1">=</span>[<span class="pl-c1">0</span>, <span class="pl-c1">2</span>])</span>
<span class="pl-k">def</span> <span class="pl-en">train_step</span>(<span class="pl-s1">params</span>, <span class="pl-s1">fixed_params</span>, <span class="pl-s1">opt_state</span>, <span class="pl-s1">images</span>, <span class="pl-s1">labels</span>, <span class="pl-s1">model</span>, <span class="pl-s1">optimizer</span>):
  <span class="pl-k">def</span> <span class="pl-en">loss_fn</span>(<span class="pl-s1">params</span>, <span class="pl-s1">fixed_params</span>, <span class="pl-s1">images</span>, <span class="pl-s1">labels</span>):
    <span class="pl-s1">logits</span> <span class="pl-c1">=</span> <span class="pl-s1">model</span>.<span class="pl-en">apply</span>({<span class="pl-s">'params'</span>: <span class="pl-en">format_params</span>(<span class="pl-s1">params</span>, <span class="pl-s1">fixed_params</span>)},
                         <span class="pl-s1">images</span>, <span class="pl-s1">train</span><span class="pl-c1">=</span><span class="pl-v">USE_DROPOUT</span>)
    <span class="pl-s1">labels</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">nn</span>.<span class="pl-en">one_hot</span>(<span class="pl-s1">labels</span>, <span class="pl-s1">logits</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>])
    <span class="pl-k">return</span> <span class="pl-c1">-</span><span class="pl-s1">jnp</span>.<span class="pl-en">mean</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">labels</span> <span class="pl-c1">*</span> <span class="pl-s1">nn</span>.<span class="pl-en">log_softmax</span>(<span class="pl-s1">logits</span>), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>))
  <span class="pl-s1">grads</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">loss_fn</span>)(<span class="pl-s1">params</span>, <span class="pl-s1">fixed_params</span>, <span class="pl-s1">images</span>, <span class="pl-s1">labels</span>)
  <span class="pl-s1">updates</span>, <span class="pl-s1">opt_state</span> <span class="pl-c1">=</span> <span class="pl-s1">optimizer</span>.<span class="pl-en">update</span>(<span class="pl-s1">grads</span>, <span class="pl-s1">opt_state</span>, <span class="pl-s1">params</span><span class="pl-c1">=</span><span class="pl-s1">params</span>)
  <span class="pl-s1">params</span> <span class="pl-c1">=</span> <span class="pl-s1">optax</span>.<span class="pl-en">apply_updates</span>(<span class="pl-s1">params</span>, <span class="pl-s1">updates</span>)
  <span class="pl-k">return</span> <span class="pl-s1">params</span>, <span class="pl-s1">opt_state</span>

<span class="pl-k">def</span> <span class="pl-en">train_loop</span>(<span class="pl-s1">paths</span>, <span class="pl-s1">ds_train</span>, <span class="pl-s1">ds_validation</span>, <span class="pl-s1">devices</span>, <span class="pl-s1">exp_config</span>):
  <span class="pl-k">global</span> <span class="pl-v">LOOP_START</span>
  <span class="pl-k">global</span> <span class="pl-v">TRAIN_PARAMS_TRANSFER</span>
  <span class="pl-k">global</span> <span class="pl-v">TRAIN</span>
  <span class="pl-s1">timing</span> <span class="pl-c1">=</span> {<span class="pl-s">'start_time'</span>: <span class="pl-s1">time</span>.<span class="pl-en">time</span>(),
            <span class="pl-s">'start_time_loop'</span>: <span class="pl-v">LOOP_START</span>}
  <span class="pl-s1">task</span> <span class="pl-c1">=</span> <span class="pl-s1">paths</span>[<span class="pl-c1">0</span>].<span class="pl-s1">task</span>
  <span class="pl-c"># The following values should be shared by all paths in this generation batch.</span>
  <span class="pl-k">for</span> <span class="pl-s1">path</span> <span class="pl-c1">in</span> <span class="pl-s1">paths</span>:
    <span class="pl-k">assert</span> <span class="pl-s1">task</span> <span class="pl-c1">==</span> <span class="pl-s1">path</span>.<span class="pl-s1">task</span>
    <span class="pl-k">assert</span> <span class="pl-s1">paths</span>[<span class="pl-c1">0</span>].<span class="pl-s1">hparams</span>[<span class="pl-s">'ds_image_size'</span>] <span class="pl-c1">==</span> <span class="pl-s1">path</span>.<span class="pl-s1">hparams</span>[<span class="pl-s">'ds_image_size'</span>]

  <span class="pl-s1">gc</span>.<span class="pl-en">collect</span>()

  <span class="pl-c"># Compile.</span>
  <span class="pl-s1">compile_train_batches_arr</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put_replicated</span>(
      <span class="pl-en">get_sample_batch</span>(
        <span class="pl-s1">paths</span>[<span class="pl-c1">0</span>].<span class="pl-s1">hparams</span>[<span class="pl-s">'ds_image_size'</span>],
        <span class="pl-s1">task</span>.<span class="pl-s1">train_batch_size</span>),
      <span class="pl-s1">devices</span>)
  <span class="pl-s1">compile_eval_batches_arr</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put_replicated</span>(
      <span class="pl-en">get_sample_batch</span>(
          <span class="pl-s1">paths</span>[<span class="pl-c1">0</span>].<span class="pl-s1">hparams</span>[<span class="pl-s">'ds_image_size'</span>],
          <span class="pl-s1">task</span>.<span class="pl-s1">validation_batch_size</span>),
      <span class="pl-s1">devices</span>)

  <span class="pl-k">for</span> <span class="pl-s1">p_id</span>, <span class="pl-s1">path</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">paths</span>):
    <span class="pl-k">if</span> <span class="pl-v">VERBOSE</span>:
      <span class="pl-en">print</span>(<span class="pl-s">'Parent'</span>)
      <span class="pl-en">print</span>(<span class="pl-en">prp</span>(<span class="pl-s1">path</span>.<span class="pl-s1">parent</span>))
      <span class="pl-en">print</span>(<span class="pl-en">prp</span>(<span class="pl-s1">path</span>))
    <span class="pl-s1">path</span>.<span class="pl-s1">device_id</span> <span class="pl-c1">=</span> <span class="pl-s1">p_id</span> <span class="pl-c1">%</span> <span class="pl-en">len</span>(<span class="pl-s1">devices</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">device</span> <span class="pl-c1">=</span> <span class="pl-s1">devices</span>[<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>]
    <span class="pl-en">print</span>(<span class="pl-s">"path:"</span>, <span class="pl-s1">p_id</span>, <span class="pl-s">"device:"</span>, <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">path</span>.<span class="pl-en">get_optimizer</span>()
    <span class="pl-s1">path</span>.<span class="pl-s1">optimizer_init_fn</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">path</span>.<span class="pl-s1">optimizer</span>.<span class="pl-s1">init</span>, <span class="pl-s1">device</span><span class="pl-c1">=</span><span class="pl-s1">path</span>.<span class="pl-s1">device</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">best_params_local</span> <span class="pl-c1">=</span> <span class="pl-c1">None</span>
    <span class="pl-s1">path</span>.<span class="pl-s1">best_opt_state_local</span> <span class="pl-c1">=</span> <span class="pl-c1">None</span>
    <span class="pl-s1">path</span>.<span class="pl-s1">best_quality</span> <span class="pl-c1">=</span> <span class="pl-c1">None</span>
    <span class="pl-s1">path</span>.<span class="pl-s1">best_score</span> <span class="pl-c1">=</span> <span class="pl-s1">path</span>.<span class="pl-s1">parent</span>.<span class="pl-en">score</span>() <span class="pl-k">if</span> <span class="pl-s1">path</span>.<span class="pl-s1">task</span> <span class="pl-c1">is</span> <span class="pl-s1">path</span>.<span class="pl-s1">parent</span>.<span class="pl-s1">task</span> <span class="pl-k">else</span> <span class="pl-c1">-</span><span class="pl-s1">np</span>.<span class="pl-s1">inf</span>
    <span class="pl-s1">path</span>.<span class="pl-s1">evals</span> <span class="pl-c1">=</span> []

    <span class="pl-c"># Launch parallel compilation of eval and train step functions.</span>
    <span class="pl-s1">params_local</span> <span class="pl-c1">=</span> <span class="pl-s1">path</span>.<span class="pl-en">get_trainable_params</span>()
    <span class="pl-en">check_is_local</span>(<span class="pl-s1">params_local</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_params_device</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put</span>(<span class="pl-s1">params_local</span>, <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_fixed_params_device</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put</span>(
        <span class="pl-s1">path</span>.<span class="pl-en">get_fixed_params</span>(),
        <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_train</span> <span class="pl-c1">=</span> <span class="pl-v">Thread</span>(
        <span class="pl-s1">target</span><span class="pl-c1">=</span><span class="pl-s1">train_step</span>,
        <span class="pl-s1">args</span><span class="pl-c1">=</span>(<span class="pl-s1">path</span>.<span class="pl-s1">compile_params_device</span>,
              <span class="pl-s1">path</span>.<span class="pl-s1">compile_fixed_params_device</span>,
              <span class="pl-s1">path</span>.<span class="pl-en">optimizer_init_fn</span>(<span class="pl-s1">params_local</span>),
              <span class="pl-s1">compile_train_batches_arr</span>[<span class="pl-s">'image'</span>][<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>],
              <span class="pl-s1">compile_train_batches_arr</span>[<span class="pl-s">'label'</span>][<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>],
              <span class="pl-s1">path</span>.<span class="pl-s1">model</span>,
              <span class="pl-s1">path</span>.<span class="pl-s1">optimizer</span>))
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_eval</span> <span class="pl-c1">=</span> <span class="pl-v">Thread</span>(
        <span class="pl-s1">target</span><span class="pl-c1">=</span><span class="pl-s1">eval_step</span>,
        <span class="pl-s1">args</span><span class="pl-c1">=</span>(<span class="pl-en">format_params</span>(
                  <span class="pl-s1">path</span>.<span class="pl-s1">compile_params_device</span>,
                  <span class="pl-s1">path</span>.<span class="pl-s1">compile_fixed_params_device</span>),
              <span class="pl-s1">compile_eval_batches_arr</span>[<span class="pl-s">'image'</span>][<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>],
              <span class="pl-s1">compile_eval_batches_arr</span>[<span class="pl-s">'label'</span>][<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>],
              <span class="pl-s1">path</span>.<span class="pl-s1">model</span>))
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_eval</span>.<span class="pl-en">start</span>()

  <span class="pl-k">for</span> <span class="pl-s1">path</span> <span class="pl-c1">in</span> <span class="pl-s1">paths</span>:
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_eval</span>.<span class="pl-en">join</span>()
    <span class="pl-k">del</span> <span class="pl-s1">path</span>.<span class="pl-s1">compile_eval</span>
    <span class="pl-s1">timing</span>[<span class="pl-s">'end_compile_eval'</span>] <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_train</span>.<span class="pl-en">start</span>()
  <span class="pl-k">del</span> <span class="pl-s1">compile_eval_batches_arr</span>

  <span class="pl-k">for</span> <span class="pl-s1">path</span> <span class="pl-c1">in</span> <span class="pl-s1">paths</span>:
    <span class="pl-s1">path</span>.<span class="pl-s1">compile_train</span>.<span class="pl-en">join</span>()
    <span class="pl-k">del</span> <span class="pl-s1">path</span>.<span class="pl-s1">compile_train</span>
    <span class="pl-k">del</span> <span class="pl-s1">path</span>.<span class="pl-s1">compile_params_device</span>
    <span class="pl-k">del</span> <span class="pl-s1">path</span>.<span class="pl-s1">compile_fixed_params_device</span>
    <span class="pl-s1">timing</span>[<span class="pl-s">'end_compile'</span>] <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
  <span class="pl-k">del</span> <span class="pl-s1">compile_train_batches_arr</span>

  <span class="pl-s1">gc</span>.<span class="pl-en">collect</span>()

  <span class="pl-c"># Parameter transfer.</span>
  <span class="pl-s1">transfer_start</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
  <span class="pl-k">for</span> <span class="pl-s1">path</span> <span class="pl-c1">in</span> <span class="pl-s1">paths</span>:
    <span class="pl-s1">path</span>.<span class="pl-s1">params_device</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put</span>(
        <span class="pl-s1">path</span>.<span class="pl-en">get_trainable_params</span>(),
        <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">fixed_params_device</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put</span>(
        <span class="pl-s1">path</span>.<span class="pl-en">get_fixed_params</span>(),
        <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)
    <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span> <span class="pl-c1">=</span> <span class="pl-s1">path</span>.<span class="pl-en">optimizer_init_fn</span>(<span class="pl-s1">path</span>.<span class="pl-s1">params_device</span>)
    <span class="pl-c"># Set opt state.</span>
    <span class="pl-k">for</span> <span class="pl-s1">c</span> <span class="pl-c1">in</span> <span class="pl-s1">path</span>.<span class="pl-s1">components</span>:
      <span class="pl-k">if</span> <span class="pl-s1">c</span>.<span class="pl-en">is_trainable</span>():
        <span class="pl-k">assert</span> <span class="pl-s1">c</span>.<span class="pl-s1">name</span> <span class="pl-c1">in</span> <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span>[<span class="pl-c1">1</span>][<span class="pl-c1">0</span>].<span class="pl-s1">trace</span>.<span class="pl-en">keys</span>()
        <span class="pl-k">if</span> <span class="pl-s1">c</span>.<span class="pl-s1">opt_state</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
          <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span> <span class="pl-c1">=</span> (
              <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span>[<span class="pl-c1">0</span>],
              (<span class="pl-s1">optax</span>.<span class="pl-v">TraceState</span>(
                  <span class="pl-s1">trace</span><span class="pl-c1">=</span><span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span>[<span class="pl-c1">1</span>][<span class="pl-c1">0</span>].<span class="pl-s1">trace</span>.<span class="pl-en">copy</span>(
                      {<span class="pl-s1">c</span>.<span class="pl-s1">name</span>: <span class="pl-s1">jax</span>.<span class="pl-en">device_put</span>(<span class="pl-s1">c</span>.<span class="pl-s1">opt_state</span>,
                                              <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)})),
               <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span>[<span class="pl-c1">1</span>][<span class="pl-c1">1</span>]
               )
          )
    <span class="pl-en">check_is_on_device</span>(<span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span>, <span class="pl-s1">path</span>.<span class="pl-s1">device</span>)

  <span class="pl-v">TRAIN_PARAMS_TRANSFER</span> <span class="pl-c1">+=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>() <span class="pl-c1">-</span> <span class="pl-s1">transfer_start</span>

  <span class="pl-s1">iter_ds_validation</span> <span class="pl-c1">=</span> <span class="pl-en">iter</span>(<span class="pl-s1">ds_validation</span>)
  <span class="pl-c"># TRAIN</span>
  <span class="pl-k">for</span> <span class="pl-s1">t_step</span>, <span class="pl-s1">train_batch</span> <span class="pl-c1">in</span> <span class="pl-en">zip</span>(
      <span class="pl-en">range</span>(<span class="pl-s1">exp_config</span>.<span class="pl-s1">num_validations_per_path_training</span>
            <span class="pl-c1">*</span> <span class="pl-s1">task</span>.<span class="pl-s1">num_train_batches_between_validations</span>),
      <span class="pl-s1">ds_train</span>,
  ):
    <span class="pl-s1">replicated_start</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
    <span class="pl-s1">train_batch_arr</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">device_put_replicated</span>(<span class="pl-s1">train_batch</span>, <span class="pl-s1">devices</span>)
    <span class="pl-v">TRAIN_PARAMS_TRANSFER</span> <span class="pl-c1">+=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>() <span class="pl-c1">-</span> <span class="pl-s1">replicated_start</span>
    <span class="pl-k">for</span> <span class="pl-s1">p_id</span>, <span class="pl-s1">path</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">paths</span>):
      <span class="pl-k">if</span> <span class="pl-s1">t_step</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span>:
        <span class="pl-s1">timing</span>[<span class="pl-s">'end_prep'</span>] <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
        <span class="pl-s1">t_step_0_time</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
      <span class="pl-s1">train_step_start</span> <span class="pl-c1">=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>()
      <span class="pl-s1">path</span>.<span class="pl-s1">params_device</span>, <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span> <span class="pl-c1">=</span> <span class="pl-en">train_step</span>(
          <span class="pl-s1">path</span>.<span class="pl-s1">params_device</span>,
          <span class="pl-s1">path</span>.<span class="pl-s1">fixed_params_device</span>,
          <span class="pl-s1">path</span>.<span class="pl-s1">opt_state_device</span>,
          <span class="pl-s1">train_batch_arr</span>[<span class="pl-s">'image'</span>][<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>],
          <span class="pl-s1">train_batch_arr</span>[<span class="pl-s">'label'</span>][<span class="pl-s1">path</span>.<span class="pl-s1">device_id</span>],
          <span class="pl-s1">path</span>.<span class="pl-s1">model</span>,
          <span class="pl-s1">path</span>.<span class="pl-s1">optimizer</span>)
      <span class="pl-v">TRAIN</span> <span class="pl-c1">+=</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>() <span class="pl-c1">-</span> <span class="pl-s1">train_step_start</span>
      <span class="pl-k">if</span> <span class="pl-s1">t_step</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span> <span class="pl-c1">and</span> <span class="pl-s1">time</span>.<span class="pl-en">time</span>() <span class="pl-c1">-</span> <span class="pl-s1">t_step_0_time</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">1</span>:
        <span class="pl-en">print</span>(<span class="pl-s">f'WARNING: First train step took: <span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">time</span>.<span class="pl-en">time</span>()<span class="pl-c1">-</span><span class="pl-s1">t_step_0_time</span>:.2f<span class="pl-kos">}</span></span> s'</span>)
    <span class="pl-k">del</span> <span class="pl-s1">train_batch</span>, <span class="pl-s1">train_batch_arr</span>

    <span class="pl-c"># EVAL</span>
    <span class="pl-c"># ...</span></pre></div>
<ul dir="auto">
<li>Full error messages/tracebacks:</li>
</ul>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Exception in thread Thread-14:
Traceback (most recent call last):
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py&quot;, line 890, in _bootstrap
    self._bootstrap_inner()
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py&quot;, line 926, in _bootstrap_inner
    self.run()
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py&quot;, line 870, in run
    self._target(*self._args, **self._kwargs)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/traceback_util.py&quot;, line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/api.py&quot;, line 476, in cache_miss
    donated_invars=donated_invars, inline=inline, keep_unused=keep_unused)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py&quot;, line 1765, in bind
    return call_bind(self, fun, *args, **params)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py&quot;, line 1781, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py&quot;, line 678, in process_call
    return primitive.impl(f, *tracers, **params)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/dispatch.py&quot;, line 185, in _xla_call_impl
    return compiled_fun(*args)
  File &quot;/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/dispatch.py&quot;, line 615, in _execute_compiled
    out_bufs_flat = compiled.execute(input_bufs_flat)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 122875791936 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:    1.43GiB
              constant allocation:         8B
        maybe_live_out allocation:  581.19MiB
     preallocated temp allocation:  114.44GiB
  preallocated temp fragmentation:  146.50MiB (0.13%)
                 total allocation:  115.86GiB
              total fragmentation:  146.53MiB (0.12%)
Peak buffers:
        Buffer 1:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 2:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 3:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 4:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 5:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 6:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================
"><pre class="notranslate">Exception <span class="pl-k">in</span> thread Thread-14:
Traceback (most recent call last):
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py<span class="pl-pds">"</span></span>, line 890, <span class="pl-k">in</span> _bootstrap
    <span class="pl-en">self._bootstrap_inner</span>()
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py<span class="pl-pds">"</span></span>, line 926, <span class="pl-k">in</span> _bootstrap_inner
    <span class="pl-en">self.run</span>()
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/threading.py<span class="pl-pds">"</span></span>, line 870, <span class="pl-k">in</span> run
    self._target(<span class="pl-k">*</span>self._args, <span class="pl-k">**</span>self._kwargs)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/traceback_util.py<span class="pl-pds">"</span></span>, line 162, <span class="pl-k">in</span> reraise_with_filtered_traceback
    <span class="pl-k">return</span> fun(<span class="pl-k">*</span>args, <span class="pl-k">**</span>kwargs)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/api.py<span class="pl-pds">"</span></span>, line 476, <span class="pl-k">in</span> cache_miss
    donated_invars=donated_invars, inline=inline, keep_unused=keep_unused)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py<span class="pl-pds">"</span></span>, line 1765, <span class="pl-k">in</span> <span class="pl-c1">bind</span>
    <span class="pl-k">return</span> call_bind(self, fun, <span class="pl-k">*</span>args, <span class="pl-k">**</span>params)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py<span class="pl-pds">"</span></span>, line 1781, <span class="pl-k">in</span> call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/core.py<span class="pl-pds">"</span></span>, line 678, <span class="pl-k">in</span> process_call
    <span class="pl-k">return</span> primitive.impl(f, <span class="pl-k">*</span>tracers, <span class="pl-k">**</span>params)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/dispatch.py<span class="pl-pds">"</span></span>, line 185, <span class="pl-k">in</span> _xla_call_impl
    <span class="pl-k">return</span> compiled_fun(<span class="pl-k">*</span>args)
  File <span class="pl-s"><span class="pl-pds">"</span>/mnt/lustre/liujun1/.conda/envs/muNet/lib/python3.7/site-packages/jax/_src/dispatch.py<span class="pl-pds">"</span></span>, line 615, <span class="pl-k">in</span> _execute_compiled
    out_bufs_flat = compiled.execute(input_bufs_flat)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory <span class="pl-k">while</span> trying to allocate 122875791936 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:    1.43GiB
              constant allocation:         8B
        maybe_live_out allocation:  581.19MiB
     preallocated temp allocation:  114.44GiB
  preallocated temp fragmentation:  146.50MiB (0.13%)
                 total allocation:  115.86GiB
              total fragmentation:  146.53MiB (0.12%)
Peak buffers:
        Buffer 1:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 2:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 3:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 4:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 5:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================

        Buffer 6:
                Size: 1.27GiB
                XLA Label: custom-call
                Shape: f32[64,16,577,577]
                ==========================
</pre></div>
<ul dir="auto">
<li>environment:<br>
python: 3.7.13<br>
jax: 0.3.14<br>
jaxlib: 0.3.14+cuda11.cudnn82</li>
</ul>