[ZhiyuLi-Nvidia]

What

Two cheaper paths on the single-threaded TransferQueue controller. Both are overrides in a controller subclass, because the controller is a Ray actor.

1. kv_retrieve_meta reads a torch tensor once per field (col_mask[col_idx]). Each read costs microseconds, which is about a quarter of a put+get cycle. The override builds the field mask with one .all(dim=0).tolist(). Results are identical.
2. FieldMeta.remove_samples rescans the shape of every remaining sample of a nested field on each clear, only to see whether the field can collapse to non-nested. The cost grows with the number of live samples, so CLEAR_META slows down late in a run. The override rescans only after the field has lost RESCAN_REMOVED_FRACTION (1/8) of its remaining samples. In a large GRPO run (about 131k rollouts per step), CLEAR_META averaged 0.3 s and rose to 1.3 s late in the run, and NOTIFY_DATA_UPDATE acknowledgements queue behind it.

tq.init starts the subclass: configure_controller_patch() runs right before it. Each override is a plain copy of the stock method. PatchedController.__init__ applies the FieldMeta override inside the controller process only, so the driver keeps stock FieldMeta. Building the subclass checks transfer_queue.__version__ against STOCK_TQ_VERSION (0.1.11, the pinned commit), so a TQ pin bump fails loudly instead of silently running a stale copy. The error says what to do: delete the module if the new TQ has the upstream change, otherwise diff the two methods and bump STOCK_TQ_VERSION.

Experiment results: kv_retrieve_meta

In-process cycle cost at the pinned TQ, without any other controller patch. One cycle is kv_retrieve_meta(create) for one new key, update_production_status with 16 fields, kv_retrieve_meta(get), and set_custom_meta. 1 torch thread, 20k live rows, 3000 cycles per run, 7 interleaved runs, median:

tree	µs per cycle	runs
head	325	321-327
head + this patch	218 (−33%)	213-223

The earlier ZMQ benchmark below also included a geometric table-growth patch, which is not in this PR or in NeMo-RL. It shows the effect on top of that patch.

ZMQ benchmark with the growth patch (96 clients, 4 requests per cycle)

One TQ controller (1 torch thread, as a num_cpus=1 Ray actor gets), real ZMQ, 96 closed-loop client processes, 20k live rows, 16 fields. Each client cycle is KV_RETRIEVE_META(create), NOTIFY_DATA_UPDATE, KV_RETRIEVE_META(get), SET_CUSTOM_META. 15 s measured window per run, 2 repetitions, uncapped 90-core node.

tree	achieved @ 8k target	controller CPU @ 4k target	CPU per request
head	4,594	90%, 99%	~236 µs
growth-only	5,620	83%, 83%	~208 µs
growth + this patch	6,469 (+15% vs growth-only)	73%, 74%	~184 µs

This run used the earlier mask form (sum(dim=0) == n, then .tolist()). The PR uses .all(dim=0).tolist(), which gives identical results on 0/1 status.

Experiment results: FieldMeta.remove_samples

One nested field with all-distinct shapes (the worst case), 16-row clears, 2000 clears per size, TQ at the pinned commit:

Live samples	Stock	This PR
20k	0.59 ms	0.010 ms
100k	10.3 ms	0.018 ms
300k	50.8 ms	0.003 ms*

* 2000 clears × 16 rows removes 32k samples, which is under 268k ÷ 8, so this row never reaches a rescan. Its steady state is about 0.02 ms, like the 100k row.

These are per field, and a recipe has several nested fields. A rescan still costs O(live samples), but it runs once per eighth of the samples removed, so the average cost per clear stays constant. A field with few samples still rescans on every clear, because one removal is already an eighth of what is left.

Behavior change

A nested field whose shapes have become uniform can stay nested until the next rescan. Readers handle both forms: BatchMeta.get_shapes uses per_sample_shapes when present and the single shape otherwise.

FieldMeta is pickled by save_checkpoint, so removed_since_scan becomes part of a controller checkpoint. It is a class attribute with default 0, so checkpoints load in both directions.

Tests

tests/unit/data_plane/test_tq_controller_patch.py has 8 tests:

• The installed TQ is STOCK_TQ_VERSION, and another version is rejected.
• kv_retrieve_meta returns the same indexes and fields as stock on partial-field, missing-key, and empty-key cases, and with create=True on new and mixed keys (same key mapping afterwards).
• remove_samples keeps a field nested until an eighth of its samples are gone, and serves the same per-sample shapes as stock after every clear.
• The patched controller carries the kv_retrieve_meta override, starting it installs FieldMeta.remove_samples, and calling configure_controller_patch() twice wraps once.

The tests ran against TQ at the pinned commit. ruff 0.9.9, pyrefly 0.24.2, and detect-secrets pass on the changed files.

Notes for review

• Files: adapters/tq_controller_patch.py (new), a 2-line hook in adapters/transfer_queue.py, pyrefly.toml, and the test.
• The copy logs through TQ's get_logger, so TQ_LOGGING_LEVEL controls its lines like stock.
• The overrides copy about 100 lines of TQ code. The remove_samples change is on perf/throttle-nested-shape-rescan. Delete this module once TQ takes both changes.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ZhiyuLi-Nvidia]

/ok to test ed27504

[zyzhou5]

LGTM!

[ZhiyuLi-Nvidia]

/ok to test fc9191b

[terrykong]

Thanks for the careful benchmarks and the clear write-up. This review ran both overrides against the stock TQ methods at the pinned commit: the field mask on empty, missing and repeated keys, 300 random clear/add sequences for remove_samples, a real Ray actor, and controller checkpoints loaded in both directions. The results matched every time. No correctness problems found. The CI run on fc9191b tested the current head, and L0 runs all 6 new tests.

Inline:

1. (C1) Add tests for the copied create=True branch and for calling configure_controller_patch() twice. Please add these in this PR.
2. (C2) Open the upstream TQ PR, link it as the condition for deleting the module, and reword the version error.
3. (C3) Use TQ's get_logger in the copy. Low priority.

PR description — please update:

• It still describes a sha256 check of each stock method and two hash tests. The head commit fc9191b replaced both with the transfer_queue.__version__ check. It also says four commits; there are six.
• The "+41% vs head" row includes a geometric-growth patch that is not in this PR or in NeMo-RL. Please add a "head + this patch" row (the in-process profile is enough).
• Please footnote that the 300k remove_samples row never reaches a rescan: 2000 clears × 16 rows = 32k removed, which is under 268k ÷ 8. So its steady state is about 0.02 ms, not 0.003 ms. The conclusion still holds.

On the naming question: tq_controller_patch.py fits better. Both overrides run on the controller, the entry point is already configure_controller_patch, and it matches tq_mooncake_checkpoint.py.

CI: gb200_L1_Functional_Tests_SFT failed because the runner lost contact with the server, so no test ran. "CI quality check" failed only because of that job. A rerun of that job should clear both.

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4769-cheaper-controller.html

Generated by Claude Code

[terrykong]

generally lgtm module the nits

[ZhiyuLi-Nvidia]

/ok to test 612cdfc