[Core] Add recurring D2H window scheduling - #131
Conversation
zaoxing
left a comment
zaoxing
left a comment
There was a problem hiding this comment.
Review of the recurring D2H window scheduling change. 10 inline comments below — 4 correctness, 1 efficiency, 1 altitude, 2 simplification, 1 reuse, 1 docs.
The two I'd prioritise before merge:
- Late-issue overruns wipe learned capacity (
recurring_d2h_grant_controller.cpp:342) — the controller already tracksclean_open/missed_open_occurrencebut only uses it to suppress timing evidence; an overrun after a drain stall still resetsmax_safe_as a "capacity shift". create_record_runtime()afterclose()(engine.py:264) — the relaxed guard plusclose()not clearing_ring_configlets a closed engine build a sink-less record ring and re-activate the transport.
Scope note: the third_party/DMI-Megatron-Integration submodule bump was not reviewed.
| decision.prior_min_unsafe = min_unsafe_; | ||
|
|
||
| if (!max_safe_.has_value() && !min_unsafe_.has_value()) { | ||
| decision.byte_limit = std::numeric_limits<uint64_t>::max(); |
There was a problem hiding this comment.
[altitude] First grant is unbounded, and an overrun discards its timing evidence
With no bounds the policy grants UINT64_MAX (everything that fits staging); when that overruns, complete() resets the timing observation, so the controller learns only B=huge and halves from there. From a multi-GiB backlog that is ~log2(backlog/capacity) consecutive overruns into the TX-critical regions the feature exists to protect (4 GiB → 2 → 1 → 512 → … → 4 MiB ≈ 10 occurrences).
The overrunning transfer's completion time already gives a first-order estimate (bytes × window_ns / transfer_ns) that is thrown away. Consider capping the initial grant (configured seed or staging_free / k) and/or using the overrun's timing as the initial B.
There was a problem hiding this comment.
Not fixed — I left this one to you deliberately, since it turns on a decision a reviewer should not make unilaterally.
The capping half changes public surface and tuning defaults: either a new RecurringD2HWindowConfig seed field (bindings, validation, and both docs) or a staging_free / k with an invented k. Choosing k wrong makes the feature worse rather than better, and that call is yours.
On the second half — seeding the initial B from the overrun's timing — I do not think it is sound, and would rather say so than implement it. For an overrun the window duration is unknown; all that is known is window_ns < transfer_ns, so bytes × window_ns / transfer_ns has no measured numerator. And B wants a value known unsafe, which attempt.bytes already is, more tightly than any estimate derived from it. Happy to be corrected if you meant something different.
One correction to the cost estimate, which may affect how you weigh this: the search does not start from UINT64_MAX in practice. The drain clamps the first transfer to what fits staging, so min_unsafe_ becomes the staged bytes and the halving count is log2(staging capacity / true capacity), not log2(backlog / capacity). Still worth bounding, but the 4 GiB → 4 MiB ≈ 10 occurrences figure assumes a 4 GiB staging ring.
|
@XbzOnGit Also a quick question: how was the parameter determined? window = dmi.RecurringD2HWindowConfig() ring = dmi.RingConfig() engine = dmi.MonitoringEngine( |
Samfisheryu
left a comment
Samfisheryu
left a comment
There was a problem hiding this comment.
Reviewed the DMI changes and the actual pinned Megatron integration at b9b2c49 / Megatron 19d1999, including the opt-in boundary and existing Ring FIFO behavior.
Pushed two small fixes in a544600: the new progress allocation used obsolete CUDA API signatures and broke the CUDA 13 build even with windows disabled; close() also retained the config and allowed a closed/deferred engine to create a record runtime. The latter now has both regression cases.
Requesting changes for the remaining runtime behavior, not documentation: I reproduced Alan's late-issue capacity reset and late-CLOSE estimate inflation in their existing threads. The new inline comment covers validation leaving the training window controller active and causing terminal fallback.
Validation: 1,129 CPU tests and 213 pinned-integration tests passed; native controller tests passed 203 assertions; native Ring tests passed 147 assertions (the multi-device ownership case was skipped with one visible GPU). Ten existing single-GPU reference/ClickHouse tests passed. An additional CUDA Graph test passed 64 replays plus warmup, reading all 65 tensor payloads back byte-for-byte. The failing validation reproduction also preserved all 10 tensor payloads: it breaks window scheduling, not tensor contents. This is not a full multi-GPU Megatron training run.
| @@ -1 +1 @@ | |||
| Subproject commit b0e141f0339095caf67ead3853d909b6265f0d3f | |||
| Subproject commit b9b2c49c2a0e09f3b49332cf5f630ecc4ad40d5f | |||
There was a problem hiding this comment.
[P2] Suspend training windows during validation, then resume them
The pinned Megatron schedule stops publishing boundaries for forward_only, but the integration runtime leaves the native training pattern active. Validation records then fill the Ring while the last training window is closed. Capacity flushes count against the training policy; after three, it permanently enters ENABLED_FALLBACK, including when training resumes.
Two-part reproduction:
- Add this test to the pinned integration's tests/test_megatron_d2h_window_schedule.py, using its existing schedule_harness (real schedule, mocked compute/P2P):
def test_train_then_validation_suspends_windows(schedule_harness):
run, events, runtime, records = schedule_harness
run(2, 0, 2)
assert records.counter == 4
events.clear()
run(2, 0, 2, forward_only=True)
assert records.counter == 4
assert not any(e[0] in ('open', 'close', 'define') for e in events)
assert not runtime.d2h_windows_active # fails: still True- Independently on a real GPU Ring/reference sink: enable windows with period=2, windows=[(1,2)], 64 KiB payload/staging rings, and fallback threshold=3. Publish OPEN, wait for ACTIVE, publish CLOSE, then capture ten 32 KiB tensors without further markers, matching forward-only scheduling. Flush and read back.
evaluation: mode=D2HWindowMode.ENABLED_FALLBACK forced_flushes=3
[d2h_window] terminal fallback count=3 threshold=3
All 10 payloads were byte-exact. The problem is that a normal validation phase permanently disables the feature. Please explicitly suspend/resume the native window policy around forward-only phases, using normal batching without charging these flushes as training-window failures. Merely stopping boundary markers is insufficient.
There was a problem hiding this comment.
Not fixed here, and I do not think it can be from this repository: the change you are asking for is in DMI-Megatron-Integration's schedule_runtime.py, a separate repo. This PR only moves the submodule pointer.
Agreed on the diagnosis and on the shape of the fix — explicitly suspending and resuming the native window policy around forward-only phases, rather than only stopping boundary markers. The capacity flushes a validation phase produces genuinely charge against the training window's policy, and three of them are terminal.
One thing that is fixable here: the PR description says it pins DMI-Megatron-Integration f416d73, but the branch records b9b2c49 — the revision you actually reviewed. The description needs updating.
a544600 to
dd5627b
Compare
The three D2H window translation units each carried their own copy of the same check_cuda/check_launch helper. Move one inline pair into ring/cuda_check.h so a failure reports the same shape everywhere. The CUDA 13 fix landed alongside it replaced the device-ordinal overloads of cudaMemAdvise/cudaMemPrefetchAsync with the cudaMemLocation form unconditionally. CUDA 13 does declare only that form, but it publishes no _v2 alias to spell it on older toolkits and docs/install.md still supports CUDA 12.x, so select on CUDART_VERSION rather than breaking those builds.
Two observations the controller made after a stall were read as if the drain had been watching the whole time. An overrun was fed to the policy unconditionally, and an overrun at or below max_safe_ is treated as a downward capacity shift: max_safe_ and the timing estimate are dropped and halving resumes. But a transfer issued after the drain reached the window late ran against a truncated window, so its overrun says nothing about capacity. A converged 64 MiB window that took a forced flush ending just before CLOSE lost its bound and re-searched from 32 MiB, under-filling for several occurrences and doing it again on the next stall. Overrun evidence is now gated on a clean occurrence; success still counts, since bytes that fit in a shortened window fit in a full one. CLOSE was observed at the first poll where the counter had reached it, however much later that poll was, and window_ns was measured to that poll. A 10 ms window whose close was seen 205 ms late estimated 64 * 205 / 5 = 2,624 MiB and granted it -- a guaranteed overrun into the region the feature exists to protect. The estimate now requires the counter to sit exactly at absolute_close. Neither the counter nor the clock shows afterwards that the drain was away, so the drain reports the stall itself: note_observation_gap() before a forced full flush and before an ordinary batching flush marks the observations that straddle the gap unusable. Two smaller items in the same path. The per-activation "active version" line went to stderr unconditionally, so an integration redefining its pattern per phase spammed every rank in production; it is now behind the debug logger with the other informational window lines. And do_window_decision walked all of scanned_ on every poll even though pending_bytes_ already holds the aligned sum, holding the lock prepare_step()/reserve_record() need; the walk now runs only when the backlog cannot fit whole.
The loop acknowledges only the newest requested pause generation, but every earlier waiter's predicate is acknowledged >= generation, so two concurrent pause_after_flush_and_wait() callers each left holding a token while only the newer one could resume. resume() then dropped the older token silently, the loop stayed parked on pause_resumed_generation_, and every later force_flush_and_wait() or reserve_record() hung with nothing pointing at the cause. Only one pause is outstanding now: a second requester queues on pause_entry_cv_ until the first resumes, so two tokens can never exist. A token that is not the outstanding one -- stale, or already resumed -- throws instead of returning, so misuse fails where the mistake is rather than as an unrelated hang later. record_drain_failure() wakes entry waiters so a dead drain reports itself instead of blocking them. DrainPauseControl is a generic interface, and today only the window subsystem's version-reuse path calls it. That path's cleanup resume now swallows a rejection: it runs while unwinding, and must not replace the failure the caller needs to see.
The registry hardcoded 0 for D2HWindowProgressKind.PACKED_VERSION_COUNTER, so adding a second kind or reordering the values would silently bind the wrong kernel. Key it by the bound enum instead, imported inside the branch so a module-level attribute access does not force the extension load at import time. docs/config.md's table listed seven of the struct's eight fields, leaving out capacity_flush_count_reset_interval_periods -- which this feature adds, validates as positive, exposes in the bindings and already documents in docs/integration-api-v1.md. The two documents described different field sets for the same struct.
3 participants
Purpose
Add an opt-in D2H policy for frameworks that know recurring execution regions where offload traffic is less likely to contend with TX-critical work. The framework supplies only a repeating window pattern and ordered boundaries; DMI-core learns how many complete-record bytes each window can carry.
The existing GPU ring, drain stream, record reconstruction, P2P thread, and sink continue to move and publish the data. The feature is disabled by default and applies only to the HookPointV1 record path.
Configuration
The integration defines sorted, non-overlapping half-open windows within a positive period
T:Window
(begin, end)in repetitionnis[n*T + begin, n*T + end). The integration then publishes each ordered boundary on its framework CUDA stream:enableddefaults toFalse. Enabled configurations require positive probe intervals and fallback threshold. Pattern definition returnsFalseonly after terminal fallback rejects further definitions. Debug mode logs real definitions, transfers, outcomes, warnings, and fallback.Megatron exposes the feature through
--dmi-recurring-d2h-windows. It supports eager, non-interleaved 1F1B with PP>1 and the standard pipeline communicator. It requiresbatch_p2p_sync=False; when both synchronous P2P and windowing are requested, the integration prints a warning and disables windowing. PP=1 and interleaved/VPP execution retain ordinary batching; multi-module pipelines are rejected when the option is enabled.batch_p2p_syncis a workaround for incorrect asynchronous batched-P2P ordering in older PyTorch releases. The coalesced NCCL work/stream-ordering fix landed in PyTorch 2.1 through pytorch/pytorch#98793. This Megatron checkout already requirestorch>=2.6.0, and the validation below used PyTorch2.13.0+cu129. An older affected stack that must keepbatch_p2p_sync=Truetherefore retains ordinary DMI batching instead of enabling recurring windows.Algorithm
Each window position learns:
A: largest proven-safe transfer;B: smallest proven-unsafe transfer; andE: clean same-window timing estimate.The
BINARY_ADAPTIVEpolicy starts with all currently transferable records. With onlyB, it halves; with both bounds, it uses their midpoint; with onlyA, it usesAor an allowed largerE. Completion before CLOSE raisesA; completion at or after CLOSE lowersB.If a transfer at or below
Aoverruns, the controller treats it as a downward capacity shift, clears stale safe/timing evidence, and resumes halving. Timing evidence is used only when OPEN observation, issue, completion, and CLOSE observation belong to the same occurrence.Records are never split. A first record larger than the target may be probed as one unit. Repeated failed minimum-record and timing-revalidation probes use increasing occurrence cooldowns, preventing a failed probe every recurrence. Indivisible-record convergence stalls at the current safe prefix and emits one warning; record chunking is outside this PR.
An overrun crossing later boundaries updates only its originating window and resumes from current progress. Repeated eligible capacity flushes trigger terminal fallback to ordinary batching; isolated evidence ages out after the configured number of periods. Lifecycle, shutdown, and oversized-record flushes do not count.
Implementation details
advance_boundary()enqueues a captureable GPU operation on the framework stream. It increments a packed 16-bit version / 48-bit counter and publishes it to CPU-visible memory without a host callback or GPU/CPU synchronization.Forced flush handling remains first in the drain loop. In an active window, the controller admits at most one transfer per occurrence. The drain packs the largest FIFO prefix of whole ready records that fits, rechecks the version and occurrence, issues D2H on the existing copy stream, and reports actual bytes and the completion boundary. Disabled, no-pattern, missed-window, and fallback states proceed to the existing
should_flush()batching path.Pattern definitions are versioned and pending definitions remain ordered. Completion from an older version cannot update current learning state. At exceptional version exhaustion, DMI synchronizes the framework stream, flushes and pauses the drain worker, clears version state, enqueues the reset and new definition, then resumes.
Compatibility and scope
The current Megatron path continues to pass its native
DMXHostEnginethroughstorage_backend="auto"and creates the record runtime without an explicitrecord_sink. It therefore retains the ClickHouse tensor-bytes path rather than selecting the capture/object-storage sink.This PR does not infer windows inside DMI-core, coordinate learning across ranks/streams/engines, split records, or claim that a logical window proves physical PCIe/NIC idleness.
Pinned integration revisions:
f416d731aecebdValidation
main: 28 passed.git diff --check: passed.A 13-iteration Qwen3-1.7B-shaped PP=2 run with
batch_p2p_sync=Falsecompared ordinary batching and recurring windows while offloading forward hidden states. Both modes stored exactly 728 logical records and 3,053,453,312 payload bytes, with zero bidirectional identity differences.For placement analysis, we conservatively treat compute windows as TX-critical. In large-scale training, these regions may also carry cross-node distributed traffic such as EP communication, so D2H traffic should avoid them.
Trace-measured D2H overlap with TX-critical windows fell by 96.25%. The controller recorded 34 successes and four overruns without minimum-record probes or terminal fallback. The short window converged to a 4 MiB grant after its larger grants overran; the long window remained successful.
Under Nsight, iterations 3--13 averaged 479.791 ms with ordinary batching and 470.591 ms with recurring windows (-1.92%) in this single run. This validates traffic placement and downward reconvergence.