<h2 dir="auto">Environment info</h2>

<ul dir="auto">
<li><code class="notranslate">transformers</code> version: 4.9.0.dev0</li>
<li>Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29</li>
<li>Python version: 3.8.10</li>
<li>PyTorch version (GPU?): not installed (NA)</li>
<li>Tensorflow version (GPU?): 2.5.0 (False)</li>
<li>Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)</li>
<li>Jax version: 0.2.17</li>
<li>JaxLib version: 0.1.68</li>
</ul>
<h3 dir="auto">Who can help</h3>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/patrickvonplaten/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/patrickvonplaten">@patrickvonplaten</a></p>
<h2 dir="auto">Information</h2>
<p dir="auto">Trying to run the experimental GPTNeo Flax script. Are getting the following error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="07/17/2021 16:08:11 - INFO - __main__ - ***** Running training *****
07/17/2021 16:08:11 - INFO - __main__ -   Num examples = 2852257
07/17/2021 16:08:11 - INFO - __main__ -   Num Epochs = 10
07/17/2021 16:08:11 - INFO - __main__ -   Instantaneous batch size per device = 3
07/17/2021 16:08:11 - INFO - __main__ -   Total train batch size (w. parallel &amp; distributed) = 24
07/17/2021 16:08:11 - INFO - __main__ -   Total optimization steps = 1188440
Epoch ... (1/10):   0%|                                                                                                                                                                                        | 0/10 [00:00&lt;?, ?it/sF0717 16:08:46.411695   76098 array.h:414] Check failed: n &lt; sizes_size                                                                                                                                     | 0/118844 [00:00&lt;?, ?it/s]
*** Check failure stack trace: ***
    @     0x7f4a22c7f347  (unknown)
    @     0x7f4a22c7ded4  (unknown)
    @     0x7f4a22c7d9c3  (unknown)
    @     0x7f4a22c7fcc9  (unknown)
    @     0x7f4a1e8e7eee  (unknown)
    @     0x7f4a1e87ab2f  (unknown)
    @     0x7f4a1e878cc2  (unknown)
    @     0x7f4a223fddb4  (unknown)
    @     0x7f4a223ff212  (unknown)
    @     0x7f4a223fce23  (unknown)
    @     0x7f4a1885856f  (unknown)
    @     0x7f4a1e8a3248  (unknown)
    @     0x7f4a1e8a4d2b  (unknown)
    @     0x7f4a1e3f202b  (unknown)
    @     0x7f4a1e8e3001  (unknown)
    @     0x7f4a1e8e0d6a  (unknown)
    @     0x7f4a1e8e08bd  (unknown)
    @     0x7f4a1e8e3001  (unknown)
    @     0x7f4a1e8e0d6a  (unknown)
    @     0x7f4a1e8e08bd  (unknown)
    @     0x7f4a1df5f13f  (unknown)
    @     0x7f4a1df5a52e  (unknown)
    @     0x7f4a1df64292  (unknown)
    @     0x7f4a1df71ffd  (unknown)
    @     0x7f4a1db5c6b6  (unknown)
    @     0x7f4a1db5c014  TpuCompiler_Compile
    @     0x7f4a28dcf956  xla::(anonymous namespace)::TpuCompiler::Compile()
    @     0x7f4a2657f0d4  xla::Service::BuildExecutables()
    @     0x7f4a265751a0  xla::LocalService::CompileExecutables()
    @     0x7f4a264b9e07  xla::LocalClient::Compile()
    @     0x7f4a264942a0  xla::PjRtStreamExecutorClient::Compile()
    @     0x7f4a2408f152  xla::PyClient::Compile()
    @     0x7f4a23e095e2  pybind11::detail::argument_loader&lt;&gt;::call_impl&lt;&gt;()
    @     0x7f4a23e09a51  pybind11::cpp_function::initialize&lt;&gt;()::{lambda()#3}::operator()()
    @     0x7f4a23df0460  pybind11::cpp_function::dispatcher()
    @           0x5f2cc9  PyCFunction_Call
https://symbolize.stripped_domain/r/?trace=7f4a22c7f347,7f4a22c7ded3,7f4a22c7d9c2,7f4a22c7fcc8,7f4a1e8e7eed,7f4a1e87ab2e,7f4a1e878cc1,7f4a223fddb3,7f4a223ff211,7f4a223fce22,7f4a1885856e,7f4a1e8a3247,7f4a1e8a4d2a,7f4a1e3f202a,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1df5f13e,7f4a1df5a52d,7f4a1df64291,7f4a1df71ffc,7f4a1db5c6b5,7f4a1db5c013,7f4a28dcf955,7f4a2657f0d3,7f4a2657519f,7f4a264b9e06,7f4a2649429f,7f4a2408f151,7f4a23e095e1,7f4a23e09a50,7f4a23df045f,5f2cc8&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7f4a2378b000-7f4a2a67dfc0,2a762cd764e70bc90ae4c7f9747c08d7:7f4a15d2d000-7f4a22fae280 
https://symbolize.stripped_domain/r/?trace=7f4acedc218b,7f4acedc220f,7f4a22c7f487,7f4a22c7ded3,7f4a22c7d9c2,7f4a22c7fcc8,7f4a1e8e7eed,7f4a1e87ab2e,7f4a1e878cc1,7f4a223fddb3,7f4a223ff211,7f4a223fce22,7f4a1885856e,7f4a1e8a3247,7f4a1e8a4d2a,7f4a1e3f202a,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1df5f13e,7f4a1df5a52d,7f4a1df64291,7f4a1df71ffc,7f4a1db5c6b5,7f4a1db5c013,7f4a28dcf955,7f4a2657f0d3,7f4a2657519f,7f4a264b9e06,7f4a2649429f&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7f4a2378b000-7f4a2a67dfc0,2a762cd764e70bc90ae4c7f9747c08d7:7f4a15d2d000-7f4a22fae280 
*** SIGABRT received by PID 76098 (TID 76098) on cpu 46 from PID 76098; ***
E0717 16:08:46.484046   76098 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0717 16:08:46.484074   76098 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0717 16:08:46.484099   76098 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0717 16:08:46.484107   76098 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0717 16:08:46.484121   76098 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0717 16:08:46.484133   76098 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0717 16:08:46.484139   76098 coredump_hook.cc:525] RAW: Discarding core.
F0717 16:08:46.411695   76098 array.h:414] Check failed: n &lt; sizes_size 
E0717 16:08:46.761921   76098 process_state.cc:771] RAW: Raising signal 6 with default behavior"><pre class="notranslate"><code class="notranslate">07/17/2021 16:08:11 - INFO - __main__ - ***** Running training *****
07/17/2021 16:08:11 - INFO - __main__ -   Num examples = 2852257
07/17/2021 16:08:11 - INFO - __main__ -   Num Epochs = 10
07/17/2021 16:08:11 - INFO - __main__ -   Instantaneous batch size per device = 3
07/17/2021 16:08:11 - INFO - __main__ -   Total train batch size (w. parallel &amp; distributed) = 24
07/17/2021 16:08:11 - INFO - __main__ -   Total optimization steps = 1188440
Epoch ... (1/10):   0%|                                                                                                                                                                                        | 0/10 [00:00&lt;?, ?it/sF0717 16:08:46.411695   76098 array.h:414] Check failed: n &lt; sizes_size                                                                                                                                     | 0/118844 [00:00&lt;?, ?it/s]
*** Check failure stack trace: ***
    @     0x7f4a22c7f347  (unknown)
    @     0x7f4a22c7ded4  (unknown)
    @     0x7f4a22c7d9c3  (unknown)
    @     0x7f4a22c7fcc9  (unknown)
    @     0x7f4a1e8e7eee  (unknown)
    @     0x7f4a1e87ab2f  (unknown)
    @     0x7f4a1e878cc2  (unknown)
    @     0x7f4a223fddb4  (unknown)
    @     0x7f4a223ff212  (unknown)
    @     0x7f4a223fce23  (unknown)
    @     0x7f4a1885856f  (unknown)
    @     0x7f4a1e8a3248  (unknown)
    @     0x7f4a1e8a4d2b  (unknown)
    @     0x7f4a1e3f202b  (unknown)
    @     0x7f4a1e8e3001  (unknown)
    @     0x7f4a1e8e0d6a  (unknown)
    @     0x7f4a1e8e08bd  (unknown)
    @     0x7f4a1e8e3001  (unknown)
    @     0x7f4a1e8e0d6a  (unknown)
    @     0x7f4a1e8e08bd  (unknown)
    @     0x7f4a1df5f13f  (unknown)
    @     0x7f4a1df5a52e  (unknown)
    @     0x7f4a1df64292  (unknown)
    @     0x7f4a1df71ffd  (unknown)
    @     0x7f4a1db5c6b6  (unknown)
    @     0x7f4a1db5c014  TpuCompiler_Compile
    @     0x7f4a28dcf956  xla::(anonymous namespace)::TpuCompiler::Compile()
    @     0x7f4a2657f0d4  xla::Service::BuildExecutables()
    @     0x7f4a265751a0  xla::LocalService::CompileExecutables()
    @     0x7f4a264b9e07  xla::LocalClient::Compile()
    @     0x7f4a264942a0  xla::PjRtStreamExecutorClient::Compile()
    @     0x7f4a2408f152  xla::PyClient::Compile()
    @     0x7f4a23e095e2  pybind11::detail::argument_loader&lt;&gt;::call_impl&lt;&gt;()
    @     0x7f4a23e09a51  pybind11::cpp_function::initialize&lt;&gt;()::{lambda()#3}::operator()()
    @     0x7f4a23df0460  pybind11::cpp_function::dispatcher()
    @           0x5f2cc9  PyCFunction_Call
https://symbolize.stripped_domain/r/?trace=7f4a22c7f347,7f4a22c7ded3,7f4a22c7d9c2,7f4a22c7fcc8,7f4a1e8e7eed,7f4a1e87ab2e,7f4a1e878cc1,7f4a223fddb3,7f4a223ff211,7f4a223fce22,7f4a1885856e,7f4a1e8a3247,7f4a1e8a4d2a,7f4a1e3f202a,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1df5f13e,7f4a1df5a52d,7f4a1df64291,7f4a1df71ffc,7f4a1db5c6b5,7f4a1db5c013,7f4a28dcf955,7f4a2657f0d3,7f4a2657519f,7f4a264b9e06,7f4a2649429f,7f4a2408f151,7f4a23e095e1,7f4a23e09a50,7f4a23df045f,5f2cc8&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7f4a2378b000-7f4a2a67dfc0,2a762cd764e70bc90ae4c7f9747c08d7:7f4a15d2d000-7f4a22fae280 
https://symbolize.stripped_domain/r/?trace=7f4acedc218b,7f4acedc220f,7f4a22c7f487,7f4a22c7ded3,7f4a22c7d9c2,7f4a22c7fcc8,7f4a1e8e7eed,7f4a1e87ab2e,7f4a1e878cc1,7f4a223fddb3,7f4a223ff211,7f4a223fce22,7f4a1885856e,7f4a1e8a3247,7f4a1e8a4d2a,7f4a1e3f202a,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1e8e3000,7f4a1e8e0d69,7f4a1e8e08bc,7f4a1df5f13e,7f4a1df5a52d,7f4a1df64291,7f4a1df71ffc,7f4a1db5c6b5,7f4a1db5c013,7f4a28dcf955,7f4a2657f0d3,7f4a2657519f,7f4a264b9e06,7f4a2649429f&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7f4a2378b000-7f4a2a67dfc0,2a762cd764e70bc90ae4c7f9747c08d7:7f4a15d2d000-7f4a22fae280 
*** SIGABRT received by PID 76098 (TID 76098) on cpu 46 from PID 76098; ***
E0717 16:08:46.484046   76098 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0717 16:08:46.484074   76098 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0717 16:08:46.484099   76098 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0717 16:08:46.484107   76098 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0717 16:08:46.484121   76098 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0717 16:08:46.484133   76098 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0717 16:08:46.484139   76098 coredump_hook.cc:525] RAW: Discarding core.
F0717 16:08:46.411695   76098 array.h:414] Check failed: n &lt; sizes_size 
E0717 16:08:46.761921   76098 process_state.cc:771] RAW: Raising signal 6 with default behavior
</code></pre></div>
<h2 dir="auto">Environment info</h2>
<ul dir="auto">
<li><code class="notranslate">transformers</code> version: 4.9.0.dev0</li>
<li>Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29</li>
<li>Python version: 3.8.10</li>
<li>PyTorch version (GPU?): 1.9.0+cpu (False)</li>
<li>Tensorflow version (GPU?): 2.7.0-dev20210705 (False)</li>
<li>Flax version: 0.3.4 (tpu)</li>
<li>Jax version: 0.2.16</li>
<li>JaxLib version: 0.1.68</li>
<li>Using distributed or parallel set-up in script?: Yes</li>
</ul>
<h3 dir="auto">Who can help</h3>
<p dir="auto">examples/research_projects/jax/model_parallel<br>
<a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/patil-suraj/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/patil-suraj">@patil-suraj</a></p>
<h2 dir="auto">Information</h2>
<p dir="auto">Model I am using GPTNeo-1.3B (for instance the one with resized to multiple of 8 embedding can be found <a href="https://huggingface.co/flax-community/gpt-neo-1.3B-resized-embed" rel="nofollow">here</a>)</p>
<p dir="auto">The problem arises when using:</p>
<ul dir="auto">
<li>
<p dir="auto">the official example scripts</p>
</li>
<li>
<p dir="auto">my own modified scripts:<br>
Same error is observed with <a href="https://github.com/ncoop57/gpt-code-clippy/blob/main/run_clm_mp_apps.py">customized script</a></p>
</li>
</ul>
<h2 dir="auto">To reproduce</h2>
<p dir="auto">Run the command below in <code class="notranslate">examples/research_projects/jax-projects/model parallel</code> folder in cloned tarnsformers repo:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="python run_clm_mp.py \
    --model_name_or_path flax-community/gpt-neo-1.3B-resized-embed  \
    --tokenizer_name gpt2 \
    --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
    --do_train  --do_eval \
    --block_size 1024 \
    --num_train_epochs 5 \
    --learning_rate 4e-6 \
    --per_device_train_batch_size 3 --per_device_eval_batch_size 3 \
    --overwrite_output_dir --output_dir ~/tmp/flax-clm \
    --cache_dir ~/datasets_cache/wikitext --dtype bfloat16 \
    --logging_steps 96 --eval_steps 96"><pre class="notranslate"><code class="notranslate">python run_clm_mp.py \
    --model_name_or_path flax-community/gpt-neo-1.3B-resized-embed  \
    --tokenizer_name gpt2 \
    --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
    --do_train  --do_eval \
    --block_size 1024 \
    --num_train_epochs 5 \
    --learning_rate 4e-6 \
    --per_device_train_batch_size 3 --per_device_eval_batch_size 3 \
    --overwrite_output_dir --output_dir ~/tmp/flax-clm \
    --cache_dir ~/datasets_cache/wikitext --dtype bfloat16 \
    --logging_steps 96 --eval_steps 96
