Allow linear and matmul to be consumed by nvFuser by default

notebooks/writing_a_trace_transform_cpu_offloading.ipynb

@@ -472,7 +472,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -507,7 +507,7 @@
     "model, args, kwargs = get_model_and_args()\n",
     "\n",
     "# Check against the vanilla `thunder.jit` model\n",
-    "expected = thunder.jit(model)(*args, **kwargs)\n",
+    "expected = thunder.jit(model, nv_enable_linear=False)(*args, **kwargs)\n",
     "\n",
     "grad_output = torch.randn_like(expected)\n",
     "expected_grads = torch.autograd.grad(expected, model.parameters(), grad_output)\n",
@@ -519,7 +519,7 @@
     "    expected_cpu = expected.to(\"cpu\")\n",
     "    expected_grads_cpu = tree_map(lambda t: t.to(\"cpu\"), expected_grads)\n",
     "\n",
-    "jmodel = thunder.jit(model, transforms=[CPUOffloading()])\n",
+    "jmodel = thunder.jit(model, nv_enable_linear=False, transforms=[CPUOffloading()])\n",
     "\n",
     "actual = jmodel(*args, **kwargs)\n",
     "\n",
@@ -642,7 +642,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -656,8 +656,8 @@
    "source": [
     "model, args, kwargs = get_model_and_args()\n",
     "\n",
-    "measurement_thunder = benchmark(thunder.jit(model), model, args, kwargs)\n",
-    "measurement_thunder_offload = benchmark(thunder.jit(model, transforms=[CPUOffloading()]), model, args, kwargs)\n",
+    "measurement_thunder = benchmark(thunder.jit(model, nv_enable_linear=False), model, args, kwargs)\n",
+    "measurement_thunder_offload = benchmark(thunder.jit(model, nv_enable_linear=False, transforms=[CPUOffloading()]), model, args, kwargs)\n",
     "\n",
     "del model, args, kwargs\n",
     "\n",

thunder/executors/nvfuserex_impl.py

@@ -2285,6 +2285,10 @@ def map_inside_replacement(x: Any) -> None:
 
 def _linear_check(a: TensorProxy, b: TensorProxy, bias: TensorProxy | None) -> bool:
     enable_linear: None | bool = get_compile_option("nv_enable_linear", "Enable nvFuser linear.")
+    # Enable linear by default if nvFuser version is 0.2.23 or later
+    # Because nvFuser 0.2.23 has a fix for https://github.com/NVIDIA/Fuser/issues/3369
+    if enable_linear is None and nvfuser_version() >= LooseVersion("0.2.23"):
+        enable_linear = True
     if not enable_linear:
         return False
     # Verify linear inputs and bias (optional) are supported tensors.
@@ -2315,7 +2319,10 @@ def _matmul_check(
     b: TensorProxy,
 ) -> bool:
     enable_matmul: None | bool = get_compile_option("nv_enable_matmul", "Enable nvFuser matmul.")
-
+    # Enable matmul by default if nvFuser version is 0.2.23 or later
+    # Because nvFuser 0.2.23 has a fix for https://github.com/NVIDIA/Fuser/issues/3369
+    if enable_matmul is None and nvfuser_version() >= LooseVersion("0.2.23"):
+        enable_matmul = True
     if not enable_matmul or not are_supported_tensors(a, b):
         return False
     return True

thunder/tests/test_networks.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 from functools import partial
 import warnings
+from looseversion import LooseVersion
 
 import pytest
 import torch
@@ -21,6 +22,7 @@
 )
 import thunder.tests.nanogpt_model as nanogpt_model
 import thunder.tests.hf_bart_self_attn as hf_bart_self_attn
+from thunder.executors.nvfuserex import nvfuser_version
 
 #
 # nanoGPT tests
@@ -550,7 +552,8 @@ def test_hf_llama():
     assert_close(res2, expected2, rtol=1e-1, atol=1e-1)
 
     # changes this to fewer as needed, the goal is to not have too many fusions
-    assert len(get_fusion_symbols(thunder.last_traces(jm)[-1])) == 7
+    num_of_nvfusion = 5 if nvfuser_version() >= LooseVersion("0.2.23") else 7
+    assert len(get_fusion_symbols(thunder.last_traces(jm)[-1])) == num_of_nvfusion
 
 
 @requiresCUDA

thunder/tests/test_nvfuser.py

@@ -356,6 +356,7 @@ def test_cse_rematerialization(executor, device, _):
         disable_torch_autograd=True,
         executors=executor.executors_list(),
         nv_enable_bookend=False,
+        nv_enable_linear=False,
     )
     compiled_func(x, y)