liger-kernel-nightly 0.5.10.dev20250611215839__tar.gz → 0.5.10.dev20250613212111__tar.gz

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_cpo_loss.py

@@ -147,7 +147,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_cpo_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_cross_entropy.py

@@ -109,7 +109,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_cross_entropy,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_distill_jsd_loss.py

@@ -248,7 +248,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_jsd_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_dpo_loss.py

@@ -166,7 +166,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_dpo_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_embedding.py

@@ -112,7 +112,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_embedding,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_fused_linear_cross_entropy.py

@@ -156,7 +156,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_cross_entropy,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_fused_linear_jsd.py

@@ -246,7 +246,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_jsd,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_jsd.py

@@ -143,7 +143,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_jsd,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_args,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_kl_div.py

@@ -110,7 +110,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_kldiv,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_args,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_kto_loss.py

@@ -299,7 +299,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_kto_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_layer_norm.py

@@ -111,7 +111,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_layer_norm,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_orpo_loss.py

@@ -149,7 +149,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_orpo_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "full", "backward"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_simpo_loss.py

@@ -147,7 +147,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_simpo_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "full", "backward"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_softmax.py

@@ -124,7 +124,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_softmax,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "full", "backward"],
         metric_name="speed",
         metric_unit="ms",
         overwrite=args.overwrite,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_swiglu.py

