/home/dat/pino/lib/python3.8/site-packages/jax/lib/xla_bridge.py:382: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.
  warnings.warn(
/home/dat/pino/lib/python3.8/site-packages/jax/lib/xla_bridge.py:369: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.
  warnings.warn(
Epoch ... (1/5):   0%|                                                     | 0/5 [00:00<?, ?it/s]
Epoch ... (1/5):   0%|                                                     | 0/5 [01:59<?, ?it/s]
Traceback (most recent call last):
  File "./run_mlm_flax.py", line 725, in <module>
    state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/_src/api.py", line 1647, in f_pmapped
    out = pxla.xla_pmap(
  File "/home/dat/pino/lib/python3.8/site-packages/jax/core.py", line 1620, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/core.py", line 1551, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/core.py", line 1623, in process
    return trace.process_map(self, fun, tracers, params)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/core.py", line 606, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 624, in xla_pmap_impl
    compiled_fun, fingerprint = parallel_callable(fun, backend, axis_name, axis_size,
  File "/home/dat/pino/lib/python3.8/site-packages/jax/linear_util.py", line 262, in memoized_fun
    ans = call(fun, *args)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 899, in parallel_callable
    compiled = xla.backend_compile(backend, built, compile_options)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/interpreters/xla.py", line 360, in backend_compile
    return backend.compile(built_c, compile_options=options)
jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: Resource exhausted: Ran out of memory in memory space hbm. Used 77.35G of 15.48G hbm. Exceeded hbm capacity by 61.87G.
Total hbm usage >= 77.87G:
    reserved        530.00M
    program          77.35G
    arguments            0B
Output size 0B; shares 0B with arguments.
Program hbm requirement 77.35G:
    global           692.0K
    scoped           253.0K
    HLO temp         77.35G (97.9% utilization: Unpadded (75.71G) Padded (77.34G), 0.0% fragmentation (10.14M))
  Largest program allocations in hbm:
  1. Size: 12.29G
     Operator: op_type="dot_general" op_name="pmap(train_step)/dot_general[ dimension_numbers=(((2,), (0,)), ((), ()))\n                              precision=None\n                              preferred_element_type=None ]" source_file="/home/dat/pino/lib/python3.8/site-packages/flax/linen/linear.py" source_line=175
     Shape: f32[16,4096,50358]{1,2,0:T(8,128)}
     Unpadded size: 12.29G
     Extra memory due to padding: 512.0K (1.0x expansion)
     XLA label: %fusion.1737.remat7 = f32[16,4096,50358]{1,2,0:T(8,128)} fusion(f32[50358]{0:T(1024)} %get-tuple-element.23314, f32[768,50358,1]{0,1,2:T(8,128)} %bitcast.5719, f32[768]{0:T(1024)} %get-tuple-element.23315, f32[768]{0:T(1024)} %get-tuple-element.23316, f32[...
     Allocation type: HLO temp
     ==========================
  2. Size: 2.62G
     Operator: op_type="mul" op_name="pmap(train_step)/mul" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)}
     Unpadded size: 2.62G
     XLA label: %fusion.1803 = f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} fusion(f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} %fusion.1984, f32[16,12,28,128]{3,2,1,0:T(8,128)} %negate.34, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2407, f32[16,12,28,128]{3,2,1,0:T(8,128)} ...
     Allocation type: HLO temp
     ==========================
  3. Size: 2.62G
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(transpose(jvp(_einsum)))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                           precision=None\n                                                           preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=626
     Shape: f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)}
     Unpadded size: 2.62G
     XLA label: %fusion.1984 = f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} fusion(f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} %fusion.1986, pred[]{:T(1024)} %fusion.9668, bf16[16,12,28,384,64]{3,2,4,1,0:T(8,128)(2,1)} %fusion.2490.remat_uncompressed.remat, f32[16,12,32,128,...
     Allocation type: HLO temp
     ==========================
  4. Size: 1.31G
     Operator: op_type="div" op_name="pmap(train_step)/div" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.1816.remat2 = bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)} fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2418, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9170, f32[16,12,28,128,128]{3,4,2,1,0:T(8,128)} %fusion.3004.remat3, f32[16,12,28,12...
     Allocation type: HLO temp
     ==========================
  5. Size: 1.31G
     Operator: op_type="div" op_name="pmap(train_step)/div" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.1805.remat6 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2407, f32[16,12,28,128]{3,2,1,0:...
     Allocation type: HLO temp
     ==========================
  6. Size: 1.31G
     Operator: op_type="div" op_name="pmap(train_step)/div" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.1805.remat6 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2407, f32[16,12,28,128]{3,2,1,0:...
     Allocation type: HLO temp
     ==========================
  7. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13419 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2417, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9169, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  8. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13420 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2416, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9168, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  9. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13421 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2415, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9167, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  10. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13422 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2414, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9166, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  11. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13423 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2413, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9165, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  12. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13424 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2412, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9164, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  13. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13425 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2411, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9163, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  14. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13426 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2410, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9162, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  15. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13427 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2409, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9161, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  16. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13428 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2408, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9160, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  17. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=591
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.11600 = (f32[16,12,28,128]{3,2,1,0:T(8,128)}, f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}) fusion(s32[16,12,30,128,384]{3,4,2,1,0:T(8,128)} %get-tuple-element.25240, bf16[16,12,28,384,64]{3,2,1,0,4:T(8,128)(2,1)} %slice.28186.remat_uncompressed, f32[...
     Allocation type: HLO temp
     ==========================
  18. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=584
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.2061.remat6 = f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)} fusion(f32[16,28,128,384]{2,3,1,0:T(8,128)} %get-tuple-element.20979, bf16[16,12,28,384,64]{3,2,4,1,0:T(8,128)(2,1)} %fusion.2489.remat_uncompressed, f32[16,12,32,128,64]{3,2,4,1,0:T(8,128)} %...
     Allocation type: HLO temp
     ==========================
  19. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=591
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.11584 = (f32[16,12,28,128]{3,2,1,0:T(8,128)}, f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}) fusion(s32[16,12,30,128,384]{3,4,2,1,0:T(8,128)} %get-tuple-element.25199, bf16[16,12,28,384,64]{3,2,1,0,4:T(8,128)(2,1)} %slice.28115, f32[16,12,32,128,64]{3,...
     Allocation type: HLO temp
     ==========================
  20. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=591
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.11592 = (f32[16,12,28,128]{3,2,1,0:T(8,128)}, f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}) fusion(s32[16,12,30,128,384]{3,4,2,1,0:T(8,128)} %get-tuple-element.25220, bf16[16,12,28,384,64]{3,2,1,0,4:T(8,128)(2,1)} %slice.28151.remat_uncompressed, f32[...
     Allocation type: HLO temp
     ==========================
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
--------------------
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "./run_mlm_flax.py", line 725, in <module>
    state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
  File "/home/dat/pino/lib/python3.8/site-packages/jax/interpreters/xla.py", line 360, in backend_compile
    return backend.compile(built_c, compile_options=options)
RuntimeError: Resource exhausted: Ran out of memory in memory space hbm. Used 77.35G of 15.48G hbm. Exceeded hbm capacity by 61.87G.
Total hbm usage >= 77.87G:
    reserved        530.00M
    program          77.35G
    arguments            0B
Output size 0B; shares 0B with arguments.
Program hbm requirement 77.35G:
    global           692.0K
    scoped           253.0K
    HLO temp         77.35G (97.9% utilization: Unpadded (75.71G) Padded (77.34G), 0.0% fragmentation (10.14M))
  Largest program allocations in hbm:
  1. Size: 12.29G
     Operator: op_type="dot_general" op_name="pmap(train_step)/dot_general[ dimension_numbers=(((2,), (0,)), ((), ()))\n                              precision=None\n                              preferred_element_type=None ]" source_file="/home/dat/pino/lib/python3.8/site-packages/flax/linen/linear.py" source_line=175
     Shape: f32[16,4096,50358]{1,2,0:T(8,128)}
     Unpadded size: 12.29G
     Extra memory due to padding: 512.0K (1.0x expansion)
     XLA label: %fusion.1737.remat7 = f32[16,4096,50358]{1,2,0:T(8,128)} fusion(f32[50358]{0:T(1024)} %get-tuple-element.23314, f32[768,50358,1]{0,1,2:T(8,128)} %bitcast.5719, f32[768]{0:T(1024)} %get-tuple-element.23315, f32[768]{0:T(1024)} %get-tuple-element.23316, f32[...
     Allocation type: HLO temp
     ==========================
  2. Size: 2.62G
     Operator: op_type="mul" op_name="pmap(train_step)/mul" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)}
     Unpadded size: 2.62G
     XLA label: %fusion.1803 = f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} fusion(f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} %fusion.1984, f32[16,12,28,128]{3,2,1,0:T(8,128)} %negate.34, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2407, f32[16,12,28,128]{3,2,1,0:T(8,128)} ...
     Allocation type: HLO temp
     ==========================
  3. Size: 2.62G
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(transpose(jvp(_einsum)))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                           precision=None\n                                                           preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=626
     Shape: f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)}
     Unpadded size: 2.62G
     XLA label: %fusion.1984 = f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} fusion(f32[16,12,28,128,1024]{3,4,2,1,0:T(8,128)} %fusion.1986, pred[]{:T(1024)} %fusion.9668, bf16[16,12,28,384,64]{3,2,4,1,0:T(8,128)(2,1)} %fusion.2490.remat_uncompressed.remat, f32[16,12,32,128,...
     Allocation type: HLO temp
     ==========================
  4. Size: 1.31G
     Operator: op_type="div" op_name="pmap(train_step)/div" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.1816.remat2 = bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)} fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2418, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9170, f32[16,12,28,128,128]{3,4,2,1,0:T(8,128)} %fusion.3004.remat3, f32[16,12,28,12...
     Allocation type: HLO temp
     ==========================
  5. Size: 1.31G
     Operator: op_type="div" op_name="pmap(train_step)/div" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.1805.remat6 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2407, f32[16,12,28,128]{3,2,1,0:...
     Allocation type: HLO temp
     ==========================
  6. Size: 1.31G
     Operator: op_type="div" op_name="pmap(train_step)/div" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=619
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.1805.remat6 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2407, f32[16,12,28,128]{3,2,1,0:...
     Allocation type: HLO temp
     ==========================
  7. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13419 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2417, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9169, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  8. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13420 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2416, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9168, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  9. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13421 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2415, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9167, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  10. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13422 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2414, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9166, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  11. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13423 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2413, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9165, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  12. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13424 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2412, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9164, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  13. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13425 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2411, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9163, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  14. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13426 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2410, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9162, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  15. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13427 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2409, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9161, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  16. Size: 1.31G
     Shape: bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}
     Unpadded size: 1.31G
     XLA label: %fusion.13428 = (bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}, bf16[16,12,28,128,1024]{3,4,2,1,0:T(8,128)(2,1)}) fusion(f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.2408, f32[16,12,28,128]{3,2,1,0:T(8,128)} %fusion.9160, f32[16,12,28,128,128]{3,4,2,1,0:...
     Allocation type: HLO temp
     ==========================
  17. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=591
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.11600 = (f32[16,12,28,128]{3,2,1,0:T(8,128)}, f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}) fusion(s32[16,12,30,128,384]{3,4,2,1,0:T(8,128)} %get-tuple-element.25240, bf16[16,12,28,384,64]{3,2,1,0,4:T(8,128)(2,1)} %slice.28186.remat_uncompressed, f32[...
     Allocation type: HLO temp
     ==========================
  18. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=584
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.2061.remat6 = f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)} fusion(f32[16,28,128,384]{2,3,1,0:T(8,128)} %get-tuple-element.20979, bf16[16,12,28,384,64]{3,2,4,1,0:T(8,128)(2,1)} %fusion.2489.remat_uncompressed, f32[16,12,32,128,64]{3,2,4,1,0:T(8,128)} %...
     Allocation type: HLO temp
     ==========================
  19. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=591
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.11584 = (f32[16,12,28,128]{3,2,1,0:T(8,128)}, f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}) fusion(s32[16,12,30,128,384]{3,4,2,1,0:T(8,128)} %get-tuple-element.25199, bf16[16,12,28,384,64]{3,2,1,0,4:T(8,128)(2,1)} %slice.28115, f32[16,12,32,128,64]{3,...
     Allocation type: HLO temp
     ==========================
  20. Size: 1008.00M
     Operator: op_type="dot_general" op_name="pmap(train_step)/jit(jvp(_einsum))/dot_general[ dimension_numbers=(((4,), (4,)), ((0, 1, 2), (0, 1, 2)))\n                                                precision=None\n                                                preferred_element_type=None ]" source_file="/home/dat/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py" source_line=591
     Shape: f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}
     Unpadded size: 1008.00M
     XLA label: %fusion.11592 = (f32[16,12,28,128]{3,2,1,0:T(8,128)}, f32[16,12,28,128,384]{3,4,2,1,0:T(8,128)}) fusion(s32[16,12,30,128,384]{3,4,2,1,0:T(8,128)} %get-tuple-element.25220, bf16[16,12,28,384,64]{3,2,1,0,4:T(8,128)(2,1)} %slice.28151.remat_uncompressed, f32[...
     Allocation type: HLO temp
     ==========================