Text Classification on GLUE on TPU using Jax/Flax : BigBird
Author: Jeevesh8
Pull request: no

The notebook for Text Classification on GLUE tasks, provided [here](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb) can do with some updates:
1. It has a missing ``from flax import traverse_utils`` which can be added.
2. Replace ``gradient_transformation`` with ``adamw(1e-7)`` in ``TrainState.create()``.
3. Probably a note to ``pip install sentencepiece`` since ``RoBERTA`` etc. models which are available in Flax at this point, use it.

Also, I couldn't get the notebook to run for ``google/bigbird-roberta-base`` with ``batch_size=1`` and the default task(``cola``).
I got the following error:
```
Shapes of inputs: (8, 1, 128) (8, 1, 128) (8, 1)
---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
<ipython-input-32-cb9e92a30675> in <module>()
      7         print(batch['input_ids'].shape, batch['attention_mask'].shape, batch['labels'].shape)
----> 8         state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
      9         progress_bar_train.update(1)

67 frames
/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    142     try:
--> 143       return fun(*args, **kwargs)
    144     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in f_pmapped(*args, **kwargs)
   1658         name=flat_fun.__name__, donated_invars=tuple(donated_invars),
-> 1659         global_arg_shapes=tuple(global_arg_shapes_flat))
   1660     return tree_unflatten(out_tree(), out)

/usr/local/lib/python3.7/dist-packages/jax/core.py in bind(self, fun, *args, **params)
   1623     assert len(params['in_axes']) == len(args)
-> 1624     return call_bind(self, fun, *args, **params)
   1625 

/usr/local/lib/python3.7/dist-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1555   with maybe_new_sublevel(top_trace):
-> 1556     outs = primitive.process(top_trace, fun, tracers, params)
   1557   return map(full_lower, apply_todos(env_trace_todo(), outs))

/usr/local/lib/python3.7/dist-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1626   def process(self, trace, fun, tracers, params):
-> 1627     return trace.process_map(self, fun, tracers, params)
   1628 

/usr/local/lib/python3.7/dist-packages/jax/core.py in process_call(self, primitive, f, tracers, params)
    608   def process_call(self, primitive, f, tracers, params):
--> 609     return primitive.impl(f, *tracers, **params)
    610   process_map = process_call

/usr/local/lib/python3.7/dist-packages/jax/interpreters/pxla.py in xla_pmap_impl(fun, backend, axis_name, axis_size, global_axis_size, devices, name, in_axes, out_axes_thunk, donated_invars, global_arg_shapes, *args)
    622                                    donated_invars, global_arg_shapes,
--> 623                                    *abstract_args)
    624   # Don't re-abstractify args unless logging is enabled for performance.

/usr/local/lib/python3.7/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
    261     else:
--> 262       ans = call(fun, *args)
    263       cache[key] = (ans, fun.stores)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/pxla.py in parallel_callable(fun, backend_name, axis_name, axis_size, global_axis_size, devices, name, in_axes, out_axes_thunk, donated_invars, global_arg_shapes, *avals)
    698   with core.extend_axis_env(axis_name, global_axis_size, None):  # type: ignore
--> 699     jaxpr, out_sharded_avals, consts = pe.trace_to_jaxpr_final(fun, global_sharded_avals, transform_name="pmap")
    700   jaxpr = xla.apply_outfeed_rewriter(jaxpr)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals, transform_name)
   1208     main.jaxpr_stack = ()  # type: ignore
-> 1209     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1210     del fun, main

/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1187     in_tracers = map(trace.new_arg, in_avals)
-> 1188     ans = fun.call_wrapped(*in_tracers)
   1189     out_tracers = map(trace.full_raise, ans)

/usr/local/lib/python3.7/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:

<ipython-input-24-f522714f3451> in train_step(state, batch, dropout_rng)
     10     grad_function = jax.value_and_grad(loss_function)
---> 11     loss, grad = grad_function(state.params)
     12     grad = jax.lax.pmean(grad, "batch")

/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    142     try:
--> 143       return fun(*args, **kwargs)
    144     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in value_and_grad_f(*args, **kwargs)
    886     if not has_aux:
--> 887       ans, vjp_py = _vjp(f_partial, *dyn_args)
    888     else:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in _vjp(fun, has_aux, *primals)
   1965     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 1966     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1967     out_tree = out_tree()

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    113   if not has_aux:
--> 114     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    115   else:

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
    100   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
--> 101   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
    102   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate)
    497     fun = trace_to_subjaxpr(fun, main, instantiate)
--> 498     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    499     assert not env

/usr/local/lib/python3.7/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:

<ipython-input-24-f522714f3451> in loss_function(params)
      5     def loss_function(params):
----> 6         logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
      7         loss = state.loss_function(logits, targets)

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, input_ids, attention_mask, token_type_ids, position_ids, params, dropout_rng, train, output_attentions, output_hidden_states, return_dict)
   1420             return_dict,
-> 1421             rngs=rngs,
   1422         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in apply(self, variables, rngs, method, mutable, capture_intermediates, *args, **kwargs)
    938         mutable=mutable, capture_intermediates=capture_intermediates
--> 939     )(variables, *args, **kwargs, rngs=rngs)
    940 

/usr/local/lib/python3.7/dist-packages/flax/core/scope.py in wrapper(variables, rngs, *args, **kwargs)
    686     with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
--> 687       y = fn(root, *args, **kwargs)
    688     if mutable is not False:

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in scope_fn(scope, *args, **kwargs)
   1177     try:
-> 1178       return fn(module.clone(parent=scope), *args, **kwargs)
   1179     finally:

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic, output_attentions, output_hidden_states, return_dict)
   1697             output_hidden_states=output_hidden_states,
-> 1698             return_dict=return_dict,
   1699         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic, output_attentions, output_hidden_states, return_dict)
   1458             output_hidden_states=output_hidden_states,
-> 1459             return_dict=return_dict,
   1460         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, hidden_states, attention_mask, deterministic, output_attentions, output_hidden_states, return_dict)
   1265             output_hidden_states=output_hidden_states,
-> 1266             return_dict=return_dict,
   1267         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, hidden_states, attention_mask, deterministic, output_attentions, output_hidden_states, return_dict)
   1221             layer_outputs = layer(
-> 1222                 hidden_states, attention_mask, deterministic=deterministic, output_attentions=output_attentions
   1223             )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, hidden_states, attention_mask, deterministic, output_attentions)
   1179         attention_outputs = self.attention(
-> 1180             hidden_states, attention_mask, deterministic=deterministic, output_attentions=output_attentions
   1181         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, hidden_states, attention_mask, deterministic, output_attentions)
   1113         attn_outputs = self.self(
-> 1114             hidden_states, attention_mask, deterministic=deterministic, output_attentions=output_attentions
   1115         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in __call__(self, hidden_states, attention_mask, deterministic, output_attentions)
    360             plan_num_rand_blocks=None,
--> 361             output_attentions=output_attentions,
    362         )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in bigbird_block_sparse_attention(self, query_layer, key_layer, value_layer, band_mask, from_mask, to_mask, from_blocked_mask, to_blocked_mask, n_heads, head_size, plan_from_length, plan_num_rand_blocks, output_attentions)
    476                 plan_from_length=plan_from_length,
--> 477                 plan_num_rand_blocks=plan_num_rand_blocks,
    478             )

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:

/usr/local/lib/python3.7/dist-packages/transformers/models/big_bird/modeling_flax_big_bird.py in _bigbird_block_rand_mask_with_head(self, from_seq_length, to_seq_length, from_block_size, to_block_size, num_heads, plan_from_length, plan_num_rand_blocks, window_block_left, window_block_right, global_block_top, global_block_bottom, global_block_left, global_block_right)
   1004                         global_block_left=global_block_left,
-> 1005                         global_block_right=global_block_right,
   1006                     )

UnfilteredStackTrace: ValueError: could not broadcast input array from shape (0) into shape (3)

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
```
### Who can help
@patrickvonplaten @vasudevgupta7  @patil-suraj 


