# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

from typing import Callable, Dict, Iterable, List, Tuple, Union
import pytest
import copy

import torch
from transformer_engine.pytorch import (
    DotProductAttention,
    GroupedLinear,
    LayerNormLinear,
    LayerNormMLP,
    Linear,
    MultiheadAttention,
    TransformerLayer,
    autocast,
    quantized_model_init,
    make_graphed_callables,
    is_fp8_available,
    is_fp8_block_scaling_available,
    is_mxfp8_available,
    is_nvfp4_available,
    is_bf16_available,
)
from transformer_engine.pytorch.quantization import FP8GlobalStateManager
import transformer_engine.pytorch.ops as te_ops
from transformer_engine.common import recipe
from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override

# Check if FP8 is supported.
fp8_available = is_fp8_available()
fp8_block_scaling_available = is_fp8_block_scaling_available()
mxfp8_available = is_mxfp8_available()
nvfp4_available = is_nvfp4_available()

# Reset RNG states.
reset_rng_states()

model_configs = {
    "small": ModelConfig(2, 32, 2, 32),
}


def nvfp4_vanilla():
    nvfp4_recipe = recipe.NVFP4BlockScaling()
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams()
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
    return nvfp4_recipe


def nvfp4_rht_and_2d_quantization():
    nvfp4_recipe = recipe.NVFP4BlockScaling()
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams(
        random_hadamard_transform=True, fp4_2d_quantization=False
    )
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(
        random_hadamard_transform=False, fp4_2d_quantization=True
    )
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams(
        random_hadamard_transform=True, fp4_2d_quantization=False
    )
    return nvfp4_recipe


def nvfp4_row_scaled():
    nvfp4_recipe = recipe.NVFP4BlockScaling(
        disable_rht=True,
        disable_stochastic_rounding=True,
        disable_2d_quantization=True,
        row_scaled_activation=True,
        backward_override="dequantized",
    )
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams()
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
    return nvfp4_recipe


def nvfp4_4over6():
    nvfp4_recipe = recipe.NVFP4BlockScaling(
        disable_rht=True,
        disable_stochastic_rounding=True,
        nvfp4_4over6="all",
    )
    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True)
    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
    return nvfp4_recipe


def check_rht_usage(recipe: recipe.Recipe) -> bool:
    # if using RHT, we can only support bf16
    # check fp4_quant_fwd_inp, fp4_quant_fwd_weight, fp4_quant_bwd_grad
    if recipe.nvfp4():
        if (
            recipe.fp4_quant_fwd_inp.random_hadamard_transform
            or recipe.fp4_quant_fwd_weight.random_hadamard_transform
            or recipe.fp4_quant_bwd_grad.random_hadamard_transform
        ):
            return True
    return False


def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> bool:
    supported_input_dtypes = []
    if recipe.nvfp4():
        supported_input_dtypes.append(torch.bfloat16)
        # if not using RHT, we can add fp32 as well
    if not check_rht_usage(recipe):
        supported_input_dtypes.append(torch.float32)
    return supported_input_dtypes


fp8_recipes = []
if mxfp8_available:
    fp8_recipes.append(recipe.MXFP8BlockScaling())
if nvfp4_available:
    fp8_recipes.append(nvfp4_rht_and_2d_quantization())
    fp8_recipes.append(nvfp4_row_scaled())
    fp8_recipes.append(nvfp4_4over6())
if fp8_block_scaling_available:
    fp8_recipes.append(recipe.Float8BlockScaling())
if fp8_available:
    fp8_recipes.append(recipe.Float8CurrentScaling())
    fp8_recipes.append(recipe.DelayedScaling())

# Supported data types
dtypes: List[torch.dtype] = [torch.float32, torch.float16]
if is_bf16_available():  # bf16 requires sm_80 or higher
    dtypes.append(torch.bfloat16)


@pytest.fixture(autouse=True)
def reset_global_fp8_state():
    yield
    FP8GlobalStateManager.reset()


def assert_all_equal(l1: List[torch.Tensor], l2: List[torch.Tensor], names=None) -> bool:
    """Check that two lists of tensors match exactly."""
    assert len(l1) == len(l2), "Unequal number of outputs."
    failure_message = "Output mismatches in:"
    failed_tensors = []
    for i, (t1, t2) in enumerate(zip(l1, l2)):
        if not torch.equal(t1, t2):
            failure_message += "\n    "
            if names is None:
                failure_message += f"tensor at idx={i}"
            else:
                failure_message += names[i]
            failed_tensors.append((t1, t2))
    if failed_tensors:
        print(failure_message)
        t1, t2 = failed_tensors[0]
        torch.testing.assert_close(t1, t2, rtol=0, atol=0)


