Skip to content

Unify LoRA target selection with HF model layouts - #3318

Merged
yueming-yuan merged 95 commits into
mainfrom
yueming/hf-lora-targets
Sep 23, 2026
Merged

yueming-yuan merged 95 commits into
mainfrom
yueming/hf-lora-targets

Conversation

@yueming-yuan

@yueming-yuan yueming-yuan commented Sep 18, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #3620 (LoRA arguments/utilities and HF config relocation only). Review this PR against that base for the HF target-selection changes.

ci-sglang-pr: #40242

Use HF module names as the source of truth for LoRA targets. Builds on #3291 and SGLang #40242.

  • Share HF target groups through --target-modules attn,mlp,unembed; ordinary LoRA can mix groups with explicit HF targets, while Tinker accepts only whole groups. Exclude DSA indexers.
  • Convert HF → checkpoint → Megatron names at the backend boundary; reject unsupported partial fused targets.
  • Build adapter configs from exported names; preserve Inkling's native adapter format.

Review fixes

  • Emit accepted canonical aliases for fully selected expert FC1 targets; grouped-expert adapters retain their fused layout.

  • Allow fully selected GDN input projections in canonical mode; reject partial expert/GDN selections.

  • Replace Tinker training flags with --target-modules attn,mlp,unembed; default to all three groups and derive SDK flags from the selection. Reject explicit module names and exclusions in Tinker.

  • Exclude packed gate_up_proj in both GLM launchers when KEEP_MOE_LORA=0.

  • Apply scoped HF exclusions to resolved module names before retaining the final target selection; ordinary HF selection does not initialize Bridge.

At ba82dee23a, pre-commit run --all-files and git diff --check pass. A CPU A/B probe against 62b70f94fe, using the real launcher functions and Transformers meta-models, distinguishes both bugs from their fixes; default launcher commands remain unchanged. Added CPU cases cover both launchers and all-layer/single-layer exclusions. The full local launcher pytest gate was blocked during collection by ModuleNotFoundError: No module named 'sglang'; no new GPU validation was performed.

Model validation and Multi-LoRA support fixes are tracked in #3584.

At 53e27ae330, a CPU A/B probe against 294fe27e766 confirms that five explicit-target/exclusion cases change from accepted to rejected, while all eight training-group flag combinations preserve their behavior. pre-commit run --all-files and git diff --check pass. No regression test files were added, as requested; no GPU validation was run for this Tinker argument fix.

At c005f80efd, CPU A/B probes against 53e27ae330 reproduce and distinguish the canonical expert/GDN initialization failures for Qwen3-MoE, GPT-OSS, Qwen3-Next, and Qwen3.5 text. Probes use native Transformers 5.12.1 meta-model names and the canonical matcher method from pinned Bridge 2e09c234a3272285140224d0d698593418b55ba5. Standard LoRA outputs and dense canonical split selections are preserved; partial expert/GDN selections remain rejected. pre-commit run --all-files and git diff --check pass. No regression test files were added, as requested. Full GPU injection, training, and adapter export were not validated by these CPU probes.

Kimi K3 integration

Define Kimi K3 HF defaults to match its existing native training recipe. The launcher now uses all-linear instead of maintaining its own target list; omitting targets or selecting attn,mlp produces the same selection. Explicit targets can still omit routed-expert w2, and the native backend rejects unsupported layouts. Shared-outer adapter names drop the expert-index wildcard; SGLang #40242 resolves exported w1/w2/w3 names to its LoRA pool modules.

Validation: an extracted argument-path A/B probe reproduced the original undefined-layout error; the updated defaults were compared against the previous launcher's exact target set for K3, K3-4layer, and K3-4layer-64experts. The HF loader/model introspection were stubbed; these checks do not exercise backend initialization. Kimi K3 manual launcher snapshot checks: 10 passed. No new regression tests or GPU validation.

Point the image at radixark/Megatron-Bridge@bridge-new (upstream main @ 8ec5c9cdc
plus the still-open #6069 cherry-pick) instead of the legacy bridge branch, so CI
can surface what breaks against the current Megatron-LM miles-main pin.
…ridge-new

The bridge-mode paths break against radixark/Megatron-Bridge@bridge-new
(= NVIDIA upstream main); every bridge-mode e2e in #3291 fails on the first
weight sync with `ValueError: not enough values to unpack (expected 3, got 2)`.

* hf_weight_iterator_bridge: upstream Bridge yields the two-field
  HFWeightTuple unless asked for source names; pass
  `with_megatron_names=True` when the export method accepts it, and read the
  third field as a tuple of 0/1/N Megatron parameter names
  (`HFSourcedWeightTuple.megatron_param_names`) instead of a single string.
  Padding and quantization rules use the first source; a tensor without a
  Megatron source (HF-only passthrough) is passed through unquantized. The
  bucketing groupby now keys on the full source tuple, so a packed
  grouped-expert tensor stays one unit. The radixark `bridge` branch (no
  keyword, single-name third field) keeps working.
* lora/utils: export_adapter_weights may yield 2 or 3 fields; only name and
  weight are used. Pass `share_expert_adapters=False` explicitly: upstream
  LoRA defaults it to True, which would silently switch miles to one shared
  adapter per expert group.
* model_provider / lora/bridge: `provider.dsa_attention_backend` no longer
  exists on Megatron-Bridge main; map `--dsa-attention-backend` onto
  Megatron-Core's native `dsa_kernel_backend` ("tilelang" -> "tilelang",
  "megatron" -> "none"), honouring an explicit `--dsa-kernel-backend`.
