liger-kernel-nightly 0.6.2.dev20251014053719__tar.gz → 0.6.2.dev20251016055812__tar.gz

liger_kernel_nightly-0.6.2.dev20251016055812/.github/workflows/docs.yml

@@ -0,0 +1,64 @@
+name: Publish documentation
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - 'docs/**'
+      - 'mkdocs.yml'
+
+permissions:
+  contents: write
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Configure Git Credentials
+        run: |
+          git config user.name github-actions[bot]
+          git config user.email 41898282+github-actions[bot]@users.noreply.github.com
+      - uses: actions/setup-python@v5
+        with:
+          python-version: 3.x
+      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 
+      - uses: actions/cache@v4
+        with:
+          key: mkdocs-material-${{ env.cache_id }}
+          path: .cache
+          restore-keys: |
+            mkdocs-material-
+      - run: pip install mkdocs-material mkdocstrings[python]
+      # ====== Backup the benchmarks from gh-pages ======
+      # This is necessary because the benchmarks are not included in the documentation build process.
+      # So we need to backup the benchmarks from gh-pages and restore them after the documentation is built.
+      - name: Backup benchmarks from gh-pages
+        run: |
+          git fetch origin gh-pages
+          # create worktree bound to local gh-pages, tracking origin/gh-pages
+          git branch -f gh-pages origin/gh-pages || true
+          mkdir -p ghp && git worktree add ghp gh-pages || true
+          if [ -d ghp/benchmarks ]; then
+            tar -C ghp -czf /tmp/benchmarks.tgz benchmarks
+          fi
+          # IMPORTANT: remove worktree so gh-pages isn't checked out anywhere
+          git worktree remove ghp --force || true
+          echo "Backed up benchmarks from gh-pages"
+      # ====== Deploy the documentation ======
+      - name: Deploy documentation
+        run: mkdocs gh-deploy --force
+      # ====== Restore the benchmarks onto gh-pages ======
+      # This is necessary because the benchmarks are not included in the documentation build process.
+      # So we need to restore the benchmarks onto gh-pages after the documentation is built.
+      - name: Restore benchmarks onto gh-pages
+        run: |
+          # Refresh remote tracking and recreate a clean worktree
+          git fetch origin gh-pages
+          git worktree add -B gh-pages ghp origin/gh-pages
+          if [ -f /tmp/benchmarks.tgz ]; then
+            tar -C ghp -xzf /tmp/benchmarks.tgz
+            git -C ghp add -A
+            git -C ghp commit -m "Restore benchmarks after gh-deploy" || echo "No changes"
+            git -C ghp push origin gh-pages
+          fi
+          git worktree remove ghp --force || true

{liger_kernel_nightly-0.6.2.dev20251014053719 → liger_kernel_nightly-0.6.2.dev20251016055812}/PKG-INFO

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

{liger_kernel_nightly-0.6.2.dev20251014053719 → liger_kernel_nightly-0.6.2.dev20251016055812}/docs/acknowledgement.md

@@ -9,7 +9,6 @@
 We referenced or used the following projects:
 
 
-
 | # | Project                                                                                      | Description                                                                             | Location                                                                                                                         | License                                                                              |
 |---|----------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------|
 | 1 | [Unsloth](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/unsloth/kernels/utils.py#L43)                              | `calculate_settings` to determine block size and warp; We reuse it for Norm and MLP     | [Liger Kernel Utils](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/utils.py#L23) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |

{liger_kernel_nightly-0.6.2.dev20251014053719 → liger_kernel_nightly-0.6.2.dev20251016055812}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.2.dev20251014053719"
+version = "0.6.2.dev20251016055812"
 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.6.2.dev20251014053719 → liger_kernel_nightly-0.6.2.dev20251016055812}/src/liger_kernel/ops/cross_entropy.py