def generate_data(
    model_config: ModelConfig,
    dtype: torch.dtype,
    warmup: bool = False,
    requires_grad: bool = True,
) -> torch.Tensor:
    """Generate synthetic data."""
    gen_func = torch.ones if warmup else torch.randn
    return gen_func(
        model_config.max_seqlen_q,
        model_config.batch_size,
        model_config.hidden_size,
        device="cuda",
        requires_grad=requires_grad,
        dtype=dtype,
    )


def get_outputs(
    model: torch.nn.Module,
    output: Union[torch.Tensor, Iterable[torch.Tensor]],
) -> List[torch.Tensor]:
    """Return grads and params for comparsion."""
    values = []
    for param in model.parameters():
        values.append(param)
        if param.grad is not None:
            values.append(param.grad)
    if isinstance(output, torch.Tensor):
        values.append(output)
    else:
        values.extend(output)
    return values


def reset_graphs(
    graphed_callables: Union[Callable, Tuple[Callable, ...], Dict[Tuple[int, int], Callable]],
) -> None:
    """Reset CUDA graphs."""
    if isinstance(graphed_callables, tuple) or isinstance(graphed_callables, list):
        for callable in graphed_callables:
            callable.reset()
    elif isinstance(graphed_callables, dict):
        for callable in graphed_callables.values():
            callable.reset()
    else:
        graphed_callables.reset()


class _Sequential(torch.nn.Sequential):
    """Sequential model that forwards keyword arguments to modules"""

    def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor:
        x = input_
        for module in self:
            x = module(x, **kwargs)
        return x


class _GroupedLinearWrapper(torch.nn.Module):
    """Adapt GroupedLinear to the [seqlen, batch, hidden] data used in this test.

    GroupedLinear expects a 2D `[total_tokens, hidden]` input plus an `m_splits`
    list, so this wrapper flattens the leading dims, splits the tokens evenly
    across the GEMMs, and restores the original shape. It also forwards
    `is_first_microbatch` so the FP8 weight-caching path is exercised under CUDA
    graphs.
    """

    def __init__(self, hidden_size: int, num_gemms: int, params_dtype: torch.dtype) -> None:
        super().__init__()
        self.num_gemms = num_gemms
        self.grouped_linear = GroupedLinear(
            num_gemms,
            hidden_size,
            hidden_size,
            device="cuda",
            params_dtype=params_dtype,
        )

    def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor:
        seqlen, batch, hidden = input_.shape
        x = input_.reshape(seqlen * batch, hidden)
        total_tokens = x.shape[0]
        assert total_tokens % self.num_gemms == 0, "tokens must split evenly across GEMMs"
        split = total_tokens // self.num_gemms
        m_splits = [split] * self.num_gemms
        out = self.grouped_linear(x, m_splits, **kwargs)
        return out.reshape(seqlen, batch, hidden)


# Supported modules
_test_cuda_graphs_modules: List[str] = [
    # Put linear first to test the case where the cuda context might not be set in
    # creating TMA descriptor for MXFP8 quantization.
    "linear",
    "transformer",
    "layernorm_mlp_nocheckpoint",
    "layernorm_mlp_checkpoint",
    "layernorm_linear",
    "mha",
    "linear_op",
]


