TpuEmbeddingEngine_WriteParameters not available in this library.

I followed all of the instructions in the training guide but when I run the device_train script, I get this error:

2022-02-23 07:56:56.271731: F external/org_tensorflow/tensorflow/core/tpu/tpu_library_init_fns.inc:104] TpuEmbeddingEngine_WriteParameters not available in this library.

This is my exact command for the training process:

python3 device_train.py --config=configs/6B.json --tune-model-path=gs://nnrap/step_383500

[mrseeker]

Okay, I fixed the issue this way:

• Use version v2-alpha instead of the tpu-vm-tf-2.8.0
• Install the package “python3-venv”
• Build & start a virtualenv environment (python3 -m venv venv and source venv/bin/activate)
• Install wheel (needed for some)
• Install the requirements
• Install Jax (as told by Google: pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html)
• Reinstall tensorflow because of numpy (pip install -U tensorflow)

This seemed to fix a LOT of my issues I was having, and its now working 😃

[mosmos6]

@mrseeker Your solution fixed my problem too. Thank you for sharing it!