# Running BigBird on TPUs **URL:** **Author:** ![Neel-Gupta](https://avatars.discourse-cdn.com/v4/letter/n/f1d935/32.png) [@Neel-Gupta](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FNeel-Gupta)%5C **Post date:** [June 24, 2021, 6:57pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Frunning-bigbird-on-tpus%2F7289%2F1 "2021-06-24T18:57:33Z") This is in continuation of my previous issue filed in this forum ([here](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbug-confirmation-bigbirdlm-not-able-to-use-flax%2F6959%2F11)). I wanted to try and run the `run_mlm_flax` [script](https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_mlm_flax.py) with a Colab TPU v2, but I am facing an error:- ```auto Module Name: WARNING:root:TPU has started up successfully with version pytorch-1.9 Traceback (most recent call last): File "xla_spawn.py", line 87, in main() File "xla_spawn.py", line 83, in main xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores) AttributeError: module 'run_mlm_flax' has no attribute '_mp_fn' ``` To my understanding, the script for spawning processes in each TPU core required a Huggingface `Trainer` object to work with. However, we apparently do not have a Trainer script for MLM/CLM for Bigbird. Compounding the difficulties is the whole `Flax-noTrainer-TPU` combination which I do not think `BigBird` has been tested out fully. Is there any possible approach to run this on TPUs? ---
**Author:** ![vasudevgupta](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/vasudevgupta/32/1251_2.png) [@vasudevgupta](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fvasudevgupta)%5C **Post date:** [June 24, 2021, 7:20pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Frunning-bigbird-on-tpus%2F7289%2F2 "2021-06-24T19:20:27Z")
Hey, I think you are mixing flax & pytorch scripts. how can you get pytorch error in flax script otherwise? FlaxBigBird works perfectly on cloud TPU. You can refer [this script](https://github.com/vasudevgupta7/bigbird/blob/main/src/bigbird_flax.py) also. Also Currently 🤗Trainer is for only PyTorch & Tensorflow. ---
**Author:** ![Neel-Gupta](https://avatars.discourse-cdn.com/v4/letter/n/f1d935/32.png) [@Neel-Gupta](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FNeel-Gupta)%5C **Post date:** [June 24, 2021, 7:29pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Frunning-bigbird-on-tpus%2F7289%2F3 "2021-06-24T19:29:58Z")
_sigh_, it’s the third time that I have had to change my scripts again. Why is that with HF half the time you have to change your entire code? The `Trainer` API was a gamechanger, until it doesn’t work with some models 😞 Why do we need to actually use the `Flax` version of BigBird to run it on TPUs? I am not asking you give me a 2-line code solution, but the least HuggingFace could strive to do is to not be _sooo_ fragmented as it is now. At this point, `Transformers` has become so bloated with different tasks, changing APIs, outdated notebooks that its a nightmare (I honestly can’t even comphrend how Open source devs even contribute here). Honestly, it would rather be easier if it be married to a single framework like pytorch - allowing simple `DataLoaders` to interface directly with Huggingface Models and be trained easily in a pytorch-style manner with XLA spawn. I think Huggingface tries to take control of too many parts of the expereince - `Datasets` for making datasets, and it doesn’t always mesh well. Or perhaps I am just too dumb, and probably unfair to expect pytorch like simplicity from such a huge lib. ---
**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C **Post date:** [June 24, 2021, 7:37pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Frunning-bigbird-on-tpus%2F7289%2F4 "2021-06-24T19:37:17Z")
> [@Neel-Gupta](#): > > Why do we need to actually use the `Flax` version of BigBird to run it on TPUs? As I have said in the other topic, you don’t. Why are _you_ using an example located in examples/ **flax** / if you don’t want to use FLAX? There is really no need for the rant after, it’s completely inappropriate when people are replying to you for free on their own time. ---
**Author:** ![Neel-Gupta](https://avatars.discourse-cdn.com/v4/letter/n/f1d935/32.png) [@Neel-Gupta](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FNeel-Gupta)%5C **Post date:** [June 24, 2021, 8:17pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Frunning-bigbird-on-tpus%2F7289%2F5 "2021-06-24T20:17:37Z")
> [@sgugger](#): > > As I have said in the other topic, you don’t. Why are _you_ using an example located in examples/ **flax** / if you don’t want to use FLAX? It was Patrick Von Platen’s suggestion and PR merge for using the Flax version with BigBird. I asked the same thing, but I didn’t get a reply from him as to what was the reason. > [@sgugger](#): > > There is really no need for the rant after, it’s completely inappropriate when people are replying to you for free on their own time. Honestly, if you really think that asking basic questions about the direction of HF is a rant then that clearly shows the attitude of the company as a whole.