def _test_cuda_graphs(
    *,
    graph_mode: str,
    module: str,
    model_config: ModelConfig,
    num_layers: int,
    dtype: torch.dtype,
    fp8: bool,
    fp8_params: bool,
    fp8_weight_caching: bool,
    fp8_recipe: recipe.Recipe,
) -> List[torch.Tensor]:
    """Helper function for CUDA graph test."""
    reset_rng_states()
    FP8GlobalStateManager.reset()

    # Operation-based API does not support FP8 weight caching.
    if module == "linear_op":
        fp8_weight_caching = False

    # Create modules.
    with quantized_model_init(enabled=fp8_params, recipe=fp8_recipe):
        if module == "transformer":
            modules = [
                TransformerLayer(
                    model_config.hidden_size,
                    model_config.hidden_size,
                    model_config.num_heads,
                    hidden_dropout=0.0,
                    attention_dropout=0.0,
                    fuse_qkv_params=True,
                    params_dtype=dtype,
                )
                for _ in range(num_layers)
            ]
        elif module == "layernorm_mlp_nocheckpoint":
            modules = [
                LayerNormMLP(
                    model_config.hidden_size,
                    model_config.hidden_size,
                    params_dtype=dtype,
                    checkpoint=False,
                )
                for _ in range(num_layers)
            ]
        elif module == "layernorm_mlp_checkpoint":
            modules = [
                LayerNormMLP(
                    model_config.hidden_size,
                    model_config.hidden_size,
                    params_dtype=dtype,
                    checkpoint=True,
                )
                for _ in range(num_layers)
            ]
        elif module == "layernorm_linear":
            modules = [
                LayerNormLinear(
                    model_config.hidden_size,
                    model_config.hidden_size,
                    params_dtype=dtype,
                )
                for _ in range(num_layers)
            ]
        elif module == "mha":
            modules = [
                MultiheadAttention(
                    model_config.hidden_size,
                    model_config.num_heads,
                    attention_dropout=0.0,
                    params_dtype=dtype,
                    fuse_qkv_params=True,
                )
                for _ in range(num_layers)
            ]
        elif module == "linear":
            modules = [
                Linear(
                    model_config.hidden_size,
                    model_config.hidden_size,
                    device="cuda",
                    params_dtype=dtype,
                )
                for _ in range(num_layers)
            ]
        elif module == "grouped_linear":
            modules = [
                _GroupedLinearWrapper(
                    model_config.hidden_size,
                    num_gemms=2,
                    params_dtype=dtype,
                )
                for _ in range(num_layers)
            ]
        elif module == "linear_op":
            modules = [
                te_ops.Sequential(
                    te_ops.Linear(
                        model_config.hidden_size,
                        model_config.hidden_size,
                        dtype=dtype,
                    ),
                )
                for _ in range(num_layers)
            ]
        else:
            raise ValueError(f"Unknown module type ({module})")

        # Initialize gradient buffers.
        for module in modules:
            for param in module.parameters():
                param.grad = torch.empty_like(param)

        # Generate model and wrap API to return graphed version.
        if graph_mode == "full":
            # Graph entire model at once.
            model = torch.nn.Sequential(*modules)
            model = make_graphed_callables(
                model,
                (generate_data(model_config, dtype, warmup=True),),
                num_warmup_iters=10,
                enabled=fp8,
                cache_quantized_params=fp8_weight_caching,
                recipe=fp8_recipe,
            )
        elif graph_mode == "individual":
            # Graph individual modules.
            modules = [
                make_graphed_callables(
                    module,
                    (generate_data(model_config, dtype, warmup=True),),
                    num_warmup_iters=10,
                    enabled=fp8,
                    cache_quantized_params=fp8_weight_caching,
                    recipe=fp8_recipe,
                )
                for module in modules
            ]
            model = _Sequential(*modules)
        else:
            model = _Sequential(*modules)

    # Optimizer.
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # Training steps.
    for _ in range(3):
        optimizer.zero_grad(set_to_none=False)
        for grad_accumulation_step in range(2):
            input_ = generate_data(model_config, dtype)
            grad_output = generate_data(model_config, dtype, requires_grad=False)
            with autocast(enabled=fp8, recipe=fp8_recipe):
                kwargs = {}
                if fp8_weight_caching:
                    kwargs["is_first_microbatch"] = grad_accumulation_step == 0
                output = model(input_, **kwargs)
            output.backward(grad_output)
        optimizer.step()

    outputs = get_outputs(model, output)
    if graph_mode == "full":
        reset_graphs(model)
    elif graph_mode == "individual":
        reset_graphs(modules)
    return outputs


