feat(refit): add RL trainer client and Megatron adapter - #616
Conversation
ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (20)
💤 Files with no reviewable changes (2)
🚧 Files skipped from review as they are similar to previous changes (7)
WalkthroughThe change adds RL trainer refit APIs, protobuf and gRPC bindings, manifest publication and retrieval, Megatron training and inference support, environment configuration, compatibility exports, documentation, and tests. ChangesRL refit protocol
Trainer publication flow
Megatron support
Validation and documentation
Estimated code review effort: 4 (Complex) | ~60 minutes Mergeability Score: 🟠 High · up to The PR adds rank-local trainer staging and publishing, but the Megatron IN_PLACE path can still leave a trainer rank waiting indefinitely, while the new package layout can make imports fail depending on order; malformed tensor geometry is also accepted. These issues can cause hangs or runtime failures, so the PR is not merge-ready until the liveness and import risks are fixed or explicitly accepted. 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}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py (2)
164-173: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReport missing QKV extras with the tensor name.
Lines 166 to 168 index
item.extrasdirectly. A missing key raisesKeyError('head_dim')without the tensor name. Every other validation in this module raisesValueErrorprefixed withitem.name. Align this path so a publisher misconfiguration identifies the offending tensor.♻️ Proposed change
- head_dim = int(item.extras["head_dim"]) - q_heads_local = int(item.extras["num_heads_local"]) - kv_heads_local = int(item.extras["num_kv_heads_local"]) + required = ("head_dim", "num_heads_local", "num_kv_heads_local") + missing = [key for key in required if key not in item.extras] + if missing: + raise ValueError(f"{item.name}: QKV aliasing requires extras {missing}") + head_dim = int(item.extras["head_dim"]) + q_heads_local = int(item.extras["num_heads_local"]) + kv_heads_local = int(item.extras["num_kv_heads_local"])🤖 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_rl/train/engines/megatron/aliases.py` around lines 164 - 173, Update the QKV validation path around item.extras access to catch missing head_dim, num_heads_local, or num_kv_heads_local metadata and raise a ValueError prefixed with item.name, preserving the existing validation behavior for present extras.
142-147: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCross-check the derived gated
full_shapeagainstitem.global_shape.
_build_gated_aliasesderivesfull_shapefrom the local tensor shape and the source group size. It never compares the result with the declareditem.global_shape. The non-fused axes therefore come from local geometry only, and a caller that declares an inconsistentglobal_shapegets no error. The single-name path at line 255 usesitem.global_shapedirectly, so the two paths trust different sources of truth.Add an assertion that the derived fused extent agrees with
item.global_shape[axis] // 2and that the remaining dimensions match.♻️ Proposed validation
source_rank, source_size = _source_rank_and_size(item, axis) full_shape = list(item.tensor.shape) full_shape[axis] = half * source_size + expected = list(int(dim) for dim in item.global_shape) + expected[axis] //= 2 + if expected != [int(dim) for dim in full_shape]: + raise ValueError( + f"{item.name}: derived gate/up shape {tuple(full_shape)} disagrees " + f"with declared global shape {item.global_shape}" + )🤖 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_rl/train/engines/megatron/aliases.py` around lines 142 - 147, Update _build_gated_aliases after deriving full_shape to validate it against item.global_shape: require the fused axis extent to equal item.global_shape[axis] // 2 and all non-fused dimensions to match exactly. Raise an assertion or equivalent validation error before constructing shard_range when the declared global shape is inconsistent.modelexpress_client/python/tests/test_refit_trainer_client.py (1)
96-96: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace the fixed sleep with a bounded poll to avoid a flaky renewal assertion.
Line 96 sleeps 0.4 s and line 130 asserts
service.registration_count >= 2. That couples the test to the renewal interval derived fromregistration_ttl_seconds=1. If the renewal interval is at or above 0.4 s, or if the CI machine is loaded, the second registration does not arrive and the test fails. Poll until the count reaches 2 with a generous deadline.♻️ Proposed change
- time.sleep(0.4) + deadline = time.monotonic() + 10.0 + while service.registration_count < 2 and time.monotonic() < deadline: + time.sleep(0.02) + assert service.registration_count >= 2Then drop the assertion at line 130.
🤖 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_refit_trainer_client.py` at line 96, Replace the fixed time.sleep(0.4) in the renewal test with a bounded polling loop that waits until service.registration_count reaches 2, using a generous deadline and short polling interval. Preserve timeout protection, then remove the separate registration_count >= 2 assertion because the poll should enforce the condition.modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/receiver.py (1)
5-7: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftThe two packages now import each other.
This module imports from
modelexpress. At the same timemodelexpress_client/python/modelexpress/refit/reshard/megatron_aliases.pyimports frommodelexpress_rl. The dependency edges therefore run in both directions between the two top-level packages. That makes the import order significant and it can produce partially initialized modules if either side later imports at package__init__level.Pick one direction. Keep the shared implementations in a single owning package and let the other package re-export only.
🤖 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_rl/inference/reshard/megatron/receiver.py` around lines 5 - 7, Remove the bidirectional package dependency by choosing one owning package for Megatron reshard implementations. Update MegatronReshardReceiver and the related megatron_aliases re-export so only the non-owning package imports and re-exports from the owning package; ensure no module under either top-level package imports back in the opposite direction.
🤖 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/modelexpress_rl/client.py`:
- Around line 242-245: Update the publication bookkeeping around
self._published_shards and StagedWeightVersionShard.publish() so every
successfully published staging allocation remains retained for the version,
including repeated publications with different buffer_owner objects;
alternatively reject duplicate publications for the same version and source slot
before replacing the existing handle. Preserve the existing lifecycle behavior
that releases retained buffers only during later eviction/release.
In `@modelexpress_client/python/modelexpress_rl/refit_pb2_grpc.py`:
- Line 24: Update the proto generation flow used by generate_proto.sh so both
generated gRPC stubs avoid the redundant f-string prefix causing the F541 lint
error; alternatively configure Ruff to exclude the generated protobuf modules.
Do not modify only refit_pb2_grpc.py or other checked-in generated output, and
ensure regeneration produces lint-clean stubs consistently.
In
`@modelexpress_client/python/modelexpress_rl/train/engines/megatron/adapter.py`:
- Around line 89-99: The source_reuse_ready fence in
MegatronTrainerAdapter.stage_shard must not wait on an unretained, unset Event:
either retain and signal the event during version retirement, or make the fence
callable raise NotImplementedError until that lifecycle exists. In
modelexpress_client/python/modelexpress_rl/train/engines/megatron/adapter.py
lines 89-99, implement the chosen behavior; in
modelexpress_client/python/tests/test_refit_megatron_adapter.py lines 108-131,
add a staged.source_reuse_ready assertion matching that behavior, following the
existing trainer-client fence test pattern.
Apply the same fix in
`@modelexpress_client/python/tests/test_refit_megatron_adapter.py` around lines
108 - 131: The test currently omits the reuse-fence wait, allowing the blocking
defect to pass unnoticed.
---
Nitpick comments:
In
`@modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/receiver.py`:
- Around line 5-7: Remove the bidirectional package dependency by choosing one
owning package for Megatron reshard implementations. Update
MegatronReshardReceiver and the related megatron_aliases re-export so only the
non-owning package imports and re-exports from the owning package; ensure no
module under either top-level package imports back in the opposite direction.
In
`@modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py`:
- Around line 164-173: Update the QKV validation path around item.extras access
to catch missing head_dim, num_heads_local, or num_kv_heads_local metadata and
raise a ValueError prefixed with item.name, preserving the existing validation
behavior for present extras.
- Around line 142-147: Update _build_gated_aliases after deriving full_shape to
validate it against item.global_shape: require the fused axis extent to equal
item.global_shape[axis] // 2 and all non-fused dimensions to match exactly.
Raise an assertion or equivalent validation error before constructing
shard_range when the declared global shape is inconsistent.
In `@modelexpress_client/python/tests/test_refit_trainer_client.py`:
- Line 96: Replace the fixed time.sleep(0.4) in the renewal test with a bounded
polling loop that waits until service.registration_count reaches 2, using a
generous deadline and short polling interval. Preserve timeout protection, then
remove the separate registration_count >= 2 assertion because the poll should
enforce the condition.
🪄 Autofix
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: c501863d-c0c2-4501-affe-903295b080bf
📒 Files selected for processing (25)
modelexpress_client/python/README.mdmodelexpress_client/python/generate_proto.shmodelexpress_client/python/modelexpress/refit/reshard/megatron_aliases.pymodelexpress_client/python/modelexpress/refit/reshard/megatron_publisher.pymodelexpress_client/python/modelexpress_rl/__init__.pymodelexpress_client/python/modelexpress_rl/client.pymodelexpress_client/python/modelexpress_rl/inference/__init__.pymodelexpress_client/python/modelexpress_rl/inference/reshard/__init__.pymodelexpress_client/python/modelexpress_rl/inference/reshard/megatron/__init__.pymodelexpress_client/python/modelexpress_rl/inference/reshard/megatron/layout.pymodelexpress_client/python/modelexpress_rl/inference/reshard/megatron/receiver.pymodelexpress_client/python/modelexpress_rl/refit_pb2.pymodelexpress_client/python/modelexpress_rl/refit_pb2_grpc.pymodelexpress_client/python/modelexpress_rl/train/__init__.pymodelexpress_client/python/modelexpress_rl/train/adapter.pymodelexpress_client/python/modelexpress_rl/train/engines/__init__.pymodelexpress_client/python/modelexpress_rl/train/engines/megatron/__init__.pymodelexpress_client/python/modelexpress_rl/train/engines/megatron/adapter.pymodelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.pymodelexpress_client/python/modelexpress_rl/train/engines/megatron/publisher.pymodelexpress_client/python/modelexpress_rl/train/manifest.pymodelexpress_client/python/tests/test_refit_megatron_adapter.pymodelexpress_client/python/tests/test_refit_trainer_client.pymodelexpress_client/python/tests/test_reshard_megatron.pymodelexpress_common/proto/refit.proto
b1e373d to
23c64e4
Compare
23c64e4 to
dd3f97a
Compare
dd3f97a to
cfaab7a
Compare
cfaab7a to
a1e3a1a
Compare
a1e3a1a to
949fe25
Compare
8893dd1 to
c1d4f31
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
Overall, this is a strong refactor. Separating the version lifecycle in |
195e138 to
8f9437d
Compare
|
Review follow-up: commit |
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
8f9437d to
3bcc6ba
Compare
4 participants
Summary
ModelExpressTrainerClientlifecycle for staging and publishing versioned trainer shardsTrainerEngineAdaptercontract and rank-local manifest servicemodelexpress_rl/train/...andmodelexpress_rl/inference/...namespacesWhy
The Redis-backed
RefitServiceprovides the version and shard control plane, but RL trainer actors still need a framework-facing client that connects their native engine tensors to those APIs. This PR adds that client boundary without wiring ModelExpress into a specific RL framework.The intended lifecycle is:
WeightVersion.WeightVersionShardmetadata and serves its immutable manifest.READYafter all expected source slots are published.source_reuse_ready; until then the Megatron IN_PLACE implementation fails that fence explicitly instead of claiming the source buffers are reusable.API and implementation notes
ModelExpressTrainerClient.initialize()resolves RL deployment policy frommodelexpress_rl.envsand shared connectivity settings frommodelexpress.envs, constructs the selected trainer adapter internally, and maintains worker registration.TrainerEngineAdapter.agent_name,nixl_metadata, and metadata-listener port protocol used by publication.initialize()remains the supported setup path.worker_endpointidentifies the trainer-side manifest service, whileserver_urlidentifies the central ModelExpress control plane.stage_shard()delegates engine-specific tensor discovery and staging toTrainerEngineAdapter.StagedWeightVersionShard.publish()is idempotent and publishes only control metadata; weight bytes remain on trainer-owned buffers.MegatronTrainerAdaptercurrently supportsIN_PLACEplusFULL_TENSOR; it derives the source slot from the initialized Megatron global rank and the NIXL metadata endpoint from the shared worker host plus the manager's listen port.modelexpress_rl; temporary re-exports at the priormodelexpress.refit.reshardpaths preserve NeMo-RL #3632 compatibility while the shared reshard core remains engine-neutral.Validation
200 passedacross the environment, trainer-client, Megatron, and refit/reshard Python suite on the rebased, single-commit head4 passedagainst a real Redis 7 backend in theRefitServicegRPC lifecycle suite, covering automatic readiness, replacement publishers, leases, release, shard deletion, and worker expiryModelExpressTrainerClientand Megatron adapter against the Rust ModelExpress server and Redis backend: the trainer registered, published its shard, served its manifest throughRefitWorkerService, and advanced the version fromSTAGINGtoREADYRELEASING, and deleted the source shardF541check, andgit diff --check origin/main...HEADpassesNot included
Summary by CodeRabbit