feat(refit): native->HF conversion capture + minimal-read generator

[tanushriya910]

What

Makes the refactored vLLM generator (#648) pull minimal weight slices by default and map a trainer-native layout into the HF loader on receive. This folds two pieces onto the #648 base:

• Native to HF conversion capture. A trainer that publishes native (non-HF) weights, e.g. an MoE with stacked experts, now converts to HF before the vLLM load-weights capture. Shared capture machinery (build_lazy_weights / capture_weights / convert_source_weights) plus an expert-parallel capture fix. Conversion is identity when convert_native_to_hf is None, so dense models are unaffected. (feat(refit): map a native trainer layout into the vLLM loader on receive #660)
• select/split view-ops in the slice plan, so a captured view that is a slice of a larger tensor resolves to the right source region. (feat(refit): view-ops slice planning for reshard #672)

How it works

• Minimal reads by default. Each rank pulls only the weight slices it needs (plan_transfer), not whole tensors. Verification needs a complete tensor to hash, which minimal reads never build, so it is gated behind MX_RESHARD_PUBLISH_DIGEST. Turn it on and any source that a slice read cannot fully recover is pulled whole and digest-checked. Leave it off, the default, and reads stay minimal with no verification.
• Capture runs on the live model. Params are reverted to their bf16 load-time skeletons via a layerwise reload, captured, then kernel tensors are restored without finalizing. This replaces the meta-twin capture.

Testing

• 37 unit tests pass (test_reshard_refit_geometry.py, test_reshard_refit_slice_plan.py), including the added select/split view-op and conversion-capture coverage.
• Validated bit-exact end to end on Qwen3-30B-A3B (MoE, trainer FSDP ep=8, vLLM tp=4 with expert parallel) driven by prime-rl. Step-0 snapshot to randomize to transfer verify returned max_abs_diff=0 (post-randomize was 35.56). The minimal-slice read path dominated the transfer (segmented reads far outnumbered full-pulls), exercising both the conversion capture and the view-ops on a real MoE.

Notes for reviewers

• MX_RESHARD_PUBLISH_DIGEST stays False by default. Minimal reads are the normal path; full-pull plus digest verification is the opt-in verification mode.
• Base is feat(refit): add RL generator client and staged vLLM installation #648, already on main.

Summary by CodeRabbit

• New Features

  • Added support for converting trainer-native weights into compatible Hugging Face and vLLM layouts.
  • Improved model weight capture, including shared recording and renamed source weights.
  • Added support for more tensor slicing patterns and mixture-of-experts layouts.

• Performance

  • Reduced staged-transfer reads by default and added detailed transfer and reconstruction timing.

• Reliability

  • Added validation for unsupported tensor conversions and improved digest-based transfer verification.
  • Published shard digests to improve integrity checking during transfers.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9484481e-6a78-4f7f-b0fd-4c9a1c6bbf30

📥 Commits

Reviewing files that changed from the base of the PR and between f4fa625 and 3711653.

📒 Files selected for processing (13)

• modelexpress_client/python/modelexpress/refit/reshard/__init__.py
• modelexpress_client/python/modelexpress/refit/reshard/geometry.py
• modelexpress_client/python/modelexpress/refit/reshard/slice_plan.py
• modelexpress_client/python/modelexpress_rl/inference/engines/vllm/__init__.py
• modelexpress_client/python/modelexpress_rl/inference/engines/vllm/adapter.py
• modelexpress_client/python/modelexpress_rl/inference/engines/vllm/context.py
• modelexpress_client/python/modelexpress_rl/inference/engines/vllm/installer.py
• modelexpress_client/python/modelexpress_rl/inference/nixl_staged_transfer.py
• modelexpress_client/python/modelexpress_rl/train/engines/fsdp/publisher.py
• modelexpress_client/python/tests/test_refit_nixl_staged_transfer.py
• modelexpress_client/python/tests/test_refit_vllm_adapter.py
• modelexpress_client/python/tests/test_reshard_refit_geometry.py
• modelexpress_client/python/tests/test_reshard_refit_slice_plan.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The PR extends reshard capture with source conversion and shared lazy weights, supports rank-reducing view plans, updates vLLM live-model capture, and adds digest-aware staged transfer planning, verification, logging, and metrics.

Changes

Reshard pipeline

Layer / File(s)	Summary
Lazy capture and source conversion modelexpress_client/python/modelexpress/refit/reshard/*, modelexpress_client/python/tests/test_reshard_refit_geometry.py	Adds public lazy-weight helpers, shared-recorder validation, source-name preservation, native-to-Hugging Face conversion, and MoE reconstruction coverage.
Rank-reducing slice reconstruction modelexpress_client/python/modelexpress/refit/reshard/slice_plan.py, modelexpress_client/python/tests/test_reshard_refit_slice_plan.py	Supports select, unbind, and split view plans. Rank-increasing views remain rejected.
vLLM live-model capture modelexpress_client/python/modelexpress_rl/inference/engines/vllm/{context.py,adapter.py,installer.py}, modelexpress_client/python/tests/test_refit_vllm_adapter.py	Passes the converter callback through vLLM context and adapter construction. Captures converted weights from the live model with layerwise reload.
Digest-aware staged transfer modelexpress_client/python/modelexpress_rl/inference/nixl_staged_transfer.py, modelexpress_client/python/modelexpress_rl/train/engines/fsdp/publisher.py, modelexpress_client/python/tests/test_refit_nixl_staged_transfer.py	Uses minimal reads by default, enables full reads and verification in digest mode, records reconstruction timing, and publishes shard digests.

Estimated code review effort: 4 (Complex) | ~60 minutes

Merge Risk: ⚪ Minimal · up to 37116

This change is merge-ready after normal checks and review; no actionable merge-blocking risk remains.

Poem

A rabbit checks each lazy thread

Views fold neatly where tensors tread
MoE shards hop through plans so bright
Digests guard the transfer right
vLLM reloads by live-model light

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.17% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 47 functions across 13 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main changes: native-to-HF conversion capture and minimal-read generation for refit.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}