liger-kernel-nightly 0.4.2.dev20241122175637__tar.gz → 0.4.2.dev20241128223426__tar.gz

{liger_kernel_nightly-0.4.2.dev20241122175637/src/liger_kernel_nightly.egg-info → liger_kernel_nightly-0.4.2.dev20241128223426}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.4.2.dev20241122175637
+Version: 0.4.2.dev20241128223426
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.4.2.dev20241122175637 → liger_kernel_nightly-0.4.2.dev20241128223426}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.4.2.dev20241122175637"
+version = "0.4.2.dev20241128223426"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.4.2.dev20241122175637 → liger_kernel_nightly-0.4.2.dev20241128223426}/src/liger_kernel/ops/layer_norm.py

@@ -180,8 +180,13 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD):
     dY = dY.view(-1, dim)
     n_rows, n_cols = dY.shape
 
+    sm_count = 1
+    if X.device.type == "cuda":
+        sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+    elif X.device.type == "xpu":
+        sm_count = torch.xpu.get_device_properties(X.device).gpu_subslice_count
+
     DX = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
-    sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     _DW = torch.empty((sm_count, n_cols), dtype=W.dtype, device=W.device)
     _DB = torch.empty((sm_count, n_cols), dtype=W.dtype, device=W.device)
 

{liger_kernel_nightly-0.4.2.dev20241122175637 → liger_kernel_nightly-0.4.2.dev20241128223426}/src/liger_kernel/ops/rms_norm.py

@@ -264,6 +264,7 @@ def rms_norm_backward(
     dY = dY.view(-1, dim)
     n_rows, n_cols = dY.shape
 
+    sm_count = 1
     if X.device.type == "cuda":
         sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
     elif X.device.type == "xpu":

{liger_kernel_nightly-0.4.2.dev20241122175637 → liger_kernel_nightly-0.4.2.dev20241128223426}/src/liger_kernel/ops/utils.py

@@ -20,6 +20,8 @@ import triton
 import triton.language as tl
 from packaging.version import Version
 
+from liger_kernel.utils import infer_device
+
 
 def is_hip() -> bool:
     return torch.version.hip is not None
@@ -69,10 +71,11 @@ def compare_version(package: str, operator: Callable, target: str):
 
 
 def get_amp_custom_fwd_bwd() -> Callable:
+    device = infer_device()
     if compare_version("torch", operator.ge, "2.4.0"):
         return (
-            functools.partial(torch.amp.custom_fwd, device_type="cuda"),
-            functools.partial(torch.amp.custom_bwd, device_type="cuda"),
+            functools.partial(torch.amp.custom_fwd, device_type=device),
+            functools.partial(torch.amp.custom_bwd, device_type=device),
         )
     return torch.cuda.amp.custom_fwd, torch.cuda.amp.custom_bwd
 

liger_kernel_nightly-0.4.2.dev20241128223426/src/liger_kernel/utils.py

@@ -0,0 +1,13 @@
+import torch
+
+
+def infer_device():
+    """
+    Get current device name based on available devices
+    """
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.xpu.is_available():
+        return "xpu"
+    else:
+        return "cpu"

{liger_kernel_nightly-0.4.2.dev20241122175637 → liger_kernel_nightly-0.4.2.dev20241128223426/src/liger_kernel_nightly.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.4.2.dev20241122175637
+Version: 0.4.2.dev20241128223426
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.4.2.dev20241122175637 → liger_kernel_nightly-0.4.2.dev20241128223426}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -2,7 +2,9 @@ LICENSE
 NOTICE
 README.md
 pyproject.toml
+src/liger_kernel/__init__.py
 src/liger_kernel/env_report.py
+src/liger_kernel/utils.py
 src/liger_kernel/chunked_loss/__init__.py
 src/liger_kernel/chunked_loss/cpo_loss.py
 src/liger_kernel/chunked_loss/dpo_loss.py