From 69e3892c27c487bf93d9d044670e9b1ca4f6e642 Mon Sep 17 00:00:00 2001
From: Zchen <161216199+ZH-CEN@users.noreply.github.com>
Date: Sun, 12 Oct 2025 22:32:12 +0800
Subject: [PATCH] =?UTF-8?q?tpu=20=E5=A4=9A=E7=BA=BF=E7=A8=8B=E7=BC=96?=
 =?UTF-8?q?=E8=AF=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 TPU_ISSUES_RECORD.md              | 71 ++++++++++++++++++++++++++++++-
 model_training_nnn/rnn_trainer.py | 13 ++++++
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/TPU_ISSUES_RECORD.md b/TPU_ISSUES_RECORD.md
index 74e6503..d4a94a0 100644
--- a/TPU_ISSUES_RECORD.md
+++ b/TPU_ISSUES_RECORD.md
@@ -218,6 +218,8 @@ All TPU training issues have been systematically identified and fixed:
 ✅ **Problem 2**: `batch_sampler=None` error → Fixed with custom collate_fn + batch_size=1
 ✅ **Problem 3**: Data type mismatch (dataset) → Fixed bf16 conversion + padding preservation
 ✅ **Problem 4**: Data type mismatch (model) → Fixed all parameter initialization with explicit bf16 dtype
+✅ **Problem 5**: Memory exhaustion → Fixed with batch_size=32 + gradient_accumulation_steps=2
+✅ **Problem 6**: Training hang logging → Added progress message for XLA compilation wait
 
 **The solution addresses dtype consistency at ALL levels**:
 - Input data loading: `.to(torch.bfloat16)`
@@ -251,10 +253,77 @@ num_dataloader_workers: 0  # TPU compatibility
 - Gradient accumulation maintains training stability
 - Effective batch size unchanged: 2 steps × 32 = 64 samples
 
+## CPU Usage During TPU Training (2025-10-12 16:00)
+
+**高CPU使用率是正常的TPU训练行为**
+
+### 问题描述
+用户观察到CPU使用率达到100%，询问是什么操作以及是否可以使用多个CPU核心。
+
+### 技术解释
+**正常行为**: TPU训练期间100% CPU使用率是预期的，原因如下：
+
+1. **XLA编译**: PyTorch XLA需要CPU进行图编译和优化
+2. **数据预处理**: CPU负责数据加载、增强和转换
+3. **主机-TPU通信**: CPU管理与TPU的数据传输
+4. **分布式协调**: 多TPU核心的同步需要CPU协调
+
+### 当前设置分析
+- `num_dataloader_workers: 0` - 为TPU兼容性禁用多进程数据加载
+- `gradient_accumulation_steps: 2` - CPU需要管理梯度累积
+- 687M参数模型 - 大模型增加CPU开销
+
+### 多核心使用
+**数据加载器工作进程已禁用**原因：
+```yaml
+num_dataloader_workers: 0 # set to 0 for TPU to avoid multiprocessing issues
+```
+
+TPU训练建议保持`num_workers=0`因为：
+- TPU与多进程数据加载存在兼容性问题
+- XLA编译已经能充分利用CPU资源
+- 避免进程间通信开销
+
+### 优化建议
+1. **保持当前设置** - `num_workers=0`是TPU最佳实践
+2. **监控系统资源** - 确保有足够RAM支持XLA编译
+3. **耐心等待编译** - 首个batch编译需5-15分钟，之后会加速
+
+**结论**: 100% CPU使用率表明系统正在进行正常的TPU训练操作，无需担心。
+
+### XLA编译优化 (2025-10-12 16:15)
+
+**问题**: XLA编译只使用单线程，浪费了多核CPU资源
+
+**解决方案**: 在`rnn_trainer.py`中添加XLA多线程优化配置：
+
+```python
+# XLA multi-threading optimization for faster compilation
+import torch_xla.core.xla_model as xm
+if xm.get_xla_supported_devices():
+    # Enable XLA multi-threading for compilation speedup
+    os.environ.setdefault('XLA_FLAGS',
+        '--xla_cpu_multi_thread_eigen=true ' +
+        '--xla_cpu_enable_fast_math=true ' +
+        f'--xla_force_host_platform_device_count={os.cpu_count()}'
+    )
+    # Set PyTorch XLA threading
+    os.environ.setdefault('PYTORCH_XLA_COMPILATION_THREADS', str(os.cpu_count()))
+```
+
+**效果**:
+- `--xla_cpu_multi_thread_eigen=true`: 启用CPU多线程Eigen库
+- `--xla_cpu_enable_fast_math=true`: 启用快速数学优化
+- `--xla_force_host_platform_device_count`: 利用所有CPU核心
+- `PYTORCH_XLA_COMPILATION_THREADS`: 设置PyTorch XLA编译线程数
+
+**预期改进**: XLA图编译时间从5-15分钟缩短到2-8分钟
+
 ## Lessons Learned
 - **Root Cause**: TPU XLA compiler requires strict dtype consistency across all tensors
 - **Key Insight**: `torch.eye()` and `torch.zeros()` default to f32 - must explicitly specify dtype
 - **Documentation**: Record issues immediately to avoid repeated debugging cycles
 - Don't overcomplicate TPU conversion - identify systematic dtype issues
 - Read Accelerate documentation carefully for parameter placement
-- TPU memory allocation: fewer cores = less total memory
\ No newline at end of file
+- TPU memory allocation: fewer cores = less total memory
+- **CPU Usage**: 100% CPU usage during TPU training is normal and expected
\ No newline at end of file
diff --git a/model_training_nnn/rnn_trainer.py b/model_training_nnn/rnn_trainer.py
index 8d4c77c..6009717 100644
--- a/model_training_nnn/rnn_trainer.py
+++ b/model_training_nnn/rnn_trainer.py
@@ -22,6 +22,18 @@ from omegaconf import OmegaConf
 from accelerate import Accelerator, DataLoaderConfiguration
 from accelerate.utils import set_seed
 
+# XLA multi-threading optimization for faster compilation
+import torch_xla.core.xla_model as xm
+if xm.get_xla_supported_devices():
+    # Enable XLA multi-threading for compilation speedup
+    os.environ.setdefault('XLA_FLAGS',
+        '--xla_cpu_multi_thread_eigen=true ' +
+        '--xla_cpu_enable_fast_math=true ' +
+        f'--xla_force_host_platform_device_count={os.cpu_count()}'
+    )
+    # Set PyTorch XLA threading
+    os.environ.setdefault('PYTORCH_XLA_COMPILATION_THREADS', str(os.cpu_count()))
+
 torch.set_float32_matmul_precision('high') # makes float32 matmuls faster on some GPUs
 torch.backends.cudnn.deterministic = True # makes training more reproducible
 torch._dynamo.config.cache_size_limit = 64
@@ -539,6 +551,7 @@ class BrainToTextDecoder_Trainer:
         train_start_time = time.time()
 
         # train for specified number of batches
+        self.logger.info("Starting training loop - loading first batch (TPU compilation may take 5-15 minutes)...")
         for i, batch in enumerate(self.train_loader):
             
             self.model.train()