[PyTorch] [torch.compile] torch.compile support for Linear

tests/pytorch/distributed/run_layer_with_overlap.py

@@ -200,6 +200,19 @@ def _parse_args(argv=None, namespace=None):
     parser.add_argument(
         "--use-cuda-graphs", action="store_true", default=False, help="Use CUDA Graphs."
     )
+    parser.add_argument(
+        "--compile",
+        action="store_true",
+        default=False,
+        help="Wrap each layer in torch.compile (tests Userbuffers on the compiled path).",
+    )
+    parser.add_argument(
+        "--compile-mode",
+        type=str,
+        default="default",
+        choices=["default", "reduce-overhead"],
+        help="torch.compile mode used when --compile is set.",
+    )
     parser.add_argument(
         "--ub-cfg", type=str, default=None, help="Optional TP config yaml file input."
     )
@@ -285,6 +298,9 @@ def _parse_args(argv=None, namespace=None):
     )
     args = parser.parse_args(argv, namespace)
 
+    if args.compile and args.use_cuda_graphs:
+        parser.error("--compile and --use-cuda-graphs are mutually exclusive.")
+
     if args.use_cuda_graphs and args.layer_type in [te.MultiheadAttention, te.TransformerLayer]:
         warnings.warn(f"{args.layer_type.__name__} does not support CUDA Graphs!")
         args.use_cuda_graphs = False
@@ -535,6 +551,17 @@ def run_fwd_bwd(model, x):
             loss.backward()
         return out
 
+    if opts.compile:
+        for i, layer in enumerate(test_model.layers):
+            # Static shapes; dynamic-shape coverage lives in tests/pytorch/test_torch_compile.py.
+            test_model.layers[i] = torch.compile(
+                layer, fullgraph=True, mode=opts.compile_mode, dynamic=False
+            )
+        dist_print(
+            f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...",
+            debug=True,
+        )
+
     torch_rng_state = torch.get_rng_state()
     cuda_rng_state = torch.cuda.get_rng_state(torch.device(f"cuda:{LOCAL_RANK}"))
     if opts.use_cuda_graphs:

tests/pytorch/distributed/run_numerics.py

@@ -310,22 +310,40 @@ def _copy_params(model_distributed, model_single):
 
 
 def _apply_models(
-    model_single_node, model_distributed, input_single_node, input_distributed, **kwargs
+    model_single_node,
+    model_distributed,
+    input_single_node,
+    input_distributed,
+    use_compile=False,
+    compile_mode="default",
+    **kwargs,
 ):
     _alloc_main_grad(model_single_node, model_distributed)  # for fuse_wgrad_accumulation=True
     input_single_node.requires_grad_()
     input_distributed.requires_grad_()
+    forward_single_node = model_single_node
+    forward_distributed = model_distributed
+    if use_compile:
+        # Reset the compile cache so parametrized cases don't trip recompile_limit.
+        torch._dynamo.reset()
+        # Static shapes; dynamic-shape coverage lives in tests/pytorch/test_torch_compile.py.
+        forward_single_node = torch.compile(
+            model_single_node, fullgraph=True, mode=compile_mode, dynamic=False
+        )
+        forward_distributed = torch.compile(
+            model_distributed, fullgraph=True, mode=compile_mode, dynamic=False
+        )
     with te.autocast(
         enabled=QUANTIZATION is not None,
         recipe=quantization_recipe(),
     ):
-        output_single_node = model_single_node(input_single_node, **kwargs)
+        output_single_node = forward_single_node(input_single_node, **kwargs)
     with te.autocast(
         enabled=QUANTIZATION is not None,
         recipe=quantization_recipe(),
         amax_reduction_group=NCCL_WORLD,
     ):
-        output_distributed = model_distributed(input_distributed, **kwargs)
+        output_distributed = forward_distributed(input_distributed, **kwargs)
     return output_single_node, output_distributed
 
 
@@ -641,12 +659,20 @@ def test_quantized_all_gather():
 #                   Linear                 #
 ############################################
 @run_distributed_test()
