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

fix

Browse Source
This commit is contained in:
Dan Saunders
2025-09-17 18:53:07 -04:00
parent d57b9c67c2
commit e62979d11d
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -71,7 +71,7 @@ def _call_grouped_mm(
lengths = torch.tensor( lengths = torch.tensor(
[a.shape[0] for a in As2], device=device, dtype=torch.int32 [a.shape[0] for a in As2], device=device, dtype=torch.int32
) )
offsets = torch.cumsum(lengths, dim=0) offsets = torch.cumsum(lengths, dim=0).to(torch.int32)
Y_cat = torch._grouped_mm( Y_cat = torch._grouped_mm(
torch.cat(As2, dim=0), torch.stack(Bs2, dim=0), offsets torch.cat(As2, dim=0), torch.stack(Bs2, dim=0), offsets
) )