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

tests/pytorch/distributed/run_layer_with_overlap.py

@@ -18,6 +18,11 @@
 import torch
 import torch.distributed as dist
 
+try:
+    from torch._dynamo.utils import counters as dynamo_counters
+except ImportError:  # pragma: no cover
+    dynamo_counters = None
+
 import transformer_engine.pytorch as te
 from transformer_engine.common.recipe import (
     DelayedScaling,
@@ -200,6 +205,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 +303,13 @@ 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; to test"
+            " torch.compile with CUDA graphs use --compile --compile-mode"
+            " reduce-overhead."
+        )
+
     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 +560,14 @@ def run_fwd_bwd(model, x):
             loss.backward()
         return out
 
+    if opts.compile:
+        for i, layer in enumerate(test_model.layers):
+            test_model.layers[i] = torch.compile(layer, fullgraph=True, mode=opts.compile_mode)
+        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:
@@ -545,7 +578,18 @@ def run_fwd_bwd(model, x):
         if not opts.benchmark:
             del test_graph
     else:
+        if opts.compile and opts.compile_mode == "reduce-overhead":
+            # Warm up so the measured run below replays captured CUDA graphs.
+            for _ in range(2):
+                torch.compiler.cudagraph_mark_step_begin()
+                run_fwd_bwd(test_model, test_x)
+                test_model.zero_grad(set_to_none=True)
+                test_x.grad = None
+            torch.compiler.cudagraph_mark_step_begin()
         test_out = run_fwd_bwd(test_model, test_x)
+        if opts.compile and opts.compile_mode == "reduce-overhead" and dynamo_counters is not None:
+            skips = dynamo_counters["inductor"]["cudagraph_skips"]
+            assert skips == 0, f"reduce-overhead fell back to eager: {skips} cudagraph skip(s)"
     test_grads = [test_out, test_x.grad]
     names = ["output", "input.grad"]
     for test_name, test_param in test_model.named_parameters():

tests/pytorch/distributed/run_numerics.py

@@ -292,6 +292,15 @@ def _check_gradients(model_distributed, model_single, main_grad_check=False):
         assert not bool(numerics_failed.item())
 
 
+def _check_input_grads(input_single_node, input_distributed, parallel_mode, sequence_parallel):
+    grad_d = input_distributed.grad
+    if parallel_mode == "row":
+        grad_d = _gather(grad_d, dim=1)
+    elif sequence_parallel:
+        grad_d = _gather(grad_d, dim=0)
+    _check_outputs(input_single_node.grad, grad_d)
+
+
 def _copy_params(model_distributed, model_single):
     for dist_param, single_param in zip(model_distributed.parameters(), model_single.parameters()):
         with torch.no_grad():
@@ -310,22 +319,35 @@ 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()
+        forward_single_node = torch.compile(model_single_node, fullgraph=True, mode=compile_mode)
+        forward_distributed = torch.compile(model_distributed, fullgraph=True, mode=compile_mode)
     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 +663,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 +726,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:
@@ -728,6 +763,8 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs):
             main_grad_check=("fuse_wgrad_accumulation" in kwargs),
         )
 
+    _check_input_grads(input_single_node, input_distributed, parallel_mode, sequence_parallel)
+
 
 def test_linear():
     """Run linear layer tests with various configurations."""
@@ -740,12 +777,20 @@ 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
+        if kwargs.get("use_compile", False) and QUANTIZATION == "fp8":
             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,45 @@ def test_layers_with_overlap_bf16(
     )
 
 
+@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"])
+@pytest.mark.parametrize(
+    "quantization",
+    [None, "fp8_current_scaling", "mxfp8"],
+    ids=["bf16", "fp8_current_scaling", "mxfp8"],
+)
+@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, quantization, compile_mode
+):
+    """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile,
+    checked numerically against the eager, non-overlap reference."""
+    if quantization is not None and linear_parallel_mode == "row":
+        pytest.skip(
+            "FP8 row-parallel UB forces differentiable fp8_output, unsupported under compile."
+        )
+    _run_layer_with_overlap(
+        te.Linear.__name__,
+        linear_parallel_mode,
+        overlap_rs_dgrad,
+        quantization is not None,
+        quantization,
+        use_compile=True,
+        compile_mode=compile_mode,
+    )
+
+
 @pytest.mark.parametrize("use_cublasmp", (False, True))
 @pytest.mark.parametrize(
     "quantization",

tests/pytorch/test_hybrid_quantization.py

@@ -496,7 +496,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(