Issue 1:
Environment info

transformers version: 4.9.0.dev0

Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29

Python version: 3.8.10

PyTorch version (GPU?): not installed (NA)

Tensorflow version (GPU?): 2.5.0 (False)

Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)

Jax version: 0.2.17

JaxLib version: 0.1.68

Who can help

@patrickvonplaten

Information

Trying to run the experimental GPTNeo Flax script. Are getting the following error:

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 & distributed) = 24
07/17/2021 16:08:11 - INFO - __main__ - Total optimization steps = 1188440
Epoch ... (1/10): 0%| | 0/10 [00:00<?, ?it/sF0717 16:08:46.411695 76098 array.h:414] Check failed: n < sizes_size | 0/118844 [00:00<?, ?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<>::call_impl<>()
@ 0x7f4a23e09a51 pybind11::cpp_function::initialize<>()::{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&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&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 < sizes_size
E0717 16:08:46.761921 76098 process_state.cc:771] RAW: Raising signal 6 with default behavior

Issue 2:
Environment info

transformers version: 4.9.0.dev0

Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29

Python version: 3.8.10

PyTorch version (GPU?): 1.9.0+cpu (False)

Tensorflow version (GPU?): 2.7.0-dev20210705 (False)

Flax version: 0.3.4 (tpu)

Jax version: 0.2.16

JaxLib version: 0.1.68

Using distributed or parallel set-up in script?: Yes

Who can help

examples/research_projects/jax/model_parallel

@patil-suraj

Information

Model I am using GPTNeo-1.3B (for instance the one with resized to multiple of 8 embedding can be found here)

The problem arises when using:

the official example scripts

my own modified scripts:

Same error is observed with customized script

To reproduce

Run the command below in examples/research_projects/jax-projects/model parallel folder in cloned tarnsformers repo:

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

Stack trace:

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 & distributed) = 24
07/16/2021 13:59:21 - INFO - __main__ - Total optimization steps = 480
Epoch ... (1/5): 0%| | 0/5 [00:00<?, ?it/sF0716 13:59:49.611617 14290 array.h:414] Check failed: n < sizes_size | 0/96 [00:00<?, ?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<>::call_impl<>()
@ 0x7efd6e1baa51 pybind11::cpp_function::initialize<>()::{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&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&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 < sizes_size
E0716 13:59:49.953522 14290 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)

Expected behavior

Training in model parallel mode.