[Feat] Integrate Megatron-LM core_v0.17.1

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "third_party/vllm-integration"]
 	path = third_party/vllm-integration
 	url = https://github.com/ProjectDMX/DMI-vLLM-Integration.git
+[submodule "third_party/DMI-Megatron-Integration"]
+	path = third_party/DMI-Megatron-Integration
+	url = https://github.com/ProjectDMX/DMI-Megatron-Integration.git

README.md

@@ -13,7 +13,7 @@
 > collaborators to explore downstream applications built on DMI such as **interpretability**, **speculative decoding**,
 > **hallucination analysis**, **distillation**, **activation steering**, and beyond. If you're interested, please [contact us](mailto:ynn1999@umd.edu,sixianx@umd.edu,zaoxing@umd.edu).
 
-> **Project Status — research preview.** DMI supports HuggingFace and official vLLM integrations across 20+ dense and MoE model families, including Qwen3, Llama, Gemma, Mistral, GPT-OSS, Phi, and Granite, with early support for emerging architectures such as Qwen3.6, Llama 4, DeepSeek V4 Flash, GLM-5.2, Kimi K3, and MiniMax-M2.7. SGLang support is on the way. APIs may change. Contributions, bug reports, and feature requests are welcome.
+> **Project Status — research preview.** DMI supports HuggingFace and official vLLM integrations across 20+ dense and MoE model families, including Qwen3, Llama, Gemma, Mistral, GPT-OSS, Phi, and Granite, with early support for emerging architectures such as Qwen3.6, Llama 4, DeepSeek V4 Flash, GLM-5.2, Kimi K3, and MiniMax-M2.7, plus Megatron-LM training. SGLang support is on the way. APIs may change. Contributions, bug reports, and feature requests are welcome.
 
 > **👀Technical Report Available:** https://arxiv.org/abs/2605.11093
 
@@ -23,8 +23,8 @@
 We are working to make DMI useful across more backends, more models, and more
 stages of the model lifecycle.
 
-- **What's next** — **SGLang**, multimodal models, a pluggable observability storage
-  stack, and end-to-end observability from inference to large-scale training with **Megatron**.
+- **What's next** — **SGLang**, multimodal models, and a pluggable observability
+  storage stack.
 - **From observation to action** — Low-latency streaming/pluggable APIs enables more downstream applications like online monitoring, activation steering,
     distillation, and speculative decoding.
 - **Broader PCIe-aware scheduling** — Extend DMI's serving-first drain governor
@@ -33,13 +33,13 @@ stages of the model lifecycle.
 
 ## About
 
-**DMI is a full-feature observability layer for LLM inference.** It gives real-time access to
+**DMI is a full-feature observability layer for LLM inference and training.** It gives real-time access to
 *any* internal model state — residual streams, attention patterns, MLP outputs,
-KV-cache slices, logits — during real serving, with minimal overhead and without
-forking the inference engine.
+KV-cache slices, logits — during inference or training, with minimal overhead.
 
-Right now, DMI works in **HuggingFace Transformers** and **vLLM** out of the box, captures
-internal tensors through CUDA-Graph–compatible hooks, and streams them off the
+Right now, DMI supports inference through **HuggingFace Transformers** and
+**vLLM**, and training through **Megatron-LM**. It captures
+internal tensors through CUDA-Graph–compatible hooks and streams them off the
 GPU via a dedicated ring buffer to a host-side drain that pushes into a
 queryable store (or drops them, for transport-only profiling).
 
@@ -65,8 +65,9 @@ That's the gap DMI fills.
 - **`Ring²`** — GPU↔CPU co-designed staging. A dedicated GPU-side payload ring
   isolates captured tensors from the KV-cache memory pool; an on-host meta ring
   is drained asynchronously.
