[[["易于理解","easyToUnderstand","thumb-up"],["解决了我的问题","solvedMyProblem","thumb-up"],["其他","otherUp","thumb-up"]],[["很难理解","hardToUnderstand","thumb-down"],["信息或示例代码不正确","incorrectInformationOrSampleCode","thumb-down"],["没有我需要的信息/示例","missingTheInformationSamplesINeed","thumb-down"],["翻译问题","translationIssue","thumb-down"],["其他","otherDown","thumb-down"]],["最后更新时间 (UTC)：2025-09-04。"],[],[],null,["# Troubleshooting JAX - TPU\n=========================\n\nThis guide provides pointers to JAX troubleshooting information to\nhelp you identify and resolve problems you might encounter while training\nJAX models on Cloud TPU.\n\nFor a more general guide to\ngetting started with Cloud TPU, see the\n[JAX quickstart](/tpu/docs/run-calculation-jax).\n| **Note:** If you aren't able to resolve your issue using this guide, see [Getting Support](/tpu/docs/getting-support) for further assistance.\n\nGeneral JAX issues\n------------------\n\nIf you run into issues while developing your training model or\ntraining with JAX, see the\n[JAX FAQ.](https://jax.readthedocs.io/en/latest/faq.html)\n\nFor more general programming\nerrors you might encounter when writing a training application with JAX, see\n[JAX Errors.](https://jax.readthedocs.io/en/latest/errors.html)\n\nProfiling JAX performance\n-------------------------\n\nYou can understand how your TPU resources are being utilized using the\ntools described in\n[Profiling JAX performance.](https://jax.readthedocs.io/en/latest/profiling.html)\n\nTroubleshooting memory issues\n-----------------------------\n\n| **Note:** TPU memory ([high bandwidth memory](/tpu/docs/system-architecture-tpu-vm#versions)) is shared across TensorCores on each TPU chip. This enables efficient coordination between the on-chip TensorCores.\n\nYou can monitor how the memory is used with the [JAX Device Memory Profiler](https://jax.readthedocs.io/en/latest/device_memory_profiling.html),\nbut you cannot directly manage how it is used.\n\nThe Device Memory Profiler can be used to:\n\n- Figure out [which arrays and executables are in TPU memory](https://jax.readthedocs.io/en/latest/device_memory_profiling.html#understanding-how-a-jax-program-is-using-gpu-or-tpu-memory) at a given time, or\n- [Track down memory leaks.](https://jax.readthedocs.io/en/latest/device_memory_profiling.html#debugging-memory-leaks)\n\nYou cannot specify how TPU memory is allocated for specific operations.\nFor more information on JAX-specific TPU performance issues, see\n[Performance Notes for using TPUs with JAX](/tpu/docs/run-calculation-jax#performance_notes).\n\nTroubleshooting TPU issues\n--------------------------\n\n### How can I verify that the TPU is running?\n\n**Details**\n\nEverything will be run on the TPU as long as JAX doesn't print\n\"No GPU/TPU found, falling back to CPU.\"\n\nYou can verify the TPU is active by either looking at `jax.devices()`, where\nyou should see several TPU devices displayed, or verify\nprogrammatically with: `assert jax.devices()[0].platform == 'tpu'`.\n\n### RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.\n\n**Details**\n\nThis runtime error message and/or finding the following in `/tmp/tpu_logs/tpu_driver.WARNING` on the TPU VM:\n`W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx`\ncan indicate that you are running the wrong TPU VM version.\n\nVerify that you are running the [current JAX runtime version](/tpu/docs/supported-tpu-versions#jax) and retry."]]
