Mutable Computational Graphs in Pure Julia: Performing "Hot Surgery" on a LLaMA Transformer

Specific Domains Machine Learning
1

Hi everyone, I am Abdallah Khemais. As an Instructor of Applied Mathematics and AI at ISITCOM Sousse Tunisia, with over two decades of experience, my work focuses on the intersection of formal mathematical theory and high-performance system design.

This drive to bridge abstract mathematics with low-level execution led me to create NeuroDSL, a reactive computational graph framework written in pure Julia. Today, I want to share a visual demonstration of how we solved a major bottleneck in dynamic network architectures.

In mechanistic interpretability and dynamic routing, modifying a network mid-training usually forces a complete recompilation of the computational graph. By treating the DAG as a persistent, mutable data structure, NeuroDSL bypasses this entirely.

In this demonstration, we visualize a “hot surgery” on the network—inserting a full residual block dynamically without halting the training loop or losing optimizer states.

Here is the architectural breakdown of what happens under the hood:

1. The Pre-Operative Architecture (Steps 0 to 150)

We initialize a compact, 3-layer LLaMA-style Transformer optimized for induction head tracking.

  • Dimensions: 64-dimensional embeddings, 128 hidden dimension, and 4 attention heads.

  • Initial Topology: The framework compiles an exact DAG consisting of 162 computational nodes.

  • Execution: The forward pass flows directly from the second layer (layer_2_out) into the consumers of the third layer (layer_3_norm1_out and layer_3_res1). The average loss stabilizes around 3.08.

2. The Topological Mutation (Step 150)

At step 150, we execute the insert_block! function. We do not rebuild a new graph; we mutate the existing one in-place.

  • Exact Identity Injection: A new LlamaBlock is injected directly after layer_2_out. The projections are initialized to the exact identity, ensuring the mathematical mapping is completely undisturbed at the moment of insertion.

  • Graph Expansion: The DAG instantly expands from 162 to 212 nodes, successfully absorbing 50 new computational nodes and 9 new learnable parameters.

  • State Preservation: This is the critical systems challenge. The AdamW optimizer perfectly preserves its momentum states (m_1 and m_2) for the original 162 nodes, while dynamically allocating fresh buffers exclusively for the 9 new parameters.

3. Lazy Invalidation & Historical Tracing

The interactive viewer tracks this mutation across time. It does not simply hide zeroed-out values; it physically alters the displayed topology based on the epoch cursor.

  • Before Step 150: The grafted sub-graph (prefixed surgery_layer_2_out_...) does not exist in the UI. If you inspect a future parameter, the tracer rigorously reports: “did not exist yet at this epoch” instead of falling back to a null value. The original graph edges remain strictly intact.

  • After Step 150: The new block materializes, the original edges are gracefully replaced by the graft, and the deeper network resumes training, pushing the loss down to 1.54.

By relying on strict topological ordering and local invalidation, we completely bypass the recompilation overhead.

(Note: As a new user on this forum, the anti-spam system prevents me from posting clickable links or embedding the video directly. You can find the resources by removing the spaces in the URLs below)

:laptop: Source Code: github . com / nevermind78 / NeuroDSL

:movie_camera: Video Demo: youtu . be / ZWD1dujm0Sk

:link: Interactive Graph Surgery: nevermind78 . github . io / NeuroDSL / graph_surgery_viz . html

Feel free to play with the interactive HTML viewer to inspect the tensor shapes, gradients, and edge mutations step-by-step. I would love to connect with anyone working on Enzyme integration, SciML, or differentiable programming!

4 Likes