diff --git a/.gitmodules b/.gitmodules index 376efaee3..7cdf51793 100644 --- a/.gitmodules +++ b/.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 diff --git a/README.md b/README.md index 9270290e4..e24d07645 100644 --- a/README.md +++ b/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,6 +153,7 @@ 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 @@ -158,7 +161,7 @@ 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. diff --git a/docs/install.md b/docs/install.md index 925ed8e2a..5d2dfd741 100644 --- a/docs/install.md +++ b/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 diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index 52b3ffd98..a1f5279c9 100644 --- a/docs/integration-api-v1.md +++ b/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 diff --git a/docs/megatron.md b/docs/megatron.md new file mode 100644 index 000000000..215804b0a --- /dev/null +++ b/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 \ + \ + --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. diff --git a/native/csrc/ring/producer.cu b/native/csrc/ring/producer.cu index 6600e7638..179f66621 100644 --- a/native/csrc/ring/producer.cu +++ b/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; diff --git a/tests/native/ring/test_producer.cu b/tests/native/ring/test_producer.cu index 361e0dfc1..c6668e5d4 100644 --- a/tests/native/ring/test_producer.cu +++ b/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 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; diff --git a/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index 11dcd777a..3b32d28b1 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -324,40 +324,6 @@ static void test_record_reservation_reclaims_per_entry() { CUDA_CHECK(cudaFree(dynamic_device)); } -static void test_false_gated_record_reclaims_its_full_reservation() { - banner("false-gated record publishes and reclaims a zero-byte task"); - DrainHarness harness(make_config()); - - const std::vector source = pattern(64, 173); - uint8_t* device = upload(source, harness.stream); - int32_t* gate = nullptr; - const int32_t disabled = 0; - CUDA_CHECK(cudaMalloc(&gate, sizeof(disabled))); - CUDA_CHECK(cudaMemcpyAsync(gate, &disabled, sizeof(disabled), - cudaMemcpyHostToDevice, harness.stream)); - - const uint64_t reserved = - ring::align_up(source.size(), ring::PAYLOAD_ALIGN); - harness.drain->reserve_record({{reserved, true}}); - ring::launch_record_producer_static( - harness.allocated.state(), device, source.size(), gate, 1, - harness.stream); - ring::DrainTask task = harness.flush_one(); - - EXPECT(task.tensor_total_bytes == 0); - EXPECT(task.alloc_bytes == 0); - EXPECT(harness.drain->pending_record_reclaims() == 0); - EXPECT(harness.drain->cpu_payload_head() == reserved); - harness.drain->apply_pending_record_reclaims(); - EXPECT(harness.drain->cpu_payload_head() == 0); - EXPECT(harness.drain->cpu_payload_tail_committed() == 0); - harness.drain->rethrow_record_reclaim_failure(); - harness.release(task); - - CUDA_CHECK(cudaFree(gate)); - CUDA_CHECK(cudaFree(device)); -} - static void test_timed_drain_flush_uses_request_generations() { banner("timed drain flush does not reuse another request generation"); ring::RingConfig cfg = make_config(512); @@ -598,7 +564,6 @@ int main() { test_repeated_wrap_delivery(); test_zero_byte_delivery(); test_record_reservation_reclaims_per_entry(); - test_false_gated_record_reclaims_its_full_reservation(); test_timed_drain_flush_uses_request_generations(); test_drain_worker_binds_owner_device(); test_record_flush_bounds_current_stream_prefix_wait(); diff --git a/third_party/DMI-Megatron-Integration b/third_party/DMI-Megatron-Integration new file mode 160000 index 000000000..b0e141f03 --- /dev/null +++ b/third_party/DMI-Megatron-Integration @@ -0,0 +1 @@ +Subproject commit b0e141f0339095caf67ead3853d909b6265f0d3f