</code></pre></div>
<p dir="auto">Stack trace:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="07/16/2021 13:59:13 - INFO - absl - A polynomial schedule was set with a non-positive `transition_steps` value; this results in a constant schedule with value `init_value`.
/home/arto/jenv/lib/python3.8/site-packages/jax/experimental/pjit.py:160: UserWarning: pjit is an experimental feature and probably has bugs!
  warn(&quot;pjit is an experimental feature and probably has bugs!&quot;)
07/16/2021 13:59:21 - INFO - __main__ - ***** Running training *****
07/16/2021 13:59:21 - INFO - __main__ -   Num examples = 2318
07/16/2021 13:59:21 - INFO - __main__ -   Num Epochs = 5
07/16/2021 13:59:21 - INFO - __main__ -   Instantaneous batch size per device = 3
07/16/2021 13:59:21 - INFO - __main__ -   Total train batch size (w. parallel &amp; distributed) = 24
07/16/2021 13:59:21 - INFO - __main__ -   Total optimization steps = 480
Epoch ... (1/5):   0%|                                                                                                                    | 0/5 [00:00&lt;?, ?it/sF0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size                                                                   | 0/96 [00:00&lt;?, ?it/s]
*** Check failure stack trace: ***
    @     0x7efd6d030347  (unknown)
    @     0x7efd6d02eed4  (unknown)
    @     0x7efd6d02e9c3  (unknown)
    @     0x7efd6d030cc9  (unknown)
    @     0x7efd68c98eee  (unknown)
    @     0x7efd68c2bb2f  (unknown)
    @     0x7efd68c29cc2  (unknown)
    @     0x7efd6c7aedb4  (unknown)
    @     0x7efd6c7b0212  (unknown)
    @     0x7efd6c7ade23  (unknown)
    @     0x7efd62c0956f  (unknown)
    @     0x7efd68c54248  (unknown)
    @     0x7efd68c55d2b  (unknown)
    @     0x7efd687a302b  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd6831013f  (unknown)
    @     0x7efd6830b52e  (unknown)
    @     0x7efd68315292  (unknown)
    @     0x7efd68322ffd  (unknown)
    @     0x7efd67f0d6b6  (unknown)
    @     0x7efd67f0d014  TpuCompiler_Compile
    @     0x7efd73180956  xla::(anonymous namespace)::TpuCompiler::Compile()
    @     0x7efd709300d4  xla::Service::BuildExecutables()
    @     0x7efd709261a0  xla::LocalService::CompileExecutables()
    @     0x7efd7086ae07  xla::LocalClient::Compile()
    @     0x7efd708452a0  xla::PjRtStreamExecutorClient::Compile()
    @     0x7efd6e440152  xla::PyClient::Compile()
    @     0x7efd6e1ba5e2  pybind11::detail::argument_loader&lt;&gt;::call_impl&lt;&gt;()
    @     0x7efd6e1baa51  pybind11::cpp_function::initialize&lt;&gt;()::{lambda()#3}::operator()()
    @     0x7efd6e1a1460  pybind11::cpp_function::dispatcher()
    @           0x5f2cc9  PyCFunction_Call
https://symbolize.stripped_domain/r/?trace=7efd6d030347,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f,7efd6e440151,7efd6e1ba5e1,7efd6e1baa50,7efd6e1a145f,5f2cc8&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
https://symbolize.stripped_domain/r/?trace=7eff9cd0b18b,7eff9cd0b20f,7efd6d030487,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
*** SIGABRT received by PID 14290 (TID 14290) on cpu 89 from PID 14290; ***
E0716 13:59:49.681807   14290 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0716 13:59:49.681854   14290 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0716 13:59:49.681862   14290 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0716 13:59:49.681870   14290 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0716 13:59:49.681876   14290 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0716 13:59:49.681886   14290 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0716 13:59:49.681891   14290 coredump_hook.cc:525] RAW: Discarding core.
F0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size 
E0716 13:59:49.953522   14290 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)"><pre class="notranslate"><code class="notranslate">07/16/2021 13:59:13 - INFO - absl - A polynomial schedule was set with a non-positive `transition_steps` value; this results in a constant schedule with value `init_value`.
/home/arto/jenv/lib/python3.8/site-packages/jax/experimental/pjit.py:160: UserWarning: pjit is an experimental feature and probably has bugs!
  warn("pjit is an experimental feature and probably has bugs!")
