feat(RL/post-2389): MX V2 Support in RL + WeightTransferEngine Support#349
Closed
KavinKrishnan wants to merge 6 commits into
Closed
feat(RL/post-2389): MX V2 Support in RL + WeightTransferEngine Support#349KavinKrishnan wants to merge 6 commits into
KavinKrishnan wants to merge 6 commits into
Conversation
4 tasks
5 tasks
KavinKrishnan
force-pushed
the
kavink/post-2389-phase3-4
branch
from
June 4, 2026 00:21
c6d19f3 to
1a3c4a8
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
KavinKrishnan
changed the title
4 tasks
KavinKrishnan
marked this pull request as ready for review
June 26, 2026 00:24
Contributor
Author
|
Ready for review — conflicts resolved, all tests green. Merged
Smoke:
Pre-merge state preserved on The PR description has been refreshed with the current dependency context — several downstream PRs (#421, #429, #450, ai-dynamo/dynamo#10900, jthomson04/RL#2 / #7) have been waiting on this to reach main. |
KavinKrishnan
added a commit
that referenced
this pull request
Jun 26, 2026
…uctors
The merge from origin/main pulled in two new constructors of
ModelMetadataRecord that did not exist when this branch was opened,
but the textual merge did not detect that this branch's added
``identity: Option<SourceIdentity>`` field needed to be supplied.
- modelexpress_server/src/p2p/backend/memory.rs::get_metadata
- modelexpress_server/src/p2p/service.rs test
test_get_metadata_preserves_artifact_source_status
Both default to identity: None (matches main's behavior pre-#349:
no identity round-tripped through these paths). Production identity
plumbing is unchanged and still flows through state.rs/kubernetes.rs/
redis.rs which were already correctly merged.
Contributor
WalkthroughAdds v2 registry metadata and server-side identity propagation, new training and receiving helpers, v2 wrappers and demo scripts, and benchmark tooling for local, Kubernetes, and sample-result runs. ChangesV2 transport and benchmark stack
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 17
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
🟡 Minor comments (6)
modelexpress_client/python/modelexpress/refit_receiver.py-432-437 (1)
432-437: 🗄️ Data Integrity & Integration | 🟡 MinorReject unknown descriptor dtypes in
receive_weights_scratch()
Accept only the explicitbf16/f16/f32mappings and raise on anything else instead of defaulting tobfloat16; keep thesize % elem_size == 0check so scratch buffers can’t be mis-sized.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/refit_receiver.py` around lines 432 - 437, In receive_weights_scratch(), stop defaulting unknown descriptor dtypes to torch.bfloat16: validate td.dtype against the explicit _DTYPE_MAP entries (bf16, f16, f32) and raise an error for anything else before allocating scratch_tensors. Keep the existing size % elem_size == 0 guard in the same flow so tensor sizes are only used when they exactly match the selected dtype’s element size, using the td, _DTYPE_MAP, and scratch_tensors logic as the key locations.modelexpress_client/python/tests/test_v2_shape_registry.py-118-122 (1)
118-122: 🎯 Functional Correctness | 🟡 MinorUse a sequence with duplicates here
{3, 1, 2, 5, 4, 5, 1}is deduplicated beforeencode_expert_set()runs, so this only covers already-unique input. Use a list or tuple with repeated values instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/tests/test_v2_shape_registry.py` around lines 118 - 122, The round-trip test for encode_expert_set/decode_expert_set is using a set literal, which removes duplicates before the codec is exercised. Update test_expert_set_codec_round_trip to use a sequence type such as a list or tuple with repeated values so encode_expert_set is verified against duplicate input, while keeping the same expected canonical encoded output and decoded set.Source: Linters/SAST tools
modelexpress_client/python/tests/test_vllm_weight_transfer.py-269-269 (1)
269-269: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winPrefix unused fixture unpacking to keep Ruff green.
v2/sdare unused at Line 269,sdis unused at Line 305, andv2is unused at Line 498. Rename them to_v2/_sdor unpack with_.Also applies to: 305-305, 498-498
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/tests/test_vllm_weight_transfer.py` at line 269, The test unpacking in test_vllm_weight_transfer is creating unused local variables that Ruff flags, so update the affected destructuring assignments to mark the unused items as intentionally ignored. In the test cases around the vllm_wt unpacking, rename the unused bindings like v2 and sd to underscore-prefixed names or replace them with _. Keep the meaningful variables unchanged and apply the same cleanup at each affected unpacking site in the test file.Source: Linters/SAST tools
modelexpress_client/python/tests/test_vllm_weight_transfer.py-43-43 (1)
43-43: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winRestore
MX_WEIGHT_TRANSFER_AUTOREGISTERin fixture teardown.This module-scoped fixture mutates process-wide environment and never restores it, so tests that run later can accidentally see auto-registration disabled.
Proposed fix
- os.environ["MX_WEIGHT_TRANSFER_AUTOREGISTER"] = "0" + old_autoregister = os.environ.get("MX_WEIGHT_TRANSFER_AUTOREGISTER") + os.environ["MX_WEIGHT_TRANSFER_AUTOREGISTER"] = "0" ... finally: + if old_autoregister is None: + os.environ.pop("MX_WEIGHT_TRANSFER_AUTOREGISTER", None) + else: + os.environ["MX_WEIGHT_TRANSFER_AUTOREGISTER"] = old_autoregister for name in _injected: sys.modules.pop(name, None)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/tests/test_vllm_weight_transfer.py` at line 43, The module-scoped fixture in test_vllm_weight_transfer.py sets MX_WEIGHT_TRANSFER_AUTOREGISTER but does not restore the original process environment afterward. Update the fixture teardown to save the previous value before mutating os.environ and restore or delete MX_WEIGHT_TRANSFER_AUTOREGISTER when the fixture finishes, using the fixture around the vllm weight transfer tests as the place to fix it.modelexpress_client/python/benchmarks/README.md-79-88 (1)
79-88: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the cluster-runner details in this README.
Line 81 says the Job needs 5 GPUs and Lines 86-88 point readers at
modelexpress/benchmarks/bench_elastic_scaling.py, but the supplied runner/manifest launch 1 trainer + 3 receivers (4 GPUs total) and execute the standalonebenchmarks/bench_elastic_scaling.pyoverlay. Following this README will send users to the wrong quota request and wrong path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/benchmarks/README.md` around lines 79 - 88, The cluster-runner README details are incorrect: the manifest/job uses 1 trainer + 3 receivers for 4 GPUs total, not 5, and the benchmark entrypoint is the standalone `benchmarks/bench_elastic_scaling.py` overlay, not `modelexpress/benchmarks/bench_elastic_scaling.py`. Update the GPU quota guidance in the README to match the actual job definition, and fix the harness path reference so readers are directed to the runner used by the manifest.modelexpress_client/python/benchmarks/results-20260530-105218/tree_fanout.json-162-169 (1)
162-169: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRegenerate this fixture from an actual fan-out run.
fanout_factoris still1.0andtrainer_egress_bytesequalstotal_delivered_bytes, so this sample does not demonstrate the tree_fanout behavior described inbench_elastic_scaling.py:675-748. Please re-record it from a run that actually relays some transfers through earlier receivers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/benchmarks/results-20260530-105218/tree_fanout.json` around lines 162 - 169, The tree_fanout fixture was captured from a run that does not actually show fan-out behavior, since fanout_factor remains 1.0 and trainer_egress_bytes matches total_delivered_bytes. Re-record the benchmark output using the tree_fanout path in bench_elastic_scaling.py so that some transfers are relayed through earlier receivers, then update the derived values in the tree_fanout.json fixture to reflect the real run.
🧹 Nitpick comments (2)
modelexpress_server/src/p2p/service.rs (1)
487-529: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a positive
identityround-trip assertion here.This PR's new contract is
GetMetadataResponse.identity, but the happy-path test still buildsidentity: None. A regression that dropsextra_parametersorrevisionwould still pass, even though downstream receivers now readmeta.identity.extra_parametersfor version filtering.Suggested fix
Ok(Some(ModelMetadataRecord { source_id: source_id.to_string(), worker_id: worker_id.to_string(), model_name: "my-model".to_string(), workers: vec![WorkerRecord { @@ }], published_at: 1234567890, - identity: None, + identity: Some(test_identity()), })) }); @@ assert_eq!(resp.mx_source_id, "abc123def456abcd"); assert_eq!(resp.worker_id, "worker-uuid-1"); + let identity = resp.identity.expect("identity should be present"); + assert_eq!(identity.model_name, "my-model");🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_server/src/p2p/service.rs` around lines 487 - 529, The happy-path test in test_get_metadata_found currently sets ModelMetadataRecord.identity to None, so it does not verify the new GetMetadataResponse.identity round-trip; update the mock record to include a real identity value and assert that svc.get_metadata returns the expected identity fields. Use the existing get_metadata test setup and the ModelMetadataRecord/GetMetadataResponse types to confirm extra_parameters and revision are preserved end to end.modelexpress_server/src/p2p/backend/redis.rs (1)
84-133: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd a round-trip test for
SourceAttributesJson.
From<&SourceIdentity>andto_source_identity()now have to stay in lockstep across a large field set. A small round-trip test with a fully populatedSourceIdentitywould catch the next silently dropped field before it breaksGetMetadataResponse.identity.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_server/src/p2p/backend/redis.rs` around lines 84 - 133, Add a round-trip test for SourceAttributesJson to keep From<&SourceIdentity> and to_source_identity() in sync across all fields. Create a fully populated SourceIdentity, convert it into SourceAttributesJson, convert it back with to_source_identity(), and assert the result matches the original so any future dropped or mismapped field is caught.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelexpress_client/python/benchmarks/bench_elastic_scaling.py`:
- Around line 159-168: trainer_egress_bytes in bench_elastic_scaling.py
currently infers “trainer vs replica” from source_worker_rank == 0, but
tree-fanout receivers are launched with worker_rank=0 so replica pulls are
miscounted as trainer egress. Update the accounting to use an unambiguous signal
from the receiver/cycle data (for example, a dedicated trainer-origin flag or an
actual origin rank) and then adjust the tree-fanout launch path in the receiver
setup so republished replicas are distinguishable from the trainer. Keep the
logic in trainer_egress_bytes and the tree-fanout receiver creation code
consistent so fanout_factor reflects real replica fan-out.
- Around line 400-418: The benchmark in receive_weights is unnecessarily
retaining every loaded tensor in the captured list via captured.extend, which
can accumulate full model state across versions and cause OOMs. Update the
receive path in bench_elastic_scaling.py so the load_weights callback discards
tensors immediately instead of appending them, and remove the now-unused
captured collection if nothing else consumes it. Keep the behavior centered
around ReceiverCycleResult, engine.receive_weights, and the load_weights
callback.
In `@modelexpress_client/python/benchmarks/k8s/bench-elastic.yaml`:
- Around line 36-203: The benchmark Job is still running with the default root
security context. Update the pod spec around the bench container to add explicit
securityContext settings such as runAsNonRoot, allowPrivilegeEscalation false,
and dropped capabilities, using the bench container and spec fields as the main
anchors. If you turn on readOnlyRootFilesystem, also add a writable emptyDir
mount for /tmp since the overlay/bootstrap logic in the container command writes
to /tmp and /results.
- Around line 120-158: The benchmark scenarios in the elastic k8s script are
swallowing failures because the commands in the scenario blocks are followed by
echo-only fallbacks, which can let the Job succeed despite a broken run. Update
the Scenario 1/2/3 command sequences in the harness invocation block so that
failures are logged and then propagated with a non-zero exit status, using the
existing $HARNESS execution flow and scenario labels to keep the error handling
tied to each scenario.
In `@modelexpress_client/python/modelexpress/nemo_rl_v2.py`:
- Around line 282-286: The publish path in nemo_rl_v2.py is accumulating stale
entries because the descriptor list is append-only while _registered_tensors is
keyed by name. Update the registry handling around the publisher flow that uses
_registry and _registered_tensors so each publish cycle either resets
per-version state or replaces any existing descriptor for the same tensor name
before appending. Make sure the long-lived publisher behavior in the tensor
registration logic stays consistent with the shape_registry contents and does
not retain duplicate or old descriptors across steps.
- Around line 843-848: The expert selection logic in the candidate scan returns
a trainer rank unconditionally, which can violate the required candidate-owned
superset for sharded MoE tensors. Update the selection path in the candidate
loop so ROLE_TRAINER is only chosen when it actually owns all experts required
by needed_experts_per_layer, and otherwise continue searching for a valid
candidate. Use the existing ownership/ranking data in nemo_rl_v2.py to validate
trainer ownership before returning, so the chosen rank always matches the
requested shard.
- Around line 1149-1151: The default tensor-parallel slice logic in nemo_rl_v2
should not silently truncate non-divisible axes by using axis_total //
target_layout.world_size. Update the slicing path around the chunk/t_start/t_end
calculation to validate that axis_total is evenly divisible by
target_layout.world_size before deriving the default range when target_range is
omitted. If it is not divisible, reject the request with a clear error instead
of computing a truncated slice, and keep the check localized near the existing
target_layout rank-based range computation.
- Line 286: The add_tensor path is storing the DTensor wrapper in
_registered_tensors, but NIXL publication should keep the local shard instead.
Update the add_tensor implementation in nemo_rl_v2.py so that when a DTensor is
passed for descriptor inference, the value saved into _registered_tensors[name]
is the underlying local tensor/shard rather than the DTensor object itself; use
the existing add_tensor and _registered_tensors symbols to locate and adjust the
registration logic.
In `@modelexpress_client/python/modelexpress/refit_receiver.py`:
- Around line 496-505: The direct metadata path in
receive_weights_from_metadata() is missing the same sidecar filtering already
applied in receive_weights() and receive_weights_scratch(), which can let __mx_
zero-size tensors reach NIXL and trigger the invalid (0,0,0) descriptor failure.
Update the metadata-to-NIXL flow in RefitReceiver to filter out those sidecars
before calling self._nixl.receive_from_source(), reusing the existing sidecar
check logic used by the other receive methods so all three paths behave
consistently.
- Around line 341-350: The NIXL receive tuple is being unpacked with the wrong
meaning, so tensor counts are recorded as skipped bytes. In the receive paths
that call self._nixl.receive_from_source() and build TransferStats, treat the
second return value as the matched/received tensor count, not bytes_skipped, and
update the stats fields accordingly so tensors_received reflects that value
while bytes_skipped only uses an actual skipped-bytes source if available. Apply
the same correction in all duplicated receive blocks referenced by the repeated
TransferStats construction.
In `@modelexpress_client/python/modelexpress/shape_descriptors.py`:
- Around line 228-243: The shard range computation in shape_descriptors’
local_shard_range logic assumes uniform DTensor shards, so it can publish
incorrect ranges for uneven layouts. Update the DTensor handling path in the
shape descriptor code to detect non-uniform sharding before computing start/end,
and raise an unsupported/validation error instead of emitting a descriptor when
the shard extents differ. Keep the fix localized around the _RealDTensor branch
and the local_extent/start/end calculation.
In `@modelexpress_client/python/modelexpress/training_publisher.py`:
- Around line 310-318: The shutdown path in training_publisher.py does not clear
the registration flag, so a later reinitialize can reuse stale state and skip
descriptor registration in publish_weights(). Update the shutdown() method on
the relevant publisher class to reset _registered along with _initialized after
closing _nixl and _client, so a fresh NixlTransferManager always re-registers
before publishing.
- Around line 283-308: `publish_layer()` currently overwrites
`self._mx_source_id`, so `mark_ready()` only marks the last published layer as
ready. Update `TrainingPublisher` to retain every `mx_source_id` returned during
the current step (for example, in a per-step collection) and have `mark_ready()`
iterate over all of them and call `_client.update_status()` for each one. Keep
`self._mx_source_id` available for the existing no-publish guard, but ensure all
layer publishes are transitioned to `SOURCE_STATUS_READY`.
In `@modelexpress_client/python/modelexpress/vllm_weight_transfer.py`:
- Around line 309-314: The mixed-TP slice discovery path in
vllm_weight_transfer.py is incorrectly constraining source lookup to the same
rank by default. Update the call to discover_v2_sources_for_slice in the
receiver logic so mixed-TP assembly allows cross-rank discovery unless
same-rank-only is explicitly required, and verify the default passed from
update_info.same_rank_only does not block adjacent trainer ranks for
target_tp_layout-based slices.
- Around line 360-365: The tree fan-out path is effectively disabled after
scratch receives because the receiver is initialized without model tensors,
leaving _registered_buffers empty so publish_self_as_source() has nothing to
advertise. Update the vllm_weight_transfer flow around
self._init_info.publish_self_as_replica and
self._receiver.publish_self_as_source so the receiver is populated with the
needed tensor/buffer metadata before publishing, or otherwise ensure
publish_self_as_source can register and emit replica sources for
scratch-initialized models.
- Around line 490-491: Reorder the base classes in _MxEngineForVllm so
MxWeightTransferEngine comes before _VllmWeightTransferEngineBase in the
inheritance list. This fixes the MRO so abstract methods from the vLLM base do
not keep the subclass abstract and any overlapping method names resolve to the
MX implementation first; update only the _MxEngineForVllm class definition.
In `@modelexpress_client/python/tests/test_bench_elastic_scaling.py`:
- Around line 100-125: The tree-fanout test fixtures are modeling an impossible
receiver/source shape for the current harness, so update the affected cases to
match what run_tree_fanout() actually records. Adjust the BenchResult
construction in test_trainer_egress_under_tree_fanout_can_be_less and the
related tree-fanout tests so they assert against the real unique source
identifier and receiver worker_rank values produced by the harness, rather than
hardcoding replica pulls with source_worker_rank=1 when all spawned receivers
are worker_rank=0. Keep the egress assertions aligned with the actual data shape
returned by the harness.
---
Minor comments:
In `@modelexpress_client/python/benchmarks/README.md`:
- Around line 79-88: The cluster-runner README details are incorrect: the
manifest/job uses 1 trainer + 3 receivers for 4 GPUs total, not 5, and the
benchmark entrypoint is the standalone `benchmarks/bench_elastic_scaling.py`
overlay, not `modelexpress/benchmarks/bench_elastic_scaling.py`. Update the GPU
quota guidance in the README to match the actual job definition, and fix the
harness path reference so readers are directed to the runner used by the
manifest.
In
`@modelexpress_client/python/benchmarks/results-20260530-105218/tree_fanout.json`:
- Around line 162-169: The tree_fanout fixture was captured from a run that does
not actually show fan-out behavior, since fanout_factor remains 1.0 and
trainer_egress_bytes matches total_delivered_bytes. Re-record the benchmark
output using the tree_fanout path in bench_elastic_scaling.py so that some
transfers are relayed through earlier receivers, then update the derived values
in the tree_fanout.json fixture to reflect the real run.
In `@modelexpress_client/python/modelexpress/refit_receiver.py`:
- Around line 432-437: In receive_weights_scratch(), stop defaulting unknown
descriptor dtypes to torch.bfloat16: validate td.dtype against the explicit
_DTYPE_MAP entries (bf16, f16, f32) and raise an error for anything else before
allocating scratch_tensors. Keep the existing size % elem_size == 0 guard in the
same flow so tensor sizes are only used when they exactly match the selected
dtype’s element size, using the td, _DTYPE_MAP, and scratch_tensors logic as the
key locations.
In `@modelexpress_client/python/tests/test_v2_shape_registry.py`:
- Around line 118-122: The round-trip test for
encode_expert_set/decode_expert_set is using a set literal, which removes
duplicates before the codec is exercised. Update
test_expert_set_codec_round_trip to use a sequence type such as a list or tuple
with repeated values so encode_expert_set is verified against duplicate input,
while keeping the same expected canonical encoded output and decoded set.
In `@modelexpress_client/python/tests/test_vllm_weight_transfer.py`:
- Line 269: The test unpacking in test_vllm_weight_transfer is creating unused
local variables that Ruff flags, so update the affected destructuring
assignments to mark the unused items as intentionally ignored. In the test cases
around the vllm_wt unpacking, rename the unused bindings like v2 and sd to
underscore-prefixed names or replace them with _. Keep the meaningful variables
unchanged and apply the same cleanup at each affected unpacking site in the test
file.
- Line 43: The module-scoped fixture in test_vllm_weight_transfer.py sets
MX_WEIGHT_TRANSFER_AUTOREGISTER but does not restore the original process
environment afterward. Update the fixture teardown to save the previous value
before mutating os.environ and restore or delete MX_WEIGHT_TRANSFER_AUTOREGISTER
when the fixture finishes, using the fixture around the vllm weight transfer
tests as the place to fix it.
---
Nitpick comments:
In `@modelexpress_server/src/p2p/backend/redis.rs`:
- Around line 84-133: Add a round-trip test for SourceAttributesJson to keep
From<&SourceIdentity> and to_source_identity() in sync across all fields. Create
a fully populated SourceIdentity, convert it into SourceAttributesJson, convert
it back with to_source_identity(), and assert the result matches the original so
any future dropped or mismapped field is caught.
In `@modelexpress_server/src/p2p/service.rs`:
- Around line 487-529: The happy-path test in test_get_metadata_found currently
sets ModelMetadataRecord.identity to None, so it does not verify the new
GetMetadataResponse.identity round-trip; update the mock record to include a
real identity value and assert that svc.get_metadata returns the expected
identity fields. Use the existing get_metadata test setup and the
ModelMetadataRecord/GetMetadataResponse types to confirm extra_parameters and
revision are preserved end to end.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7b18aa69-5beb-4fe1-8379-17174cee0c41
⛔ Files ignored due to path filters (1)
modelexpress_client/go/gen/modelexpress/p2p/p2p.pb.gois excluded by!**/*.pb.go,!**/gen/**
📒 Files selected for processing (28)
modelexpress_client/python/benchmarks/README.mdmodelexpress_client/python/benchmarks/__init__.pymodelexpress_client/python/benchmarks/bench_elastic_scaling.pymodelexpress_client/python/benchmarks/k8s/bench-elastic.yamlmodelexpress_client/python/benchmarks/results-20260530-105218/compile_target.jsonmodelexpress_client/python/benchmarks/results-20260530-105218/elastic_scale.jsonmodelexpress_client/python/benchmarks/results-20260530-105218/tree_fanout.jsonmodelexpress_client/python/benchmarks/run_cluster_bench.shmodelexpress_client/python/modelexpress/__init__.pymodelexpress_client/python/modelexpress/nemo_rl_v2.pymodelexpress_client/python/modelexpress/p2p_pb2.pymodelexpress_client/python/modelexpress/p2p_pb2_grpc.pymodelexpress_client/python/modelexpress/refit_receiver.pymodelexpress_client/python/modelexpress/shape_descriptors.pymodelexpress_client/python/modelexpress/training_publisher.pymodelexpress_client/python/modelexpress/vllm_weight_transfer.pymodelexpress_client/python/scripts/v2_dtensor_e2e_demo.pymodelexpress_client/python/scripts/v2_moe_e2e_demo.pymodelexpress_client/python/tests/test_bench_elastic_scaling.pymodelexpress_client/python/tests/test_v2_shape_registry.pymodelexpress_client/python/tests/test_v2_source_picker.pymodelexpress_client/python/tests/test_vllm_weight_transfer.pymodelexpress_common/proto/p2p.protomodelexpress_server/src/p2p/backend.rsmodelexpress_server/src/p2p/backend/kubernetes.rsmodelexpress_server/src/p2p/backend/redis.rsmodelexpress_server/src/p2p/service.rsmodelexpress_server/src/p2p/state.rs
KavinKrishnan
force-pushed
the
kavink/post-2389-phase3-4
branch
from
June 26, 2026 00:39
5d14209 to
8da2c34
Compare
KavinKrishnan
force-pushed
the
kavink/post-2389-phase3-4
branch
from
June 26, 2026 01:11
8da2c34 to
57cae41
Compare
This comment was marked as spam.
This comment was marked as spam.
KavinKrishnan
force-pushed
the
kavink/post-2389-phase3-4
branch
from
June 26, 2026 01:14
57cae41 to
0432982
Compare
…timizations
Lands the v2 client surface and the surrounding RL workstream needed
to unblock the downstream Megatron-MX and perf PRs.
## What this contributes
### v2 client surface (this PR's primary deliverable)
* `MxV2TrainingPublisher` + `MxV2RefitReceiver` (modelexpress.nemo_rl_v2)
-- the fat-client surface for per-rank shard publish and
receiver-side multi-source assembly.
* `MxWeightTransferEngine` (modelexpress.vllm_weight_transfer) -- an
adapter implementing vLLM's upstream WeightTransferEngine ABC.
* `TensorDescriptorV2.extra_parameters` (map<string,string>) plus
`SourceIdentity.revision` (string) -- the two proto extensions that
carry all per-tensor + per-source RL metadata.
### Phase 3a -- compile_target + compile_metadata on TensorDescriptorV2
New per-tensor fields default to ``hf_raw`` / ``{}``; the wire encoder
omits them when default so existing payloads stay byte-identical.
New constants: ``COMPILE_TARGET_HF_RAW``, ``_VLLM_FUSED``,
``_DEEPGEMM_FP8``, ``_CUTLASS_FP8``, ``_TRTLLM``. New helper
``compile_target_matches(descriptor, *, allowed_targets,
required_metadata=None)`` for receiver-side filtering with whitelist
plus required-metadata-subset semantics.
### Phase 3b -- compile_target_filter on discover_v2_sources
New kwargs ``compile_target_filter`` (whitelist set) and
``required_compile_metadata`` (subset-of-every-tensor's
compile_metadata). Candidates with no v2 registry are rejected when
either filter is set; candidates with mixed compile targets are
rejected if any tensor falls outside the allowed set.
``V2SourceCandidate.compile_targets: frozenset[str]`` exposed for
caller introspection.
### Phase 4 -- Multi-source slice discovery for mixed trainer/inference TP
New types: ``TargetTPLayout``, ``SliceSource``, ``SliceCoveragePlan``.
New methods ``MxV2RefitReceiver.discover_v2_sources_for_slice`` and
``MxV2RefitReceiver.receive_via_plan`` -- planner walks v2 candidates
per tensor, intersects each publisher's local_shard_range against the
receiver's requested slice, emits the minimal candidate set covering
it; surfaces coverage gaps and shard_axis mismatches in plan.missing.
``receive_via_plan`` orchestrates per-candidate scratch RDMA pulls
and stitches via torch.cat along the shard axis.
### Proto extensions
* ``TensorDescriptorV2.extra_parameters: map<string, string>`` -- the
escape hatch downstream RL clients use for per-tensor metadata
(megatron_role, compile_target, expert_id, revision, training_step,
...). Heavily used by #429.
* ``SourceIdentity.revision: string`` -- content-addressed weight
version. Non-empty value guarantees two sources with identical
SourceIdentity have bit-identical weight bytes, enabling
decentralized modes (no central coordinator) to use mx_source_id
as a full content check.
* Redis backend ``SourceAttributesJson`` carries the new fields plus
the union of main's compatibility additions
(backend_framework_version, torch_version, cuda_version,
triton_version, gpu_arch, compile_config_digest).
## Tests
68/68 unit tests green on this branch's surface (test_v2_source_picker,
test_v2_shape_registry, test_source_id, test_types, test_vllm_adapter).
## Compatibility
Every new field has a backwards-compatible default and every new method
arg is optional. Existing demos and downstream consumers run unchanged.
## Downstream impact
This PR is the prerequisite for several PRs in flight:
* #421 -- publish_self_as_source tree fan-out fix
* #429 -- Megatron-Core MX clients
* #450 -- MX_RDMA_NIC_PIN=stripe multi-NIC mode
* ai-dynamo/dynamo#10900 -- first-time upstream port of
dynamo.vllm.mx_refit extension
* ai-dynamo/dynamo#10901 -- buffer-caching perf fix in Dynamo extension
* jthomson04/RL#2 (merged) + #7 -- NeMo RL Megatron-MX integration plus
perf fix
Once this PR lands, kavink/nemo_rl_moe (the integration branch holding
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
#421 and #429) can be retired by rebasing those PRs onto main.
KavinKrishnan
force-pushed
the
kavink/post-2389-phase3-4
branch
from
June 26, 2026 05:49
0432982 to
85dfa4e
Compare
Contributor
Author
|
Pre-review pass: addressed the anticipated reviewer concerns before they land. Force-pushed Mapping anticipated concerns to mitigations applied:
Other changes worth noting:
Diff stat on top of the prior CodeRabbit-response commit: 8 files changed (3 deletions of JSON results, 1 file move with shim, 4 in-place edits for parameterization + path updates). No semantic change to the v1/v2 client surface itself. |
Contributor
|
@KavinKrishnan I identified a NIXL registration leak in the scratch receive path while testing this branch. In a live NIXL benchmark, the receiver agent metadata grew linearly with repeated scratch receives: 9878 -> 19190 -> 28502 -> 37814 -> 47126I implemented a fix that scopes scratch tensor registrations to the duration of a single receive and explicitly deregisters them once the transfer completes. Persistent tensor registrations remain unchanged and are still released at shutdown, preserving the expected publisher metadata lifecycle. With this fix, the same benchmark remains stable and bounded: count: 21
first 10: [9882, 9878, 9878, 9878, 9878, 9878, 9878, 9878, 9878, 9878]
last 10: [9878, 9878, 9878, 9878, 9878, 9878, 9878, 9878, 9878, 9878]
first: 9882 last: 9878 delta: -4I pushed a commit to this PR branch with the fix and validation details: #457 Validation was performed using |
…arding
Adds the no-allgather refit contract that lets RL framework integrations
publish trainer-side local shards directly (via DTensor.to_local() or an
explicit placement descriptor) instead of materializing tensor.full_tensor()
before NIXL registration. The receiver side intersects per-rank ownership
against per-rank requests with a pure-function planner and drives one-sided
NIXL READs against the matching source ranks.
New modules (all under modelexpress_client/python/modelexpress):
- rl_slice_descriptors.py - Typed contract: SliceOwnership, SliceRequest,
SegmentPlan, CoveragePlan + PlanIncompleteError + QuantizationMetadataError.
Framework-agnostic; no torch dependency.
- rl_reshard_planner.py - Pure plan_coverage(sources, requests) -> CoveragePlan.
Filters by dtype + compile_target + compile_metadata; refuses zero-copy on
quantization_scope='global-required'; prefers same-rank sources for
multi-NIC routing on segmented fabrics. Framework-agnostic; no torch.
- rank_local_publisher.py - RankLocalPublisher wrapping MxTrainingPublisher.
add_dtensor() auto-derives placement from a DTensor's .placements +
.device_mesh and calls .to_local() (no allgather). add_explicit_shard()
takes a PlacementDescriptor for non-DTensor inputs (e.g. Megatron-Core).
- integrations/verl_checkpoint_engine.py - VerlMxCheckpointEngine (trainer)
and VerlMxRolloutLoader (inference). Drop-in replacement for the existing
full_tensor() publish + collective_rpc("load_weights") pull. Falls back
to MxRefitReceiver.receive_weights() when the receiver lacks the
per-segment fast path, so existing deployments stay green.
Properties:
- No global allgather on the trainer side: each FSDP rank publishes only
its local shard.
- No "rank 0 holds the full model" memory spike during publish.
- Cross-parallelism reshard is by construction (FSDP=N -> TP=M is N*M/gcd
segments per receiver, all driven by the planner).
- Same-rank routing happens at planning time - composes with the existing
multi-NIC routing policies the MX integrations already use.
- Compile-target filtering keeps staged kernel rollouts safe (a receiver
not yet upgraded to read cutlass_fp8 filters those sources out).
- Quantization that needs a global scale fails loud
(QuantizationMetadataError on quantization_scope='global-required') so
callers fall back to a full-copy install path for the affected tensors
rather than silently corrupting them.
Tests: 44 CPU-only unit tests across three files. No torch.distributed
required; all green in 0.3s on a dev box and inside a cluster Job pod.
Sibling surface to the v1 MxTrainingPublisher / MxRefitReceiver - existing
v1 callers are unaffected.
(cherry picked from commit b01a5ab of
kavink/verl-rank-to-rank-prototype; docs/RL/VERL_MX_OVERVIEW.md
references scrubbed since process docs were removed from this PR.)
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
…k harness
Adds the per-segment NIXL READ primitive that the rank-to-rank planner
drives, plus the benchmark harness that exercises the contract end-to-end.
MxRefitReceiver gets two new entry points:
- prefetch_source(mx_source_id, worker_id) -> remote_agent_name
Resolves the source's NIXL metadata once via MxClient.get_metadata,
loads it into the local NIXL agent via add_remote_agent, and caches
the resulting agent handle keyed by (source_id, worker_id). Subsequent
receive_segment calls for the same source skip the gRPC round-trip
entirely.
- receive_segment(remote_agent_name, source_addr, byte_count,
target_addr, source_device_id, timeout_seconds)
Issues a single one-sided NIXL READ for one contiguous segment. Used
by VerlMxRolloutLoader (and by other framework integrations going
forward) to drive each SegmentPlan emitted by the planner without the
bulk name-matched receive_weights path.
Together these are the data-plane half of the rank-to-rank contract; the
existing receive_weights path stays in place and is the documented
fallback for receivers that don't yet expose receive_segment.
New benchmark harness (modelexpress_client/python/benchmarks/
bench_verl_rank_to_rank.py) spins up two trainer-side publishers + two
inference-side receivers as threads on a single GPU pod, transfers a
synthetic shard-aware tensor through real NIXL against the live MX
server, checksum-verifies the bytes landed correctly, and reports
per-cycle Gbps + per-source byte distribution + savings vs the v1
receive_weights baseline. Runs both rank-to-rank and v1 in the same
process for an apples-to-apples comparison.
Tests: 9 new unit tests in test_receive_segment.py covering
prefetch_source caching + per-(source,worker) cache isolation +
not-found errors, and receive_segment's NIXL parameter wiring +
single-READ contract + handle release + error/timeout paths +
initialize() guard. Total RL contract test count is now 53 (44
contract + 9 data-plane), all CPU-only, all green in 0.32s.
(cherry picked from commit 64b20e0 of
kavink/verl-rank-to-rank-prototype; docs/RL/VERL_MX_OVERVIEW.md
reference dropped + benchmarks/__init__.py extended to describe both
bench_elastic_scaling.py and bench_verl_rank_to_rank.py.)
Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
This was referenced Jun 30, 2026
Lets a plain, already-materialized shard (e.g. a post-conversion output buffer that is no longer a DTensor) publish as a SHARD of a larger logical tensor, so receivers reassemble across ranks instead of the publisher pre-gathering with DTensor.full_tensor(). This is the ModelExpress substrate half of the prime-rl "on-the-fly shard-in without allgather" work for non-expert weights. prime-rl's ShardedSlot holds each rank's FSDP row-shard as a plain converted buffer; describe_tensor previously emitted REPLICATE for plain tensors (collapsing shards into full replicas), which forced the trainer to allgather first. NonExpertShardSpec carries the global shape + this rank's row range explicitly so the multi-source planner can gather. Changes: - shape_descriptors.NonExpertShardSpec + describe_tensor(shard_spec=...): short-circuit to a SHARD descriptor, validating local extent matches the range width and the range is within global bounds. - MxV2TrainingPublisher.add_tensor(shard_spec=...): pass-through. - Export NonExpertShardSpec from the package. - tests/test_non_expert_shard_spec.py: descriptor emission, validation errors, and a 4-rank reassembly plan-coverage check (5 tests). - benchmarks/bench_non_expert_rank_to_rank.py: Tier-2 cluster harness — N publishers each hold only their row-shard, receiver reassembles the full tensor via the multi-source planner, byte-identity verified, and it asserts the plan drew from all N ranks (no single rank held the full tensor). Full RL contract suite: 95 passing (90 prior + 5 new). The prime-rl trainer-side change (drop the GatheredSlot gather override; publish ShardedSlot shards with this spec) rides in the prime-rl repo. The prime-rl receiver hybrid (gather non-experts / same-rank experts) is a separate, larger follow-up; this harness isolates and proves the new publish + reassembly mechanism on real RDMA. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
Cycle 0 worked but cycle 1+ raced the trainers' re-publish and found zero candidates (empty plan: not fully_covered, missing=[]). Mirror the prime-rl worker's transient-retry: poll discover_v2_sources_for_slice with backoff until all N ranks' shards are covered or timeout. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
The receive_via_plan scratch path does a rendezvous that needs target-side CPU progress; in this single-process threaded harness the publisher threads block at the cross-cycle barrier during a receive, so the multi-source scratch pull hung after the first shard. Switch to the proven one-sided plan_coverage + receive_segment path (same as the verl and EP benches, which complete with blocked targets): assemble all N row-shards directly into a pre-registered full buffer, then byte-verify each rank's rows. Restores UCX_TLS to the working transport list. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
Contributor
Author
|
Superseded by the consolidated integration branch in #482, which already contains this v2-base + WeightTransferEngine work (now the native "mx" backend). Closing to reduce open-PR surface — the branch and this discussion are preserved, and this will be re-cut as a clean, small review slice (v2 client base) from #482. No work is lost. |
KavinKrishnan
added a commit
that referenced
this pull request
Jul 6, 2026
Squashed integration of the MX v2 RL weight-refit path (previously 53 commits; see closed PRs #349/#465/#450/#470 for the incremental history). Client (modelexpress_client/python/modelexpress): - nemo_rl_v2.py: MxV2TrainingPublisher / MxV2RefitReceiver, multi-source pick_megatron_slice_plans, tree fan-out (publish_self_as_source replica_uid unique-sid + discover_v2_sources prefer_replicas). - megatron_translator.py: Megatron→HF translate (qkv un-interleave, gated-MLP split, per-expert unstack). - nixl_transfer.py: sliced pull (receive_sliced_from_source), host/device memtype for pinned-CPU staging, rebind_tensors buffer caching; reconciled with the accelerator-backend abstraction on main. - ucx_utils.py: MX_RDMA_NIC_PIN=stripe multi-NIC parallelism. - rl_expert_layout.py: MoE expert-parallel wire filtering (compute_local_expert_ids). - vmm/: VMM-arena single-region registration. - engines/vllm/weight_transfer.py: MxWeightTransferEngine — native vLLM WeightTransferEngine backend registered as "mx". - engines/vllm/mdl.py: MDL (Mapped Direct Load) decoupled load-side fast-loader. Server (modelexpress_server/src/p2p) + proto: - SourceIdentity.extra_parameters round-trip + revision, additive with main's version/arch fields. Docs: MX-RL design + integration/benchmark context under docs/RL, docs/slides. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
KavinKrishnan
added a commit
that referenced
this pull request
Jul 6, 2026
Squashed integration of the MX v2 RL weight-refit path (previously 53 commits; see closed PRs #349/#465/#450/#470 for the incremental history). Client (modelexpress_client/python/modelexpress): - nemo_rl_v2.py: MxV2TrainingPublisher / MxV2RefitReceiver, multi-source pick_megatron_slice_plans, tree fan-out (publish_self_as_source replica_uid unique-sid + discover_v2_sources prefer_replicas). - megatron_translator.py: Megatron→HF translate (qkv un-interleave, gated-MLP split, per-expert unstack). - nixl_transfer.py: sliced pull (receive_sliced_from_source), host/device memtype for pinned-CPU staging, rebind_tensors buffer caching; reconciled with the accelerator-backend abstraction on main. - ucx_utils.py: MX_RDMA_NIC_PIN=stripe multi-NIC parallelism. - rl_expert_layout.py: MoE expert-parallel wire filtering (compute_local_expert_ids). - vmm/: VMM-arena single-region registration. - engines/vllm/weight_transfer.py: MxWeightTransferEngine — native vLLM WeightTransferEngine backend registered as "mx". - engines/vllm/mdl.py: MDL (Mapped Direct Load) decoupled load-side fast-loader. Server (modelexpress_server/src/p2p) + proto: - SourceIdentity.extra_parameters round-trip + revision, additive with main's version/arch fields. Docs: MX-RL design + integration/benchmark context under docs/RL, docs/slides. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
KavinKrishnan
added a commit
that referenced
this pull request
Jul 6, 2026
Squashed integration of the MX v2 RL weight-refit path (previously 53 commits; see closed PRs #349/#465/#450/#470 for the incremental history). Client (modelexpress_client/python/modelexpress): - nemo_rl_v2.py: MxV2TrainingPublisher / MxV2RefitReceiver, multi-source pick_megatron_slice_plans, tree fan-out (publish_self_as_source replica_uid unique-sid + discover_v2_sources prefer_replicas). - megatron_translator.py: Megatron→HF translate (qkv un-interleave, gated-MLP split, per-expert unstack). - nixl_transfer.py: sliced pull (receive_sliced_from_source), host/device memtype for pinned-CPU staging, rebind_tensors buffer caching; reconciled with the accelerator-backend abstraction on main. - ucx_utils.py: MX_RDMA_NIC_PIN=stripe multi-NIC parallelism. - rl_expert_layout.py: MoE expert-parallel wire filtering (compute_local_expert_ids). - vmm/: VMM-arena single-region registration. - engines/vllm/weight_transfer.py: MxWeightTransferEngine — native vLLM WeightTransferEngine backend registered as "mx". - engines/vllm/mdl.py: MDL (Mapped Direct Load) decoupled load-side fast-loader. Server (modelexpress_server/src/p2p) + proto: - SourceIdentity.extra_parameters round-trip + revision, additive with main's version/arch fields. Docs: MX-RL design + integration/benchmark context under docs/RL, docs/slides. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
KavinKrishnan
added a commit
that referenced
this pull request
Jul 6, 2026
Squashed integration of the MX v2 RL weight-refit path (previously 53 commits; see closed PRs #349/#465/#450/#470 for the incremental history). Client (modelexpress_client/python/modelexpress): - nemo_rl_v2.py: MxV2TrainingPublisher / MxV2RefitReceiver, multi-source pick_megatron_slice_plans, tree fan-out (publish_self_as_source replica_uid unique-sid + discover_v2_sources prefer_replicas). - megatron_translator.py: Megatron→HF translate (qkv un-interleave, gated-MLP split, per-expert unstack). - nixl_transfer.py: sliced pull (receive_sliced_from_source), host/device memtype for pinned-CPU staging, rebind_tensors buffer caching; reconciled with the accelerator-backend abstraction on main. - ucx_utils.py: MX_RDMA_NIC_PIN=stripe multi-NIC parallelism. - rl_expert_layout.py: MoE expert-parallel wire filtering (compute_local_expert_ids). - vmm/: VMM-arena single-region registration. - engines/vllm/weight_transfer.py: MxWeightTransferEngine — native vLLM WeightTransferEngine backend registered as "mx". - engines/vllm/mdl.py: MDL (Mapped Direct Load) decoupled load-side fast-loader. Server (modelexpress_server/src/p2p) + proto: - SourceIdentity.extra_parameters round-trip + revision, additive with main's version/arch fields. Docs: MX-RL design + integration/benchmark context under docs/RL, docs/slides. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
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.
Summary
Adds the framework-agnostic client-side surface MX needs to support RL weight refit.
The same primitives are consumed unchanged by every RL framework integration (NeMo RL, PrimeRL, verl) — each framework just provides a thin glue layer on top, no per-framework forking of the MX clients. Introduces both layers as a single deliverable, because RL needs the v2 fat-client API but the v2 layer reuses the v1 building blocks underneath:
I. v1 base (
MxTrainingPublisher/MxRefitReceiverintraining_publisher.py,refit_receiver.py)II. v2 fat-client layer on top (
MxV2TrainingPublisher/MxV2RefitReceiver/MxWeightTransferEngineinnemo_rl_v2.py,engines/vllm/weight_transfer.py)Along with four RL primitives:
MxV2TrainingPublisher.add_tensorTensorDescriptorV2.extra_parameters(new map field on the proto)SourceIdentity.revision(new string field on the proto) +MxV2RefitReceiver.discover_v2_sources(min_version=...)MxWeightTransferEngineWeightTransferEngineABC (vLLM PR #43375, native RL APIs blog) so RL frameworks drop MX in without modifying vLLM.Status
origin/main; union-merge forSourceIdentityfield set, regenerated proto stubs viagenerate_proto.shwith pinnedgrpcio-tools==1.66.2).mergeable: MERGEABLE. DCO + Python Tests + Test Suite (stable/beta/nightly) all green.Reviewer guidance
A few things worth pre-empting for reviewers coming in cold:
Why v1 + v2 are one PR
The v1 base and v2 fat-client layer are released together because:
MxV2RefitReceiverinstantiates anMxRefitReceiverinternally — the v2 layer is a wrapper around v1's NIXL register + publish + receive primitives, not a replacement.mainactually usesMxTrainingPublisher/MxRefitReceiverstandalone).The v2 layer adds the four RL primitives listed above; v1 is the substrate it sits on.
Why
MxWeightTransferEnginedoesn't live underEngineAdapterThere are two distinct "engine" abstractions in this codebase, and they target different lifecycle phases:
EngineAdapter(existing)LoadStrategyChain.run(model, ctx)engines/vllm/adapter.pyMxWeightTransferEngine(this PR)WeightTransferEngineABC (vLLM PR #43375)engines/vllm/weight_transfer.pyBoth live under
engines/vllm/so they're discoverable side-by-side, but they're independent — different contracts, different hook points, different__init__shapes.MxWeightTransferEnginedoes not extendEngineAdapterbecause it isn't going throughLoadStrategyChain; it's plugged directly into vLLM's own factory and called per-refit-cycle, not at startup.A back-compat shim at
modelexpress/vllm_weight_transfer.pyre-exports the public symbols so existing callers don't break.Why
nemo_rl_v2.pyis long (~1350 lines)The file is the fat-client API surface — one cohesive class (
MxV2RefitReceiver) with its dataclasses, slice planner, and source picker. Splitting it across modules would split the API across imports for downstream RL framework integrations that already learn one entry point. The length is intentional, not a refactor TODO.Benchmarks and demo scripts
benchmarks/bench_elastic_scaling.pyandscripts/v2_*_e2e_demo.pyare bundled here because they exercise the v2 client surface in a self-contained way — the benchmark is the validation harness for the new client API, and the demos are the canonical usage examples downstream consumers reference. The benchmark Job manifest is fully parameterized (<NAMESPACE>,<IMAGE>,<IMAGE_PULL_SECRET>); the README documents the env-var protocol.What this PR contributes
Per-tensor
compile_target+compile_metadataonTensorDescriptorV2New per-tensor fields default to
hf_raw/{}; the wire encoder omits them when default so existing payloads stay byte-identical. New constants for canonical-BF16, vLLM-fused, DeepGEMM-FP8, CUTLASS-FP8, and TRT-LLM kernel formats. New helpercompile_target_matches()for receiver-side filtering with whitelist + required-metadata-subset semantics.compile_target_filteronMxV2RefitReceiver.discover_v2_sourcesNew kwargs
compile_target_filter(whitelist set) andrequired_compile_metadata(dict that must be a subset of every tensor'scompile_metadata). Candidates with no v2 registry are rejected when either filter is set; candidates with mixed compile targets are rejected if any tensor falls outside the allowed set.Multi-source slice discovery for mixed trainer/inference TP
New types:
TargetTPLayout,SliceSource,SliceCoveragePlan. New methodsMxV2RefitReceiver.discover_v2_sources_for_sliceandMxV2RefitReceiver.receive_via_plan. The planner walks v2 candidates per tensor, intersects each publisher'slocal_shard_rangeagainst the receiver's requested slice, emits the minimal candidate set covering it; surfaces coverage gaps and shard-axis mismatches inplan.missing.receive_via_planorchestrates per-candidate scratch RDMA pulls and stitches viatorch.catalong the shard axis.Proto extensions
TensorDescriptorV2.extra_parameters: map<string, string>— escape hatch downstream RL clients use for per-tensor metadata (parallelism role, kernel format, expert ID, version, etc.). Heavily used by #429.SourceIdentity.revision: string— content-addressed weight version. Non-empty value guarantees two sources with identicalSourceIdentityhave bit-identical weight bytes, enabling decentralized modes (no central coordinator) to usemx_source_idas a full content check.SourceAttributesJsoncarries the new fields plus the union of main's compatibility additions (backend_framework_version,torch_version,cuda_version,triton_version,gpu_arch,compile_config_digest).Tests
68/68 client unit tests + cluster-tier green.
Compatibility
Every new field has a backwards-compatible default and every new method arg is optional. Existing demos and downstream consumers run unchanged. The back-compat shim at
modelexpress/vllm_weight_transfer.pyre-exports the public symbols from the new canonical location atmodelexpress.engines.vllm.weight_transfer.Downstream impact
This PR is the prerequisite for the rest of the RL workstream. Several downstream PRs already in flight import the v2 surface or
nemo_rl_v2module:publish_self_as_sourcetree fan-out fix (merged into thekavink/nemo_rl_moeintegration branch, awaiting this PR to reach main)MX_RDMA_NIC_PIN=stripemulti-NIC modedynamo.vllm.mx_refitextension; importsMxV2RefitReceiver+ Megatron modulesOnce this PR lands,
kavink/nemo_rl_moe(the integration branch holding #421 + #429) can be retired by rebasing those PRs onto main directly.