feat(data-plane): track data-plane time, latency percentiles and byte volume - #3616
Open
ZhiyuLi-Nvidia wants to merge 60 commits into
Open
feat(data-plane): track data-plane time, latency percentiles and byte volume#3616ZhiyuLi-Nvidia wants to merge 60 commits into
ZhiyuLi-Nvidia wants to merge 60 commits into
Conversation
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/data_plane_observability_metrics
branch
from
August 24, 2026 00:50
c67f613 to
48ed054
Compare
Contributor
Author
|
/ok to test 48ed054 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/data_plane_observability_metrics
branch
from
August 27, 2026 07:06
5e79a5d to
67be5a7
Compare
Contributor
Author
|
/ok to test 67be5a7 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/data_plane_observability_metrics
branch
2 times, most recently
from
August 30, 2026 15:27
9cbeea4 to
501a7fd
Compare
Contributor
Author
|
/ok to test 501a7fd |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/data_plane_observability_metrics
branch
from
September 8, 2026 23:44
15a6c26 to
4ea3bc7
Compare
Contributor
Author
|
/ok to test 4ea3bc7 |
… volume MetricsDataPlaneClient recorded per-event bytes and wall_ms but aggregated only bytes, so there was no answer to "what did the data plane cost this step" and no way to compare backends. Adds, all derived at the DataPlaneClient interface so simple and mooncake_cpu are measured identically: * total_wall_ms plus a per-op breakdown (calls, errors, wall_ms, bytes, keys) with derived mean_ms, mb_per_s and pct_of_total_ms. by_op partitions total_wall_ms exactly. * p50/p99 per op from a fixed-bucket histogram. Buckets rather than retained samples keep memory O(1) and, more importantly, make the counts additive: per-rank histograms sum into one cluster-wide distribution, which averaging per-rank percentiles cannot do. A mean hides tail latency -- exactly where MR registration churn and queueing show up. * fit_latency_bandwidth(): least-squares split of wall_ms into fixed per-request overhead vs bandwidth, answering whether an op is overhead- or bandwidth-dominated. Sufficient statistics only, so these are additive across ranks too. * Communication volume: bytes_written / bytes_read / comm_volume_bytes, derived from by_op so bytes have a single source of truth. Distinct from bytes_outstanding, which is occupancy rather than traffic. * _td_nontensor(): non-tensor payload bytes, which _td_bytes omits. TQ ships non-tensors over a separate msgpack path, so volume was undercounted by whatever metadata rides along. The fit reports its own validity. R^2 alone is not sufficient -- a chunked step function and a quadratic both fit a line above R^2 0.93 while producing a meaningless split -- so model_trustworthy also requires a non-negative intercept, which is what catches those. A degenerate case where all requests are the same size (common in RL) is reported as "unidentifiable" rather than given an arbitrary split. Byte accounting verified against ground truth: _td_bytes equals t.contiguous().nbytes (what mooncake registers and sends) for contiguous, sliced, transposed and stride-0 expanded views and across bf16/bool/int64/fp8; _td_nontensor matches msgpack.packb byte-for-byte across strings, ints, bytes, nested containers and utf-8. Off by default. Measured overhead with the wrapper enabled is 22 us per get and 96 us per put at 256 keys -- 0.04% of a 59 ms operation. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The metrics existed but nothing published them: DataPlaneStats accumulated in-process and snapshot() was never called, so a run produced no data-plane numbers. Hooks grpo_train_sync -- the trainer data_plane.enabled selects (run_grpo.py) -- to log one snapshot per step alongside the existing timing/train metrics, and surfaces the observability toggle in the exemplar configs (still off by default). Cumulative counters are differenced into per-step values. The key metric is data_plane/frac_of_step: per-op shares only say where data-plane time went, never whether it mattered against compute -- and on DSv3 the storage backend proved nearly invisible against step time, so that ratio is what decides whether optimising the data plane is worth anything. Scope: only the driver's own client is visible here. Clients are built per process (tq_policy bootstraps one; rollout actors and policy workers build their own), so this covers the trainer's reads rather than cluster-wide traffic. The stats were designed additive -- histogram buckets and regression sums both sum correctly -- so a cluster-wide gather can be layered on without changing the accounting. No-op when observability is disabled: the plain adapter has no snapshot(), so the hook returns immediately. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Four parallel review agents (reuse, simplification, efficiency, altitude) reviewed the two preceding commits. Fixes applied: Altitude/reuse (both agents, independently): delta arithmetic, unit conversion and the per-op namespace moved off the trainer onto MetricsDataPlaneClient.get_step_metrics(), mirroring VllmGeneration.get_step_metrics(). The trainer hook drops 71 -> 27 lines and the module-level _DP_PREV_SNAPSHOT global is gone -- the previous reading is the client's state, and it lives there now. The async trainer can log the same metrics with one call against its own client instead of copying ~60 lines. The enable-gate is an isinstance check rather than a getattr probe on a method name. Efficiency: by_op.setdefault(op, OpStats()) built a throwaway OpStats and its 16-bucket histogram on *every* op before discarding it (0.62 us/op, ~2.8% of a 22 us get); replaced with get-then-insert. Efficiency (bug): the _estimate_encoded_bytes node budget did not bound the walk. Recursive calls returned 0 once exhausted, but the enclosing sum() generators kept iterating, so cost was O(size) not O(max_nodes) -- 88 ms on a 1M-element list. Containers now break on exhaustion: 1.88 ms at 1M elements, and the docstring no longer claims a bound it did not provide. Simplification: _td_nontensor merged into _td_bytes so non-tensor bytes are counted in a *single* traversal -- wiring them in as a second walk would have doubled the per-put cost over a structure holding hundreds of keys. This also gives the msgpack estimator its first caller: non-tensor payload now counts toward comm_volume_bytes instead of being silently omitted. Dropped OpStats.ok_calls (identically calls - errors) and fit["negative_intercept"] (identically fixed_ms < 0), and removed defensive .get() on keys snapshot() always sets. Verified: byte math still matches msgpack.packb exactly (tensor + non-tensor), per-step deltas correct, no-op when observability is off. Skipped, with reasons: TQ's _quantile_from_cumulative is a private staticmethod on a Ray actor taking a live prometheus Histogram, so not importable; TensorDictBase.bytes() cannot replace _td_bytes now that it also counts non-tensor leaves; the ~10 us default log_event callback and _record_put's O(n_keys) inserts are pre-existing and out of this diff's scope; cluster-wide aggregation across the other three per-process clients is a follow-up, not a cleanup. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The metrics wrapper sits on the hot path of every transfer, so its cost is a design constraint. Measured against a no-op inner client at 256 keys / 18 MB, per put/get: tensors only 62us / 17us -> 28us / 7us (2.2x / 2.4x) + NonTensorData metadata 83us / 49us -> 48us / 27us (1.7x / 1.8x) + 256-row message_log 5124us / 5057us -> 57us / 34us (91x / 149x) Four changes get there, none of which alter the numbers reported (byte accounting is bit-identical on every shape both versions walk the same way; the NonTensorStack estimate moves by 0.7%): - `_td_bytes` walks with `items()` instead of `keys()` + `get()`. `get()` re-resolves each nested key from the root, which measured 2.7x the cost of the traversal `items()` already performs. - `NonTensorData` / `NonTensorStack` are matched by type rather than by `hasattr`. Both are tensorclasses whose attribute misses fall through a `__getattr__` costing ~2.8us per probe. - A `NonTensorStack` is extrapolated from a strided 4-row sample instead of `tolist()`-ing every row. This is the 91x case: a 256-row message_log stack cost 5.1 ms per put, i.e. 8.7% of a 59 ms operation, which made the metrics more expensive than most of what they measured. - `_emit` builds its event dict only when a callback is registered, and `_record_put` drops the per-key remainder spreading (an `enumerate` and a compare per key to move at most one byte each). The non-tensor size estimator now dispatches on exact type instead of an isinstance chain the common leaves sat 5-7 branches deep in, and charges its node budget per container element rather than per visited node. Also adds `data_plane.observability.verify_tensor_hash` (off by default): each put records a per-row `torch.hash_tensor` fingerprint, each get re-checks it, so a tensor that changes between wire-in and wire-out is reported as `hash/mismatches` instead of being trained on. Fingerprints are per row, so a 256-row put read back as eight shards of 32 still reconciles, and rows this process never wrote are counted under `rows_unverified` rather than as clean. XOR reduction is commutative, so a single `hash_tensor` call would report a mis-sharded read that swaps two rows of `arange`-shaped data as clean. Dtype and per-row shape go in as a salt and a second reduction over the row's even-index sub-multiset breaks the permutation symmetry; a boolean mask, where XOR sees only parity, is the documented remaining blind spot. The check costs ~1.2 ms per 18 MB on each side, hence off by default. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…call
Verifying the guard against the real wire layout found it was not guarding
anything. `column_io` packs every per-token field through
`codec.pack_jagged_fields` before `put_samples`, so `input_ids`,
`generation_logprobs`, `token_mask` and `advantages` all arrive as jagged
nested tensors — which `_row_fingerprints` skipped. It reported
`fields_skipped=6, mismatches=0`: a clean bill of health for a payload it
had never looked at. Eight of ten injected corruptions went undetected.
Nested leaves are now fingerprinted by padding to a rectangle and reducing
per row. The padding is free rather than approximate: XOR is its own
identity on zero and a zero pad bitcasts to a zero 64-bit word for every
dtype on this wire (verified int64/int32/bf16/fp32/bool), so a jagged put
reconciles against the dense read `_from_wire` hands back when rows happen
to be uniform.
Also cuts the fingerprint back to a single `torch.hash_tensor` call plus a
dtype salt. The previous second, strided reduction doubled the cost to buy
detection of a within-row element permutation — a rearrangement nothing on
this wire performs. Rows are compared per sample id, so a swap *between*
rows never needed it. All ten corruption modes are caught either way.
Verification, against pack_jagged_fields in and _from_wire out:
false positives 0 over a 500-row randomized soak, every shard grouping
from 1 to 256, reversed id order, field subsets, delta
writes and overwritten fields
false negatives 0 of 10 — single-element change per dtype, truncated
row, zeroed row, bf16->fp32, wrong sample, swapped rows
`hash/fields_skipped` is now logged per step, because that counter at
`0 mismatches` is exactly what this bug looked like from the outside.
Dropping the jagged special case in `_td_bytes` sped the metrics path up
as well: `numel()` already reports a nested tensor's total element count
and costs half of reaching through `.values()`. On the real jagged payload
(256 ragged rows, 12 MB) put goes 201us -> 99us and get 162us -> 76us,
where the earlier dense-payload benchmark had missed this path entirely.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…angle
Per-row fingerprints for jagged leaves meant padding each field out to a
rectangle before hashing, and that rectangle is sized by the longest row,
not by the payload. On a realistically ragged batch it was 3.5x the real
data and cost 13x more than hashing the values buffer (5378us vs 405us) —
a bill that grows with how ragged the batch is, for a field the digest
only has to answer "did anything change" about.
Now split by what each layout can do cheaply:
rectangular leaf one digest per row (hash_tensor(dim=1)) — names the
sample that diverged, survives shard reads
jagged leaf one digest over the values buffer, XORed per row with
that row's length, marked batch-scoped
verify_tensor_hash on a 12 MB jagged batch: 8.0ms -> 2.4ms per put, and
the cost now tracks real bytes instead of padding waste.
The scheme has to follow the field, not the layout in hand. A field packed
jagged comes back *dense* whenever its rows are uniform, because
`_from_wire` densifies those — so choosing per layout made the two sides
compute different things and every row of a uniform batch reported a
mismatch (940 false alarms across the soak). `_check_hashes` now replays
the scheme recorded at put; the values buffer of a uniform jagged field
and the dense tensor it densifies into are the same elements in the same
order, so the digests agree by construction. A shard read of a
batch-scoped field reports unverified rather than mismatching.
Also bitcasts every leaf to a same-width integer before hashing.
`hash_tensor` has no float8 kernel — it raises NotImplementedError on
float8_e4m3fn, which propagated straight out of put_samples and would have
taken the transfer down with it. The view is free and the dtype still
enters through the salt.
Re-verified against the same injected corruptions, zero false alarms:
caught single-element change in every dtype, truncated row, zeroed
row, bf16->fp32, row served from the wrong sample
clean 500-row soak, shard groupings 1..256, reversed id order,
field subsets, delta writes, overwritten fields
limit mis-shard on a jagged field is caught only when the swapped
rows differ in length: 58/60 on ragged data, never on a
uniform-length batch. Rectangular fields catch it always.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Asking what the hash check actually catches, per leaf type, turned up a third silent hole. `_check_hashes` drops a field whose recorded scheme cannot be reproduced by the read, which is correct — reporting it as a row-level mismatch would cry wolf on every sharded fetch — but it dropped it without counting, so the report still read clean. It fires for real divergence, not just shard reads: a rectangular put that comes back jagged (truncating one row makes the batch ragged) has no comparable digest, so a genuine truncation was being swallowed. That case now lands in `hash/fields_skipped`, which is already logged per step. This is the same shape as the original bug — a field quietly stops being compared and zero mismatches reads as a clean bill of health. Every path that declines to check something now increments a counter. Also documents that detection is not attribution. A jagged digest covers the whole values buffer, so any change moves every row's value: it says "this batch is wrong", never "this sample is wrong". Measured on an 8-row batch, one element changed in u3: jagged leaf caught, flags all 8 rows rectangular leaf caught, names u3 `pack_jagged_fields` leaves ~94% of the payload jagged, so batch-level is the normal resolution and the README now says so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Cleanup pass over the metrics/verification work. No behaviour change: the corruption matrix still catches every mode it did, the false-alarm soak is still clean, and put/get stay at 99us/77us on the jagged payload. observability.py 1119 -> 1036 lines, prose:code 0.62 -> 0.58 net -96 lines - Reverts the exact-type `_ENCODERS` dispatch table to the isinstance chain it replaced. The table only helped exact-type hits while making every subclass and unknown type walk 13 isinstance checks instead of 8 — and `_NONTENSOR_STACK_SAMPLES` had already collapsed the call count the table was optimising per-call cost for. Removes 10 module-level names and ~45 lines. - Shares `_pop_partition_keys` between the byte accounting and the fingerprint store, which were running near-identical teardown 170 lines apart in lockstep by inspection only. `_drop_hashes` disappears. - Collapses `_batch_scoped_rows` (partition -> field -> row count, every int in an inner dict identical) to `_batch_scope` (partition -> (field names, row count)). - One `_as_list` for the three copy-pasted `sample_ids` coercions. - Logs `hash/rows_recorded`, which was write-only outside tests. - Tests reuse the existing `_hash_client()` helper. Trims docstrings that had grown into commit messages: benchmark numbers measured on one machine that silently rot, and prose duplicated verbatim from the README added in the same series. The traps stay — NonTensorData .tolist() broadcasting, NonTensorStack subclassing LazyStackedTensorDict, setdefault's eager default — since those are what a reader would trip on. Two review findings were raised and not acted on: - Making jagged digests row-local would delete the batch-scoping concept, but the claim that batch-scoping loses coverage on sharded reads does not hold: the sharded get_samples is in worker_mixin, a process that never wrote those rows and so reports rows_unverified either way. The driver reads back the same full meta it wrote. Row-local slicing also measured slower than the padding it would replace. - `_as_int_view` duplicates weight_transfer_sparse_codec.integer_view, which cannot be imported here: it pulls in ray via the vllm config chain, and data_plane is imported eagerly and tested in a Ray-free venv. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
A nested tensor's ``nbytes``, ``numel`` and ``element_size`` all route through ``__torch_function__`` and measure ~16 us per leaf; the packed values buffer answers the same question in ~0.1 us. Every per-token field on this wire is nested, so that was paid four or five times per transfer and dominated the whole wrapper. Against the PR head, on 256 ragged rows / 12 MB: put 205 us -> 37 us (5.5x) get 168 us -> 15 us (11x) The public attribute is not an escape: ``v.nbytes`` on a nested tensor measured 16.1 us, indistinguishable from the arithmetic it would replace, because it dispatches the same way. Guarded twice, because a wrong byte count is worse than a slow one — it feeds ``max_bytes_per_key_seen``, which exists to spot wire regressions, so an overcount reads as the very thing it is watching for. An absent ``_values`` falls back to the public attribute. A buffer holding more elements than the offsets describe means the tensor views a larger allocation, which ``torch.nested.narrow`` produces; nothing in nemo_rl builds one (all four construction sites use the compacting ``as_nested_tensor``) but ``_td_bytes`` is handed whatever a caller passes. Both paths have a regression test. Byte accounting verified unchanged across contiguous, sliced, transposed, stride-0 expanded, bf16/bool/int64/float8, jagged, old strided nested and narrow-view leaves. Also documents that `enabled: true` installs `log_event` by default, so turning metrics on currently emits one logger.info per data-plane op. Left as-is since the config documents that default, but it makes the lazy event-dict build in `_emit` dead unless a caller passes on_event=None. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Cleanup pass over the nested byte-count fast path. No behaviour change — byte accounting is identical, the corruption matrix is unchanged, and put/get stay at 37us/15us. The guard comment was wrong in a way that invited breaking it. It said the first check exists in case `_values` is "absent", which reads as an argument for `if buf is None`. Measured: `_values` is a bound *method* on every dense tensor, never None, so `buf is None` would send every dense leaf down the private path. The check is on the type and the docstring now says why. Also: - `_estimate_encoded_bytes` still answered "bytes of a tensor" with `numel() * element_size()`, the exact expression this series replaced one function over. It calls `_tensor_bytes` now, so there is one answer to the question and a nested tensor reaching that branch does not pay the dispatch the fast path exists to avoid. - `_td_bytes`'s docstring still described leaves as `numel * element_size`. Arithmetically true, greppably stale. - Dropped the benchmark figures and the "all four construction sites" count from the docstring: the numbers are in README.md and a site count rots the next time a site is added. Kept the narrow-view trap. - Tests reuse the existing `_jagged()` helper, hoisted above its callers. Two findings were raised and deliberately left: - `int(offsets[-1])` is 73% of the nested-leaf cost, but eleven measured alternatives were all slower except `off[-1].item()`, which wins 0.07us on a 14us function. Not worth the churn. Worth re-measuring on a GPU box though: if `_offsets` is ever on CUDA that line is a D2H sync, which would dominate. No path in nemo_rl puts device tensors on this wire today (`kv_first_write` takes `final_batch_cpu`). - Reading torch privates is the established altitude here — `codec.py` already does `getattr(item, "_non_tensordict", None)` with the same shape one module over — and the risk that `_values` changes meaning is pinned by a unit test rather than a version comment. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Turns `data_plane.observability.enabled` on in the exemplar config. It only engages when `data_plane.enabled` is true, which is itself false by default, so this changes nothing for runs that don't use the data plane — and for runs that do, the cost is below what can be measured. Against a live TransferQueue, three configs of the same workload interleaved (256x2048, 5.4 MB jagged per step, 15 iterations, paired per-iteration deltas so machine drift cancels): metrics off 182.9 ms/step metrics on 183.4 ms/step +0.44% run-to-run spread is 4.4 ms The wrapper itself is 37 us/put and 15 us/get. `verify_tensor_hash` stays off: it costs ~4% and, until fingerprints travel in `KVBatchMeta.tags`, it verifies almost nothing across the process split it would need to cross. Removes the default `log_event` sink, which this change makes mandatory rather than optional. `on_event` fired on every single transfer, so defaulting metrics on would have defaulted a `logger.info` per data-plane op with it (measured 13 us/op with INFO enabled, versus 1.5 us with no sink). The metrics surface is `get_step_metrics()`, which the trainer logs once a step under the `data_plane/` prefix and which fans out to whatever backends the run enabled. A per-op hook is now opt-in via `observability.callback`; `log_event` is still exported for it. This also makes the lazy event-dict build in `_emit` live rather than dead code — with no sink registered, the per-op dict is never allocated. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
`_record_put` was 56% of the put path — a 256-iteration Python loop giving every key an even share of the batch total. There was never a per-key truth to keep: the split was even going in, and a subset clear released the mean coming out. One total plus one live-key set says the same thing and lets `set.update` do the per-key work in C. _record_put 18.6 us -> 3.0 us (6.3x, 256 keys) put path 35.2 us -> 20.0 us Against the PR head that is 199.5 us -> 20.0 us per put and 160.0 -> 15.2 per get. Bytes are released pro rata on clear, and clearing the last key releases the remainder exactly, so a partition reconciles to zero however it is chopped up — verified end to end through live TransferQueue (`bytes_outstanding` back to 0.00 MB, byte totals unchanged). Also trims the per-step metric set from 8 series per op tag to 5. Dropped `mean_ms`, which is exactly `wall_s / calls`, and the absolute `fixed_overhead_ms` / `bandwidth_mb_s`, which are cumulative regressions that barely move per step — `overhead_frac` carries the actionable half (is this op overhead- or bandwidth-bound). `snapshot()` still returns all of it for a one-off inspection; this only changes what is logged every step. At the ~8 op tags a real run produces that is ~70 series down to ~46. Verified in wandb: the metrics reach `Logger.log_metrics(prefix="data_plane")` and fan out to every enabled backend. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Three defects visible the moment the metrics reached a W&B chart. Mixed units on sibling series. `wall_s` sat next to `p99_ms` on the same axis, so a step read `0.008` beside `24.85` and looked like a data-plane bug rather than an axis one. Volumes had it worse one dimension over: a realistic step moved `0.00017` GB. Every duration is now `_ms` and every volume `_mb`, with a test asserting no `_s`/`_gb` series can come back. A flat p50. The latency histogram is cumulative and never reset, so the per-step percentile was a lifetime figure that stopped moving as soon as the distribution settled. A p99 pinned to bucket edges. With the handful of calls an op makes in one step, a percentile off a 15-bucket histogram is geometry, not data: one sample in `(10, 25]` always yields `10 + 15*0.99 = 24.85`, which is exactly the constant that showed up. The bouncing was the sample changing bucket. Per-step percentiles are replaced by `max_ms` — exact, no quantisation, meaningful at one call, and scoped to the step. Scoping matters: the first cut of this used a lifetime max, which is monotonic and goes flat the moment the worst call has been seen, reproducing the defect it replaced. `get_step_metrics` now zeroes the window after reporting. The percentiles stay in `snapshot()`, where the cumulative sample count justifies them. Per-op series go 5 -> 4: calls, wall_ms, max_ms, overhead_frac. Verified against live TransferQueue with runs in the grpo-dev-zhiyul W&B group; `max_ms` now varies step to step while the lifetime worst is still readable from `snapshot()`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Listing the emitted series showed the two scopes were not the same shape, and one of the gaps mattered. `step/hash/*` was emitted only by the driver path, but `_log_data_plane_metrics` prefers the cluster path whenever the fan-out reaches more than one process -- which is every real run. So with `verify_tensor_hash` on, `mismatches` never reached the logger: the guard ran, counted, and logged ERROR lines, while the series anyone would chart or alert on was absent. A guard whose findings are not reported is not a guard. Both paths now go through one `_hash_deltas`, which emits nothing when the guard never ran -- five always-zero series would read as "checked, nothing wrong" rather than "not checked". `step/self/*` was cluster-only for the same reason, so the single-process fallback silently lacked the number that says what observability cost. `step/bytes_written_mb` / `step/bytes_read_mb` went the other way: computed on both paths, dropped by `headline_series`, absent from the table, read by nobody. Removed -- per-op `mb` already splits the same traffic finer, since put's is the write volume and get's is the read volume. Verified: 184 unit tests, 48/48 metric audit checks, 23/23 hash-guard checks. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…olume
A real GRPO run on one GB200 node found that none of these series ever
reached wandb. `_log_data_plane_metrics` ran *after* the
`step_finished=True` log that commits the wandb step, and anything logged
against a committed step is dropped -- so every value was computed,
printed to stdout, and silently discarded. Reading a finished run back out
of the wandb API is what caught it: 85 distinct keys, zero `data_plane/*`.
Every synthetic harness missed it because they call `log_metrics` with no
preceding commit. Now logged before the commit, with a test that asserts
the source order, since the drop happens inside wandb where a fake logger
cannot see it.
Also adds `step/volume_mb/by_op/{op}`. `comm_volume_mb` alone hides the
asymmetry that matters: on that run get moved 20.8 MB against put's
2.7 MB, because every DP rank fetches its shard once for the logprob pass
and again for the train pass. The total cannot say that, and the direction
is what tells you whether to look at reads or writes.
Both are transfers rather than data size -- a byte written and later read
counts on both sides, and each process's transfers are summed -- which the
README now states, along with the scope caveat that makes put read small:
the rollout actor builds its own client, is not on the worker group, and so
its `kv_first_write` of the whole batch is in neither figure.
Verified on a live 5-process run: frac_of_step 1.3% of a 15.4 s step, get
66% of data-plane time, and the breakdown table logged per step.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
A real 5-process GRPO run reported 3584 mismatches per step, identically, on a job that trained fine. All of them were the guard's own doing: zero row-level mismatches, every one from the "written uniform, read ragged" branch. Root cause is the per-field scheme record. Once per-row fields became keys in that mapping, handing the whole mapping to ``_row_fingerprints`` as ``batch_scoped_fields`` forced *every* field onto the batch-scoped path on read, where the comparison then failed because the recorded value is a width rather than a row count. Only the scoped names go through now. That branch was also too harsh on its own terms. Writing a shard with uniform rows and reading it back inside a batch whose other rows differ in length is the normal shape of the pipeline, not a corruption, so failing the whole field cannot be right. The row *lengths* remain comparable, and a row that changed length is a real divergence -- so the width is recorded and checked per row, and the content it can no longer compare counts as an abstention in ``fields_skipped`` rather than being reported as wrong. Verified on a live 5-process run: 1536 rows checked per step, 0 mismatches and 0 abstentions over six consecutive steps. Plus 188 unit tests, 23/23 guard checks and 48/48 audit checks, with a regression test reproducing the false positive directly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…t costs Two complaints against this guard, both fair, both fixable. "It cries wolf." Twice now it has reported every row of every field wrong, identically, every step, and both times the cause was its own bookkeeping rather than the wire. Neither was diagnosable from the log: the digests alone say nothing about which scheme was used or what shape the read was. Per-sample lines now carry the scheme, the row index within the read, and the row length. And a step whose mismatches reach `rows_checked` is called out in the log as more likely a bug in the check than in the wire, because that is what the shape means -- a reader should not have to work that out from a raw count that looks plausible in isolation. "It costs 2.4 ms." That figure was a microbenchmark. Measured against a live TransferQueue moving 23.6 MB per step it is 10.2 ms, +11% of data-plane time -- and under 0.1% of a real GRPO step, where the data plane is a few percent of the whole. The cost is billed to `step/self/overhead_ms` along with the rest of the measurement, so it shows up on the dashboard instead of being quoted from a benchmark run elsewhere. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Four review passes over today's commits. What they found, and what changed:
The tokenizer workaround was at the wrong depth *and* pulled the wrong
lever. `get_tokenizer` is one of four tokenizer constructors, and the
uncovered ones include a live decode path (`trtllm_http_server` calls
`decode` on a tokenizer built in `trtllm_worker_async`). Worse, the defect
is a *missing attribute*, and disabling `clean_up_tokenization_spaces`
instead silently changes decoded text for every non-BPE tokenizer --
transformers only skips cleanup for BPE (`type(model).__name__ == "BPE"`),
so Unigram and WordPiece models were quietly losing it. Moved into
`_patch_transformers_tokenizer_class_set`, which already exists for exactly
this -- a version-pinned transformers defect that must reach every
tokenizer however built -- and which carries the version assert and removal
TODO this change had neither of. Supplying the class-level default
transformers itself uses fixes the crash and changes nobody's output.
`_RESERVED_NAMESPACES` was a deny-list against an open set: any later
`step/<x>/<field>` series would have become a phantom row in the breakdown
table until someone remembered to extend it. Per-op detail now lives under
`step/by_op/<op>/<field>` and is recognised by what it is. Those keys are
table-only -- `headline_series` drops them -- so nothing charted changes.
The per-op assembly was copy-pasted between the two step-metric paths,
reintroducing one level up the drift the shared helpers exist to prevent;
it is now one `_op_series`. Volume was published under two keys; it keeps
the `by_op` one. The write scheme was a positional `("scoped", n)` /
`("rows", width)` tuple whose two arms meant different things in the same
slot, under an annotation that said `int` -- now a `_WriteScheme` NamedTuple
that names both. `row_lens` no longer allocates n_rows copies of one width
on the uniform path, and the uniformity test is a `count` rather than a set
build.
Deepest of them: wandb accepts a log against a committed step, returns
cleanly, and drops it. That silently cost this feature every series it had.
Fixing the call order left the contract as a comment in one of three call
sites, so `WandbLogger` now warns when it happens, for every caller.
Verified: 191 unit tests, 23/23 guard checks, 48/48 audit checks, validate
ALL PASS.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The /simplify pass fixed three things that are real but are not this PR's subject, and carrying them here makes the diff argue for changes a reviewer did not come to review. Out, kept as `tokenizer-decode-fix.patch` for its own PR: the `TokenizersBackend` decode workaround in `nemo_rl/models/policy/__init__.py`. It fixes a transformers defect, not a data-plane one; it happens to be what unblocked the e2e run for this feature, which is not the same as belonging to it. Out, to be raised separately: the `WandbLogger` warning on writes to an already-committed step. Shared logging infrastructure used by ppo, grpo and sft, changed on the strength of a bug found in one caller. The ordering fix in `grpo_sync.py` already prevents the failure here; making the Logger say so for every caller is a wider argument than this diff should make. Out: the `-tq_simple-obs` recipe, which existed to launch my own verification run and is wired into no test suite, so by the repo's own convention it is not a recipe. `log_table` stays -- the breakdown table cannot be logged without it. Verified: 190 unit tests, 23/23 guard checks, 48/48 audit checks. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
CI's lint job runs pyrefly, which I had not run locally -- ruff alone passes on all three of these. `_hash_deltas` builds counter deltas, which are ints, under a signature promising floats. `get_data_plane_snapshot` returns the result of a ``getattr``-ed callable, so its type is `object`, not the `dict[str, Any] | None` it declares. And in `factory`, dropping the `log_event` default left `on_event` as the argument pyrefly objects to, while the existing ignore sat on the constructor line rather than the argument's own. Verified: pyrefly clean across `nemo_rl/data_plane/`, 222 unit tests, 23/23 guard checks, 48/48 audit checks. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The lint job runs three ruff hooks; the second is isort (--select I), which I had not run -- the import block in the observability tests grew a symbol at a time and ended up unsorted. Formatter and linter alone pass on it. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
torch has no bitwise_xor CUDA kernel for UInt64, so XOR-ing the salt into the digest tensor while it is still on the device raises NotImplementedError. Any backend whose get returns device tensors hits it -- TransferQueue register mode under GDR does, and it took the step down at the first verified get. The batch-scoped branch a few lines below already salts after .tolist(); this makes the per-row branch agree. The digests move to the host for comparison either way, so nothing is paid for it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The single-controller path has its own step loop and never called _log_data_plane_metrics, so enabling observability there built the metrics client, paid for its counters on every op, and logged nothing. An empty dashboard reads the same as a data plane that cost nothing, which is the worse of the two failures. Driver scope only, and the prefix says so: this client issues the advantage stage's get and the post-train clear, while the bulk traffic is the trainer and generation workers' own clients in their own processes. comm_volume_mb here is well under what the job moved. grpo_sync prefers a cluster view by fanning out over its policy worker group, but nothing implements collect_data_plane_snapshots yet, so that path also falls through to the driver's counters alone -- there is no cluster view to mirror until one exists. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Rewrite tests/unit/data_plane/test_observability.py (72 -> 50 tests), add a mock e2e test for the data-plane logging path in test_grpo.py, and make the hash guard and metrics panel non-raising so neither can fail a training step. Enable verify_tensor_hash on all 19 TQ nightlies and gate them on hash counters being zero. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Enable data_plane.observability.verify_tensor_hash on all 24 nightlies that
run the data plane, and assert step/hash/{mismatches,guard_failures} are zero
via check_metrics. The guard only counts and logs -- it never raises -- so CI
was the only place a finding could stop anything, and nothing asserted on it.
The assertion accepts either scope: grpo_sync logs data_plane/cluster/* once
the worker fan-out lands, single_controller only ever logs data_plane/driver/*.
Wrappers skip the assertion under TEST_DRYRUN: the delegated base runs in a
subshell, so common.env's dryrun exit does not reach the wrapper tail and
check_metrics would fail on a metrics.json that was never created.
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The affine model was misspecified. It fits wall_ms ~ fixed_ms + n_bytes / bandwidth, which assumes per-request overhead is constant -- but real overhead scales with payload: msgpack encode, the .contiguous() copy taken before the pointer, RDMA memory registration, staging-buffer memcpy. Those costs get split arbitrarily between the intercept and the slope, so overhead_ms understates and bandwidth_mb_s overstates by an unknown margin. More size variation does not fix that; it just makes a wrong split look stable. It also usually declined to answer. GRPO writes the same batch shape every step, so request sizes are near-uniform and the identifiability gate (coefficient of variation < 0.05) returns regime=unidentifiable. Observed on a real RDMA-vs-SimpleStorage pair: put emitted no split on either arm, and get emitted one on RDMA but not on SimpleStorage -- the same workload producing the metric on one transport and not the other, which is useless for comparing them. A 5-step run logged 19 data_plane series and none came from the fit. And it was not free: three float multiply-accumulates per op on the hot path (sum_bytes_sq, sum_bytes_ms, sum_ms_sq) plus ok_wall_ms, retained per rank and merged across processes, to feed a model that mostly abstained. Removes fit_latency_bandwidth, _latency_split, those four accumulators and their merge entries, percent_of_dataplane/by_cause/*, and the two table columns. step/self/overhead_ms is unrelated and stays -- it is a direct measurement of what observability itself cost. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Padding cost was invisible to the per-op metrics, and asymmetrically so.
pack_jagged_fields runs in the caller before put_samples is entered, so it
never reached by_op at all; _from_wire runs inside the adapter's get_samples,
where it was billed as transport. Both are CPU work proportional to payload,
so a transport comparison was crediting the wire for densification.
Adds step/codec/{pack_ms,unpack_ms}, measured with utils.timer and drained by
snapshot() into cumulative counters so they merge across processes and
difference per step through the paths every other counter already uses.
Kept out of total_wall_ms: folding CPU work into it would redefine
frac_of_step and percent_of_dataplane to mean something other than time spent
in the data plane.
The drain is gated on reset_step_window for the same reason step_max_ms is --
it is destructive, so an inspection snapshot must not delete time the step
reader is about to report.
Adds Timer.drain (pop-and-sum under one lock; reduce-then-reset drops a
concurrent record) and a should_log flag on Timer.record, whose _fmt call
otherwise dominates the cost of recording on a hot path.
Coverage matches comm_volume and is documented beside the field: only a
process that drains reports its own pad cost, and the rollout actor is not on
the worker group the fan-out reaches, so kv_first_write is omitted.
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
A real run logged step/wall_ms = 78800.9, which reads as noise beside a
674 s step clock. The rest of nemo_rl already times in seconds --
utils.timer returns perf_counter deltas, timing/setup/*_time_s, and
total_step_time, which is the very denominator frac_of_step divides by --
so the data plane was the outlier.
step/wall_ms -> step/wall_s, step/codec/{pack,unpack}_ms -> _s.
step/self/overhead_ms stays ms: at 73 ms it is sub-second, and it is a
cost-of-measurement figure read against the per-op table rather than
against the step clock. The breakdown table stays ms throughout, where
sub-millisecond percentiles are still legible -- p50_ms=0.011 beats
p50_s=1.1e-05.
No precision cost: the logger stores float32 and floating point is
scale-invariant, so /1e3 moves the decimal point rather than truncating
(measured ~1e-8 relative, against float32's ~1e-7 inherent).
The README previously asserted every duration is _ms with no exceptions,
which is now false; its Units paragraph states the split instead.
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
`verify_tensor_hash` fingerprinted each row with a bare
`torch.hash_tensor` fold. Its only implemented mode is an XOR reduction,
which is blind to anything leaving the multiset of element words alone.
Measured against the old implementation:
- a row zero-padded 8 -> 12 hashed identically (`x ^ 0 == x`)
- `(4, 8)` and `(32)` holding the same elements hashed identically
- a dtype change at equal byte width hashed identically
- a row folded against a duplicate of itself folded to zero
All are plausible codec or wire bugs, and all were silence. The recorded
`row_lens` did not cover the gap either: it held the *flattened* width,
so `(N, 4, 8)` and `(N, 32)` both recorded 32, and it was compared only
on the ragged fallback path, never on the primary one.
Mix the row's dtype and shape into the fold and the first three become
divergences. The values fold is untouched -- only the seed it combines
with is new, so the on-device reduction that made the digest cheap is
kept:
fold = torch.hash_tensor(view.reshape(n_rows, -1), dim=1)
seed = crc32(f"{dtype}|{row_shape}")
digest = fold ^ seed
The shape never travels and is never compared: one integer per row is
stored, as before. A shape change moves the seed, which moves the
digest, which surfaces as an ordinary mismatch.
Most of this diff is deletion. The fold has no ragged kernel, so ragged
leaves used a second, coarser digest over the whole values buffer, plus
a `_WriteScheme` recorded at put and replayed at get so both sides
agreed on which granularity had been used. None of it is needed: XOR is
associative and elementwise, so `hash_tensor(row)` equals
`hash_tensor(rect, dim=1)[i]`, and a ragged leaf can fold one row at a
time and still compare against a vectorized read. `_WriteScheme`,
`_batch_scope`, `_FieldDigest`, `_row_len`, `_dtype_salt`, `_as_int_view`
and the `length_only` branch all go. A shard read is now checked rather
than abstained on, and a ragged divergence names the sample rather than
the batch.
The seed's shape is the *row's*, and both layouts must agree on it: a
dense `(N, L, D)` and the jagged form whose values are `(total, D)` both
report `(L, D)`. That is what lets a field packed jagged and read back
densified reconcile. Deriving the dense row's length from its offsets
instead says `(1, L, D)` and makes every round trip a false mismatch.
Measured on a 107 MB batch of 1536 rows x 4 fields: 8 ms, against 94 ms
for the bare fold, 64 ms for a crc32 over each row's bytes and 146 ms for
blake2b. The sequential hashes cost more because they read every byte on
the host one row at a time; they also catch a within-row permutation,
which this does not. That limit is taken deliberately for the cost,
pinned by a test and recorded in README.md -- swapping `_leaf_digests`
for the crc32 form is a one-function change if a reordering bug is ever
suspected.
Also: hoist the partition lookup above the hashing so a consumer-side
get stops fingerprinting a payload it discards, and refresh the stale
`torch.hash_tensor` / "~2.4 ms per 12 MB" figures in ObservabilityConfig
and the two grpo_math_1B.yaml comments.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…mitted (#4044) Co-authored-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Rebasing onto main dropped the Mooncake 1-D promotion (the constant, the `_promote_1d_leaves` helper and its import), which is intended -- it is not needed. My conflict resolution kept the *use* of PROMOTE_1D_FIELDS inside `_from_wire` while the name it referenced was gone, which would have raised NameError on the first unpack. Caught by a symbol-level audit of the rebase rather than by the checks I ran first: parsing, conflict-marker scans and comparing the core files all passed over it, because the name is only resolved at runtime. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
cluster_step_metrics and get_step_metrics differenced the same counters, differing only in n_procs and the fan-out cost, so the arithmetic moves to one _step_metrics core. The six hash counter names were listed twice in different orders and the comm-volume sum three times; both become one definition. Both trainers hand-copied the same try/except and the same headline+table+print, so those move to metrics_never_fail_the_step and log_step_metrics, which also fixes grpo_sync's print missing flush=True. No metric changes: verified on cluster (n_procs=5) and driver (n_procs=1) paths, hash counters identical to the pre-refactor runs. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…vability wiring Three pre-commit failures from the previous commit, plus two cleanups found while reading the file: * ruff isort wanted collections.abc sorted above contextlib. * ruff-format wanted a blank line before SingleControllerActor. _log_data_plane_metrics. * pyrefly rejected row_shape at the _leaf_digests call: the jagged and dense branches bind lambdas with different capture defaults, and their union is not assignable to (int) -> tuple[int, ...]. Declaring the name before the branch pins it; the tail=tail / shape=shape defaults are deliberate and stay. * percentile_from_hist is public-named with no caller outside this module (its only user is _clamped_percentiles, and it is not re-exported from nemo_rl.data_plane), so it becomes _percentile_from_hist. * Dropped bytes_f in _emit -- assigned and never read. No behaviour change: the annotation is not evaluated under from __future__ import annotations, and the removed assignment had no reader. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/data_plane_observability_metrics
branch
from
September 9, 2026 09:58
4ea3bc7 to
2d5719b
Compare
Contributor
Author
|
/ok to test 2d5719b |
- `fold ^ seed`: `Tensor.item()` widens to `int | float`; a hash digest is integral and XOR is int-only, so cast at the use. - the two `row_shape` lambdas bound `tail`/`shape` as default arguments, widening the signature past `Callable[[int], tuple[int, ...]]`. Each is consumed by the `_leaf_digests` call in the same loop iteration, so the late-binding guard was never needed. - `td.items()`: pyrefly cannot narrow tensordict's recursive union when breaking cycles; suppressed in the module's existing idiom. pyrefly clean; tests/unit/data_plane 344 passed, 3 skipped. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Contributor
Author
|
/ok to test 05348d2 |
zyzhou5
reviewed
Sep 9, 2026
| reuse_registered_buffers: true | ||
| staging_buffer_size: 268435456 | ||
| observability: | ||
| verify_tensor_hash: true |
Contributor
There was a problem hiding this comment.
Should we turn this off for normal recipes and only turn it on in the automated test launch script?
Contributor
There was a problem hiding this comment.
wondering whether having verification to all those tests will make any obvious difference to the time spent on automated test.
| Reading a real TransferQueue step: | ||
|
|
||
| ``` | ||
| step/frac_of_step 0.074 the data plane is 7% of the step |
Contributor
There was a problem hiding this comment.
Should we clarify this might overlap with other computation so its not like 7% of the step time is spent on tq exclusively?
| view = view.view(_INT_VIEW_BY_WIDTH[view.element_size()]) | ||
| lengths = [hi - lo for lo, hi in zip(bounds, bounds[1:])] | ||
| if lengths and lengths.count(lengths[0]) == len(lengths): | ||
| folds = torch.hash_tensor(view.reshape(len(lengths), -1), dim=1).tolist() |
Contributor
There was a problem hiding this comment.
Agent's suggestion: XOR cancels repeated values in pairs, so [1,1,1,1] and [0,0,0,0] produce the same fingerprint for the same dtype and shape. GRPO advantages repeat a scalar across tokens, so zeroing an even-length advantage row can pass verification.
Contributor
Author
There was a problem hiding this comment.
Good catch — the parity blindness is real, but I'd keep the XOR fold.
The false-positive rate is 0. The digest is a deterministic function of exact bits plus a dtype|row_shape salt, so a mismatch is always a real divergence, never a flake. That is the property this debug check is bought for.
The miss rate for your case is bounded by row-length parity. A row that is one scalar over n real tokens folds to v for odd n and 0 for even n — pad zeros are XOR-neutral, so it is the real token count that sets the parity, and that varies per sample:
| corrupted | missed with prob |
|---|---|
| bit flip / stale buffer / wrong shard / truncation | ~2⁻³² per fp32 row |
| advantages zeroed batch-wide | 2⁻ᴺ, N = rows (≈2⁻⁵¹² at 512/step) |
| advantages zeroed on one sample | ½ |
| permutation within a row | 1 (documented, deliberate) |
Batch-wide corruption — what a transfer-layer bug actually looks like — needs every row to be even-length at once. Single-row at 50% is the genuine hole; I judged that acceptable for a debug-only check.
If we want it closed, it is one extra reduction, not a new algorithm: fold rect.sum(dim=1) beside the XOR. Sum is commutative and associative like XOR, so hash_tensor(row) == hash_tensor(rect, dim=1)[i] still holds and the jagged↔dense reconciliation survives, but a constant row now gives n*v instead of 0.
Either way I'll extend the README note so it records duplicate cancellation next to the permutation blindness.
| digests = self._row_fingerprints(fields, sample_ids) | ||
| if not digests: | ||
| return | ||
| partition_hashes = self._hash_by_partition.setdefault(partition_id, {}) |
Contributor
There was a problem hiding this comment.
Agent's suggestion: _hash_by_partition belongs to each client, but step cleanup runs through the driver’s client. The rollout actor and policy workers therefore retain hashes for samples.
Could you doublecheck this? I think this could lead to perf issue.
Contributor
Author
There was a problem hiding this comment.
Confirmed, and it was wider than the hashes. Fixed in 47c27ca.
_record_clear only fires in the process that issues the clear, and per single_controller.py:21-36 only SC ever calls it — GenWorker and the value actor put through their own clients. So all three per-partition stores kept every uid those processes ever wrote. _keys_by_partition and _bytes_by_partition leak with verify_tensor_hash off, which is the default, so now/bytes_outstanding_mb and n_keys_outstanding also stop meaning occupancy on those processes.
The fix, in _record_put so it runs whatever the guard flag says:
live = set(self._inner.list_sample_ids(partition_id))
stale = self._keys_by_partition.get(partition_id, set()) - live
if stale:
self._record_clear(partition_id, list(stale))Throttled to one listing per 16384 rows put, so a client that clears its own writes never makes the call. list_sample_ids is metadata-only and documented for reconciliation (interfaces.py:557). Test covers both guard states: writer puts two batches, another process clears one, writer's accounting drops exactly those uids on its next put.
… sample The wire-hash guard records a fingerprint per row on put and releases it on clear_samples, but on the SingleController path GenWorker and the value actor put through their own clients while only SC clears, so their stores grew for the life of the run. Reconcile against list_sample_ids -- metadata-only, and documented for this -- on a row-count throttle, and run the release through _record_clear so the byte and key accounting follows the same uids. A client that clears normally stays below the throttle and never makes the call. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster>
… cleared Follow-up to 9452460, which hung the release off the hash guard and so only ran with verify_tensor_hash on. The leak is in the byte and key accounting, which is on by default: _record_clear only fires in the process that issues the clear, and on the SC path GenWorker and the value actor put through their own clients while only SC clears. Move the trigger to _record_put and diff against _keys_by_partition, the store every put populates. All three per-partition stores are then released by the one rule a real clear uses, on every client, whatever the guard flag says. A failing list_sample_ids no longer reaches the caller: the put already succeeded, so the release is skipped, logged once, and retried next window. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster>
The one-shot log and its flag were scaffolding around a four-line fix. A failed listing now just skips the release; the next window retries it. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster>
The guard folded a digest on put and re-folded it on get, but held the wire-in reading in the putting process. That verifies a same-process round trip only, and the transfers worth checking are not those: the rollout actor writes what the policy workers read, so every one of those rows landed in rows_unverified and mismatches stayed 0 for want of anything to compare. Hold the readings in one actor instead. Put records, a reader with no local reading fetches -- one fetch per get, not per row. Digests are 8 bytes per row per field, so that is ~100 kB against a get already moving tens of MB. The comparison stays in the reading client so HashStats keeps counting where the transfer happened. Off Ray the handle is None and the local store answers every lookup, which is what single-process callers and the unit tests do. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster>
This reverts commit 332f5ca. The actor's store is not persisted, so after load_checkpoint every read reports rows_unverified and the guard silently stops guarding -- the same failure it was added to fix, one level up. It also put a blocking ray.get on every get, serialised through one actor across every DP rank. Carrying the digest in TQ beside the row it describes does not have either property. Reverting so that lands on a clean base. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this adds
Per-step visibility into what the TransferQueue data plane costs, so "is the
data plane my bottleneck, and which part of it" is answerable from a dashboard
instead of a profiler.
Off by default at the top level (
data_plane.enabled: false); when the dataplane is on,
observability.enabled: truewraps the client inMetricsDataPlaneClient. Measured overhead is ~15-20 µs per op, and thewrapper bills itself so that cost is visible rather than asserted.
The metrics
Everything is logged under
data_plane/cluster/(the driver plus every policyworker, summed) or
data_plane/driver/when the fan-out reaches one process.step/frac_of_stepstep/percent_of_dataplane/by_op/{put,get,clear,register}step/percent_of_dataplane/by_cause/{fixed_overhead,transfer}step/volume_mb/by_op/{get,put}step/wall_ms,step/comm_volume_mbnow/bytes_outstanding_mbstep/self/{overhead_ms,frac}step/hash/*verify_tensor_hashonPlus a per-op breakdown table (
data_plane/{cluster,driver}/breakdown),ordered worst-first, carrying the detail that would be noise as 32 line charts:
calls, mean/max/p50/p90 ms, the overhead/transfer split, and MB.
Two denominators, deliberately.
frac_of_stepdivides by the step;percent_of_dataplanedivides by the data plane. A workload can be 43%putand still not be worth touching — you need both.
Reading the numbers correctly
Three things are easy to misread, and are documented in
nemo_rl/data_plane/README.md:both sides, and every process's transfers are summed. Correct for "what
crossed the wire"; useless for "how big was the batch".
putreads small. The rollout actor builds its own client and is not onthe worker group, so its
kv_first_writeof the whole batch is in neithervolume_mbnorcomm_volume_mb.by_causeis often absent. Request sizes in RL are frequently uniform,which makes the affine fit unidentifiable. It reports nothing rather than an
arbitrary split.
Optional wire-hash guard
verify_tensor_hash: falseby default. When on, fingerprints every tensor rowwith
torch.hash_tensoron the way in and out and reports divergence — the onefailure class nothing else notices, because a mis-sharded read trains happily
on the wrong rows. Catches a row served from the wrong sample, two rows swapped
between DP ranks, truncation, a bf16→fp32 change, and a single element in any
dtype. Costs 10.2 ms/step on a 23.6 MB step (+11% of data-plane time, under
0.1% end to end), billed to
step/self/overhead_ms.Blind spot, measured rather than assumed: a batch-scoped digest is an XOR
reduction, so it cannot see a permutation within a row (0/200 for a two-token
swap). Rows of uniform width get true per-row digests and catch it.
Verification
222 unit tests, plus 23/23 adversarial hash-guard checks (10 corruption
modes caught, zero false alarms over a 500-row soak, every shard grouping)
and 48/48 metric-audit checks against an independently computed ground truth.
A real 25-step GRPO run — Llama-3.2-1B on OpenMathInstruct-2, 1 node ×
4 GB200, TransferQueue
simple, 5 processes reporting:https://wandb.ai/nvidia/nemo-rl-dataplane-obs/runs/w316chaj
frac_of_step1.3% of a 15.4 s step ·get65% of data-plane time · 24 MBmoved per step · 38,400 rows hash-checked, 0 mismatches.
Notes for review
needs 40 (~4 observations above the rank). Below that the key is absent
rather than reporting the maximum under a percentile's name. p90 rather than
p99 because a step holds tens of calls — a p99 off 58 samples equalled the
maximum 80% of the time.
max_msis scoped to the step by being reset by its reader, since amaximum cannot be differenced out of a cumulative counter.
observability.pyalready on main.