-- **HF + vLLM integration** — no engine fork required by the user. Plug in
-  through a worker class (vLLM) or a thin generation wrapper (HF).
+- **HF, vLLM, and Megatron-LM integrations** — use a thin generation wrapper
+  for HF, a worker integration for an unmodified official vLLM installation,
+  or the Megatron-LM training integration.
 - **Configurable offloading** — capture your hidden states on GPU, stage on host,
   and stream into a queryable store; visualize from notebooks (check out the [Demo](#demo) below).
 - **Quantified overhead** — measured against vanilla HF, HF's `output_hidden_states`,
@@ -116,11 +117,12 @@ Full setup, additional results, and how to reproduce:
 ## Get started
 
 Start with the [core installation guide](docs/install.md), then choose the
-HuggingFace or vLLM path depending on the runtime you want to inspect. The
-project currently supports installation from source. Use a separate environment
-and checkout for each backend, and install only the integration you need. The
-snippet below shows the minimal vLLM entry point. The version-matched integration
-checkout connects DMI to an unmodified official vLLM installation.
+HuggingFace, vLLM, or Megatron-LM path depending on the runtime you want to
+inspect. The project currently supports installation from source. Use a
+separate environment and checkout for each backend, and install only the
+integration you need. The snippet below shows the minimal vLLM entry point. The
+integration checkout connects DMI to an unmodified official
+vLLM installation.
 
 ```python
 import os
@@ -151,14 +153,15 @@ for o in llm.generate(["The answer is"], SamplingParams(max_tokens=16)):
 | **[Core installation](docs/install.md)** | Install DMI from source and build the native backend |
 | **[HuggingFace](docs/huggingface.md)** | Run HF generation, monitored generation, and offline benchmark scripts |
 | **[vLLM](docs/vllm.md)** | Run DMI through the vLLM offline API or `vllm serve` |
+| **[Megatron-LM](docs/megatron.md)** | Run DMI during Megatron-LM training |
 
 ## Contribute
 
 DMI is an early research system from FrootLab at the University of Maryland, and
 we welcome contributions from users, researchers, and systems builders. Useful
 contributions include bug reports, documentation fixes, benchmark reproduction
 notes, new model integrations, and backend-specific improvements for
-HuggingFace or vLLM.
+HuggingFace, vLLM, or Megatron-LM.
 
 - **Questions, bugs, and feature requests.** Please open a GitHub issue with the
   model, backend, hardware, and reproduction steps when applicable.

docs/install.md

@@ -31,14 +31,15 @@ nvidia-smi
 
 ## 1. Clone the repository
 
-The repo uses three git submodules: the DMI HuggingFace integration, the
-version-matched DMI-vLLM integration, and the `clickhouse-cpp` C++ client.
-The commands below fetch all three repositories; they do not install either
-Python integration.
+The repo uses four git submodules: the DMI HuggingFace integration, the
+version-matched DMI-vLLM integration, the version-matched DMI-Megatron
+integration, and the `clickhouse-cpp` C++ client. The commands below fetch all
+four repositories; they do not install any Python integration.
 
-The command below creates one backend checkout. If you plan to use both
-backends, repeat it with distinct target directories such as `DMI-hf` and
-`DMI-vllm`; do not share one checkout between their environments.
+The command below creates one backend checkout. If you plan to use multiple
+backends, repeat it with distinct target directories such as `DMI-hf`,
+`DMI-vllm`, and `DMI-megatron`; do not share one checkout between their
+environments.
 
 ```bash
 git clone --recursive https://github.com/ProjectDMX/DMI.git
@@ -52,6 +53,7 @@ Expected submodule paths:
 
 - `third_party/transformers/` — modified HF Transformers (`gpt2_p`, `qwen3_p`, `llama_p`)
 - `third_party/vllm-integration/` — DMI integration for an unmodified official vLLM installation
+- `third_party/DMI-Megatron-Integration/` — DMI integration with its pinned Megatron-LM fork at `third_party/DMI-Megatron-Integration/third_party/megatron-lm/`
 - `third_party/clickhouse-cpp/` — ClickHouse C++ client linked into the native backend
 
 ## 2. Install ClickHouse server
@@ -94,7 +96,7 @@ If conda is not already installed, follow the
 first. Then:
 
 ```bash
-DMI_BACKEND_ENV=dmi-hf  # Example; use dmi-vllm in the vLLM checkout.
+DMI_BACKEND_ENV=dmi-hf  # Example; use dmi-vllm or dmi-megatron for those checkouts.
 conda env create -f environment.yml --name "$DMI_BACKEND_ENV"
 conda activate "$DMI_BACKEND_ENV"
 ```
@@ -212,11 +214,12 @@ ClickHouse; use the host benchmark separately with a running server.
 
 ## 6. Choose one backend
 
-Continue with either the [HuggingFace guide](huggingface.md) or the
-[vLLM guide](vllm.md). Use a separate environment and checkout for each
-backend. The HuggingFace path installs a modified Transformers checkout,
-whereas the vLLM path installs its own official dependency set; do not install
-the HuggingFace integration in the vLLM environment.
+Continue with the [HuggingFace guide](huggingface.md), [vLLM guide](vllm.md),
+or [Megatron-LM guide](megatron.md). Use a separate environment and checkout
+for each backend. The HuggingFace path installs a modified Transformers
+checkout, the vLLM path installs its own official dependency set, and the
+Megatron-LM path installs its version-matched integration and pinned fork. Do
+not mix their framework dependencies in one environment.
 
 The native extension is also environment-specific: its Python suffix, Torch
 ABI, CUDA selection, and runtime paths come from the active environment. Each

docs/integration-api-v1.md

@@ -319,13 +319,18 @@ The public operations are `create_record_runtime()`, `bind_hook()`,
 
 `RecordRuntime.bind_hook()` assigns stable output IDs and binds the hook to the
 record ring. `RecordRuntime.emit_output()` reserves and publishes one eager
-descriptor before its producer. `RecordRuntime.prepare_replay()` publishes
-fresh descriptors for an existing physical plan. Every producer occurrence
-has exactly one descriptor; a gated-off occurrence uses an empty descriptor
-and a zero-byte task.
-
-A nonempty producer descriptor must contain a `PayloadSlice`; an empty
-producer descriptor is permitted.
+descriptor before its host-selected producer occurrence.
+`RecordRuntime.prepare_replay()` publishes fresh descriptors for the
+host-selected entries in an existing physical plan. Every accepted producer
+occurrence has exactly one descriptor and one task.
+
+When a hook uses a device gate, the integration must apply the same selection
+before host reservation and descriptor publication. A gated-off occurrence is
+absent from the record stream: it has no reservation, descriptor, task,
+payload write, or ring-counter advancement. This is distinct from an accepted
+zero-row occurrence, which has a real empty descriptor and publishes its
+descriptor-backed zero-byte task. A nonempty producer descriptor must contain
+a `PayloadSlice`; an empty producer descriptor is permitted.
 
 After device-to-host transfer, the native ring pairs each descriptor with its
 owned contiguous CPU payload as a backend-neutral record envelope. A native

docs/megatron.md

@@ -0,0 +1,61 @@
+# Megatron-LM usage
+
+DMI supports Megatron-LM training through the version-matched `DMI-Megatron-Integration` submodule. That integration pins the corresponding `Megatron-LM-DMI` fork, so a recursive DMI clone provides the complete source set required for this backend.
+
+## Install the Megatron-LM backend
+
+Use a dedicated environment and DMI checkout for Megatron-LM. Do not reuse a HuggingFace or vLLM environment.
+
+The version-matched integration owns the framework-specific requirements, tested dependency configuration, Transformer Engine source selection, pinned Megatron fork installation, and import verification. Follow its [installation guide](../third_party/DMI-Megatron-Integration/docs/install.md).
+
+In a recursive DMI checkout, the integration and its Megatron fork are already present, so skip that guide's clone step. Establish the final PyTorch, CUDA, Transformer Engine, and Megatron dependency versions first, then complete the DMI [core installation](install.md) from this checkout and install the integration from `third_party/DMI-Megatron-Integration/`.
+
+The DMI native extension is built against the active PyTorch and CUDA ABI. Even if the DMI core installation was completed earlier, rebuild the extension in the final Megatron environment if a later package installation upgrades or replaces PyTorch or its NVIDIA CUDA runtime packages, or if the CUDA toolkit used for compilation changes. Installing Megatron or Transformer Engine does not by itself require a rebuild when that ABI remains unchanged.
+
+```bash
+make -C native clean
+make -C native -j
+```
+
+## Enable DMI in a training run
+
+Use the pinned fork's `pretrain_gpt.py` with the normal Megatron model, data, optimizer, and parallelism arguments. Enable DMI by appending its existing CLI options to that command:
+
+```bash
+torchrun --nproc_per_node="$NPROC_PER_NODE" \
+  third_party/DMI-Megatron-Integration/third_party/megatron-lm/pretrain_gpt.py \
+  <your normal Megatron training arguments> \
+  --dmi-enable \
+  --dmi-hook-selection router-summary \
+  --dmi-model-id "$MODEL_ID" \
+  --dmi-db-host "$DMX_DB_HOST" \
+  --dmi-db-port "${DMX_DB_PORT:-9000}" \
+  --dmi-db-database "$DMX_DB_DATABASE" \
+  --dmi-clickhouse-table dmi_training_tensors
+```
+
+`--dmi-hook-selection` accepts a comma-separated integration hook selection. When it is omitted, the integration defaults to `router-summary`. Hook-specific topology requirements are validated during startup before training proceeds.
+
+The DMI options can also be supplied through their corresponding `DMI_*` environment variables. CLI values take precedence over environment values. For example:
+
+```bash
+export DMI_ENABLE=1
+export DMI_HOOK_SELECTION=router-summary
+export DMI_MODEL_ID=my-training-run
+export DMI_DB_HOST=127.0.0.1
+export DMI_DB_PORT=9000
+export DMI_DB_DATABASE=default
+export DMI_CLICKHOUSE_TABLE=dmi_training_tensors
+```
+
+## Choose the output mode
+
+Set `--dmi-db-host` or `DMI_DB_HOST` to write training records to ClickHouse. The integration creates and writes its schema-driven training tables through DMI's public storage API. Complete the ClickHouse setup in the [core installation guide](install.md) before starting the run.
+
+Leave the database host empty for capture and transport without persistence:
+
+```bash
+export DMI_DB_HOST=
+```
+
+In either mode, use the same Megatron workload and parallelism arguments you would use without DMI. The integration is activated only when `--dmi-enable` or `DMI_ENABLE=1` is present.

native/csrc/ring/producer.cu

@@ -379,21 +379,6 @@ __device__ inline void record_copy_chunk_with_wrap(
     }
 }
 
-// A false record gate still consumes the task reservation made by the host.
-// Publish an empty entry in FIFO order without touching payload state.
-__device__ inline void record_publish_zero_byte_task(RingState& ring) {
-    if (threadIdx.x != 0) return;
-    const uint32_t finished = atomicAdd(&g_block_done_counter, 1);
-    if (finished != gridDim.x - 1) return;
-
-    const uint64_t task_head = *ring.task_head;
-    const TaskEntry entry{};
-    task_publish(ring.task_entries, ring.task_cap, task_head, entry);
-    *ring.task_head = task_head + 1;
-
-    g_block_done_counter = 0;
-}
-
 __device__ inline void record_copy_contiguous(
     uint8_t*       payload_buf,
     const TwoSpan& spans,
@@ -427,10 +412,7 @@ __global__ void record_producer_static_kernel(
     const int32_t*   emit_gate,
     int32_t          emit_value) {
     if (g_ring_null_mode) return;
-    if (!record_emit_allowed(emit_gate, emit_value)) {
-        record_publish_zero_byte_task(ring);
-        return;
-    }
+    if (!record_emit_allowed(emit_gate, emit_value)) return;
 
     const uint64_t gtid = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x;
     const uint64_t stride = uint64_t(gridDim.x) * blockDim.x;
@@ -456,10 +438,7 @@ __global__ void record_producer_prefix_kernel(
     const int32_t*   emit_gate,
     int32_t          emit_value) {
     if (g_ring_null_mode) return;
-    if (!record_emit_allowed(emit_gate, emit_value)) {
-        record_publish_zero_byte_task(ring);
-        return;
-    }
+    if (!record_emit_allowed(emit_gate, emit_value)) return;
 
     int64_t rows = *row_count_dev_ptr;
     if (rows < 0) rows = 0;
@@ -494,10 +473,7 @@ __global__ void record_producer_chunked_kernel(
     const int32_t*   emit_gate,
     int32_t          emit_value) {
     if (g_ring_null_mode) return;
-    if (!record_emit_allowed(emit_gate, emit_value)) {
-        record_publish_zero_byte_task(ring);
-        return;
-    }
+    if (!record_emit_allowed(emit_gate, emit_value)) return;
 
     __shared__ int64_t selected[PRODUCER_MAX_K];
     __shared__ int64_t prefix[PRODUCER_MAX_K + 1];
@@ -553,10 +529,7 @@ __global__ void record_producer_seq_prefix_pack_kernel(
     const int32_t*   emit_gate,
     int32_t          emit_value) {
     if (g_ring_null_mode) return;
-    if (!record_emit_allowed(emit_gate, emit_value)) {
-        record_publish_zero_byte_task(ring);
-        return;
-    }
+    if (!record_emit_allowed(emit_gate, emit_value)) return;
     (void)valid_count_dev_ptr;
 
     int64_t encoded_rows = valid_prefix_sum_dev_ptr[batch];
@@ -604,10 +577,7 @@ __global__ void record_producer_segmented_pack_kernel(
     const int32_t*   emit_gate,
     int32_t          emit_value) {
     if (g_ring_null_mode) return;
-    if (!record_emit_allowed(emit_gate, emit_value)) {
-        record_publish_zero_byte_task(ring);
-        return;
-    }
+    if (!record_emit_allowed(emit_gate, emit_value)) return;
 
     const uint64_t input_rows = feature_bytes == 0
         ? 0 : nbytes_upper / feature_bytes;

tests/native/ring/test_producer.cu

@@ -405,8 +405,8 @@ static void test_record_sequence_and_segmented_pack() {
     CUDA_CHECK(cudaFree(device));
 }
 
-static void test_record_device_gate_publishes_zero_byte_entries() {
-    banner("false record gates publish zero-byte entries in FIFO order");
+static void test_record_device_gate_rejects_without_publication() {
+    banner("false record gates leave ring state unchanged");
     ring::AllocatedRing allocated(make_config());
     allocated.init();
     ring::RingState& state = allocated.state();
@@ -434,10 +434,11 @@ static void test_record_device_gate_publishes_zero_byte_entries() {
     CUDA_CHECK(cudaDeviceSynchronize());
 
     constexpr uint64_t gated_variants = 5;
+    EXPECT(*state.task_head == 0);
     for (uint64_t sequence = 0; sequence < gated_variants; ++sequence) {
-        expect_entry(state, sequence, 0);
+        EXPECT(state.task_entries[sequence].ready_seq ==
+               ring::READY_SEQ_SENTINEL);
     }
-    EXPECT(*state.task_head == gated_variants);
     EXPECT(*state.payload_head == initial_payload_head);
     EXPECT(*state.actual_bytes_counter == initial_actual_bytes);
     std::vector<uint8_t> payload_after(state.payload_cap);
@@ -451,9 +452,9 @@ static void test_record_device_gate_publishes_zero_byte_entries() {
     ring::launch_record_producer_static(
         state, device, source.size(), gate, 1);
     CUDA_CHECK(cudaDeviceSynchronize());
-    expect_entry(state, gated_variants, source.size());
-    EXPECT(read_payload(state, state.task_entries[gated_variants]) == source);
-    EXPECT(*state.task_head == gated_variants + 1);
+    expect_entry(state, 0, source.size());
+    EXPECT(read_payload(state, state.task_entries[0]) == source);
+    EXPECT(*state.task_head == 1);
     EXPECT(*state.payload_head == initial_payload_head +
            ring::align_up(source.size(), ring::PAYLOAD_ALIGN));
     EXPECT(*state.actual_bytes_counter == initial_actual_bytes + source.size());
@@ -477,7 +478,7 @@ int main() {
     test_record_chunked_compacts_payload_head();
     test_record_chunked_unaligned_boundaries();
     test_record_sequence_and_segmented_pack();
-    test_record_device_gate_publishes_zero_byte_entries();
+    test_record_device_gate_rejects_without_publication();
 
     std::printf("Results: %d passed, %d failed\n", g_pass, g_fail);
     return g_fail == 0 ? 0 : 1;