Comments:
I think it is probably not valid to use ``block_sparse`` attention, with ``num_random_blocks=3`` and ``block_size=64`` with a sequence length of just ``128``(which is the length for ``cola``).

The notebook works fine upon changing the ``attention_type`` to ``original_full``, as it should be for such short sequence length. Sorry, my bad.

Probably can raise a better error though; when trying to run the model with invalid combination of smaller sequence lengths, with the default ``num_random_blocks`` and ``block_size``. 

Shall I open a pr? For the notebook or better error message, in case you guys don't have time for one?

Thanks for the proposed changes and the very in-detail issue! Feel free to open a PR :-) 

Regarding the Big Bird bug, maybe @vasudevgupta7 had an idea? :-)

Hey @Jeevesh8, yes block - sparse attention is not supported for this configuration. Your seqlen must be greater than (5 + 2*num_random_blocks)*block size. 

Also it’s recommended to use original_full (instead of block sparse) till 1024 seqlen. 

You can find other constraints in the end of this blog post: https://huggingface.co/blog/big-bird


Let me know if you get into any other troubles. Would be happy to help. 

Thank you @vasudevgupta7 shall I add error messages corresponding to the conditions in your implementation?
Also, @patrickvonplaten there seems to be another possible bug in the Big Bird Tokenizer:
```python
!pip install transformers sentencepiece
from transformers import BigBirdTokenizer
tokenizer = BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")
print(tokenizer.decode(tokenizer.eos_token_id), tokenizer.decode(tokenizer.bos_token_id))
```
The above code produce ``
 
