TPU out of memory (OOM) with flax train a language model GPT2
Author: sarataylor2000
Pull request: no

I'm training the GPT2 on the Flax and TPU i get TPU out of memory error while i have enough memory and without filling the memory it says out of memory like this:
** jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESO**URCE_EXHAUSTED: Ran out of memory in memory space hbm. Used 17.91G of 15.48G hbm. Exceeded hbm capacity by 2.43G.**

Im training the below code:
the orginal repository use 64 batch size it works but when i use 256 batch size i get OOM TPU error while i have much more memory 

Github page of below code:
**https://github.com/huggingface/transformers/tree/main/examples/flax/language-modeling**
**python run_clm_flax.py \
    --output_dir="./norwegian-gpt2" \
    --model_type="gpt2" \
    --config_name="./norwegian-gpt2" \
    --tokenizer_name="./norwegian-gpt2" \
    --dataset_name="oscar" \
    --dataset_config_name="unshuffled_deduplicated_no" \
    --do_train --do_eval \
    --block_size="512" \
    --per_device_train_batch_size="256" \
    --per_device_eval_batch_size="256" \
    --learning_rate="5e-3" --warmup_steps="1000" \
    --adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
    --overwrite_output_dir \
    --num_train_epochs="20" \
    --logging_steps="500" \
    --save_steps="2500" \
    --eval_steps="2500" \
    --push_to_hub**

in the github i found a similar erro:
**link: https://github.com/google/flax/discussions/1690**
it propose to implement the model on cup then transfer it to TPU but it did not explain where to transfer to cpu and when back to TPU:

here is the explanation to transfer on cup first:

**This is quite odd for sure. Fragmentation and being close to the limit in terms of memory could off course result in errors that appear almost randomly. One thing you could try is to initialize the model on CPU jax.jit(model.init, backend="cpu") The params are moved to TPU automatically during training or during replication of the state (eg jax_utils.replicate)**


Here is the FULL  error i get while training on TPU:

