AI-Hypercomputer / AI-Hypercomputer/maxtext
bug: Numeric divergence between the mHC layer and its Pallas kernel
未關閉
bug
- 主要語言
- Python
- 星號
- 2.4k
- 分支
- 607
- 平均合併
- 2 天 19 小時
- 30 天內合併 PR
- 158
描述
### Bug report
DeepSeek v4. The reference path had its own mapping helper that cast alpha_scale and
beta down to self.dtype before computing sigmoid(alpha_scale * h + beta),
while the Pallas path called common.compute_sigmoid_gate, which keeps the
scale and bias in float32. With dtype: bfloat16 the two branches therefore
gated on different numbers, so toggling use_mhc_pallas_kernel changed the
layer's output beyond accumulation noise.
Fixed in PR:
https://github.com/AI-Hypercomputer/maxtext/pull/5179
### Logs/Output
_No response_
### Environment Information
_No response_
### Additional Context
_No response_
貢獻指南
評估
這個 Issue 還沒有評估資料。