Carry symbolic shapes through compile custom ops

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."
     )
@@ -485,6 +498,9 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False):
         torch.testing.assert_close(test_param, ref_param, rtol=0.0, atol=0.0)
     dist_print("Copied parameters from test model to reference model...", debug=True)
 
+    if opts.compile and opts.use_cuda_graphs:
+        raise ValueError("--compile and --use-cuda-graphs are mutually exclusive.")
+
     # Fp8 recipe setup
     fp8_format = Format.HYBRID
     fp8_recipe = None
@@ -535,6 +551,18 @@ def run_fwd_bwd(model, x):
             loss.backward()
         return out
 
+    if opts.compile:
+        for i, layer in enumerate(test_model.layers):
+            # dynamic=False for now: symbolic shapes would land in an OpaqueValueBundle
+            # op arg whose hash chokes on non-nested SymInt (see run_numerics).
+            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,45 @@ 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:
+        # Each parametrized case compiles the same module.forward code object with
+        # a different shape/recipe; with dynamic=False those guards accumulate and
+        # eventually trip Dynamo's recompile_limit. Reset so every case starts from
+        # a clean compile cache (mirrors the single-GPU torch.compile tests).
+        torch._dynamo.reset()
+        # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle
+        # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static
+        # shapes (recompile per shape) until the bundle handles symbolic shapes.
+        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 +664,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 +727,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,13 +776,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:
             continue
+        # debug instrumentation forces the eager fallback, so compile is a no-op there.
+        if kwargs.get("use_compile", False) and NVTE_TEST_NVINSPECT_ENABLED:
+            continue
         for parallel_mode in ["column", "row"]:
             for sequence_parallel in [False, True]:
                 _test_linear(parallel_mode, sequence_parallel, **kwargs)

tests/pytorch/distributed/test_comm_gemm_overlap.py

@@ -111,6 +111,8 @@ def _run_layer_with_overlap(
     quantization,
     num_layers=1,
     use_cublasmp=False,
+    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 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,40 @@ 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).
+
+    Userbuffers is expected to stay on Linear's compiled custom-op path (the
+    collective lives inside the opaque op), so this checks that torch.compile +
+    Userbuffers stays numerically correct against the eager, non-overlap reference.
+    ``compile_mode="reduce-overhead"`` additionally exercises CUDA-graph trees on
+    top of the Userbuffers collectives.
+    """
+    _run_layer_with_overlap(
+        te.Linear.__name__,
+        linear_parallel_mode,
+        overlap_rs_dgrad,
+        False,
+        None,
+        compile=True,
+        compile_mode=compile_mode,
+    )
+
+
 @pytest.mark.parametrize("use_cublasmp", (False, True))
 @pytest.mark.parametrize(
     "quantization",

tests/pytorch/test_sanity.py

@@ -1154,6 +1154,41 @@ def test_quantized_model_init_high_precision_init_val():
     ), "clear_high_precision_init_val() not work"
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("move", ["cuda", "cpu", "half"])
+def test_quantized_param_attrs_survive_apply(move):
+    """Attributes attached to a quantized parameter survive nn.Module._apply.
+
+    Quantized parameters implement the flatten protocol, so ``_apply`` moves them
+    with ``swap_tensors``, which exchanges the parameter's whole ``__dict__``.
+    Anything attached from the outside rides out on the discarded tensor unless
+    the module re-attaches it.
+    """
+    with quantized_model_init(preserve_high_precision_init_val=True):
+        model = Linear(64, 64)
+
+    weight = model.weight
+    expected = weight.get_high_precision_init_val()
+    weight.probe_attr = "attached-from-outside"
+
+    if move == "cuda":
+        model = model.cuda()
+    elif move == "cpu":
+        model = model.cpu()
+    else:
+        model = model.half()
+
+    weight = model.weight
+    assert hasattr(weight, "get_high_precision_init_val"), f"accessor lost by .{move}()"
+    assert hasattr(weight, "clear_high_precision_init_val"), f"accessor lost by .{move}()"
+    torch.testing.assert_close(weight.get_high_precision_init_val(), expected, rtol=0, atol=0)
+    assert weight.probe_attr == "attached-from-outside", f"custom attr lost by .{move}()"
+
+    # The accessor must read the surviving parameter, not the discarded one.
+    weight.clear_high_precision_init_val()
+    assert weight.get_high_precision_init_val() is None
+
+
 @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
 def test_grouped_linear_single_param_preserves_high_precision_init(monkeypatch):
     """Grouped MXFP8 and discrete weights produce identical FP32 master initialization."""