[Paper] Morphling: Fast, Fused, and Flexible GNN Training at Scale

Source: arXiv - 2512.01678v1

Overview

Morphling is a domain‑specific code generator that turns high‑level GNN models into highly tuned implementations for CPUs, GPUs, and distributed clusters. By fusing irregular graph traversals with dense matrix math and adapting to the sparsity of the data at runtime, it delivers order‑of‑magnitude speedups over popular libraries such as PyTorch Geometric and DGL.

Key Contributions

• Architecture‑aware code synthesis – Generates separate OpenMP, CUDA, and MPI kernels from a single GNN description, exploiting the strengths of each hardware platform.
• Fused graph‑matrix pipelines – Eliminates the costly intermediate buffers that plague existing frameworks, improving cache locality and reducing memory traffic.
• Sparsity‑aware runtime – Dynamically chooses dense or sparse execution paths based on feature‑level statistics, skipping work on zero‑valued entries.
• Portable primitive library – A curated set of low‑level, architecture‑specific building blocks (e.g., vectorized scatter‑add, warp‑level reductions) that can be reused across models.
• Comprehensive evaluation – Benchmarks on 11 real‑world graphs (varying size, density, and feature dimensions) showing up to 66× speedup and 15× lower peak memory usage.

Methodology

1. High‑level specification – Users write their GNN model in a PyTorch‑like DSL (e.g., message‑passing layers, aggregation functions).
2. Intermediate representation (IR) – Morphling parses the DSL into an IR that separates graph‑centric operations (edge‑wise message passing) from dense linear algebra (feature transforms).
3. Backend specialization – The IR is fed to a code‑generation engine that selects the appropriate primitive implementations for the target platform:
   • CPU (OpenMP) – Vectorized loops with cache‑blocking and NUMA‑aware thread placement.
   • GPU (CUDA) – Warp‑level cooperative kernels, shared‑memory tiling, and fused kernels that combine edge scatter/gather with matrix multiplication.
   • Distributed (MPI) – Partition‑aware data layouts and halo‑exchange routines that keep communication overhead low.
4. Sparsity profiling – At the start of each epoch, Morphling samples feature tensors to estimate the proportion of zeros. If sparsity crosses a configurable threshold, it swaps in sparse kernels (CSR/CSC formats) for the dense path; otherwise it stays dense.
5. Compilation & execution – The generated C++/CUDA code is compiled just‑in‑time (JIT) and linked back to the Python front‑end, allowing seamless integration with existing training pipelines.

Results & Findings

• Throughput: Training epochs finish in seconds rather than minutes for most benchmark datasets.
• Memory: By fusing kernels and using compact layouts (e.g., packed edge lists + column‑major feature matrices), the peak resident set size drops dramatically, enabling graphs with > 100 M edges on a single 32 GB GPU.
• Scalability: The MPI backend scales near‑linearly up to 8 nodes for the largest datasets, confirming that the code generator respects data locality and minimizes inter‑node traffic.

Practical Implications

• Faster prototyping – Data scientists can iterate on GNN architectures without waiting hours for each training run, accelerating research and product development cycles.
• Cost savings – The 10‑20× speedup translates directly into lower cloud compute bills; the reduced memory footprint allows larger models to run on commodity hardware.
• Edge & production deployment – The ability to generate CPU‑only kernels means GNN inference can be embedded in services where GPUs are unavailable (e.g., recommendation engines, fraud detection).
• Framework‑agnostic integration – Since Morphling outputs standard C++/CUDA libraries, existing PyTorch or TensorFlow pipelines can link against them with minimal code changes.
• Future‑proofing – The modular primitive library can be extended to emerging accelerators (e.g., TPUs, Habana) by adding new backend implementations without rewriting the high‑level model.

Limitations & Future Work

• Model coverage – Current support focuses on common message‑passing GNNs (GCN, GraphSAGE, GAT). More exotic operators (e.g., subgraph pooling, attention over edges) require additional primitives.
• Static sparsity thresholds – The runtime sparsity heuristic is simple; adaptive learning‑based policies could further improve the dense/sparse decision.
• Compilation overhead – JIT compilation adds a one‑time cost of several seconds, which is negligible for long training runs but noticeable for quick experiments.
• Distributed fault tolerance – The MPI backend assumes a stable cluster; integrating checkpoint‑restart mechanisms would make it more robust for production workloads.

Morphling demonstrates that a carefully engineered, architecture‑aware code synthesis pipeline can turn the “slow and memory‑hungry” reputation of GNN training on its head, opening the door for large‑scale graph AI on everyday hardware.

Authors

• Anubhab
• Rupesh Nasre

Paper Information

• arXiv ID: 2512.01678v1
• Categories: cs.LG, cs.DC, cs.PL
• Published: December 1, 2025
• PDF: Download PDF