@@ -161,7 +161,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_swiglu,
-        kernel_operation_modes=["forward"],
+        kernel_operation_modes=["forward", "full", "backward"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/benchmark_tvd.py

@@ -126,7 +126,7 @@ if __name__ == "__main__":
 
     run_benchmarks(
         bench_test_fn=bench_speed_tvd,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "full", "backward"],
         metric_name="speed",
         metric_unit="ms",
         **common_args,

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/benchmark/scripts/utils.py

@@ -235,7 +235,7 @@ def update_benchmark_data_csv(
                     pass
             else:
                 existing_data_dict[row_key] = row_dict
-
+    os.makedirs(os.path.dirname(filename_abs_path), exist_ok=True)
     with open(filename_abs_path, mode="w", newline="") as file:
         writer = csv.DictWriter(file, fieldnames=fieldnames)
         writer.writeheader()

liger_kernel_nightly-0.5.10.dev20250613212111/dev/modal/benchmarks.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+
+import modal
+
+ROOT_PATH = Path(__file__).parent.parent.parent
+REMOTE_ROOT_PATH = "/root/liger-kernel"
+PYTHON_VERSION = "3.12"
+
+image = modal.Image.debian_slim(python_version=PYTHON_VERSION).pip_install("uv")
+
+app = modal.App("liger_benchmarks", image=image)
+
+# mount: add local files to the remote container
+repo = image.add_local_dir(ROOT_PATH, remote_path=REMOTE_ROOT_PATH)
+
+
+@app.function(gpu="H100", image=repo, timeout=60 * 45)
+def liger_benchmarks():
+    import os
+    import subprocess
+
+    subprocess.run(
+        ["uv pip install -e '.[dev]' --system"],
+        check=True,
+        shell=True,
+        cwd=REMOTE_ROOT_PATH,
+    )
+    subprocess.run(["make run-benchmarks"], check=True, shell=True, cwd=REMOTE_ROOT_PATH)
+
+    file_path = Path(REMOTE_ROOT_PATH) / "benchmark" / "data" / "all_benchmark_data.csv"
+    print(f"Checking if file exists at: {file_path}")
+    print(f"File exists: {os.path.exists(file_path)}")
+
+    if not os.path.exists(file_path):
+        print("Listing directory contents:")
+        data_dir = file_path.parent
+        if os.path.exists(data_dir):
+            print(f"Contents of {data_dir}:")
+            print(os.listdir(data_dir))
+        else:
+            print(f"Data directory {data_dir} does not exist")
+        raise FileNotFoundError(f"Benchmark data file not found at {file_path}")
+
+    with open(file_path, "rb") as f:
+        data = f.read()
+        print(f"Successfully read {len(data)} bytes of data")
+        return data
+
+
+@app.local_entrypoint()
+def main():
+    try:
+        # Run the benchmarks and get the data
+        print("Starting benchmark run...")
+        benchmark_data = liger_benchmarks.remote()
+
+        if not benchmark_data:
+            raise ValueError("No data received from remote function")
+
+        # Save the data locally
+        local_data_path = ROOT_PATH / "benchmark" / "data" / "all_benchmark_data.csv"
+        print(f"Attempting to save data to: {local_data_path}")
+
+        local_data_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(local_data_path, "wb") as f:
+            f.write(benchmark_data)
+
+        print(f"Successfully saved {len(benchmark_data)} bytes to: {local_data_path}")
+
+    except Exception as e:
+        print(f"Error occurred: {str(e)}")
+        raise

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.5.10.dev20250611215839"
+version = "0.5.10.dev20250613212111"
 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.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/src/liger_kernel_nightly.egg-info/PKG-INFO

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

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/test/convergence/bf16/test_mini_models.py

@@ -38,6 +38,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_moe
 from test.utils import DEFAULT_DATASET_PATH
 from test.utils import MiniModelConfig
 from test.utils import assert_verbose_allclose
+from test.utils import get_logprobs
+from test.utils import get_topk
 from test.utils import revert_liger_kernel_to_gemma
 from test.utils import revert_liger_kernel_to_gemma2
 from test.utils import revert_liger_kernel_to_gemma3_text
@@ -851,17 +853,17 @@ def run_mini_model(
         eval_output = model(**eval_batch)
     print(f"Eval Loss: {eval_output.loss.item()}")
     loss_list.append(eval_output.loss.item())
-
+    topk_logprobs = get_topk(get_logprobs(eval_output.logits))
     MINI_MODEL_SETUPS[model_name].liger_kernel_patch_revert_func(**revert_kwargs)
     return {
         "loss": loss_list,
-        "logits": eval_output.logits,
+        "topk_logprobs": topk_logprobs.values,
         "model": model,
     }
 
 
 @pytest.mark.parametrize(
-    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logits_atol, logits_rtol, param_atol, param_rtol",
+    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
     [
         pytest.param(
             "mini_llama3",
@@ -884,7 +886,7 @@ def run_mini_model(
             1e-3,
             1e-2,
             1e-1,
-            1e-2,
+            1e-1,
             1e-2,
             1e-2,
             marks=[
@@ -902,7 +904,7 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1,  # 1e-1
+            1e-1,  # 1e-1
             1e-1,  # 1e-2
             1e-2,
             1e-2,
@@ -972,7 +974,7 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1,  # 1e-1
+            1e-1,  # 1e-1
             1e-1,  # 1e-2
             1e-2,
             1e-2,
@@ -1111,8 +1113,8 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1e-1,
             1e-2,
+            1e-1,
             1e-2,
             1e-2,
             marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
@@ -1124,8 +1126,8 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1e-1,
             1e-2,
+            1e-1,
             1e-2,
             1e-2,
             marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
@@ -1153,8 +1155,8 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1e-1,
-            1e-2,
+            3e-1,
+            4e-1,
             1e-2,
             1e-2,
             marks=[
@@ -1174,8 +1176,8 @@ def test_mini_model(
     dtype,
     loss_atol,
     loss_rtol,
-    logits_atol,
-    logits_rtol,
+    logprobs_atol,
+    logprobs_rtol,
     param_atol,
     param_rtol,
 ):
@@ -1193,13 +1195,13 @@ def test_mini_model(
         rtol=loss_rtol,
     )
 
-    # Compare the logits from evaluation step
-    if expected_output["logits"] is not None and actual_output["logits"] is not None:
+    # Compare the topk logprobs from evaluation step
+    if expected_output["topk_logprobs"] is not None and actual_output["topk_logprobs"] is not None:
         assert_verbose_allclose(
-            expected_output["logits"],
-            actual_output["logits"],
-            atol=logits_atol,
-            rtol=logits_rtol,
+            expected_output["topk_logprobs"],
+            actual_output["topk_logprobs"],
+            atol=logprobs_atol,
+            rtol=logprobs_rtol,
         )
 
     # Compare the params from the last step

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/test/convergence/bf16/test_mini_models_multimodal.py

@@ -20,6 +20,8 @@ from test.utils import FAKE_CONFIGS_PATH
 from test.utils import UNTOKENIZED_DATASET_PATH
 from test.utils import MiniModelConfig
 from test.utils import assert_verbose_allclose
+from test.utils import get_logprobs
+from test.utils import get_topk
 from test.utils import is_torchvision_available
 from test.utils import load_image_processing_config
 from test.utils import load_processor_config
@@ -764,13 +766,17 @@ def run_mini_model_multimodal(
 
         print(f"Step {i}, Loss: {output.loss.item()}")
         loss_list.append(output.loss.item())
-
+    topk_logprobs = get_topk(get_logprobs(output.logits))
     MINI_MODEL_SETUPS[model_name].liger_kernel_patch_revert_func(**revert_kwargs)
-    return {"loss": loss_list, "logits": output.logits, "model": model}
+    return {
+        "loss": loss_list,
+        "topk_logprobs": topk_logprobs.values,
+        "model": model,
+    }
 
 
 @pytest.mark.parametrize(
-    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logits_atol, logits_rtol, param_atol, param_rtol",
+    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
     [
         pytest.param(
             "mini_qwen2_vl",
@@ -917,8 +923,8 @@ def test_mini_model_multimodal(
     dtype,
     loss_atol,
     loss_rtol,
-    logits_atol,
-    logits_rtol,
+    logprobs_atol,
+    logprobs_rtol,
     param_atol,
     param_rtol,
 ):
@@ -937,12 +943,12 @@ def test_mini_model_multimodal(
         rtol=loss_rtol,
     )
 
-    # Compare the logits from the last step
+    # Compare the topk logprobs from evaluation step
     assert_verbose_allclose(
-        expected_output["logits"],
-        actual_output["logits"],
-        atol=logits_atol,
-        rtol=logits_rtol,
+        expected_output["topk_logprobs"],
+        actual_output["topk_logprobs"],
+        atol=logprobs_atol,
+        rtol=logprobs_rtol,
     )
 
     # Compare the params from the last step

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/test/convergence/bf16/test_mini_models_with_logits.py

@@ -38,6 +38,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_moe
 from test.utils import DEFAULT_DATASET_PATH
 from test.utils import MiniModelConfig
 from test.utils import assert_verbose_allclose
+from test.utils import get_logprobs
+from test.utils import get_topk
 from test.utils import revert_liger_kernel_to_gemma
 from test.utils import revert_liger_kernel_to_gemma2
 from test.utils import revert_liger_kernel_to_gemma3_text
@@ -842,12 +844,17 @@ def run_mini_model(
         print(f"Step {i}, Loss: {output.loss.item()}")
         loss_list.append(output.loss.item())
 
+    topk_logprobs = get_topk(get_logprobs(output.logits))
     MINI_MODEL_SETUPS[model_name].liger_kernel_patch_revert_func(**revert_kwargs)
-    return {"loss": loss_list, "logits": output.logits, "model": model}
+    return {
+        "loss": loss_list,
+        "topk_logprobs": topk_logprobs.values,
+        "model": model,
+    }
 
 
 @pytest.mark.parametrize(
-    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logits_atol, logits_rtol, param_atol, param_rtol",
+    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
     [
         pytest.param(
             "mini_llama3",
@@ -1058,8 +1065,8 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1e-1,
             1e-2,
+            1e-1,
             1e-2,
             1e-2,
             marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
@@ -1071,8 +1078,8 @@ def run_mini_model(
             torch.bfloat16,
             1e-3,
             1e-2,
-            1e-1,
             1e-2,
+            1e-1,
             1e-2,
             1e-2,
             marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
@@ -1159,8 +1166,8 @@ def test_mini_model(
     dtype,
     loss_atol,
     loss_rtol,
-    logits_atol,
-    logits_rtol,
+    logprobs_atol,
+    logprobs_rtol,
     param_atol,
     param_rtol,
 ):
@@ -1180,12 +1187,12 @@ def test_mini_model(
 
     # No logits are materialized
     # import pdb; pdb.set_trace()
-    # Compare the logits from the last step
+    # Compare the topk logprobs from evaluation step
     assert_verbose_allclose(
-        expected_output["logits"],
-        actual_output["logits"],
-        atol=logits_atol,
-        rtol=logits_rtol,
+        expected_output["topk_logprobs"],
+        actual_output["topk_logprobs"],
+        atol=logprobs_atol,
+        rtol=logprobs_rtol,
     )
 
     # Compare the params from the last step

{liger_kernel_nightly-0.5.10.dev20250611215839 → liger_kernel_nightly-0.5.10.dev20250613212111}/test/convergence/fp32/test_mini_models.py

@@ -38,6 +38,8 @@ from liger_kernel.transformers import apply_liger_kernel_to_qwen3_moe
 from test.utils import DEFAULT_DATASET_PATH
 from test.utils import MiniModelConfig
 from test.utils import assert_verbose_allclose
+from test.utils import get_logprobs
+from test.utils import get_topk
 from test.utils import revert_liger_kernel_to_gemma
 from test.utils import revert_liger_kernel_to_gemma2
 from test.utils import revert_liger_kernel_to_gemma3_text
@@ -849,17 +851,17 @@ def run_mini_model(
         eval_output = model(**eval_batch)
     print(f"Eval Loss: {eval_output.loss.item()}")
     loss_list.append(eval_output.loss.item())
-
+    topk_logprobs = get_topk(get_logprobs(eval_output.logits))
     MINI_MODEL_SETUPS[model_name].liger_kernel_patch_revert_func(**revert_kwargs)
     return {
         "loss": loss_list,
-        "logits": eval_output.logits,
+        "topk_logprobs": topk_logprobs.values,
         "model": model,
     }
 
 
 @pytest.mark.parametrize(
-    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logits_atol, logits_rtol, param_atol, param_rtol",
+    "model_name, num_steps, lr, dtype, loss_atol, loss_rtol, logprobs_atol, logprobs_rtol, param_atol, param_rtol",
     [
         ("mini_llama3", 32, 1e-4, torch.float32, 1e-8, 2e-5, 1e-4, 1e-5, 5e-3, 1e-5),
         pytest.param(
@@ -1013,7 +1015,7 @@ def run_mini_model(
         # TODO: mixtral is flaky so disable the test for now
         # ("mini_mixtral", 32, 1e-4, torch.float32, 5e-4, 1e-4, 5e-3, 1e-5, 1e-2, 1e-5),
         # Gemma 1.1 and 2 has more tolerance because currently, the kernel is not a perfect match (casts are not done the same way)
-        ("mini_gemma1", 32, 1e-4, torch.float32, 1e-8, 1e-4, 5e-3, 1e-5, 5e-3, 1e-5),
+        ("mini_gemma1", 32, 1e-4, torch.float32, 1e-8, 1e-4, 5e-3, 1e-2, 5e-3, 1e-5),
         ("mini_gemma1.1", 32, 1e-4, torch.float32, 1e-8, 1e-4, 5e-3, 1e-5, 5e-3, 1e-5),
         ("mini_gemma2", 32, 1e-4, torch.float32, 1e-8, 1e-4, 5e-3, 1e-5, 5e-3, 1e-5),
         pytest.param(
@@ -1041,8 +1043,8 @@ def test_mini_model(
     dtype,
     loss_atol,
     loss_rtol,
-    logits_atol,
-    logits_rtol,
+    logprobs_atol,
+    logprobs_rtol,
     param_atol,
     param_rtol,
 ):
@@ -1060,13 +1062,13 @@ def test_mini_model(
         rtol=loss_rtol,
     )
 
-    # Compare the logits from evaluation step
-    if expected_output["logits"] is not None and actual_output["logits"] is not None:
+    # Compare the topk logprobs from evaluation step
+    if expected_output["topk_logprobs"] is not None and actual_output["topk_logprobs"] is not None:
         assert_verbose_allclose(
-            expected_output["logits"],
-            actual_output["logits"],
-            atol=logits_atol,
-            rtol=logits_rtol,
+            expected_output["topk_logprobs"],
+            actual_output["topk_logprobs"],
+            atol=logprobs_atol,
+            rtol=logprobs_rtol,
         )
 
     # Compare the params from the last step