403 lines
18 KiB
Python
403 lines
18 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
TPU训练内存监控工具 - 专注于训练过程中的实时内存和MXU监控
|
||
适用于TPU v5e-8环境
|
||
"""
|
||
|
||
import tensorflow as tf
|
||
import time
|
||
import numpy as np
|
||
|
||
def monitor_tpu_during_training():
|
||
"""训练过程中的TPU实时内存和MXU监控"""
|
||
print("📊 TPU训练实时监控工具")
|
||
print("=" * 50)
|
||
|
||
# 获取TPU设备
|
||
try:
|
||
tpu_devices = tf.config.list_logical_devices('TPU')
|
||
print(f"📍 发现TPU设备: {len(tpu_devices)}个")
|
||
if not tpu_devices:
|
||
print("❌ 未发现TPU设备")
|
||
return
|
||
except Exception as e:
|
||
print(f"❌ 无法检测TPU设备: {e}")
|
||
return
|
||
|
||
def get_detailed_memory_snapshot():
|
||
"""获取详细的内存快照，包含所有核心信息"""
|
||
snapshot = {}
|
||
total_current = 0
|
||
total_peak = 0
|
||
active_cores = 0
|
||
core_details = []
|
||
|
||
for i, device in enumerate(tpu_devices):
|
||
try:
|
||
memory_info = tf.config.experimental.get_memory_info(device.name)
|
||
if memory_info and 'current' in memory_info:
|
||
current_mb = memory_info['current'] // (1024 * 1024)
|
||
peak_mb = memory_info.get('peak', memory_info['current']) // (1024 * 1024)
|
||
|
||
if current_mb > 1: # >1MB算活跃
|
||
active_cores += 1
|
||
total_current += current_mb
|
||
total_peak += peak_mb
|
||
core_details.append(f"Core{i}:{current_mb}MB")
|
||
|
||
snapshot[f'core_{i}'] = {
|
||
'current': current_mb,
|
||
'peak': peak_mb,
|
||
'device': device.name
|
||
}
|
||
else:
|
||
snapshot[f'core_{i}'] = {'current': 0, 'peak': 0, 'device': device.name}
|
||
except Exception as e:
|
||
snapshot[f'core_{i}'] = {'current': 0, 'peak': 0, 'device': device.name, 'error': str(e)}
|
||
|
||
snapshot['summary'] = {
|
||
'total_current': total_current,
|
||
'total_peak': total_peak,
|
||
'active_cores': active_cores,
|
||
'total_cores': len(tpu_devices),
|
||
'core_details': core_details
|
||
}
|
||
return snapshot
|
||
|
||
def test_mxu_performance():
|
||
"""测试MXU性能和计算能力"""
|
||
print("\n🧮 MXU计算性能测试:")
|
||
|
||
mxu_results = []
|
||
try:
|
||
with tf.device(tpu_devices[0].name):
|
||
# 测试不同规模的矩阵运算
|
||
test_configs = [
|
||
(2000, "2K×2K", tf.bfloat16),
|
||
(4000, "4K×4K", tf.bfloat16),
|
||
(6000, "6K×6K", tf.bfloat16),
|
||
]
|
||
|
||
for size, desc, dtype in test_configs:
|
||
try:
|
||
# 获取测试前内存
|
||
pre_mem = get_detailed_memory_snapshot()
|
||
|
||
start_time = time.time()
|
||
|
||
# 创建矩阵并执行MXU密集型运算
|
||
matrix_a = tf.random.normal([size, size], dtype=dtype)
|
||
matrix_b = tf.random.normal([size, size], dtype=dtype)
|
||
|
||
@tf.function
|
||
def mxu_operation():
|
||
# 连续矩阵运算，充分使用MXU
|
||
result = tf.matmul(matrix_a, matrix_b)
|
||
result = tf.matmul(result, matrix_a)
|
||
return tf.reduce_sum(result)
|
||
|
||
result = mxu_operation()
|
||
# 使用result确保计算被执行
|
||
_ = result.numpy()
|
||
end_time = time.time()
|
||
|
||
# 获取测试后内存
|
||
post_mem = get_detailed_memory_snapshot()
|
||
|
||
duration = end_time - start_time
|
||
# 计算FLOPS (两次矩阵乘法)
|
||
flops = 2 * (2 * size**3)
|
||
tflops = flops / duration / 1e12
|
||
|
||
memory_used = post_mem['summary']['total_current'] - pre_mem['summary']['total_current']
|
||
|
||
print(f" {desc} ({dtype.name}): {duration:.3f}s, {tflops:.1f}TFLOPS, 内存+{memory_used}MB")
|
||
|
||
mxu_results.append({
|
||
'size': size,
|
||
'tflops': tflops,
|
||
'duration': duration,
|
||
'memory_used': memory_used
|
||
})
|
||
|
||
except Exception as e:
|
||
print(f" {desc}: 测试失败 - {str(e)[:50]}")
|
||
|
||
# MXU性能分析
|
||
if mxu_results:
|
||
max_tflops = max(r['tflops'] for r in mxu_results)
|
||
total_memory = sum(r['memory_used'] for r in mxu_results if r['memory_used'] > 0)
|
||
|
||
# TPU v5e-8单核理论性能
|
||
theoretical_tflops = 275 # bf16峰值性能
|
||
efficiency = (max_tflops / theoretical_tflops) * 100
|
||
|
||
print(f"\n 📊 MXU性能汇总:")
|
||
print(f" 峰值性能: {max_tflops:.1f} TFLOPS")
|
||
print(f" 理论峰值: {theoretical_tflops} TFLOPS")
|
||
print(f" MXU效率: {efficiency:.1f}%")
|
||
print(f" 计算内存占用: {total_memory}MB")
|
||
|
||
if efficiency > 80:
|
||
status = "🟢 优秀"
|
||
elif efficiency > 50:
|
||
status = "🟡 良好"
|
||
elif efficiency > 20:
|
||
status = "🟠 中等"
|
||
else:
|
||
status = "🔴 需优化"
|
||
|
||
print(f" 性能评级: {status}")
|
||
|
||
except Exception as e:
|
||
print(f" MXU测试失败: {e}")
|
||
|
||
try:
|
||
print("🎯 开始TPU训练监控...")
|
||
|
||
# 1. 获取初始状态
|
||
print("\n📸 初始TPU状态:")
|
||
baseline_snapshot = get_detailed_memory_snapshot()
|
||
|
||
print(f" 总内存使用: {baseline_snapshot['summary']['total_current']}MB")
|
||
print(f" 活跃核心: {baseline_snapshot['summary']['active_cores']}/{baseline_snapshot['summary']['total_cores']}")
|
||
|
||
# 显示各核心详细状态
|
||
for i in range(len(tpu_devices)):
|
||
core = baseline_snapshot[f'core_{i}']
|
||
if core['current'] > 0 or core['peak'] > 0:
|
||
print(f" Core{i}: 当前{core['current']}MB, 峰值{core['peak']}MB")
|
||
|
||
# 2. MXU性能基准测试
|
||
test_mxu_performance()
|
||
|
||
# 3. 创建分布式策略 - 使用项目验证的TPU初始化代码
|
||
print(f"\n🔄 使用项目标准TPU初始化...")
|
||
try:
|
||
# 使用项目里验证过的TPU初始化代码
|
||
# 禁用GPU避免冲突
|
||
try:
|
||
tf.config.set_visible_devices([], 'GPU')
|
||
print("🚫 GPU已禁用，避免CUDA冲突")
|
||
except:
|
||
pass
|
||
|
||
# 使用标准的TPU初始化流程
|
||
print("🚀 使用官方TensorFlow TPU初始化...")
|
||
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
|
||
tf.config.experimental_connect_to_cluster(resolver)
|
||
tf.tpu.experimental.initialize_tpu_system(resolver)
|
||
|
||
# 验证TPU设备
|
||
tpu_devices_check = tf.config.list_logical_devices('TPU')
|
||
print(f"✅ TPU设备验证: 发现 {len(tpu_devices_check)} 个设备")
|
||
|
||
# 创建TPU策略
|
||
strategy = tf.distribute.TPUStrategy(resolver)
|
||
print(f"✅ 成功创建TPU策略: {strategy.num_replicas_in_sync}个副本")
|
||
use_distributed = True
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ 分布式策略失败: {str(e)[:80]}")
|
||
print(" 将使用单设备模拟")
|
||
use_distributed = False
|
||
|
||
# 4. 模拟Brain-to-Text训练场景
|
||
print(f"\n🧠 模拟Brain-to-Text训练场景...")
|
||
|
||
if use_distributed:
|
||
# 分布式训练模拟
|
||
with strategy.scope():
|
||
print("📦 创建分布式模型参数...")
|
||
|
||
# 创建接近真实Brain-to-Text模型的参数 (修复维度匹配)
|
||
model_components = {
|
||
# GRU层权重：第一层接收512维输入，后续层接收256维
|
||
'gru_layer_0': tf.Variable(tf.random.normal([512, 256]), name='gru_0'),
|
||
'gru_layer_1': tf.Variable(tf.random.normal([256, 256]), name='gru_1'),
|
||
'gru_layer_2': tf.Variable(tf.random.normal([256, 256]), name='gru_2'),
|
||
'output_projection': tf.Variable(tf.random.normal([256, 41]), name='output'),
|
||
# 添加day-specific层模拟 (输入512维，输出512维)
|
||
'day_weights': [tf.Variable(tf.random.normal([512, 512]), name=f'day_{i}') for i in range(8)]
|
||
}
|
||
|
||
# 检查模型加载后内存
|
||
after_model = get_detailed_memory_snapshot()
|
||
model_memory = after_model['summary']['total_current'] - baseline_snapshot['summary']['total_current']
|
||
print(f"🧠 模型加载完成: +{model_memory}MB, {after_model['summary']['active_cores']}个活跃核心")
|
||
|
||
# 训练循环模拟
|
||
print(f"\n🔄 开始训练循环监控...")
|
||
|
||
for step in range(10):
|
||
step_start_time = time.time()
|
||
|
||
@tf.function
|
||
def distributed_training_step():
|
||
# 模拟真实训练数据大小
|
||
batch_size = 32
|
||
seq_length = 1000
|
||
features = 512
|
||
|
||
# 输入数据
|
||
neural_data = tf.random.normal([batch_size, seq_length, features])
|
||
targets = tf.random.uniform([batch_size, seq_length], maxval=41, dtype=tf.int32)
|
||
|
||
# 模拟前向传播
|
||
x = neural_data
|
||
|
||
# Day-specific transformation (简化版本避免复杂的维度操作)
|
||
# 模拟day-specific变换：对每个时间步应用相同变换
|
||
day_weight = model_components['day_weights'][0] # 简化：使用第一个day权重
|
||
# 对最后一个维度进行变换: [batch, seq, 512] @ [512, 512] -> [batch, seq, 512]
|
||
x = tf.matmul(x, day_weight)
|
||
|
||
# 为CTC损失添加目标使用（模拟）
|
||
target_length = tf.reduce_sum(tf.cast(targets > 0, tf.int32), axis=1)
|
||
# 简化的CTC相关计算
|
||
batch_loss_weight = tf.reduce_mean(tf.cast(target_length, tf.float32))
|
||
|
||
# GRU layers simulation
|
||
for i in range(3):
|
||
layer_name = f'gru_layer_{i}'
|
||
weight = model_components[layer_name]
|
||
|
||
# 处理张量维度：第一层从3D输入，后续层从2D输入
|
||
if i == 0:
|
||
# 第一层：取最后时间步 [batch, seq, features] -> [batch, features]
|
||
if len(x.shape) == 3:
|
||
x = x[:, -1, :] # 取最后时间步
|
||
x = tf.nn.tanh(tf.matmul(x, weight))
|
||
else:
|
||
# 后续层：直接处理2D张量 [batch, features] -> [batch, features]
|
||
x = tf.nn.tanh(tf.matmul(x, weight))
|
||
|
||
# 输出投影
|
||
logits = tf.matmul(x, model_components['output_projection'])
|
||
|
||
# CTC loss模拟（使用batch_loss_weight作为权重）
|
||
base_loss = tf.reduce_mean(tf.square(logits))
|
||
loss = base_loss * batch_loss_weight
|
||
|
||
return loss
|
||
|
||
# 执行训练步骤
|
||
per_replica_loss = strategy.run(distributed_training_step)
|
||
# 聚合分布式结果
|
||
loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, axis=None)
|
||
step_duration = time.time() - step_start_time
|
||
|
||
# 获取当前内存状态
|
||
current_snapshot = get_detailed_memory_snapshot()
|
||
step_memory = current_snapshot['summary']['total_current']
|
||
memory_delta = step_memory - baseline_snapshot['summary']['total_current']
|
||
|
||
# 显示详细训练状态
|
||
active_cores_info = f"({', '.join(current_snapshot['summary']['core_details'])})" if current_snapshot['summary']['core_details'] else "(无活跃)"
|
||
|
||
print(f" Step {step:2d}: loss={float(loss.numpy()):.4f}, "
|
||
f"时间={step_duration:.3f}s, "
|
||
f"内存={step_memory}MB(+{memory_delta}), "
|
||
f"活跃={current_snapshot['summary']['active_cores']}/{current_snapshot['summary']['total_cores']} {active_cores_info}")
|
||
|
||
# 每5步显示峰值内存
|
||
if step % 5 == 0:
|
||
peak_info = f"峰值: {current_snapshot['summary']['total_peak']}MB"
|
||
print(f" {peak_info}")
|
||
|
||
time.sleep(0.2) # 短暂暂停观察
|
||
|
||
else:
|
||
# 单设备训练模拟（改进版）
|
||
print("🔸 单设备训练模拟...")
|
||
|
||
with tf.device(tpu_devices[0].name):
|
||
# 创建较小的模型参数
|
||
simple_weights = tf.Variable(tf.random.normal([512, 256]), name='simple_net')
|
||
|
||
for step in range(8):
|
||
step_start = time.time()
|
||
|
||
# 创建较大的数据批次
|
||
batch_data = tf.random.normal([64, 1000, 512]) # 增大batch size
|
||
|
||
# 模拟计算密集型操作
|
||
@tf.function
|
||
def compute_step():
|
||
x = tf.reshape(batch_data, [-1, 512])
|
||
result = tf.matmul(x, simple_weights)
|
||
result = tf.nn.relu(result)
|
||
return tf.reduce_mean(result)
|
||
|
||
result = compute_step()
|
||
step_duration = time.time() - step_start
|
||
|
||
# 获取内存状态
|
||
snapshot = get_detailed_memory_snapshot()
|
||
memory_change = snapshot['summary']['total_current'] - baseline_snapshot['summary']['total_current']
|
||
|
||
print(f" Step {step}: result={result.numpy():.4f}, "
|
||
f"时间={step_duration:.3f}s, "
|
||
f"内存变化=+{memory_change}MB, "
|
||
f"峰值={snapshot['summary']['total_peak']}MB")
|
||
|
||
# 5. 最终分析报告
|
||
final_snapshot = get_detailed_memory_snapshot()
|
||
total_growth = final_snapshot['summary']['total_current'] - baseline_snapshot['summary']['total_current']
|
||
peak_usage = final_snapshot['summary']['total_peak']
|
||
|
||
print(f"\n📈 训练监控报告:")
|
||
print(f" 总内存增长: +{total_growth}MB")
|
||
print(f" 峰值内存使用: {peak_usage}MB ({peak_usage/1024:.2f}GB)")
|
||
print(f" 最终活跃核心: {final_snapshot['summary']['active_cores']}/{final_snapshot['summary']['total_cores']}")
|
||
|
||
# 各核心最终状态
|
||
print(f" 各核心最终状态:")
|
||
has_changes = False
|
||
for i in range(len(tpu_devices)):
|
||
final_core = final_snapshot[f'core_{i}']
|
||
baseline_core = baseline_snapshot[f'core_{i}']
|
||
current_change = final_core['current'] - baseline_core['current']
|
||
peak_change = final_core['peak'] - baseline_core['peak']
|
||
|
||
if current_change != 0 or peak_change != 0:
|
||
has_changes = True
|
||
print(f" Core{i}: 当前{final_core['current']}MB(+{current_change}), 峰值{final_core['peak']}MB(+{peak_change})")
|
||
|
||
if not has_changes:
|
||
print(f" 所有核心内存无明显变化")
|
||
|
||
# 分布式使用分析
|
||
if final_snapshot['summary']['active_cores'] == 1:
|
||
print(f"\n⚠️ 分布式问题诊断:")
|
||
print(f" 只有1个核心活跃，其他7个核心空闲")
|
||
print(f" 可能原因: TPU策略配置问题或模型未正确分布")
|
||
print(f" 建议: 检查分布式策略和模型分片")
|
||
elif final_snapshot['summary']['active_cores'] > 4:
|
||
print(f"\n✅ 分布式状态良好:")
|
||
print(f" {final_snapshot['summary']['active_cores']}个核心活跃，多核心并行工作正常")
|
||
else:
|
||
print(f"\n🟡 分布式部分工作:")
|
||
print(f" {final_snapshot['summary']['active_cores']}个核心活跃，可能存在负载不均衡")
|
||
|
||
print("✅ TPU训练监控完成")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 训练监控失败: {e}")
|
||
import traceback
|
||
print(f"详细错误: {traceback.format_exc()[:300]}")
|
||
|
||
if __name__ == "__main__":
|
||
print("🚀 TPU训练内存监控工具")
|
||
print("专注于训练过程中的实时内存和性能监控")
|
||
print("适用于TPU v5e-8环境")
|
||
print()
|
||
|
||
monitor_tpu_during_training()
|
||
|
||
print(f"\n🎯 监控要点总结:")
|
||
print(f" 1. 确认所有8个TPU核心是否活跃")
|
||
print(f" 2. 监控内存增长模式和峰值使用")
|
||
print(f" 3. 检测MXU计算性能和效率")
|
||
print(f" 4. 验证分布式策略是否正常工作")
|
||
print(f" 5. 识别可能的内存泄漏或性能瓶颈") |