KenteaKe11 commited on
Commit ·
fe9c6f7
1
Parent(s): 99eb3d1
KenT-300M: EMA memory with learned per-layer alpha, 1 WNN per layer, 24 layers, fp32 EMA
Browse files
model.py
CHANGED
|
@@ -150,6 +150,12 @@ def _affine_scan(a: torch.Tensor, b: torch.Tensor):
|
|
| 150 |
|
| 151 |
|
| 152 |
def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 153 |
"""H_t = (1 - a m_t) H_{t-1} + a m_t u_t, tính song song cho mọi t.
|
| 154 |
|
| 155 |
u: (B,T,N,D) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D) fp32
|
|
|
|
| 150 |
|
| 151 |
|
| 152 |
def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
|
| 153 |
+
"""EMA song song, luôn chạy fp32 (tắt autocast để matmul không bị ép về fp16)."""
|
| 154 |
+
with torch.autocast(device_type=u.device.type, enabled=False):
|
| 155 |
+
return _ema_scan_fp32(u.float(), mask, h0.float(), a.float(), block)
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def _ema_scan_fp32(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
|
| 159 |
"""H_t = (1 - a m_t) H_{t-1} + a m_t u_t, tính song song cho mọi t.
|
| 160 |
|
| 161 |
u: (B,T,N,D) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D) fp32
|