如何查看Google TPU的VRAM使用情况？

Colab环境下的查看方式 #

• 内置监控面板：确认运行时已切换为TPU后，点击Colab界面右侧的「资源」标签，可直接查看TPU的实时内存占用、设备利用率等数据。
• TensorFlow代码查询：

import tensorflow as tf

# 连接并初始化TPU集群
tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)

# 获取所有TPU逻辑设备并打印内存信息
devices = tf.config.list_logical_devices('TPU')
for idx, device in enumerate(devices):
    mem_info = tf.config.experimental.get_memory_info(device)
    print(f"TPU Core {idx}:")
    print(f"  当前已用内存: {mem_info['current'] / (1024**3):.2f} GB")
    print(f"  峰值内存占用: {mem_info['peak'] / (1024**3):.2f} GB")

GCP TPU虚拟机环境下的查看方式 #

• 终端命令tpu-stats：直接在VM的终端执行该命令，会输出TPU的整体内存使用、各核心内存状态、设备负载等详细统计数据。
• PyTorch XLA代码查询：

import torch_xla.core.xla_model as xm

# 获取TPU设备
device = xm.xla_device()
# 查询内存信息
mem_stats = xm.get_memory_info(device)
used_gb = mem_stats['used'] / (1024**3)
free_gb = mem_stats['free'] / (1024**3)

print(f"TPU已用内存: {used_gb:.2f} GB")
print(f"TPU剩余内存: {free_gb:.2f} GB")

注意：TPU采用分布式内存架构，内存分散在各个计算核心中，部分工具会分别展示每个核心的内存情况，部分会提供汇总数据。

内容的提问来源于stack exchange，提问作者Seungwoo Ryu