Initial commit: E4B-MarkBase model integration with passing tests
CI / build-and-test (push) Has been cancelled
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
This commit is contained in:
@@ -0,0 +1,131 @@
|
||||
# TEXT Generation優化計畫
|
||||
|
||||
## 問題分析
|
||||
- 單個 forward pass 有 11 個 waitUntilCompleted() 呼叫
|
||||
- 每個呼叫阻塞 CPU 等待 GPU 完成
|
||||
- E4B: 11.3秒/token, 12B: 5.8秒/token(應該 <1秒)
|
||||
|
||||
## 根本原因
|
||||
```
|
||||
目前流程:
|
||||
Embedding → wait()
|
||||
Scale → wait()
|
||||
PerLayer → wait()
|
||||
Layer 0 → wait()
|
||||
Layer 1 → wait()
|
||||
...
|
||||
Layer 42 → wait()
|
||||
LM Head → wait()
|
||||
Readback → wait()
|
||||
|
||||
總共: 11+ 次同步等待
|
||||
```
|
||||
|
||||
## 優化策略
|
||||
|
||||
### 1. Batch Commands(最高優先)
|
||||
```swift
|
||||
// 優化後:
|
||||
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
||||
|
||||
// 所有操作加入同一個 command buffer
|
||||
try dequantizeRow(...) // 不等待
|
||||
try scaleBuffer(...) // 不等待
|
||||
for layer in layers { // 不等待
|
||||
try layer.forward(...)
|
||||
}
|
||||
try lmHead.forward(...) // 不等待
|
||||
|
||||
// 最後才等待
|
||||
cmdBuf.commit()
|
||||
cmdBuf.waitUntilCompleted() // 只等待一次!
|
||||
```
|
||||
|
||||
**預期改善:**
|
||||
- 從 11次等待 → 1次等待
|
||||
- 減少 GPU-CPU同步開銷
|
||||
- 預估速度提升 10倍以上
|
||||
|
||||
### 2. 移除 Per-Layer同步(次要)
|
||||
```swift
|
||||
// Line 1120-1135: Per-layer norm loop有同步等待
|
||||
for layerIdx in 0..<numHiddenLayers {
|
||||
try rmsNorm(...) // ← 應該批次處理
|
||||
cmdBufNorm.commit()
|
||||
cmdBufNorm.waitUntilCompleted() // ← 移除這個
|
||||
}
|
||||
```
|
||||
|
||||
**優化:**
|
||||
- 用單個 kernel批次處理所有 layers的 norm
|
||||
- 或用 separate command buffer不等待
|
||||
|
||||
### 3. Forward Pass重構
|
||||
```swift
|
||||
// 新結構:
|
||||
public func forwardOptimized(tokenId: Int, position: Int) throws -> [Float] {
|
||||
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
||||
|
||||
// 所有 GPU操作批次執行
|
||||
try embeddingPhase(cmdBuf) // Embedding + Scale + PerLayer
|
||||
try layersPhase(cmdBuf) // All layers in one batch
|
||||
try lmHeadPhase(cmdBuf) // LM Head + Final Norm
|
||||
|
||||
cmdBuf.commit()
|
||||
cmdBuf.waitUntilCompleted() // 只等待一次
|
||||
|
||||
return readbackLogits()
|
||||
}
|
||||
```
|
||||
|
||||
### 4. Kernel Fusion(進階)
|
||||
```swift
|
||||
// 合併多個操作成單個 kernel
|
||||
kernel void embedding_scale_norm(
|
||||
device float* embedding,
|
||||
device float* scale,
|
||||
device float* norm_weight,
|
||||
device float* output,
|
||||
uint id [[thread_position_in_grid]]
|
||||
) {
|
||||
// 一次執行:dequantize + scale + norm
|
||||
float val = dequantize(embedding[id]);
|
||||
val *= scale[0];
|
||||
val = rms_norm(val, norm_weight[id]);
|
||||
output[id] = val;
|
||||
}
|
||||
```
|
||||
|
||||
## 實作步驟
|
||||
|
||||
### Step 1: 修改 Model.swift forward function
|
||||
1. 移除所有中間的 waitUntilCompleted()
|
||||
2. 只在最後保留一個 waitUntilCompleted()
|
||||
3. 所有操作加入同一個 command buffer
|
||||
|
||||
### Step 2: 測試驗證
|
||||
1. 確保數值正確(無NaN)
|
||||
2. 測量時間改善
|
||||
|
||||
### Step 3: Kernel Fusion(optional)
|
||||
1. 建立組合 kernel
|
||||
2. 進一步減少 kernel launch overhead
|
||||
|
||||
## 預期成果
|
||||
|
||||
| Metric | 目前 | 優化後(預估) |改善倍數 |
|
||||
|--------|------|---------------|---------|
|
||||
| E4B token生成 | 11.3秒 | ~1秒 | **10倍** |
|
||||
| 12B token生成 | 5.8秒 | ~0.5秒 | **10倍** |
|
||||
| waitUntilCompleted呼叫 | 11次 | 1次 | **11倍** |
|
||||
| GPU-CPU同步 | 高 | 低 | **顯著** |
|
||||
|
||||
## 限制
|
||||
- 需要修改大量 forward pass邏輯
|
||||
- 需要確保數值穩定性
|
||||
- 可能影響 debugging能力
|
||||
|
||||
## 替代方案(如果 batch困難)
|
||||
1. 使用 asynchronous completion handlers
|
||||
2. Pipeline多個 forward passes
|
||||
3. Pre-compute KV cache for common tokens
|
||||
Reference in New Issue
Block a user