AI-Hypercomputer / AI-Hypercomputer/maxtext

bug: Numeric divergence between the mHC layer and its Pallas kernel

Ouverte
#5,195 2 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
bug
Langage dominant
Python
Étoiles
2.4k
Forks
607
Merge moyen
2 j 19 h
PR mergées (30 j)
158

Description

### 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_

Guide de contribution

Ouvrir le guide de contribution

Évaluation

Cette issue n'a pas encore été évaluée.

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.