Skip to content

feat(deepseek-v4): add native LoRA construction and conversion bridge - #2772

Draft
zhoutong-hai wants to merge 2 commits into
radixark:mainfrom
zhoutong-hai:codex/deepseek-v4-lora-bridge
Draft

feat(deepseek-v4): add native LoRA construction and conversion bridge#2772
zhoutong-hai wants to merge 2 commits into
radixark:mainfrom
zhoutong-hai:codex/deepseek-v4-lora-bridge

Conversation

@zhoutong-hai

@zhoutong-hai zhoutong-hai commented Aug 26, 2026

Copy link
Copy Markdown

Depends on #2771. The grouped wo_a execution path must land before this PR enables and wraps that target.

Problem

Current Miles pins Megatron-Bridge at 7f0fb345, which does not register DeepseekV4ForCausalLM. Letting Bridge select a provider therefore cannot construct the current native V4 graph. Treating it as V3 or using incomplete inferred mappings is worse: a strict current-main load exposed thousands of unmapped or missing checkpoint parameters.

V4 also owns both a main-attention wq_b and a nested DSA indexer linear_wq_b. The generic MLA alias maps bare wq_b to the indexer, so an unqualified V4 treatment can wrap the wrong module.

Finally, a V4 portable export is valid only as an adapter against the pinned native base. The conversion-only bridge cannot safely claim to emit a complete merged base, and a failed adapter conversion must not leave a success marker or strand peer ranks in a barrier.

Fix

  • Detect the exact V4 architecture and construct it through the Miles native provider/spec.
  • Apply PEFT in a pre-wrap hook, before DDP construction.
  • Qualify the five main V4 attention targets so wq_b cannot also match the nested indexer.
  • Register a conversion-only Bridge for the current native graph, with explicit attention mappings and explicit replicated/column/automatic semantics for hyper-connections, compressors, routing tables, attention sink, and indexer weights.
  • Keep portable V4 output adapter-only against the pinned base.
  • Make required adapter export fail closed on empty or failed Bridge conversion.
  • Collect writer/export failures across all ranks before raising, so rank-local I/O failure cannot leave peers at the trailing barrier.
  • Write .complete only after a valid adapter is present.
  • Reject unsupported native V4 Multi-LoRA instead of silently falling back.

This does not duplicate the MXFP4 ingestion work in #2717. It is scoped to current main and its pinned Bridge/native V4 graph. If the large Megatron bump in #2706 lands first, the mapping layer will need a focused rebase because that PR intentionally renames and relayouts the V4 graph. The general PEFT-export work in #2580 is related; this PR adds the V4 requirement that portable export cannot remain best-effort.

Validation

  • 8 focused tests cover mapping ownership, native-provider dispatch, main/indexer target qualification, adapter-only layout, exact architecture selection, empty export, completion-marker failure, and two-rank writer-failure propagation.
  • The writer-failure case confirms both ranks receive the same rank-0 I/O error instead of deadlocking.
  • Tests run against the exact pinned Megatron-Bridge and Miles Megatron sources; unavailable GPU-only packages were stubbed only for import.
  • Ruff and diff checks pass.
  • The underlying current-main graph and explicit mappings were exercised on a 32-rank V4 LoRA run: strict base load reached zero unmapped/missing warnings, then completed finite updates and trainer-to-rollout publications.

Remaining qualification

A successful adapter-only artifact still needs to be loaded back against the pinned base before portable export is called complete. This PR remains draft until that round trip is recorded and #2771 lands.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

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