07/16/2021 13:59:21 - INFO - __main__ - ***** Running training *****
07/16/2021 13:59:21 - INFO - __main__ -   Num examples = 2318
07/16/2021 13:59:21 - INFO - __main__ -   Num Epochs = 5
07/16/2021 13:59:21 - INFO - __main__ -   Instantaneous batch size per device = 3
07/16/2021 13:59:21 - INFO - __main__ -   Total train batch size (w. parallel &amp; distributed) = 24
07/16/2021 13:59:21 - INFO - __main__ -   Total optimization steps = 480
Epoch ... (1/5):   0%|                                                                                                                    | 0/5 [00:00&lt;?, ?it/sF0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size                                                                   | 0/96 [00:00&lt;?, ?it/s]
*** Check failure stack trace: ***
    @     0x7efd6d030347  (unknown)
    @     0x7efd6d02eed4  (unknown)
    @     0x7efd6d02e9c3  (unknown)
    @     0x7efd6d030cc9  (unknown)
    @     0x7efd68c98eee  (unknown)
    @     0x7efd68c2bb2f  (unknown)
    @     0x7efd68c29cc2  (unknown)
    @     0x7efd6c7aedb4  (unknown)
    @     0x7efd6c7b0212  (unknown)
    @     0x7efd6c7ade23  (unknown)
    @     0x7efd62c0956f  (unknown)
    @     0x7efd68c54248  (unknown)
    @     0x7efd68c55d2b  (unknown)
    @     0x7efd687a302b  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd6831013f  (unknown)
    @     0x7efd6830b52e  (unknown)
    @     0x7efd68315292  (unknown)
    @     0x7efd68322ffd  (unknown)
    @     0x7efd67f0d6b6  (unknown)
    @     0x7efd67f0d014  TpuCompiler_Compile
    @     0x7efd73180956  xla::(anonymous namespace)::TpuCompiler::Compile()
    @     0x7efd709300d4  xla::Service::BuildExecutables()
    @     0x7efd709261a0  xla::LocalService::CompileExecutables()
    @     0x7efd7086ae07  xla::LocalClient::Compile()
    @     0x7efd708452a0  xla::PjRtStreamExecutorClient::Compile()
    @     0x7efd6e440152  xla::PyClient::Compile()
    @     0x7efd6e1ba5e2  pybind11::detail::argument_loader&lt;&gt;::call_impl&lt;&gt;()
    @     0x7efd6e1baa51  pybind11::cpp_function::initialize&lt;&gt;()::{lambda()#3}::operator()()
    @     0x7efd6e1a1460  pybind11::cpp_function::dispatcher()
    @           0x5f2cc9  PyCFunction_Call
https://symbolize.stripped_domain/r/?trace=7efd6d030347,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f,7efd6e440151,7efd6e1ba5e1,7efd6e1baa50,7efd6e1a145f,5f2cc8&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
https://symbolize.stripped_domain/r/?trace=7eff9cd0b18b,7eff9cd0b20f,7efd6d030487,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
*** SIGABRT received by PID 14290 (TID 14290) on cpu 89 from PID 14290; ***
E0716 13:59:49.681807   14290 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0716 13:59:49.681854   14290 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0716 13:59:49.681862   14290 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0716 13:59:49.681870   14290 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0716 13:59:49.681876   14290 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0716 13:59:49.681886   14290 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0716 13:59:49.681891   14290 coredump_hook.cc:525] RAW: Discarding core.
F0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size 
E0716 13:59:49.953522   14290 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">Training in model parallel mode.</p>