@pytest.mark.parametrize("module", _test_cuda_graphs_modules)
@pytest.mark.parametrize("dtype", dtypes)
@pytest.mark.parametrize("fp8_params", (False, True))
@pytest.mark.parametrize("fp8_recipe", fp8_recipes + [None], ids=recipe_id)
@pytest.mark.parametrize("backward_override", (None, "high_precision", "dequantized"))
def test_make_graphed_callables(
    *,
    module: str,
    model_config: str = "small",
    num_layers: int = 3,
    dtype: torch.dtype,
    fp8_params: bool,
    fp8_recipe: recipe.Recipe,
    backward_override: str,
    fp8_weight_caching: bool = False,
) -> None:

    fp8 = fp8_recipe is not None

    skip_unsupported_backward_override(module, fp8_recipe, backward_override)
    if fp8:
        fp8_recipe = copy.deepcopy(fp8_recipe)
        fp8_recipe.backward_override = backward_override

    if fp8_params and not fp8:
        pytest.skip("FP8 needed for FP8 parameters.")
    if fp8_weight_caching and not fp8:
        pytest.skip("FP8 needed for FP8 parameters.")
    if fp8 and (fp8_recipe.float8_block_scaling() or fp8_recipe.nvfp4()) and module == "linear_op":
        pytest.skip(
            f"Module not yet supported for {fp8_recipe.__class__.__name__} with CUDA graphs"
        )
    if fp8 and fp8_recipe.nvfp4():
        if getattr(fp8_recipe, "row_scaled_activation", False) and module == "mha":
            pytest.skip("Row-scaled NVFP4 CUDA graph coverage applies to GEMM modules.")
        if dtype not in get_nvfp4_inp_supported_dtypes(fp8_recipe, dtype):
            pytest.skip(
                f"Input dtype {dtype} not supported for NVFP4 Recipe"
                f" {fp8_recipe.__class__.__name__}"
            )
        if fp8_params:
            pytest.skip("NVFP4 params not supported")
    if (
        fp8
        and fp8_recipe.delayed()
        and torch.cuda.get_device_capability() >= (10, 0)
        and module == "layernorm_mlp_checkpoint"
    ):
        pytest.skip(
            "CUDA graphs not supported for LayerNormMLP "
            "with checkpoint=True, SM>=10, "
            "and DelayedScaling recipe"
        )

    # Run model with different CUDA graph settings.
    model_config = model_configs[model_config]
    kwargs = dict(
        module=module,
        model_config=model_config,
        num_layers=num_layers,
        dtype=dtype,
        fp8=fp8,
        fp8_params=fp8_params,
        fp8_weight_caching=fp8_weight_caching,
        fp8_recipe=fp8_recipe,
    )
    # Put graphed callables first to test the case where the cuda context might not be set in
    # creating TMA descriptor for MXFP8 quantization.
    graph_outputs_mode1 = _test_cuda_graphs(graph_mode="full", **kwargs)
    graph_outputs_mode2 = _test_cuda_graphs(graph_mode="individual", **kwargs)
    outputs = _test_cuda_graphs(graph_mode="none", **kwargs)

    # Check that results match.
    assert_all_equal(outputs, graph_outputs_mode1)
    assert_all_equal(outputs, graph_outputs_mode2)


_test_make_graphed_callables_with_fp8_weight_caching_modules = [
    "transformer",
    "layernorm_mlp_nocheckpoint",
    "layernorm_mlp_checkpoint",
    "layernorm_linear",
    "linear",
    "mha",
]


@pytest.mark.parametrize(
    "module",
    _test_make_graphed_callables_with_fp8_weight_caching_modules,
)
@pytest.mark.parametrize("dtype", dtypes)
@pytest.mark.parametrize("fp8_params", (False, True))
@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id)
@pytest.mark.parametrize("backward_override", (None, "high_precision", "dequantized"))
def test_make_graphed_callables_with_fp8_weight_caching(
    *,
    module: str,
    dtype: torch.dtype,
    fp8_params: bool,
    fp8_recipe: recipe.Recipe,
    backward_override: str,
) -> None:
    test_make_graphed_callables(
        module=module,
        dtype=dtype,
        fp8_params=fp8_params,
        fp8_recipe=fp8_recipe,
        backward_override=backward_override,
        fp8_weight_caching=True,
    )


# Per-tensor FP8 recipes that support GroupedLinear FP8 weight caching.
_grouped_linear_fp8_weight_caching_recipes = []
if fp8_available:
    _grouped_linear_fp8_weight_caching_recipes.append(recipe.DelayedScaling())
    _grouped_linear_fp8_weight_caching_recipes.append(recipe.Float8CurrentScaling())


@pytest.mark.skipif(not fp8_available, reason="FP8 is not supported")
@pytest.mark.parametrize("dtype", dtypes)
@pytest.mark.parametrize("fp8_params", (False, True))
@pytest.mark.parametrize("fp8_recipe", _grouped_linear_fp8_weight_caching_recipes, ids=recipe_id)
def test_make_graphed_callables_grouped_linear_with_fp8_weight_caching(
    *,
    dtype: torch.dtype,
    fp8_params: bool,
    fp8_recipe: recipe.Recipe,
    model_config: str = "small",
    num_layers: int = 3,
) -> None:
    """GroupedLinear must thread `is_first_microbatch` into the FP8 weight-update
    skip tensor under CUDA graphs.

    With `fp8_weight_caching` enabled, the graphed and non-graphed runs only match
    when `skip_fp8_weight_update` is propagated for every microbatch. Before the
    fix, GroupedLinear hardcoded it to `None`, so the cached FP8 weights diverged
    from the eager reference. This regresses if that propagation is dropped again.
    """
    config = model_configs[model_config]
    kwargs = dict(
        module="grouped_linear",
        model_config=config,
        num_layers=num_layers,
        dtype=dtype,
        fp8=True,
        fp8_params=fp8_params,
        fp8_weight_caching=True,
        fp8_recipe=fp8_recipe,
    )
    graph_outputs_full = _test_cuda_graphs(graph_mode="full", **kwargs)
    graph_outputs_individual = _test_cuda_graphs(graph_mode="individual", **kwargs)
    outputs = _test_cuda_graphs(graph_mode="none", **kwargs)

    assert_all_equal(outputs, graph_outputs_full)
    assert_all_equal(outputs, graph_outputs_individual)