`` in output. Is it supposed to be so? I was expecting it to be ``
 
`` . Are the pre-trained weights provided on hub, corresponding to this or the opposite way? Or they don't use any of those tokens? And just use ``[CLS]`` and ``[SEP]`` ?

@Jeevesh8 feel free to create PR for adding error messages as @patrickvonplaten also approved on those in his comment above.

Regarding tokenizer, you are right. Tokenizer code looks absolutely fine, but tokenizer_config (in Hub) is wrong (this happened because tokenizer_config was not updated after the PR got reviewed earlier). I will fix BigBird tokenizer_config from Hub.  @patrickvonplaten please approve that once. 

Also, i think we are relying on [SEP], [CLS] mainly.

Sorry to trouble again @vasudevgupta7 , but why is it hard-coded [here](https://github.com/huggingface/transformers/blob/0085e712ddf80fa5cd5f355498fe7f13b839eafa/src/transformers/models/big_bird/modeling_flax_big_bird.py#L460) that ``last_idx`` be 1024 ? Shouldn't we be randomly choosing 3 blocks from all over the sequence ?

Also, is ``[CLS]`` token attending to every other token in the implementation? I am trying to find it through code base, but let me know, if you get the time to reply here. Also, wanted to know if I add multiple ``[CLS]`` tokens, will they all attend globally?

Regarding `last_idx`, this is how original bigbird attention has been implemented ([see this](https://github.com/google-research/bigbird/blob/db06498ec8804c6438111938d8654b66ddaccd5d/bigbird/core/attention.py#L743)). Not sure, why bigbird authors kept 1024. We kept same number to keep HF implementation similar to original.

Only first & last blocks (i.e. first 64 & last 64 tokens by default) are global & choosing of global tokens is independent of CLS token.  Now, since `[CLS]` token lies in first block (since usually added as 1st token), it will be global and will attend all other tokens.

Thank you for the quick reply @vasudevgupta7 . 

I have opened an issue in there original repo [here](https://github.com/google-research/bigbird/issues/18) . It will probably shed some light on the topic of the 1024.

Also, what little I gathered from the code, if the last 64 tokens has mostly ``<pad>``, then they won't be attending globally. Am I right on this?

I think that 1024 may be there to avoid this attention to and from pads.. although not sure.

Yeah, last block will not get attended if last 64 tokens are `<pad>`. To mitigate that issue, you can probably increase `block_size` so that overall tokens to attend increases or remains same.

Or probably use dynamic batching and padding with similar sequence length elements batched together.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.