Support transfomers v5

.github/workflows/nvi-ci.yml

@@ -46,7 +46,7 @@ jobs:
     - name: Run checkstyle
       run: make checkstyle
 
-  tests:
+  correctness-tests:
     runs-on: ubuntu-latest
     needs: [checkstyle]
     env:
@@ -69,15 +69,14 @@ jobs:
 
     - name: Run tests
       run: |
-        modal run dev.modal.tests
+        modal run dev.modal.tests::liger_correctness_tests
 
-  tests-bwd:
+  convergence-tests:
     runs-on: ubuntu-latest
     needs: [checkstyle]
     env:
       MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
       MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
-      REBUILD_IMAGE: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
 
     steps:
     - name: Checkout code
@@ -95,4 +94,109 @@ jobs:
 
     - name: Run tests
       run: |
-        modal run dev.modal.tests_bwd
+        modal run dev.modal.tests::liger_convergence_tests
+
+
+  correctness-tests-with-transformers-4-52-0:
+    runs-on: ubuntu-latest
+    needs: [checkstyle]
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v6
+
+    - name: Set up Python
+      uses: actions/setup-python@v6
+      with:
+        python-version: '3.10'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install modal
+
+    - name: Run tests
+      run: |
+        modal run dev.modal.tests::liger_oldest_v4_correctness_tests
+
+
+
+  convergence-tests-with-transformers-4-52-0:
+    runs-on: ubuntu-latest
+    needs: [checkstyle]
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v6
+
+    - name: Set up Python
+      uses: actions/setup-python@v6
+      with:
+        python-version: '3.10'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install modal
+
+    - name: Run tests
+      run: |
+        modal run dev.modal.tests::liger_oldest_v4_convergence_tests
+
+  correctness-tests-with-transformers-4-57-6:
+    runs-on: ubuntu-latest
+    needs: [checkstyle]
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v6
+
+    - name: Set up Python
+      uses: actions/setup-python@v6
+      with:
+        python-version: '3.10'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install modal
+
+    - name: Run tests
+      run: |
+        modal run dev.modal.tests::liger_latest_v4_correctness_tests
+
+
+
+  convergence-tests-with-transformers-4-57-6:
+    runs-on: ubuntu-latest
+    needs: [checkstyle]
+    env:
+      MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
+      MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v6
+
+    - name: Set up Python
+      uses: actions/setup-python@v6
+      with:
+        python-version: '3.10'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install modal
+
+    - name: Run tests
+      run: |
+        modal run dev.modal.tests::liger_latest_v4_convergence_tests

benchmark/scripts/benchmark_llama4_rope.py

@@ -40,8 +40,6 @@ def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRu
         num_key_value_heads=num_kv_heads,
         head_dim=head_dim,
         max_position_embeddings=seq_len,
-        rope_theta=10000.0,
-        rope_scaling=None,  # Use default rope type
     )
 
     rotary_emb = transformers_version_dispatch(
@@ -134,8 +132,6 @@ def bench_memory_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkR
         num_key_value_heads=num_kv_heads,
         head_dim=head_dim,
         max_position_embeddings=seq_len,
-        rope_theta=10000.0,
-        rope_scaling=None,  # Use default rope type
     )
 
     rotary_emb = transformers_version_dispatch(

dev/modal/tests.py

@@ -6,6 +6,8 @@
 REMOTE_ROOT_PATH = "/root/liger-kernel"
 PYTHON_VERSION = "3.12"
 
+OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION = "4.52.0"
+
 image = modal.Image.debian_slim(python_version=PYTHON_VERSION).pip_install("uv")
 
 app = modal.App("liger_tests", image=image)
@@ -15,7 +17,77 @@
 
 
 @app.function(gpu="H100!", image=repo, timeout=90 * 60)
-def liger_tests():
+def liger_correctness_tests():
+    import subprocess
+
+    subprocess.run(
+        ["uv pip install -e '.[dev]' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
+
+
+@app.function(gpu="H100!", image=repo, timeout=90 * 60)
+def liger_convergence_tests():
+    import subprocess
+
+    subprocess.run(
+        ["uv pip install -e '.[dev]' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
+
+
+oldest_v4_app = modal.App("liger_oldest_v4_tests", image=image)  # 4.52.0
+
+
+@oldest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
+def liger_oldest_v4_correctness_tests():
+    import subprocess
+
+    subprocess.run(
+        ["uv pip install -e '.[dev]' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(
+        [f"uv pip install 'transformers=={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
+
+
+@oldest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
+def liger_oldest_v4_convergence_tests():
+    import subprocess
+
+    subprocess.run(
+        ["uv pip install -e '.[dev]' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(
+        [f"uv pip install 'transformers=={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
+
+
+latest_v4_app = modal.App("liger_latest_v4_tests", image=image)  # 4.57.6
+
+
+@latest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
+def liger_latest_v4_correctness_tests():
     import subprocess
 
     subprocess.run(
@@ -24,5 +96,29 @@ def liger_tests():
         shell=True,
         cwd=REMOTE_ROOT_PATH,
     )
+    subprocess.run(
+        [f"uv pip install 'transformers>={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}, <5.0.0' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
     subprocess.run(["make test"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
+
+
+@latest_v4_app.function(gpu="H100!", image=repo, timeout=90 * 60)
+def liger_latest_v4_convergence_tests():
+    import subprocess
+
+    subprocess.run(
+        ["uv pip install -e '.[dev]' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(
+        [f"uv pip install 'transformers>={OLDEST_SUPPORTED_TRANSFORMERS_V4_VERSION}, <5.0.0' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
     subprocess.run(["make test-convergence"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)

setup.py

@@ -33,7 +33,7 @@ def get_optional_dependencies():
     """Get optional dependency groups."""
     return {
         "dev": [
-            "transformers>=4.52.0, <5.0.0",
+            "transformers>=4.52.0",
             "matplotlib>=3.7.2",
             "ruff>=0.12.0",
             "pytest>=7.1.2",

src/liger_kernel/transformers/__init__.py

@@ -21,6 +21,7 @@
 from liger_kernel.transformers.softmax import LigerSoftmax  # noqa: F401
 from liger_kernel.transformers.sparsemax import LigerSparsemax  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP  # noqa: F401
+from liger_kernel.transformers.swiglu import LigerExperts  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerPhi3SwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerQwen3MoeSwiGLUMLP  # noqa: F401
 from liger_kernel.transformers.swiglu import LigerSwiGLUMLP  # noqa: F401