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

hardening

Browse Source
This commit is contained in:
Dan Saunders
2025-09-15 19:41:10 -04:00
parent f6ed8ddc01
commit fef47a5b7c
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -16,11 +16,13 @@ def available() -> bool:
ver = tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:2]) ver = tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:2])
if ver < (2, 8): if ver < (2, 8):
return False return False
# Check for aten grouped mm ops # Require Hopper+ (SM90) per torch error message and check op presence
return hasattr(torch.ops, "aten") and ( if not torch.cuda.is_available():
hasattr(torch.ops.aten, "_grouped_mm") return False
or hasattr(torch.ops.aten, "_scaled_grouped_mm") major, minor = torch.cuda.get_device_capability()
) if major < 9:
return False
return hasattr(torch.ops, "aten") and hasattr(torch.ops.aten, "_grouped_mm")
except Exception: except Exception:
return False return False