@@ -414,6 +414,8 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
         Returns:
         tuple: A tuple with the compouted losses with respect to loss and z loss. The elements are tensors or None.
         """
+        input_requires_grad = _input.requires_grad
+
         loss, z_loss, _input = cross_entropy_forward(
             _input,
             target,
@@ -428,7 +430,8 @@ class LigerCrossEntropyFunction(torch.autograd.Function):
         # TODO: investigation
         # If we don't detach the _input tensor, the memory will double
         # Not sure why but seems that there will be a time both grad and value exist but in different location
-        ctx.save_for_backward(_input.detach())
+        if input_requires_grad:
+            ctx.save_for_backward(_input.detach())
         ctx.return_z_loss = return_z_loss
 
         return loss, z_loss

{liger_kernel_nightly-0.6.2.dev20251014053719 → liger_kernel_nightly-0.6.2.dev20251016055812}/src/liger_kernel/ops/fused_linear_cross_entropy.py

@@ -31,6 +31,8 @@ def fused_linear_cross_entropy_forward(
     assert isinstance(return_z_loss, bool), f"return_z_loss must be True or False. Got: {return_z_loss}"
     device = _input.device
 
+    input_requires_grad = _input.requires_grad
+
     # inputs have shape: BT x H
     # materialized activations will have shape: BT x V
     # the increase in memory = BT x V
@@ -49,12 +51,13 @@ def fused_linear_cross_entropy_forward(
     grad_input = torch.zeros_like(_input, device=device)
 
     # we use fp32 for loss and gradients accumulator
-    if accum_dtype is None:
-        grad_weight = torch.zeros_like(weight, device=device) if weight.requires_grad else None
-        grad_bias = torch.zeros_like(bias, device=device) if bias is not None else None
-    else:
-        grad_weight = torch.zeros_like(weight, dtype=accum_dtype, device=device) if weight.requires_grad else None
-        grad_bias = torch.zeros_like(bias, dtype=accum_dtype, device=device) if bias is not None else None
+    if input_requires_grad:
+        if accum_dtype is None:
+            grad_weight = torch.zeros_like(weight, device=device) if weight.requires_grad else None
+            grad_bias = torch.zeros_like(bias, device=device) if bias is not None else None
+        else:
+            grad_weight = torch.zeros_like(weight, dtype=accum_dtype, device=device) if weight.requires_grad else None
+            grad_bias = torch.zeros_like(bias, dtype=accum_dtype, device=device) if bias is not None else None
 
     loss_1d = torch.zeros(BT, dtype=torch.float32, device=device)
     z_loss_1d = torch.zeros(BT, dtype=_input.dtype, device=_input.device) if return_z_loss else None
@@ -150,7 +153,7 @@ def fused_linear_cross_entropy_forward(
             RETURN_Z_LOSS=return_z_loss,
             HAS_WEIGHT=True if ce_weight is not None else False,
             HAS_SOFTCAPPING=True if softcap is not None else False,
-            HAS_GRADIENTS=_input.requires_grad,
+            HAS_GRADIENTS=input_requires_grad,
             BLOCK_SIZE=BLOCK_SIZE,
             num_warps=32 if not is_hip() else 16,
         )
@@ -172,12 +175,13 @@ def fused_linear_cross_entropy_forward(
             scaling_factors_expanded = scaling_factors.unsqueeze(-1)  # chunk_size x 1
             grad_logits_chunk = grad_logits_chunk * scaling_factors_expanded
 
-        grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
+        if input_requires_grad:
+            grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
 
-        if grad_weight is not None and _input.requires_grad:
+        if grad_weight is not None and input_requires_grad:
             grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()
 
-        if bias is not None and _input.requires_grad:
+        if bias is not None and input_requires_grad:
             torch.add(
                 input=grad_bias,
                 other=grad_logits_chunk.sum(dim=0),

{liger_kernel_nightly-0.6.2.dev20251014053719 → liger_kernel_nightly-0.6.2.dev20251016055812}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

liger_kernel_nightly-0.6.2.dev20251014053719/.github/workflows/docs.yml

@@ -1,33 +0,0 @@
-name: Publish documentation
-on:
-  push:
-    branches:
-      - main
-    paths:
-      - 'docs/**'
-      - 'mkdocs.yml'
-
-permissions:
-  contents: write
-jobs:
-  deploy:
-    if: False
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Configure Git Credentials
-        run: |
-          git config user.name github-actions[bot]
-          git config user.email 41898282+github-actions[bot]@users.noreply.github.com
-      - uses: actions/setup-python@v5
-        with:
-          python-version: 3.x
-      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 
-      - uses: actions/cache@v4
-        with:
-          key: mkdocs-material-${{ env.cache_id }}
-          path: .cache
-          restore-keys: |
-            mkdocs-material-
-      - run: pip install mkdocs-material mkdocstrings[python]
-      - run: mkdocs gh-deploy --force