feat(modelexpress): add S3 delta weight update backend

[nv-hwoo]

Summary

This PR adds ModelExpress as a VIME weight-update backend, with initial support for S3-based canonical
delta weight synchronization.

Instead of publishing a complete checkpoint after every training update, the trainer captures an initial
Hugging Face baseline and publishes subsequent policy updates as versioned XOR deltas through
ModelExpress.

  flowchart LR
      T[Megatron trainer] -->|canonical HF weights| MX[ModelExpress trainer client]
      MX -->|XOR delta artifacts| S3[S3-compatible storage]
      MX -->|version lifecycle| C[ModelExpress catalog]
      C --> V[vLLM ModelExpress backend]
      S3 --> V

Loading

What changed

• Adds modelexpress as a valid --update-weight-transport.
• Lazily imports the ModelExpress implementation only when that transport is selected.
• Configures vLLM to use its modelexpress weight-transfer backend.
• Converts Megatron weights into canonical Hugging Face tensor chunks.
• Distributes chunk ownership across trainer ranks so each chunk is published once.
• Automatically creates the catalog-only, READY v0 baseline before initializing rollout engines.
• Captures the initial seed checkpoint as the trainer’s delta baseline.
• Publishes later versions as:
  • S3 XOR deltas by default.
  • Periodic full Hugging Face checkpoints when full_hf_checkpoint_interval is configured.
• Drives the vLLM update lifecycle:
  • pause generation
  • flush cache
  • start update
  • install the ModelExpress version
  • finish update
  • resume generation
• Reports delta density, transferred bytes, staging time, publication time, and activation time.

Configuration

  --update-weight-transport modelexpress
  --modelexpress-config '{
    "model_name": "policy",
    "server_url": "modelexpress:8101",
    "initial_base_version_id": "policy-v0",
    "seed_checkpoint_path": "/models/policy",
    "refit_checkpoint_dir": "/tmp/mx-refit",
    "s3_uri_prefix": "s3://weight-updates/policy",
    "s3_endpoint_url": "http://minio:9000",
    "s3_region_name": "us-east-1",
    "full_hf_checkpoint_interval": 10
  }'

Scope

This initial integration supports S3-compatible object storage and the canonical XOR-delta path.
ModelExpress P2P transfer and Dynamo-specific external rollout integration are outside this PR.

Testing

• 114 targeted unit tests passed.
• Added coverage for:
  • argument parsing and validation
  • lazy updater construction
  • vLLM backend configuration
  • automatic and retryable v0 registration
  • distributed tensor-chunk ownership
  • delta and periodic full-checkpoint publication
  • update lifecycle ordering and failure handling
  • metrics aggregation
• Ruff and git diff --check pass.

E2E Validation

Ran Qwen3-30B-A3B on B200 nodes with following setup:

• 4 Trainer GPUs (TP2 / PP2 / EP2)
• 4 Rollout GPUs (TP4)
• Run 50 rollout steps
• Full HF checkpoint interval for N=25

Delta
• Versions: v1–v24, v26–v49
• Mean: 46.68s
• Median: 44.94s

Full HF checkpoint
• Versions: v25
• Mean: 545.56s
• Median: 545.56s

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34583849 | 📁 Comparing 22aabda against latest (ce92eff)

  🔍 Preview build  

45 files changed · ± 45 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...

[gemini-code-assist]

Code Review

This pull request introduces a new weight transport option, modelexpress, allowing Megatron weights to be published through ModelExpress and installed in vLLM. It adds the UpdateWeightFromModelExpress class, CLI arguments, and comprehensive unit tests. The review comments suggest several robustness improvements, including validating required configuration keys, defensively handling potential None values in the configuration, and safely handling cases where metrics might be empty.

[aoshen02]

Hi, It does not need to change the vllm code?

[nv-hwoo]

Hi @aoshen02, we use vLLM's Weight Transfer Engine plugin to define our custom extension in our ModelExpress codebase, so we don't need any direct changes in vLLM.

Also, when used with existing ModelExpress client in vLLM engine, we also support P2P RDMA fan-out across the rollout engines during weight update or cold-start (due to crash or restart) as well as S3 fallback. We don't have a clean doc yet but here are the recent PRs:

• feat(refit): add P2P-first active refit ai-dynamo/modelexpress#740
• fix(refit): reconcile distributed RL cold starts ai-dynamo/modelexpress#738
• feat(refit): add RL cold-start load strategy chain ai-dynamo/modelexpress#727
• feat(refit): add target-version delta replay ai-dynamo/modelexpress#710

[aoshen02]

Hi @aoshen02, we use vLLM's Weight Transfer Engine plugin to define our custom extension in our ModelExpress codebase, so we don't need any direct changes in vLLM.

Also, when used with existing ModelExpress client in vLLM engine, we also support P2P RDMA fan-out across the rollout engines during weight update or cold-start (due to crash or restart) as well as S3 fallback. We don't have a clean doc yet but here are the recent PRs:

• feat(refit): add P2P-first active refit ai-dynamo/modelexpress#740
• fix(refit): reconcile distributed RL cold starts ai-dynamo/modelexpress#738
• feat(refit): add RL cold-start load strategy chain ai-dynamo/modelexpress#727
• feat(refit): add target-version delta replay ai-dynamo/modelexpress#710

Please join this vllm slack channel: