Skip to content

Support MoVA models and live weight sync - #86

Draft
akhauriyash wants to merge 5 commits into
prodfrom
yash.akhauri/mova-live-sync
Draft

Support MoVA models and live weight sync#86
akhauriyash wants to merge 5 commits into
prodfrom
yash.akhauri/mova-live-sync

Conversation

@akhauriyash

Copy link
Copy Markdown
Collaborator

What changed

  • Registers and validates native Megatron MoVA arguments.
  • Builds MoVATransformerConfig and the MoVA decoder specification for actor, reference, and old-actor models.
  • Maps packed Q/G/K[/V], routed value experts, grouped norms, gates, and selection-only router-bias buffers to SGLang names.
  • Reshards Megatron input-sharded value experts into SGLang output-sharded tensors at update time.
  • Keeps MoVA value experts on attention TP rather than FFN expert parallel groups.
  • Rejects unsafe in-place updates, flushes stale routed-V caches, and optionally verifies every rollout engine's weight version.

Why

MoVA rollouts must preserve training/inference forward semantics while retaining SGLang's communication-free output-sharded value projection during decode.

Validation

  • 59 focused tests pass after merging current prod, including current K3 coverage.
  • Ruff passes on all changed model, sync, version-check, and test files.
  • Broadcast and P2P staging contracts cover all 64 value-expert shards and router-bias buffers.

Convert gated dense and routed-value attention tensors into the canonical xLLM inference layout. Keep MoVA value experts on attention TP rather than FFN expert parallelism, synchronize router bias buffers, and reject stale cache-preserving updates.
Register the MoVA and xLLM router CLI contract in Miles and select Megatron's MoVA config and heterogeneous block spec whenever value experts are enabled. Fail early on incompatible provider, checkpoint-conversion, and attention settings while preserving the ordinary Transformer provider unchanged.
Add a cheap opt-in acceptance check that queries every rollout engine after each synchronization. Keep the existing random CI sample as the default and avoid full tensor comparisons.
Include the upstream K3 estimator before pinning the acceptance image. Resolve the shared argument-test addition by retaining both MoVA and K3 coverage.
Weight-version metadata must never advance after a failed or timed-out RDMA write. Drain every submitted transfer, clear the queue reliably, and surface aggregate failures to abort the update before rollout resumes.
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