def generate_data_for_dot_product_attention(
    model_config: ModelConfig,
    dtype: torch.dtype,
    warmup: bool = False,
) -> List[torch.Tensor]:
    """Generate synthetic data for dot product attention."""
    gen_func = torch.ones if warmup else torch.randn
    return [
        gen_func(
            model_config.max_seqlen_q,
            model_config.batch_size,
            model_config.num_heads,
            model_config.kv_channels,
            device="cuda",
            requires_grad=True,
            dtype=dtype,
        )
        for _ in range(3)
    ]


def _test_cuda_graphs_with_dot_product_attention(
    *,
    with_graph: bool,
    model_config: ModelConfig,
    dtype: torch.dtype,
) -> List[torch.Tensor]:
    """Helper function for CUDA graph test."""
    reset_rng_states()
    FP8GlobalStateManager.reset()

    # Create dot product attention module.
    assert model_config.hidden_size % model_config.num_heads == 0
    model = DotProductAttention(
        model_config.num_heads,
        model_config.kv_channels,
        attention_dropout=0.0,
    )

    # Graph model if needed.
    if with_graph:
        model = make_graphed_callables(
            model,
            generate_data_for_dot_product_attention(model_config, dtype, warmup=True),
            num_warmup_iters=10,
            enabled=False,
        )

    # Forward and backward passes.
    for _ in range(3):
        inputs = generate_data_for_dot_product_attention(model_config, dtype)
        grad_output = generate_data(model_config, dtype, requires_grad=False)
        output = model(*inputs)
        output.backward(grad_output)

    outputs = get_outputs(model, output)
    if with_graph:
        reset_graphs(model)
    return outputs


@pytest.mark.parametrize("dtype", dtypes)
def test_make_graphed_callables_with_dot_product_attention(
    *,
    model_config: str = "small",
    dtype: torch.dtype,
) -> None:
    """Test CUDA graphs with dot product attention."""
    model_config = model_configs[model_config]
    kwargs = dict(model_config=model_config, dtype=dtype)
    outputs = _test_cuda_graphs_with_dot_product_attention(with_graph=False, **kwargs)
    graph_outputs = _test_cuda_graphs_with_dot_product_attention(with_graph=True, **kwargs)
    assert_all_equal(outputs, graph_outputs)


def _test_cuda_graphs_with_kwargs(
    *,
    with_graph: bool,
    model_config: ModelConfig,
    dtype: torch.dtype,
) -> List[torch.Tensor]:
    """Helper function for CUDA graph test with keyword arguments."""
    reset_rng_states()

    # Initialize model.
    model = TransformerLayer(
        model_config.hidden_size,
        model_config.hidden_size,
        model_config.num_heads,
        hidden_dropout=0.0,
        attention_dropout=0.0,
        self_attn_mask_type="arbitrary",
        fuse_qkv_params=True,
        params_dtype=dtype,
    )

    # Initialize gradient buffers.
    for param in model.parameters():
        param.grad = torch.empty_like(param)

    # Make graphed version of model if needed.
    if with_graph:
        attn_mask = torch.zeros(
            (
                model_config.batch_size,
                1,
                model_config.max_seqlen_q,
                model_config.max_seqlen_kv,
            ),
            dtype=torch.bool,
            device="cuda",
        )
        model = make_graphed_callables(
            model,
            (generate_data(model_config, dtype, warmup=True),),
            sample_kwargs=dict(attention_mask=attn_mask),
            allow_unused_input=True,
        )

    # Optimizer.
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # Training loop.
    for _ in range(3):
        optimizer.zero_grad(set_to_none=False)
        for grad_accumulation_step in range(2):
            input_ = generate_data(model_config, dtype)
            grad_output = generate_data(model_config, dtype, requires_grad=False)
            attn_mask = torch.randint(
                2,
                (
                    model_config.batch_size,
                    1,
                    model_config.max_seqlen_q,
                    model_config.max_seqlen_kv,
                ),
                dtype=torch.bool,
                device="cuda",
            )
            output = model(input_, attention_mask=attn_mask)
            output.backward(grad_output)
        optimizer.step()

    outputs = get_outputs(model, output)
    if with_graph:
        reset_graphs(model)
    return outputs


def test_make_graphed_callables_with_kwargs(
    *,
    model_config: str = "small",
    dtype: torch.dtype = torch.float32,
) -> None:
    """Test CUDA graphs with keyword arguments."""
    model_config = model_configs[model_config]
    kwargs = dict(model_config=model_config, dtype=dtype)
    outputs = _test_cuda_graphs_with_kwargs(with_graph=False, **kwargs)
    graph_outputs = _test_cuda_graphs_with_kwargs(with_graph=True, **kwargs)
    assert_all_equal(outputs, graph_outputs)