* test_tinker_cookbook_gateway: install the cookbook with --no-deps (plus the
  three deps the image lacks). Its `transformers<=5.5.4` pin downgraded the
  container's transformers 5.12.1, which breaks `import megatron.bridge`
  (Exaone 4.5 bridge needs >=5.10) for every later test in that shard.

Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…ntime_config

tests/fast/backends/megatron_utils/test_bridge_mtp_detachment.py exec()s
this function's AST in a bare namespace to avoid GPU-only Megatron imports,
so a module-level name raises NameError there. A function-local import
resolves at call time and keeps that CPU test green.

Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
…st shim

_install_bridge_pp_group_unwrap wrapped MegatronParamMapping.broadcast_obj_from_pp_rank
as (self, obj, name=None), but Bridge's signature has been (self, obj, cache_key=None)
since upstream #608. The fork's `bridge` branch never passed the keyword, so the
mismatch stayed silent; Megatron-Bridge main calls broadcast_obj_from_pp_rank(...,
cache_key=...) from the MoE expert mappings (param_mapping.py) and the grouped
MXFP8 export path (quant_bridge.py / quant_mapping.py), which raised

  TypeError: broadcast_obj_from_pp_rank() got an unexpected keyword argument 'cache_key'

in CI round 5 (test_qwen3_5_35b_a3b_lora_ci, test_nemotron_3_ultra_4layer_ci).
Forward *args/**kwargs so the shim tracks upstream's signature.

Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
@yueming-yuan yueming-yuan changed the title Define shared HF LoRA target groups Unify LoRA target selection with HF model layouts Sep 18, 2026
@yueming-yuan
yueming-yuan marked this pull request as ready for review September 18, 2026 22:24

@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 repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

Tip: disable this comment in your organization's Code Review settings.

@yueming-yuan yueming-yuan added run-ci-lora run-ci-multi-lora Run Tinker cookbook and multi-tenant gateway GPU e2e tests labels Sep 18, 2026
@yueming-yuan
yueming-yuan changed the base branch from main to bump/megatron-bridge-bridge-new September 18, 2026 23:52
@yueming-yuan
yueming-yuan force-pushed the yueming/hf-lora-targets branch from ead988a to 1fc4938 Compare September 18, 2026 23:52
Resolve duplicate moves to the existing refactored files; the resulting tree is identical to the pre-stack tree.
@yueming-yuan
yueming-yuan changed the base branch from main to yueming/lora-arguments September 23, 2026 04:52
@yueming-yuan
yueming-yuan added this pull request to stack #3627 September 23, 2026 04:57

@Zhichenzzz Zhichenzzz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTHumans

@yueming-yuan

Copy link
Copy Markdown
Collaborator Author

/rerun-test tests/e2e/megatron/model_scripts/test_kimi_k3_4layer_lora_ci.py

@github-actions

github-actions Bot commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

❌ tests/e2e/megatron/model_scripts/test_kimi_k3_4layer_lora_ci.py failed on stage-c-8-gpu-h200 in 11m24s — workflow run

The execution job failed; inspect the workflow run for the failing step.

@yueming-yuan

Copy link
Copy Markdown
Collaborator Author

/rerun-test tests/e2e/megatron/model_scripts/test_kimi_k3_4layer_lora_ci.py

@github-actions

github-actions Bot commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

✅ tests/e2e/megatron/model_scripts/test_kimi_k3_4layer_lora_ci.py passed on stage-c-8-gpu-h200 in 16m13s — workflow run

Keep the new sampling-mask import and the relocated LoRA utility import in generate_endpoint_utils.
Bring the main conflict resolution from #3620 into #3318.
Base automatically changed from yueming/lora-arguments to main September 23, 2026 22:43
Main 1a3fd73 has the same tree as the already-merged #3620 head 1d18255. Preserve the existing #3318 tree while recording the squash commit as an ancestor.
@yueming-yuan
yueming-yuan merged commit 7d1d15b into main Sep 23, 2026
11 of 16 checks passed
@yueming-yuan
yueming-yuan deleted the yueming/hf-lora-targets branch September 23, 2026 22:47
@yueming-yuan yueming-yuan mentioned this pull request Sep 24, 2026
7 of 15 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-lora run-ci-multi-lora Run Tinker cookbook and multi-tenant gateway GPU e2e tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants