This is in continuation of my previous issue filed in this forum (here 2).
I wanted to try and run the run_mlm_flax script with a Colab TPU v2, but I am facing an error:-
Module Name: <module 'run_mlm_flax' from '/content/run_mlm_flax.py'>

WARNING:root:TPU has started up successfully with version pytorch-1.9
Traceback (most recent call last):
  File "xla_spawn.py", line 87, in <module>
    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?