$\"banner\"\/$ <\/p>\n

JAX on Cloud TPUs offers highly effective acceleration for machine studying workflows. When working in distributed cloud environments, you want specialised instruments to debug your workflows, together with accessing logs, {hardware} metrics, and extra. This weblog publish serves as a sensible information to varied debugging and profiling methods.<\/p>\n

Selecting the best instrument: Core Parts and Dependencies<\/b><\/h3>\n
On the coronary heart of the system are two predominant parts that almost all debugging instruments depend upon:<\/p>\n
\n
libtpu<\/a> (which comprises libtpu.so<\/i>, the TPU Runtime): That is probably the most elementary piece of software program. It is a shared library on each Cloud TPU VM that comprises the XLA compiler, the TPU driver, and the logic for speaking with the {hardware}. Nearly each debugging instrument interacts with or is configured by way of libtpu.<\/li>\n
JAX<\/a> and jaxlib<\/a> (The Framework): JAX is the Python library the place you write your mannequin code. jaxlib is its C++ backend, which acts because the bridge to libtpu.so<\/li>\n<\/ol>\n
The connection between these parts and the debugging instruments is illustrated within the diagram beneath.<\/p>\n<\/div>\n
\n
\n
$\"relationship_diagram\"\/$ <\/p><\/div><\/div>\n
\n
Here’s a breakdown of the precise instruments, their dependencies, and the way they relate to one another.<\/p>\n<\/div>\n
\n
\n
$\"tool_table_updated\"\/$ <\/p><\/div><\/div>\n
\n
In abstract, libtpu is the central pillar that almost all debugging instruments depend on, both for configuration (logging, HLO dumps) or for querying real-time knowledge (monitoring, profiling). Different instruments, like XProf, additionally function on the Python stage to examine the state of your JAX program instantly. By understanding these relationships, you possibly can extra successfully select the best instrument for the precise challenge you might be going through.<\/p>\n
Important Logging and Diagnostic Flags for Each Workload<\/b><\/h3>\nVerbose Logging<\/h3>\n
Probably the most important step for debugging is to allow verbose logging. With out it, you might be flying blind. These flags must be thought-about on each employee of your TPU slice, to log all the things from TPU runtime setup to program execution steps with timestamps<\/p>\n<\/div>\n
\n
\n
$\"log_updated\"\/$ <\/p><\/div><\/div>\n
\n
If you wish to allow the above default flags on each TPU employee nodes, run the next command:<\/p>\n<\/div>\n
\n
gcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} \n --zone ${ZONE} --worker=all --node=all \n --command='TPU_VMODULE=slice_configuration=1,real_program_continuator=1 TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 python3 -c \"import jax; print(f\"Host {jax.process_index()}: International units: {jax.device_count()}, Native units: {jax.local_device_count()}\")\"'<\/code><\/pre>\n\n Plain textual content\n <\/p>\n<\/div>\n \nLibtpu logs are robotically generated in \/tmp\/tpu_logs\/tpu_driver.INFO on every TPU VM. This file is your floor reality for what the TPU runtime is doing. To get logs from all TPU VMs, you possibly can run the next bash script:<\/p>\n<\/div>\n \n#!\/bin\/bash \n \nTPU_NAME=\"your TPU TPU_NAME\" \nPROJECT=\"challenge on your TPU\" \nZONE=\"zone on your TPU\" \nBASE_LOG_DIR=\"path to the place you need the logs to be downloaded to\" \n \nNUM_WORKERS=$(gcloud compute tpus tpu-vm describe $TPU_NAME --zone=$ZONE --project=$PROJECT | grep tpuVmSelflink | awk -F'[:\/]' '{print $13}' | uniq | wc -l) \n \necho \"Variety of staff = $NUM_WORKERS\" \n \nfor ((i=0; i<$NUM_WORKERS; i++)) \ndo \n mkdir -p ${BASE_LOG_DIR}\/$i \n echo \"gcloud compute tpus tpu-vm scp ${TPU_NAME}:\/tmp\/tpu_logs\/* ${BASE_LOG_DIR}\/$i\/ --zone=${ZONE} --project=${PROJECT} --worker=$i\" \n echo \"Obtain logs from employee=$i\" \n gcloud compute tpus tpu-vm scp ${TPU_NAME}:\/tmp\/tpu_logs\/* ${BASE_LOG_DIR}\/$i\/ --zone=${ZONE} --project=${PROJECT} --worker=$i \nexecuted<\/code><\/pre>\n\n Plain textual content\n <\/p>\n<\/div>\n \nOn Google Colab, you possibly can set the above atmosphere variables utilizing os.environ, and entry the logs within the \u201cInformation\u201d part within the left sidebar.<\/p>\n Listed here are some instance snippets from a log file:<\/p>\n<\/div>\n \n... \nI1031 19:02:51.863599 669 b295d63588a.cc:843] Course of id 669 \nI1031 19:02:51.863609 669 b295d63588a.cc:848] Present working listing \/content material \n... \nI1031 19:02:51.863621 669 b295d63588a.cc:866] Construct instrument: Bazel, launch r4rca-2025.05.26-2 (mainline @763214608) \nI1031 19:02:51.863621 669 b295d63588a.cc:867] Construct goal: \nI1031 19:02:51.863624 669 b295d63588a.cc:874] Command line arguments: \nI1031 19:02:51.863624 669 b295d63588a.cc:876] argv[0]: '.\/tpu_driver' \n... \n 19:02:51.863784 669 init.cc:78] Distant crash gathering hook put in. \nI1031 19:02:51.863807 669 tpu_runtime_type_flags.cc:79] --tpu_use_tfrt not specified. Utilizing default worth: true \nI1031 19:02:51.873759 669 tpu_hal.cc:448] Registered plugin from module: breakpoint_debugger_server \n... \nI1031 19:02:51.879890 669 pending_event_logger.cc:896] Enabling PjRt\/TPU occasion dependency logging \nI1031 19:02:51.880524 843 device_util.cc:124] Discovered 1 TPU v5 lite chips. \n... \nI1031 19:02:53.471830 851 2a886c8_compiler_base.cc:3677] CODE_GENERATION stage length: 3.610218ms \nI1031 19:02:53.471885 851 isa_program_util_common.cc:486] (HLO module jit_add): Executable fingerprint:0cae8d08bd660ddbee7ef03654ae249ae4122b40da162a3b0ca2cd4bb4b3a19c<\/code><\/pre>\n\n Plain textual content\n <\/p>\n<\/div>\n \nTPU Monitoring Library<\/h3>\nThe TPU Monitoring Library is a strategy to programmatically acquire insights about workflow efficiency on TPU {hardware} (utilization, capability, latency, and extra). It is part of the libtpu bundle, which is robotically put in (as a dependency) with jax[tpu], so you can begin utilizing the monitoring API instantly.<\/p>\n<\/div>\n \n# Specific set up \npip istall \"jax[tpu]\" libtpu<\/code><\/pre>\n\n Shell\n <\/p>\n<\/div>\n \nYou may view all supported metrics with tpumonitoring.list_supported_metrics() and get particular metrics with tpumonitoring.get_metric<\/code>. For instance, the next snippet prints the duty_cycle knowledge and outline:<\/p>\n<\/div>\n \nfrom libtpu.sdk import tpumonitoring \n \nduty_cycle_metric = tpumonitoring.get_metric(\"duty_cycle_pct\") \nduty_cycle_data = duty_cycle_metric.knowledge \nprint(\"TPU Obligation Cycle Information:\") \nprint(f\" Description: {duty_cycle_metric.description}\") \nprint(f\" Information: {duty_cycle_data}\")<\/code><\/pre>\n\n Python\n <\/p>\n<\/div>\n \nYou’d usually combine tpumonitoring instantly in your JAX packages, throughout mannequin coaching, earlier than inference, and so on. Study extra concerning the Monitoring Library within the Cloud TPU documentation<\/a>.<\/p>\n tpu-info<\/h3>\nThe tpu-info<\/a> command-line instrument is a straightforward strategy to get a real-time view of TPU reminiscence and different utilization metrics, much like nvidia-smi for GPUs.<\/p>\n Set up on all staff and nodes<\/p>\n<\/div>\n \ngcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} \n --zone ${ZONE} --worker=all --node=all \n --command='pip set up tpu-info'<\/code><\/pre>\n\n Plain textual content\n <\/p>\n<\/div>\n\nSSH into one employee and node to test chip utilization metrics<\/p>\n<\/div>\n\ngcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} \n --zone ${ZONE} --worker=0 --node=0 \n \ntpu-info<\/code><\/pre>\n\n Plain textual content\n <\/p>\n<\/div>\n\nWhen chips are in use, course of IDs, reminiscence utilization, and responsibility cycle% will likely be displayed<\/p>\n<\/div>\n\n\n <\/p><\/div><\/div>\n\nWhen no chips are in use, the TPU VM will present no exercise<\/p>\n<\/div>\n \n\n <\/p><\/div><\/div>\n \nStudy extra about different metrics and streaming mode in the documentation<\/a>.<\/p>\n On this publish, we mentioned some TPU logging and monitoring choices. Subsequent on this sequence, we\u2019ll discover find out how to debug your JAX packages beginning with producing HLO dumps, and profiling your code with the XProf.<\/p>\n<\/div><\/div>\n\n","protected":false},"excerpt":{"rendered":" JAX on Cloud TPUs offers highly effective acceleration for machine studying workflows. When working in distributed cloud environments, you want specialised instruments to debug your workflows, together with accessing logs, {hardware} metrics, and extra. This weblog publish serves as a sensible information to varied debugging and profiling methods. Selecting the best instrument: Core Parts and […]<\/p>\n","protected":false},"author":2,"featured_media":10552,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[56],"tags":[234,3888,305,3712,78,2551,1598,213,7308],"class_list":["post-10550","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-software","tag-cloud","tag-debugging","tag-developers","tag-essential","tag-guide","tag-jax","tag-techniques","tag-tools","tag-tpus"],"_links":{"self":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/10550","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=10550"}],"version-history":[{"count":1,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/10550\/revisions"}],"predecessor-version":[{"id":10551,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/10550\/revisions\/10551"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/media\/10552"}],"wp:attachment":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=10550"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=10550"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=10550"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}