Repository navigation
Unify LoRA target selection with HF model layouts - #3318
Conversation
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>
There was a problem hiding this comment.
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.
ead988a to
1fc4938
Compare
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Resolve duplicate moves to the existing refactored files; the resulting tree is identical to the pre-stack tree.
|
/rerun-test tests/e2e/megatron/model_scripts/test_kimi_k3_4layer_lora_ci.py |
|
❌ The execution job failed; inspect the workflow run for the failing step. |
|
/rerun-test tests/e2e/megatron/model_scripts/test_kimi_k3_4layer_lora_ci.py |
|
✅ |
Keep the new sampling-mask import and the relocated LoRA utility import in generate_endpoint_utils.
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.
--target-modules attn,mlp,unembed; ordinary LoRA can mix groups with explicit HF targets, while Tinker accepts only whole groups. Exclude DSA indexers.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_projin both GLM launchers whenKEEP_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-filesandgit diff --checkpass. A CPU A/B probe against62b70f94fe, 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 byModuleNotFoundError: 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 against294fe27e766confirms 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-filesandgit diff --checkpass. No regression test files were added, as requested; no GPU validation was run for this Tinker argument fix.At
c005f80efd, CPU A/B probes against53e27ae330reproduce 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 Bridge2e09c234a3272285140224d0d698593418b55ba5. Standard LoRA outputs and dense canonical split selections are preserved; partial expert/GDN selections remain rejected.pre-commit run --all-filesandgit diff --checkpass. 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-linearinstead of maintaining its own target list; omitting targets or selectingattn,mlpproduces the same selection. Explicit targets can still omit routed-expertw2, and the native backend rejects unsupported layouts. Shared-outer adapter names drop the expert-index wildcard; SGLang #40242 resolves exportedw1/w2/w3names 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.