def test_make_graphed_callables_returns_owned_parameter_grads() -> None:
    """Parameter grads returned from graph replay must not alias static graph buffers."""
    reset_rng_states()
    model_config = model_configs["small"]
    dtype = torch.float32
    model = torch.nn.Linear(
        model_config.hidden_size,
        model_config.hidden_size,
        bias=False,
        device="cuda",
        dtype=dtype,
    )
    model = make_graphed_callables(
        model,
        (generate_data(model_config, dtype, warmup=True, requires_grad=False),),
    )

    seen_grads = []

    def save_grad(grad):
        seen_grads.append(grad)
        return grad

    hook = model.weight.register_hook(save_grad)
    try:
        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 1
        first_grad = seen_grads[0]
        first_grad_ptr = first_grad.data_ptr()
        first_grad_snapshot = first_grad.clone()

        model.zero_grad(set_to_none=True)

        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 2
        assert first_grad.data_ptr() == first_grad_ptr
        assert seen_grads[1].data_ptr() != first_grad_ptr
        torch.testing.assert_close(first_grad, first_grad_snapshot, rtol=0, atol=0)
    finally:
        hook.remove()
        reset_graphs(model)


def test_make_graphed_callables_accumulates_owned_parameter_grads() -> None:
    """Parameter grad accumulation must not reuse overwritten static graph buffers."""
    reset_rng_states()
    model_config = model_configs["small"]
    dtype = torch.float32
    model = torch.nn.Linear(
        model_config.hidden_size,
        model_config.hidden_size,
        bias=False,
        device="cuda",
        dtype=dtype,
    )
    model = make_graphed_callables(
        model,
        (generate_data(model_config, dtype, warmup=True, requires_grad=False),),
    )

    input_1 = generate_data(model_config, dtype, requires_grad=False)
    grad_1 = generate_data(model_config, dtype, requires_grad=False)
    input_2 = generate_data(model_config, dtype, requires_grad=False)
    grad_2 = generate_data(model_config, dtype, requires_grad=False)
    expected_grad = torch.einsum("...o,...i->oi", grad_1, input_1) + torch.einsum(
        "...o,...i->oi", grad_2, input_2
    )

    try:
        model.zero_grad(set_to_none=True)
        model(input_1).backward(grad_1)
        model(input_2).backward(grad_2)
        torch.testing.assert_close(model.weight.grad, expected_grad, rtol=0, atol=0)
    finally:
        reset_graphs(model)


def test_make_graphed_callables_preserves_skipped_parameter_grad_alias() -> None:
    """Delayed-wgrad parameters are excluded from returned-grad clone handling."""
    reset_rng_states()
    model_config = model_configs["small"]
    dtype = torch.float32
    model = torch.nn.Linear(
        model_config.hidden_size,
        model_config.hidden_size,
        bias=False,
        device="cuda",
        dtype=dtype,
    )
    model.weight.skip_backward_post_hook = True
    model = make_graphed_callables(
        model,
        (generate_data(model_config, dtype, warmup=True, requires_grad=False),),
    )

    seen_grads = []

    def save_grad(grad):
        seen_grads.append(grad)
        return grad

    hook = model.weight.register_hook(save_grad)
    try:
        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 1
        first_grad_ptr = seen_grads[0].data_ptr()

        model.zero_grad(set_to_none=True)

        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 2
        assert seen_grads[1].data_ptr() == first_grad_ptr
    finally:
        hook.remove()
        reset_graphs(model)


def test_make_graphed_callables_can_skip_returned_parameter_grad_clone() -> None:
    """Parameter grad clone handling can be disabled for callers that manage lifetimes."""
    reset_rng_states()
    model_config = model_configs["small"]
    dtype = torch.float32
    model = torch.nn.Linear(
        model_config.hidden_size,
        model_config.hidden_size,
        bias=False,
        device="cuda",
        dtype=dtype,
    )
    model = make_graphed_callables(
        model,
        (generate_data(model_config, dtype, warmup=True, requires_grad=False),),
        clone_param_grads_on_return=False,
    )

    seen_grads = []

    def save_grad(grad):
        seen_grads.append(grad)
        return grad

    hook = model.weight.register_hook(save_grad)
    try:
        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 1
        first_grad_ptr = seen_grads[0].data_ptr()

        model.zero_grad(set_to_none=True)

        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 2
        assert seen_grads[1].data_ptr() == first_grad_ptr
    finally:
        hook.remove()
        reset_graphs(model)