-def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs):
+def _test_linear(
+    parallel_mode=None,
+    sequence_parallel=False,
+    use_compile=False,
+    compile_mode="default",
+    **kwargs,
+):
     """Test the linear layer with specified parallel mode and sequence parallelization.
 
     Args:
         parallel_mode (str): 'row' or 'column' parallelism.
         sequence_parallel (bool): Enable sequence parallelism if True.
+        use_compile (bool): Wrap the modules in ``torch.compile`` before running.
+        compile_mode (str): ``torch.compile`` mode ("default" or "reduce-overhead").
         kwargs (dict): Additional arguments for the linear layer.
     """
     # Set parameter data type
@@ -696,7 +722,12 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs):
 
     # Apply models
     output_single_node, output_distributed = _apply_models(
-        model_single_node, model_distributed, input_single_node, input_distributed
+        model_single_node,
+        model_distributed,
+        input_single_node,
+        input_distributed,
+        use_compile=use_compile,
+        compile_mode=compile_mode,
     )
 
     if "return_bias" in kwargs:
@@ -740,12 +771,18 @@ def test_linear():
         {"params_dtype": torch.float16 if QUANTIZATION != "nvfp4" else torch.bfloat16},
         {"delay_wgrad_compute": True},
         {"save_original_input": True},
+        {"use_compile": True},
+        {"use_compile": True, "compile_mode": "reduce-overhead"},
     ]
 
     for kwargs in kwargs_list:
         if kwargs.get("save_original_input", False) and QUANTIZATION == "fp8":
             continue
-        if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED:
+        # use_compile: debug instrumentation forces the eager fallback, so
+        # compile is a no-op there.
+        if NVTE_TEST_NVINSPECT_ENABLED and (
+            kwargs.get("delay_wgrad_compute", False) or kwargs.get("use_compile", False)
+        ):
             continue
         for parallel_mode in ["column", "row"]:
             for sequence_parallel in [False, True]:

tests/pytorch/distributed/test_comm_gemm_overlap.py

@@ -111,6 +111,8 @@ def _run_layer_with_overlap(
     quantization,
     num_layers=1,
     use_cublasmp=False,
+    use_compile=False,
+    compile_mode="default",
 ):
     test_path = TEST_ROOT / "run_layer_with_overlap.py"
     test_cmd = LAUNCH_CMD + [
@@ -129,6 +131,10 @@ def _run_layer_with_overlap(
     if overlap_rs_dgrad:
         test_cmd.append("--overlap-rs-dgrad")
 
+    if use_compile:
+        test_cmd.append("--compile")
+        test_cmd.append(f"--compile-mode={compile_mode}")
+
     if fp8:
         if quantization in ("fp8_delayed_scaling", "fp8_current_scaling") and not fp8_available:
             pytest.skip(reason_for_no_fp8)
@@ -281,6 +287,34 @@ def test_layers_with_overlap_bf16(
     )
 
 
+@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"])
+@pytest.mark.parametrize(
+    "linear_parallel_mode,overlap_rs_dgrad",
+    [
+        ("row", False),
+        ("column", False),
+        ("column", True),
+    ],
+    ids=[
+        "ROW-PARALLEL",
+        "COL-PARALLEL - BULK DGRAD/WGRAD",
+        "COL-PARALLEL - DGRAD+RS",
+    ],
+)
+def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode):
+    """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16),
+    checked numerically against the eager, non-overlap reference."""
+    _run_layer_with_overlap(
+        te.Linear.__name__,
+        linear_parallel_mode,
+        overlap_rs_dgrad,
+        False,
+        None,
+        use_compile=True,
+        compile_mode=compile_mode,
+    )
+
+
 @pytest.mark.parametrize("use_cublasmp", (False, True))
 @pytest.mark.parametrize(
     "quantization",

tests/pytorch/test_hybrid_quantization.py

@@ -495,7 +495,7 @@ def test_supports_only_rowwise_all_gather_nvfp4_columnwise(self):
         ``gather_along_first_dim`` cannot operate on a columnwise-only
         NVFP4 hybrid sub-storage. ``HybridQuantizer.supports_only_rowwise_all_gather``
         must return True in this case so ``_linear_forward_impl`` /
-        ``_linear_backward`` preserve rowwise data (which NVFP4 can
+        ``_linear_backward_impl`` preserve rowwise data (which NVFP4 can
         dequantize) instead.
         """
         hq = HybridQuantizer(