Skip to content

[ckpt] feat: Add versioned weight update manifests - #5769

Draft
yaoyu-33 wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
yaoyu-33:yuya/landscape-weight-update-manifest-20260824
Draft

[ckpt] feat: Add versioned weight update manifests#5769
yaoyu-33 wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
yaoyu-33:yuya/landscape-weight-update-manifest-20260824

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds a small transport-neutral contract for versioning and verifying the existing one-shot Hugging Face-coordinate weight stream returned by AutoBridge.export_hf_weights().

WeightUpdateRecorder observes names, shapes, and dtypes while tensors flow through without retaining payloads. The producer sends a WeightUpdateManifest as the stream trailer. Before activating staged weights, the consumer validates the exact recorded inventory plus expected model, config, update mode, base version, and target version.

This is a bounded implementation under #4515 and #5371. It does not change a protected AutoBridge signature.

Supported trigger and impact

The RL integration guide already documents streaming export_hf_weights() and reusable ordered conversion tasks across TP/PP/MoE/CP. It currently leaves identity, completeness, ordering, and version visibility to downstream transport code.

Recent primary evidence shows why this boundary matters:

  • verl#7324 fixed a sharded FSDP+EP update that emitted only 32 of 64 experts.
  • TorchTitan#4236 added reusable HF/vLLM state-dict adaptation for distributed weight movement.
  • miles#2589 made serving-side version transitions explicit because they affect in-flight samples.

Changelog

  • Add WeightUpdateRecorder, which preserves a one-shot stream and retains metadata only.
  • Add immutable, versioned WeightUpdateManifest and WeightUpdateTensor types.
  • Send the manifest as a trailer, avoiding full payload buffering or a second export pass.
  • Validate exact tensor count/order/name/shape/dtype and consumer model/config/mode/base/target context.
  • Strictly reject duplicate/unknown JSON fields, non-canonical scalar types, invalid UTF-8, and digest mismatches.
  • Document the flow in current and nightly RL integration docs; released 0.4.2 docs remain unchanged.

Validation

Fail-before on unmodified base dd150c1d73ed9a25ad3ee87149cb98a0a2ca3614:

uv run python -m pytest tests/unit_tests/models/test_weight_update_manifest.py::test_manifest_round_trip_and_validate_complete_update -q
1 failed: the conversion package does not expose a versioned weight-update manifest

Pass-after for the same regression plus adjacent export/documentation contracts:

uv run python -m pytest tests/unit_tests/models/test_weight_update_manifest.py tests/unit_tests/models/test_auto_bridge.py::TestAutoBridge::test_export_hf_weights tests/unit_tests/doc_consistency/test_readme_consistency.py::test_rl_integration_get_model_calls_supply_required_keyword_only_arguments tests/unit_tests/doc_consistency/test_readme_consistency.py::test_rl_integration_builds_a_finalized_config_with_runtime_objects -q
19 passed

Additional gates:

uv run mypy --strict src/megatron/bridge/models/conversion/weight_update_manifest.py
Success: no issues found in 1 source file

git diff --check
passed

uv run pre-commit run --all-files
passed

Scope and non-goals

  • The recorder and manifest do not retain tensor payloads.
  • Completeness is checked relative to the producer's recorded inventory; this does not independently prove the producer exported every model tensor.
  • The digest detects accidental metadata corruption. It is not tensor-content hashing or authentication.
  • No NCCL, ZMQ, HTTP, storage, acknowledgement, retry, runtime activation, or full-resync framework is added.
  • No conversion mapping, checkpoint, quantization, adapter lifecycle, dependency, CI workflow, or public conversion signature is changed.

Before this PR is ready for review

  • Focused fail-before/pass-after regression recorded.
  • Focused unit and adjacent tests pass.
  • Strict mypy and all-files pre-commit pass.
  • DCO sign-off present.

Commit: 3b3549c936f77aa0c97c7c0eaa733692960f1e1e.

@copy-pr-bot

copy-pr-bot Bot commented Aug 24, 2026

Copy link
Copy Markdown

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.

@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 733a69f

@yaoyu-33
yaoyu-33 force-pushed the yuya/landscape-weight-update-manifest-20260824 branch from 733a69f to 73d022f Compare August 24, 2026 19:24
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 73d022f

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
@yaoyu-33
yaoyu-33 force-pushed the yuya/landscape-weight-update-manifest-20260824 branch from 73d022f to 3b3549c Compare August 24, 2026 19:34
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 3b3549c

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant