tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix compile

Browse Source
This commit is contained in:
Dan Saunders
2025-09-19 13:59:57 -04:00
parent ce21da9177
commit 8564961423
1 changed files with 0 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/kernels/moe/torch_grouped.py
View File

@@ -342,9 +342,6 @@ def moe_ffn_forward_grouped(
counts_i32 = assignments.to(device=device, dtype=torch.int32) counts_i32 = assignments.to(device=device, dtype=torch.int32)
offsets = torch.cumsum(counts_i32, dim=0).to(dtype=torch.int32) offsets = torch.cumsum(counts_i32, dim=0).to(dtype=torch.int32)
if not torch.is_nonzero(offsets[-1]):
zero = torch.zeros_like(x_flat)
return zero.view(bsz, seqlen, hdim), router_logits
mm_dtype = torch.bfloat16 if expert_dtype == torch.bfloat16 else expert_dtype mm_dtype = torch.bfloat16 if expert_dtype == torch.bfloat16 else expert_dtype
routed_in = routed_input.to(mm_dtype) routed_in = routed_input.to(mm_dtype)
w_gate_t = w_gate.transpose(-2, -1).to(mm_dtype) w_gate_t = w_gate.transpose(-2, -1).to(mm_dtype)