From efdb905eac3407aeae15a712142bdc2a41e89397 Mon Sep 17 00:00:00 2001 From: Michael Lazos Date: Wed, 12 Aug 2026 17:44:53 -0700 Subject: [PATCH] Enable compiler-managed CUDA graphs for UMA inference Torch compile can capture UMA compute regions, but internal neighbor generation changes the exact edge shape from one frame to the next. Those shape changes prevent reliable CUDA graph reuse even when the atom count is fixed. Expose reduce-overhead as an inference compile mode and mark prediction boundaries for CUDA graph trees. For internal NVIDIA v3 graphs, generate the neighbor list before model execution and pad its edges to a configurable bucket. Padded edges are moved beyond the cutoff so they contribute zero energy, force, and stress while keeping the compiled input shape stable. Neighbor construction remains outside capture because the released NVIDIA implementation allocates temporary buffers. The internal path is limited to CUDA, graph version 3, and non-distributed inference. External graphs can use reduce-overhead without padding. Test Plan: ``` PYTHONPATH="$PWD/src" pytest -q -s tests/core/graph/test_padded_nvidia_graph.py tests/core/models/uma/test_padded_edges.py tests/core/units/mlip_unit/test_inference_settings.py PYTHONPATH="$PWD/src" pytest -q -s tests/core/units/mlip_unit/test_predict.py::test_reduce_overhead_internal_graph_predict pre-commit run --files docs/core/common_tasks/ase_calculator.md src/fairchem/core/graph/padded_nvidia_graph.py src/fairchem/core/models/uma/escn_md.py src/fairchem/core/units/mlip_unit/api/inference.py src/fairchem/core/units/mlip_unit/predict.py tests/core/graph/test_padded_nvidia_graph.py tests/core/models/uma/test_padded_edges.py tests/core/units/mlip_unit/test_inference_settings.py tests/core/units/mlip_unit/test_predict.py ``` --- docs/core/common_tasks/ase_calculator.md | 5 +- .../core/graph/padded_nvidia_graph.py | 54 +++++++++++++ src/fairchem/core/models/uma/escn_md.py | 28 +++++-- .../core/units/mlip_unit/api/inference.py | 23 ++++++ src/fairchem/core/units/mlip_unit/predict.py | 43 ++++++++++- tests/core/graph/test_padded_nvidia_graph.py | 76 +++++++++++++++++++ tests/core/models/uma/test_padded_edges.py | 49 ++++++++++++ .../mlip_unit/test_inference_settings.py | 40 ++++++++++ tests/core/units/mlip_unit/test_predict.py | 74 ++++++++++++++++++ 9 files changed, 383 insertions(+), 9 deletions(-) create mode 100644 src/fairchem/core/graph/padded_nvidia_graph.py create mode 100644 tests/core/graph/test_padded_nvidia_graph.py create mode 100644 tests/core/models/uma/test_padded_edges.py diff --git a/docs/core/common_tasks/ase_calculator.md b/docs/core/common_tasks/ase_calculator.md index 4a9611e06e..079c96821e 100644 --- a/docs/core/common_tasks/ase_calculator.md +++ b/docs/core/common_tasks/ase_calculator.md @@ -92,8 +92,10 @@ The advanced user might quickly see that **default**, **batch**, and **turbo** m | activation_checkpointing | this uses a custom chunked activation checkpointing algorithm and allows significant savings in memory for a small inference speed penalty. If you are predicting on systems >1000 atoms, we recommend keeping this on. However, if you want the absolute fastest inference possible for small systems, you can turn this off | | merge_mole | This is useful in long rollout applications where the system composition stays constant. By pre-merge the MoLE weights, we can save both memory and compute. | | compile | This uses torch.compile to significantly speed up computation. Due to the way pytorch traces the internal graph, it requires a long compile time during the first iteration and can even recompile anytime it detected a significant change in input dimensions. It is not recommended if you are computing frequently on very different atomic systems. | +| compile_mode | Set to `"reduce-overhead"` to use compiler-managed CUDA graphs. With internal graph generation this requires version 3, a CUDA device, and no graph parallelism. It is intended for repeated evaluations of similarly sized systems. | | external_graph_gen | Only use this if you want to use an external graph generator. This should be rarely used except for development | | internal_graph_gen_version | currently we support v2[default], an internal implementation that is better suited for parallelism and v3 the neighborlist from Nvidia Alchemi library which is faster for single gpu operations. | +| internal_graph_edge_bucket_size | Number of edges per padded bucket when `compile_mode="reduce-overhead"` uses internal graph version 3. Larger buckets reduce shape changes but perform more work on masked edges. Default 1024. | | edge_chunk_size | Experimental. Used for padding edge sizes. This helps reduce re-compilations from torch compile, default to None | | use_quaternion_wigner | enable quaternion-based Wigner D matrix computation. If false we fall back to euler-angle based rotations. default True. | | base_precision_dtype | governs the main precision type of the computation, default to FP32, FP64 is also supported | @@ -109,8 +111,9 @@ settings = InferenceSettings( activation_checkpointing=False, merge_mole=True, compile=True, + compile_mode="reduce-overhead", external_graph_gen=False, - internal_graph_gen_version=2, + internal_graph_gen_version=3, ) predictor = pretrained_mlip.get_predict_unit( diff --git a/src/fairchem/core/graph/padded_nvidia_graph.py b/src/fairchem/core/graph/padded_nvidia_graph.py new file mode 100644 index 0000000000..bf2d6619f2 --- /dev/null +++ b/src/fairchem/core/graph/padded_nvidia_graph.py @@ -0,0 +1,54 @@ +"""Bucketed NVIDIA neighbor graphs for compiled inference.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from fairchem.core.common import gp_utils +from fairchem.core.graph.radius_graph_pbc_nvidia import radius_graph_pbc_nvidia + + +class PaddedNvidiaGraphGenerator: + """Generate v3 neighbor graphs and replace data's edges with a padded bucket.""" + + def __init__(self, settings, backbone): + if gp_utils.initialized(): + raise ValueError( + "padded NVIDIA graph generation does not support graph parallelism" + ) + if not getattr(backbone, "supports_padded_edges", False): + raise ValueError( + f"{type(backbone).__name__} does not support padded neighbor edges" + ) + if torch.device(next(backbone.parameters()).device).type != "cuda": + raise ValueError("padded NVIDIA graph generation requires CUDA") + self.cutoff = float(backbone.cutoff) + self.max_neighbors = int(backbone.max_neighbors) + self.enforce_max_neighbors_strictly = bool( + backbone.enforce_max_neighbors_strictly + ) + self.edge_bucket_size = settings.internal_graph_edge_bucket_size + + def generate(self, data): + edge_index, cell_offsets, neighbors = radius_graph_pbc_nvidia( + data, + self.cutoff, + self.max_neighbors, + self.enforce_max_neighbors_strictly, + pbc=data.pbc, + ) + num_edges = edge_index.shape[1] + edge_capacity = ( + (num_edges + self.edge_bucket_size - 1) // self.edge_bucket_size + ) * self.edge_bucket_size + pad = edge_capacity - num_edges + + data.edge_index = F.pad(edge_index, (0, pad)) + data.cell_offsets = F.pad(cell_offsets.to(data.pos.dtype), (0, 0, 0, pad)) + data.nedges = neighbors.clone() + data.nedges[-1] += pad + data.edge_valid_mask = ( + torch.arange(edge_capacity, device=edge_index.device) < num_edges + ) + return data diff --git a/src/fairchem/core/models/uma/escn_md.py b/src/fairchem/core/models/uma/escn_md.py index b70d0c4a37..0d0c175df3 100644 --- a/src/fairchem/core/models/uma/escn_md.py +++ b/src/fairchem/core/models/uma/escn_md.py @@ -278,6 +278,8 @@ def resolve_dataset_mapping( @registry.register_model("escnmd_backbone") class eSCNMDBackbone(nn.Module, MOLEInterface): + supports_padded_edges = True + def __init__( self, max_num_elements: int = 100, @@ -663,10 +665,12 @@ def _generate_graph(self, data_dict): # Compute shifts from cell offsets if len(data_dict["natoms"]) == 1: - # Single system: use matmul (compile-friendly, no data-dependent ops) - shifts = data_dict["cell_offsets"].to( - data_dict["cell"].dtype - ) @ data_dict["cell"].squeeze(0) + offsets = data_dict["cell_offsets"].to(data_dict["cell"].dtype) + if "edge_valid_mask" in data_dict: + cell_per_edge = data_dict["cell"].expand(offsets.shape[0], -1, -1) + shifts = torch.bmm(offsets.unsqueeze(1), cell_per_edge).squeeze(1) + else: + shifts = offsets @ data_dict["cell"].squeeze(0) else: # Batched: need repeat_interleave for variable edges per system cell_per_edge = data_dict["cell"].repeat_interleave( @@ -684,6 +688,13 @@ def _generate_graph(self, data_dict): - data_dict["pos"][data_dict["edge_index"][1]] + shifts ) # [n_edges, 3] + edge_valid_mask = data_dict.get("edge_valid_mask", None) + if edge_valid_mask is not None: + padding_vec = edge_distance_vec.new_zeros(3) + padding_vec[0].fill_(self.cutoff + 1.0) + edge_distance_vec = torch.where( + edge_valid_mask.unsqueeze(1), edge_distance_vec, padding_vec + ) # pylint: disable=E1102 edge_distance = torch.linalg.norm( edge_distance_vec, dim=-1, keepdim=False @@ -952,7 +963,14 @@ def build_inference_settings(cls, settings: InferenceSettings) -> dict: if settings.edge_chunk_size is not None: overrides["edge_chunk_size"] = settings.edge_chunk_size if settings.external_graph_gen is not None: - overrides["otf_graph"] = not settings.external_graph_gen + padded_internal_graph = ( + settings.compile_mode == "reduce-overhead" + and not settings.external_graph_gen + and settings.internal_graph_gen_version == 3 + ) + overrides["otf_graph"] = not ( + settings.external_graph_gen or padded_internal_graph + ) if settings.internal_graph_gen_version is not None: overrides["radius_pbc_version"] = settings.internal_graph_gen_version if settings.use_quaternion_wigner is not None: diff --git a/src/fairchem/core/units/mlip_unit/api/inference.py b/src/fairchem/core/units/mlip_unit/api/inference.py index eadfe02f0c..9271819bd5 100644 --- a/src/fairchem/core/units/mlip_unit/api/inference.py +++ b/src/fairchem/core/units/mlip_unit/api/inference.py @@ -201,7 +201,30 @@ class InferenceSettings: # VRAM but allow bigger systems; reduce if you run into OOM errors. max_atoms: int | None = None + # Optional torch.compile mode. "reduce-overhead" enables Inductor CUDA graphs. + compile_mode: str | None = None + + # Edge bucket used by padded NVIDIA internal graph generation. + internal_graph_edge_bucket_size: int = 1024 + def __post_init__(self): + if self.compile_mode not in (None, "reduce-overhead"): + raise ValueError("compile_mode must be None or 'reduce-overhead'") + if self.compile_mode is not None and not self.compile: + raise ValueError("compile_mode requires compile=True") + if ( + type(self.internal_graph_edge_bucket_size) is not int + or self.internal_graph_edge_bucket_size < 1 + ): + raise ValueError("internal_graph_edge_bucket_size must be positive") + if ( + self.compile_mode == "reduce-overhead" + and not self.external_graph_gen + and self.internal_graph_gen_version != 3 + ): + raise ValueError( + "internal reduce-overhead requires internal_graph_gen_version=3" + ) if isinstance(self.base_precision_dtype, str): self.base_precision_dtype = getattr(torch, self.base_precision_dtype) assert ( diff --git a/src/fairchem/core/units/mlip_unit/predict.py b/src/fairchem/core/units/mlip_unit/predict.py index d9a4dbbad1..7a031997d3 100644 --- a/src/fairchem/core/units/mlip_unit/predict.py +++ b/src/fairchem/core/units/mlip_unit/predict.py @@ -38,6 +38,7 @@ ) from fairchem.core.components.batch_server import get_app_handle_with_retry from fairchem.core.datasets.atomic_data import AtomicData, warn_if_upcasting +from fairchem.core.graph.padded_nvidia_graph import PaddedNvidiaGraphGenerator from fairchem.core.models.uma.nn.execution_backends import ( ExecutionMode, maybe_update_settings_backend, @@ -150,6 +151,7 @@ def __init__( self.lazy_model_intialized = False self.assert_on_nans = assert_on_nans self._warned_upcast = False + self._padded_graph_generator = None if self.model.module.backbone.regress_config.direct_forces: logging.warning( @@ -448,6 +450,8 @@ def _fall_back_from_fast_path( # 1. change flags self.inference_settings.merge_mole = False self.inference_settings.compile = False + self.inference_settings.compile_mode = None + self._padded_graph_generator = None if self.inference_settings.execution_mode == ExecutionMode.UMAS_FAST_GPU: self.inference_settings.execution_mode = ExecutionMode.GENERAL self.lazy_model_intialized = False @@ -493,9 +497,17 @@ def predict( if single_atom_result is not None: return single_atom_result + if self.inference_settings.compile_mode == "reduce-overhead": + torch.compiler.cudagraph_mark_step_begin() + + use_padded_internal_graph = ( + self.inference_settings.compile_mode == "reduce-overhead" + and not self.inference_settings.external_graph_gen + ) + # Regular model prediction path # this needs to be .clone() to avoid issues with graph parallel modifying this data with MOLE - data_device = data.to(self.device).clone() + data_device = data.clone().to(self.device) dtype = self.inference_settings.base_precision_dtype if not self._warned_upcast: @@ -505,6 +517,12 @@ def predict( data_device[key] = val.to(dtype) backbone = self.model.module.backbone + if use_padded_internal_graph: + if self._padded_graph_generator is None: + self._padded_graph_generator = PaddedNvidiaGraphGenerator( + self.inference_settings, backbone + ) + data_device = self._padded_graph_generator.generate(data_device) _prepare_inference_gradients(backbone, data_device) # Model handles any per-prediction checks (e.g., MOLE consistency) @@ -514,13 +532,28 @@ def predict( self._fall_back_from_fast_path(error) self._lazy_init(data) self.model.module.on_predict_check(data_device) - return self._run_inference(data_device, undo_element_references) def _lazy_init(self, data: AtomicData) -> None: """ Lazy initialization on first predict call. """ + if ( + self.inference_settings.compile_mode == "reduce-overhead" + and not self.inference_settings.external_graph_gen + ): + if gp_utils.initialized(): + raise ValueError( + "internal reduce-overhead does not support graph parallelism" + ) + if not getattr(self.model.module.backbone, "supports_padded_edges", False): + raise ValueError( + "internal reduce-overhead is not supported by " + f"{type(self.model.module.backbone).__name__}" + ) + if torch.device(self._requested_device).type != "cuda": + raise ValueError("internal reduce-overhead requires CUDA") + # Model handles its own preparation (MOLE merge, eval mode, etc.) self.model.module.prepare_for_inference(data, self.inference_settings) # Inference differentiates outputs with respect to inputs, not weights. @@ -539,7 +572,11 @@ def _lazy_init(self, data: AtomicData) -> None: # The model's scalars are fixed at inference, so this skips dynamo's # TensorifyScalarRestartAnalysis retrace during compile. torch._dynamo.config.specialize_float = True - self.model = torch.compile(self.model, dynamic=True) + self.model = torch.compile( + self.model, + dynamic=True, + mode=self.inference_settings.compile_mode, + ) self.lazy_model_intialized = True diff --git a/tests/core/graph/test_padded_nvidia_graph.py b/tests/core/graph/test_padded_nvidia_graph.py new file mode 100644 index 0000000000..264e55f8c2 --- /dev/null +++ b/tests/core/graph/test_padded_nvidia_graph.py @@ -0,0 +1,76 @@ +"""Tests for bucketed NVIDIA neighbor graphs.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +import fairchem.core.graph.padded_nvidia_graph as padded_graph +from fairchem.core.graph.padded_nvidia_graph import PaddedNvidiaGraphGenerator + + +class _GraphData(SimpleNamespace): + def clone(self): + return _GraphData( + **{ + key: value.clone() if torch.is_tensor(value) else value + for key, value in vars(self).items() + } + ) + + +def _generator(bucket_size=4): + generator = object.__new__(PaddedNvidiaGraphGenerator) + generator.cutoff = 5.0 + generator.max_neighbors = 32 + generator.enforce_max_neighbors_strictly = False + generator.edge_bucket_size = bucket_size + return generator + + +def test_generate_pads_edges_to_bucket(monkeypatch): + edge_index = torch.tensor([[1, 0, 2], [0, 1, 1]]) + cell_offsets = torch.arange(9, dtype=torch.float64).reshape(3, 3) + neighbors = torch.tensor([2, 1]) + + def graph(*args, **kwargs): + return edge_index, cell_offsets, neighbors + + monkeypatch.setattr(padded_graph, "radius_graph_pbc_nvidia", graph) + data = _GraphData( + pos=torch.empty(3, 3, dtype=torch.float32), + pbc=torch.ones(2, 3, dtype=torch.bool), + ) + + result = _generator().generate(data) + + assert result is data + torch.testing.assert_close(result.edge_index[:, :3], edge_index) + torch.testing.assert_close( + result.edge_index[:, 3], torch.zeros(2, dtype=torch.long) + ) + torch.testing.assert_close(result.cell_offsets[:3], cell_offsets.float()) + torch.testing.assert_close(result.cell_offsets[3], torch.zeros(3)) + torch.testing.assert_close(result.nedges, torch.tensor([2, 2])) + torch.testing.assert_close( + result.edge_valid_mask, torch.tensor([True, True, True, False]) + ) + + +def test_generate_keeps_exact_bucket_shape(monkeypatch): + edge_index = torch.tensor([[1, 0, 2, 1], [0, 1, 1, 2]]) + cell_offsets = torch.zeros(4, 3) + neighbors = torch.tensor([4]) + monkeypatch.setattr( + padded_graph, + "radius_graph_pbc_nvidia", + lambda *args, **kwargs: (edge_index, cell_offsets, neighbors), + ) + data = _GraphData(pos=torch.empty(3, 3), pbc=torch.ones(1, 3, dtype=torch.bool)) + + result = _generator().generate(data) + + torch.testing.assert_close(result.edge_index, edge_index) + torch.testing.assert_close(result.nedges, neighbors) + assert result.edge_valid_mask.all() diff --git a/tests/core/models/uma/test_padded_edges.py b/tests/core/models/uma/test_padded_edges.py new file mode 100644 index 0000000000..b0296b2f53 --- /dev/null +++ b/tests/core/models/uma/test_padded_edges.py @@ -0,0 +1,49 @@ +"""Tests for padded neighbor edges in UMA.""" + +from __future__ import annotations + +import torch + +from fairchem.core.models.uma.escn_md import eSCNMDBackbone + + +def test_precomputed_padding_is_moved_beyond_cutoff(): + backbone = object.__new__(eSCNMDBackbone) + torch.nn.Module.__init__(backbone) + backbone.cutoff = 5.0 + backbone.otf_graph = False + data = { + "pos": torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + "cell": torch.eye(3).unsqueeze(0), + "natoms": torch.tensor([2]), + "edge_index": torch.tensor([[0, 0], [1, 0]]), + "cell_offsets": torch.zeros(2, 3), + "edge_valid_mask": torch.tensor([True, False]), + } + + graph = backbone._generate_graph(data) + + torch.testing.assert_close(graph["edge_distance"][0], torch.tensor(1.0)) + torch.testing.assert_close(graph["edge_distance"][1], torch.tensor(6.0)) + torch.testing.assert_close(data["scatter_target"], torch.tensor([1, 0])) + + +def test_precomputed_periodic_shifts_match_internal_graph_math(): + backbone = object.__new__(eSCNMDBackbone) + torch.nn.Module.__init__(backbone) + backbone.cutoff = 5.0 + backbone.otf_graph = False + data = { + "pos": torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + "cell": torch.diag(torch.tensor([3.0, 4.0, 5.0])).unsqueeze(0), + "natoms": torch.tensor([2]), + "edge_index": torch.tensor([[0], [1]]), + "cell_offsets": torch.tensor([[1.0, 0.0, 0.0]]), + } + + graph = backbone._generate_graph(data) + + torch.testing.assert_close( + graph["edge_distance_vec"], torch.tensor([[2.0, 0.0, 0.0]]) + ) + torch.testing.assert_close(graph["edge_distance"], torch.tensor([2.0])) diff --git a/tests/core/units/mlip_unit/test_inference_settings.py b/tests/core/units/mlip_unit/test_inference_settings.py index 19ebde2ed8..4d5ec9d73e 100644 --- a/tests/core/units/mlip_unit/test_inference_settings.py +++ b/tests/core/units/mlip_unit/test_inference_settings.py @@ -71,6 +71,46 @@ def test_invalid_string_raises(): InferenceSettings(base_precision_dtype="int8") +def test_compile_mode_validation(): + with pytest.raises(ValueError, match="compile_mode must be"): + InferenceSettings(compile=True, compile_mode="max-autotune") + with pytest.raises(ValueError, match="requires compile=True"): + InferenceSettings(compile_mode="reduce-overhead") + + settings = InferenceSettings( + compile=True, compile_mode="reduce-overhead", external_graph_gen=True + ) + assert settings.compile_mode == "reduce-overhead" + + +def test_internal_reduce_overhead_requires_graph_v3(): + with pytest.raises(ValueError, match="internal_graph_gen_version=3"): + InferenceSettings(compile=True, compile_mode="reduce-overhead") + + +def test_reduce_overhead_settings_accept_valid_internal_graph_configuration(): + settings = InferenceSettings( + compile=True, + compile_mode="reduce-overhead", + external_graph_gen=False, + internal_graph_gen_version=3, + ) + assert settings.internal_graph_edge_bucket_size == 1024 + + +def test_reduce_overhead_fields_preserve_existing_positional_arguments(): + settings = InferenceSettings(False, True, False, False, True) + + assert settings.wigner_cuda is True + assert settings.compile_mode is None + + +@pytest.mark.parametrize("bucket_size", [0, -1, 1.5, True]) +def test_internal_graph_edge_bucket_size_must_be_positive_integer(bucket_size): + with pytest.raises(ValueError, match="must be positive"): + InferenceSettings(internal_graph_edge_bucket_size=bucket_size) + + # --- to_omegaconf --- diff --git a/tests/core/units/mlip_unit/test_predict.py b/tests/core/units/mlip_unit/test_predict.py index be5d70ce7f..7471372c62 100644 --- a/tests/core/units/mlip_unit/test_predict.py +++ b/tests/core/units/mlip_unit/test_predict.py @@ -39,6 +39,7 @@ from fairchem.core.common.gp_utils import GraphParallelConfig from fairchem.core.datasets.atomic_data import AtomicData, atomicdata_list_to_batch from fairchem.core.datasets.common_structures import get_fcc_crystal_by_num_atoms +from fairchem.core.graph.padded_nvidia_graph import PaddedNvidiaGraphGenerator from fairchem.core.models.uma.compat import UMA_1P1_MODEL_ID from fairchem.core.models.uma.nn.execution_backends import UMASFastGPUBackend from fairchem.core.units.mlip_unit import InferenceSettings, MLIPPredictUnit @@ -215,6 +216,79 @@ def test_single_dataset_predict(internal_graph_gen_version, pretrained_checkpoin ) +@pytest.mark.compile_gpu() +@pytest.mark.pretrained("uma-s-1p2") +def test_reduce_overhead_internal_graph_predict(pretrained_checkpoint, monkeypatch): + base_atoms = bulk("Cu", "fcc", a=3.6).repeat((2, 2, 2)) + expanded_atoms = base_atoms.copy() + expanded_atoms.set_cell(expanded_atoms.cell * 1.8, scale_atoms=True) + frames = [base_atoms, expanded_atoms, base_atoms] + batches = [ + atomicdata_list_to_batch([AtomicData.from_ase(atoms, task_name="omat")]) + for atoms in frames + ] + + common_settings = dict( + tf32=True, + activation_checkpointing=False, + merge_mole=True, + external_graph_gen=False, + internal_graph_gen_version=3, + execution_mode="umas_fast_gpu", + ) + eager_predictor = get_predict_unit_for_test( + pretrained_checkpoint, + device="cuda", + inference_settings=InferenceSettings(**common_settings), + ) + references = [ + { + name: value.detach().cpu() + for name, value in eager_predictor.predict(batch).items() + } + for batch in batches + ] + + edge_capacities = [] + generate = PaddedNvidiaGraphGenerator.generate + + def record_capacity(generator, data): + result = generate(generator, data) + edge_capacities.append(result.edge_index.shape[1]) + return result + + monkeypatch.setattr(PaddedNvidiaGraphGenerator, "generate", record_capacity) + compiled_predictor = get_predict_unit_for_test( + pretrained_checkpoint, + device="cuda", + inference_settings=InferenceSettings( + **common_settings, + compile=True, + compile_mode="reduce-overhead", + internal_graph_edge_bucket_size=64, + ), + ) + + predictions = [] + first_prediction = None + for batch, reference in zip(batches, references, strict=True): + prediction = compiled_predictor.predict(batch) + predictions.append(prediction) + if first_prediction is None: + first_prediction = { + name: value.detach().clone() for name, value in prediction.items() + } + for name, expected in reference.items(): + torch.testing.assert_close( + prediction[name].detach().cpu(), expected, rtol=5e-3, atol=ATOL + ) + + assert edge_capacities[0] != edge_capacities[1] + assert edge_capacities[0] == edge_capacities[2] + for name, expected in first_prediction.items(): + torch.testing.assert_close(predictions[0][name], expected, rtol=0, atol=0) + + @pytest.mark.gpu() @pytest.mark.parametrize("internal_graph_gen_version", [2, 3]) @pytest.mark.pretrained("uma-s-1p1", "uma-s-1p2")