Traceback (most recent call last):
  File "/kaggle/input/yyyyyyyyyyyyy/casual_model_unsupervised_train.py.txt", line 845, in <module>
    main()
  File "/kaggle/input/yyyyyyyyyyyyy/casual_model_unsupervised_train.py.txt", line 752, in main
    state, train_metric = p_train_step(state, batch)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/api.py", line 2026, in cache_miss
    out_tree, out_flat = f_pmapped_(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/api.py", line 1902, in pmap_f
    out = pxla.xla_pmap(
  File "/usr/local/lib/python3.8/site-packages/jax/core.py", line 1859, in bind
    return map_bind(self, fun, *args, **params)
  File "/usr/local/lib/python3.8/site-packages/jax/core.py", line 1891, in map_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/usr/local/lib/python3.8/site-packages/jax/core.py", line 1862, in process
    return trace.process_map(self, fun, tracers, params)
  File "/usr/local/lib/python3.8/site-packages/jax/core.py", line 680, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/usr/local/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 792, in xla_pmap_impl
    compiled_fun, fingerprint = parallel_callable(
  File "/usr/local/lib/python3.8/site-packages/jax/linear_util.py", line 285, in memoized_fun
    ans = call(fun, *args)
  File "/usr/local/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 823, in parallel_callable
    pmap_executable = pmap_computation.compile()
  File "/usr/local/lib/python3.8/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 1091, in compile
    self._executable = PmapExecutable.from_hlo(self._hlo, **self.compile_args)
  File "/usr/local/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 1214, in from_hlo
    compiled = dispatch.compile_or_get_cached(
  File "/usr/local/lib/python3.8/site-packages/jax/_src/dispatch.py", line 768, in compile_or_get_cached
    return backend_compile(backend, computation, compile_options)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/jax/_src/dispatch.py", line 713, in backend_compile
    return backend.compile(built_c, compile_options=options)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Ran out of memory in memory space hbm. Used 17.91G of 15.48G hbm. Exceeded hbm capacity by 2.43G.

Total hbm usage >= 18.43G:
    reserved        530.00M 
    program          17.91G 
    arguments            0B 

Output size 0B; shares 0B with arguments.

Program hbm requirement 17.91G:
    global           132.0K
    HLO temp         17.91G (99.8% utilization: Unpadded (14.75G) Padded (14.78G), 17.5% fragmentation (3.13G))

  Largest program allocations in hbm:

  1. Size: 6.14G
     Operator: op_name="pmap(train_step)/jit(main)/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/usr/local/lib/python3.8/site-packages/flax/linen/linear.py" source_line=196
     Shape: f32[64,511,50257]{1,2,0:T(8,128)}
     Unpadded size: 6.12G
     Extra memory due to padding: 13.14M (1.0x expansion)
     XLA label: fusion.3719 = fusion(get-tuple-element.1407, bitcast.1602), kind=kOutput, calls=fused_computation.2774
     Allocation type: HLO temp
     ==========================

  2. Size: 3.07G
     Operator: op_name="pmap(train_step)/jit(main)/jit(transpose(jvp(log_softmax)))/add_any" source_file="/usr/local/lib/python3.8/site-packages/optax/_src/loss.py" source_line=172
     Shape: bf16[64,511,50257]{1,2,0:T(8,128)(2,1)}
     Unpadded size: 3.06G
     Extra memory due to padding: 6.57M (1.0x expansion)
     XLA label: fusion.8 = fusion(get-tuple-element.1576, get-tuple-element.1575, slice.3481, divide.161), kind=kLoop, calls=fused_computation.8
     Allocation type: HLO temp
     ==========================


**Does any smart guy know the solution to this problem?**
thanks



Comments:
Or you just need to use a smaller batch size to avoid the OOM error. cc @sanchit-gandhi who may have other ideas.




> Or you just need to use a smaller batch size to avoid the OOM error. cc @sanchit-gandhi who may have other ideas.

thanks for the reply i really appreciate it
 smaller batch size  has some disadvantages 
1- it takes huge amount of time for training
2-it is impossible to train large models like gpt-neo 1.3B and 2.7B parameters if i set the batch size to 1 still i get the OOM error so training large models would be to tally impossible, 
any way do you think is there any reliable solution like transfer model on cpu then TPU or change the TPU memory limitation directly?


Hey @sarataylor2000!

> the orginal repository use 64 batch size

I believe this was the _effective_ batch size, which is computed as: `effective_batch_size = per_device_batch_size * num_devices`

In your script, you're setting `per_device_batch_size=256`. Supposing you're using a TPU v3-8, you have 8 TPU cores, which means you have 8 devices. This means your effective batch size is: 256 * 8 = 2048. This is a factor of 2048/64 = 32 times larger than the original repo!

If you need to use an effective batch size of 64, you can work out your per device batch size as: `per_device_batch_size = effective_batch_size / num_devices = 64 / 8 = 8`

So we only need a `per_device_batch_size=8` here!

> it propose to implement the model on cup then transfer it to TPU but it did not explain where to transfer to cpu and when back to TPU

The OOM memory we're getting with your example is happening on the `pmap` step, i.e. when we're performing training. The model is loaded up perfectly fine, so no need to change the model loading logic (transferring from CPU -> TPU)

You can try setting `dtype="bfloat16"` to load the params in bfloat16 precision and save memory on the model weights + optimiser states: https://github.com/huggingface/transformers/blob/92ce53aab859012f7714dae6d6fce7a7d701e75f/examples/flax/language-modeling/run_mlm_flax.py#L168

> Hey @sarataylor2000!
> 
> > the orginal repository use 64 batch size
> 
> I believe this was the _effective_ batch size, which is computed as: `effective_batch_size = per_device_batch_size * num_devices`
> 
> In your script, you're setting `per_device_batch_size=256`. Supposing you're using a TPU v3-8, you have 8 TPU cores, which means you have 8 devices. This means your effective batch size is: 256 * 8 = 2048. This is a factor of 2048/64 = 32 times larger than the original repo!
> 
> If you need to use an effective batch size of 64, you can work out your per device batch size as: `per_device_batch_size = effective_batch_size / num_devices = 64 / 8 = 8`
> 
> So we only need a `per_device_batch_size=8` here!
> 
> > it propose to implement the model on cup then transfer it to TPU but it did not explain where to transfer to cpu and when back to TPU
> 
> The OOM memory we're getting with your example is happening on the `pmap` step, i.e. when we're performing training. The model is loaded up perfectly fine, so no need to change the model loading logic (transferring from CPU -> TPU)
> 
> You can try setting `dtype="bfloat16"` to load the params in bfloat16 precision and save memory on the model weights + optimiser states:
> 
> https://github.com/huggingface/transformers/blob/92ce53aab859012f7714dae6d6fce7a7d701e75f/examples/flax/language-modeling/run_mlm_flax.py#L168

first of  thanks so much for reply i really appreciate it, 

1_the total batch is 256 and per devise batch is 32
2_for changing the dtype it works for models up to T5_large but doesn't work with models like T5 x-large or T5 xx-large
or gpt-nep 1.3 & 2.3 b parameter and gpt-J 8bit this big models if i set per device batch=1 NOT WORK AT ALL!
**so is there any way to prevent OOM and fine-tune this big models?**
 



ok seems no one is capable to solve this TPU OOM problem i suspend here there would be some genuine guys but all SUCKS! No one know the solution

@sarataylor2000 We do not tolerate this kind of language in this repository. You can learn more by reading our code of conduct [here](https://github.com/huggingface/transformers/blob/main/CODE_OF_CONDUCT.md). As a result, I have blocked you for seven days.
This is an opensource repository, you get all the code for free. You are not entitled to an answer or a free debugging session.

Hey @sarataylor2000,

What TPU device are you using? A v3-8? If so, it's going to be difficult running super large models like T5 XXL using `pmap`, as each TPU core only has 16GB of memory.

One thing we should definitely try is enabling gradient checkpointing:
https://github.com/huggingface/transformers/blob/21a2d900eceeded7be9edc445b56877b95eda4ca/examples/flax/language-modeling/run_mlm_flax.py#L110

If that doesn't work, we'll have to resort to some heavy engineering to make this work. This is very advanced and thus outside the scope of the `transformers` library, so I've left some pointers here:
1. Add [`scan_with_axes`](https://github.com/google/flax/blob/1f6b0949d964fbc99f8f8b9541caff54226d0a78/flax/linen/partitioning.py#L378). See [seq2seq-speech/modeling_flax_bart.py](https://github.com/sanchit-gandhi/seq2seq-speech/blob/main/models/modeling_flax_bart.py) for an example of the code changes you need to make here
2. Use `pmap` to shard the model and activations across devices, see [JAX pmap](https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html) and [bloom-jax-inference](https://github.com/huggingface/bloom-jax-inference/blob/main/bloom_inference/modeling_bloom/modeling_bloom.py)

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/main/CONTRIBUTING.md) are likely to be ignored.