def test_make_graphed_callables_snapshots_parameter_grad_clone_policy() -> None:
    """Parameter grad clone policy is fixed at capture time."""
    reset_rng_states()
    model_config = model_configs["small"]
    dtype = torch.float32
    model = torch.nn.Linear(
        model_config.hidden_size,
        model_config.hidden_size,
        bias=False,
        device="cuda",
        dtype=dtype,
    )
    model = make_graphed_callables(
        model,
        (generate_data(model_config, dtype, warmup=True, requires_grad=False),),
    )
    model.weight.skip_backward_post_hook = True

    seen_grads = []

    def save_grad(grad):
        seen_grads.append(grad)
        return grad

    hook = model.weight.register_hook(save_grad)
    try:
        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 1
        first_grad = seen_grads[0]
        first_grad_ptr = first_grad.data_ptr()
        first_grad_snapshot = first_grad.clone()

        model.zero_grad(set_to_none=True)

        output = model(generate_data(model_config, dtype, requires_grad=False))
        output.backward(generate_data(model_config, dtype, requires_grad=False))

        assert len(seen_grads) == 2
        assert seen_grads[1].data_ptr() != first_grad_ptr
        torch.testing.assert_close(first_grad, first_grad_snapshot, rtol=0, atol=0)
    finally:
        hook.remove()
        reset_graphs(model)


def _make_capture_time_hooks(
    modules: Tuple[torch.nn.Module, ...],
    records: List[Tuple[int, str]],
) -> List[Dict[str, Dict[int, Callable]]]:
    """Make capture-time hooks that record call order."""

    def make_hook(module_idx: int, hook_name: str) -> Callable:
        expected_module = modules[module_idx]

        def hook(module: torch.nn.Module) -> None:
            assert module is expected_module
            assert not torch.cuda.is_current_stream_capturing()
            records.append((module_idx, hook_name))

        return hook

    return [
        {
            "forward_pre_hooks": {0: make_hook(module_idx, "forward_pre_hooks")},
            "forward_hooks": {0: make_hook(module_idx, "forward_hooks")},
            "backward_pre_hooks": {0: make_hook(module_idx, "backward_pre_hooks")},
            "backward_hooks": {0: make_hook(module_idx, "backward_hooks")},
        }
        for module_idx in range(len(modules))
    ]


@pytest.mark.parametrize("with_order", (False, True))
def test_make_graphed_callables_with_capture_time_hooks(with_order: bool) -> None:
    """Test capture-time hooks around warmup and graph capture."""
    num_warmup_iters = 2
    modules = (
        torch.nn.Linear(8, 8, device="cuda"),
        torch.nn.Linear(8, 8, device="cuda"),
    )
    sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in modules)
    records = []
    hook_order = [
        (0, "forward_pre_hooks"),
        (0, "forward_hooks"),
        (1, "forward_pre_hooks"),
        (1, "forward_hooks"),
        (1, "backward_pre_hooks"),
        (1, "backward_hooks"),
        (0, "backward_pre_hooks"),
        (0, "backward_hooks"),
    ]

    graphed_callables = make_graphed_callables(
        modules,
        sample_args,
        num_warmup_iters=num_warmup_iters,
        _order=[1, 2, -2, -1] if with_order else None,
        capture_time_hooks=_make_capture_time_hooks(modules, records),
    )

    assert records == hook_order * (num_warmup_iters + 1)

    for graphed in graphed_callables:
        x = torch.randn(4, 8, device="cuda", requires_grad=True)
        y = graphed(x)
        y.backward(torch.ones_like(y))
    assert records == hook_order * (num_warmup_iters + 1)
    reset_graphs(graphed_callables)


