Repository navigation
Add native MoVA training and checkpoint conversion - #7
Draft
akhauriyash wants to merge 13 commits into
Draft
akhauriyash wants to merge 13 commits into
akhauriyash wants to merge 13 commits into
Conversation
Implement the xLLM-compatible MoVA forward path, heterogeneous dense-prefix layout, routed value experts, router semantics, FLOP accounting, and focused correctness tests. Keeping the model primitive separate from checkpoint conversion makes the training integration independently reviewable.
Validate and stream native or Hugging Face xLLM exports into topology-neutral MCore checkpoints, preserve the MoVA launch contract on reload, and test independent forward parity plus TP/PP/EP resharding. The converter refuses incomplete, mismatched, or overwritten sources.
Megatron FSDP materializes modules through reset_parameters. Teach the custom MoVA normalization and grouped value experts to honor meta construction and initialize correctly after sharding, avoiding multi-billion-parameter eager CUDA allocation.
Lock the three-layer dense prefix to global execution order across PP2/VPP4 chunks so future scheduling changes cannot silently alter the converted architecture.
Map the MoVA attention-input column projection onto Transformer Engine's standard QKV user-buffer role. This preserves baseline behavior and avoids disabling TP communication overlap for the shared transformer config.
Exercise the standard Transformer Engine router fusion against the independent value-projection reference and the production TP2 grouped-GEMM layout before enabling it in profiling.
Keep the checkpoint-compatible eager normalization as the default while exposing PyTorch RMSNorm as an isolated, correctness-tested performance candidate. The opt-in path preserves xLLM BF16 scale rounding and carries no new state.
xLLM evaluates its dot auxiliary objective independently on each context partition and averages replicated gradients. Restrict compatibility-mode statistics to TP while retaining CP-averaged logging, and cover the TP2-by-CP2 contract with a distributed gradient regression.
The xLLM HF exporter writes its bridge default rather than the source training coefficient. Require the native balancing type and coefficient explicitly, validate their consistency, log both values, and test the per-router mapping so checkpoint continuation cannot silently use the wrong loss scale.
MoVA conversion stores the native router-loss contract for both routed sublayers and the source RMSNorm epsilon. Restore those values under --use-checkpoint-args while preserving standard checkpoint behavior, so continuation cannot silently change the objective or normalization.
The production router now passes with_padding_mask into _apply_aux_loss. Keep the xLLM routing-map regression aligned with that callback contract so the port is testable on the current RL360 Megatron baseline.
Production Megatron restores the standard MoE router score function for non-MoVA checkpoints. Assert that inherited behavior while continuing to verify MoVA-only arguments remain isolated.
Current production MLP checkpoint merging reads low_memory_resume from Megatron global arguments. Supply and restore that normal launch default in the standalone eight-rank reshard regression.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
Why
Continuing the 36B MoVA checkpoint in Megatron requires forward-equivalent weights and model math, plus a deterministic mapping for live SGLang updates.
Validation