def _test_cuda_graphs_with_interleaved_pipeline_parallelism(
    *,
    with_graph: bool,
    model_config: ModelConfig,
    dtype: torch.dtype,
    reuse_graph_input_output_buffers: bool = False,
    clone_param_grads_on_return: bool = True,
) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
    """Simulate Megatron-LM interleaved pipeline parallelism."""
    reset_rng_states()

    # Pipeline parallel configuration.
    num_layers = 2
    num_microbatches = 3
    layer_order = [1, 2, 1, 2, -2, -1, 1, 2, -2, -1, -2, -1]

    # Initialize model.
    model = torch.nn.ModuleList(
        [
            Linear(
                model_config.hidden_size,
                model_config.hidden_size,
                params_dtype=dtype,
            )
            for _ in range(num_layers)
        ]
    )

    # Initialize gradient buffers.
    for param in model.parameters():
        param.grad = torch.empty_like(param)

    # Make graphed version of model if needed.
    layer_forwards = {
        (i % num_layers, i // num_layers): model[i % num_layers]
        for i in range(num_layers * num_microbatches)
    }
    if with_graph:
        sample_args = tuple(
            (generate_data(model_config, dtype, warmup=True),)
            for _ in range(num_layers * num_microbatches)
        )
        layer_forwards = make_graphed_callables(
            tuple(model),
            sample_args,
            allow_unused_input=True,
            _order=layer_order,
            _reuse_graph_input_output_buffers=reuse_graph_input_output_buffers,
            clone_param_grads_on_return=clone_param_grads_on_return,
        )
        layer_forwards = {
            (i // num_microbatches, i % num_microbatches): forward
            for i, forward in enumerate(layer_forwards)
        }

    # Optimizer.
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # Training loop.
    for _ in range(3):
        optimizer.zero_grad(set_to_none=False)

        # Generate data.
        inputs = {}
        grad_outputs = {}
        for layer_idx in range(num_layers):
            for microbatch_idx in range(num_microbatches):
                x = generate_data(model_config, dtype)
                dy = generate_data(model_config, dtype, requires_grad=False)
                idxs = (layer_idx, microbatch_idx)
                inputs[idxs] = x
                grad_outputs[idxs] = dy

        # Cache for layer outputs.
        outputs = {}
        output_snapshots = {} if reuse_graph_input_output_buffers else None

        def forward(layer_idx: int, microbatch_idx: int):
            """Helper function for forward steps"""
            idxs = (layer_idx, microbatch_idx)
            outputs[idxs] = layer_forwards[idxs](inputs[idxs])
            if output_snapshots is not None:
                # Reused graph output buffers are only valid until their corresponding backward.
                output_snapshots[idxs] = outputs[idxs].detach().clone()

        def backward(layer_idx: int, microbatch_idx: int):
            """Helper function for backward steps"""
            outputs[layer_idx, microbatch_idx].backward(grad_outputs[layer_idx, microbatch_idx])

        # Forward and backward steps.
        forward(0, 0)
        forward(1, 0)
        forward(0, 1)
        forward(1, 1)
        backward(1, 0)
        backward(0, 0)
        forward(0, 2)
        forward(1, 2)
        backward(1, 1)
        backward(0, 1)
        backward(1, 2)
        backward(0, 2)

        # Optimizer step.
        optimizer.step()

    output_values = output_snapshots if output_snapshots is not None else outputs
    output_values = [y for _, y in sorted(output_values.items())]
    outputs = get_outputs(model, output_values)
    final_weights = [param.detach().clone() for param in model.parameters()]
    if with_graph:
        reset_graphs(layer_forwards)
    return outputs, final_weights


def test_make_graphed_callables_with_interleaved_pipeline_parallelism(
    *,
    model_config: str = "small",
    dtype: torch.dtype = torch.float16,
) -> None:
    """Test CUDA graphs with Megatron-LM interleaved pipeline parallelism."""
    model_config = model_configs[model_config]
    kwargs = dict(model_config=model_config, dtype=dtype)
    outputs, weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism(
        with_graph=False,
        **kwargs,
    )
    graph_outputs, graph_weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism(
        with_graph=True,
        **kwargs,
    )
    assert_all_equal(outputs, graph_outputs)
    assert_all_equal(weights, graph_weights)


def test_make_graphed_callables_with_interleaved_pipeline_parallelism_reused_buffers(
    *,
    model_config: str = "small",
    dtype: torch.dtype = torch.float16,
) -> None:
    """Test CUDA graphs with reused input/output buffers."""
    model_config = model_configs[model_config]
    kwargs = dict(model_config=model_config, dtype=dtype)
    outputs, weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism(
        with_graph=False,
        **kwargs,
    )
    graph_outputs, graph_weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism(
        with_graph=True,
        reuse_graph_input_output_buffers=True,
        **kwargs,
    )
    assert_all_equal(outputs, graph_outputs)
    assert_all_equal(weights, graph_weights)


def test_make_graphed_callables_with_interleaved_pipeline_parallelism_reused_buffers_no_param_grad_clone(
    *,
    model_config: str = "small",
    dtype: torch.dtype = torch.float16,
) -> None:
    """Test reused input/output buffers when returned parameter grad clones are disabled."""
    model_config = model_configs[model_config]
    kwargs = dict(model_config=model_config, dtype=dtype)
    outputs, weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism(
        with_graph=False,
        **kwargs,
    )
    graph_outputs, graph_weights = _test_cuda_graphs_with_interleaved_pipeline_parallelism(
        with_graph=True,
        reuse_graph_input_output_buffers=True,
        clone_param_grads_on_return=False,
        **kwargs,
    )
    assert_all_equal(outputs, graph_outputs)
    assert_all_equal(weights, graph_weights)
