fix(ckpt): resolve tied embeddings for builder-backed models at PP > 1

[kamran-nvidia]

What does this PR do ?

Fixes conversion of tied-embedding models built through the builder path (get_model_config() → get_model()) at pipeline parallelism > 1. Raised by @cuichenx on #6379 and #6401; the same bug affects Qwen3 on main since #6365.

The bug: the conversion helpers read share_embeddings_and_output_weights from the model's .config. On the builder path that is only the nested TransformerConfig; the flag lives on the outer ModelConfig, so the helpers fell back to "untied". At PP > 1 the last stage owns a separate output_layer.weight, which conversion then tried to convert. The post-import broadcast that copies the input embeddings into that output layer was also skipped. On a tied Qwen3 checkpoint at PP=2 (current main):

ValueError: Hugging Face checkpoint is missing mapped parameter(s):
  output_layer.weight -> lm_head.weight

Affected: scripts/conversion GPU import/export with --pp > 1, and AutoBridge.get_model() / load_hf_weights() at PP > 1 for families with USE_MODEL_CONFIG_FOR_CONVERSION = True (Qwen3 on main, plus the families migrating in #6379, #6383, #6392 and #6401). PP=1 and AutoBridge.import_ckpt() (single-process) never allocate the duplicate output parameter, and the provider path keeps the flag on .config, so the existing L0 PP=2 roundtrip (provider path) did not catch it.

Changelog

• MegatronModelBridge._share_embeddings_and_output_weights() now takes the unwrapped model instead of its config:
  • If model.config declares the flag (provider-built models and VLM wrapper configs), it is used as before, so the provider path is unchanged. Some VLM wrappers set their own module attribute to False while the config says tied; the config still wins.
  • Otherwise, which is the builder path, it falls back to the flag Megatron Core stores on the built model, checking language_model first because it owns the output layer.
• All callers in model_bridge.py and quant_bridge.py pass the unwrapped model. The helper is private, and the existing one-argument test stubs keep working.
• Tests:
  • Unit: helper resolution, including config-wins-over-wrapper and the language-model fallback, and build_conversion_tasks with the flag on the model.
  • Functional (L0 Qwen3 group, 2 GPUs): test_qwen3_tied_embeddings_builder_pp builds the tied toy Qwen3 at PP=2 through the builder path. It checks that the last-stage output_layer.weight equals the HF embed_tokens, and that the export reproduces every HF tensor with no lm_head.weight.

Verification

GB200, NeMo 26.10 container, Megatron Core a98aed553, tiny random-weight tied checkpoints, TP=1 / PP=2:

Check	Before	After
Qwen3, builder load + export (main)	ValueError (missing lm_head.weight)	output layer = embeddings; 46/46 tensors exact
Qwen3, scripts/conversion GPU import + export	ValueError	bit-exact, 46 tensors
Qwen2, builder load + export (with #6379)	ValueError (missing lm_head.weight)	output layer = embeddings; 50/50 tensors exact
EXAONE 4.0, builder load + export (with #6401)	ValueError (no mapping for output_layer.weight)	output layer = embeddings; 46/46 tensors exact

• Unit tests: 596 passed, 1 skipped across test_model_bridge.py, test_model_bridge_lora.py, test_fp8_param_export.py, quantization/test_quant_bridge.py, test_gemma3_vl_bridge.py, test_auto_bridge.py, test_gpt_model_config.py and models/qwen.
• The new functional test could not run locally: its fixture downloads the Qwen3 tokenizer and the compute node has no Hub access. The script it runs is the one used for the table above.
• ruff check / ruff format --check: clean (ruff 0.16.7; pre-commit itself not run).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• refactor(model): build Qwen2 and MiMo from builder-backed model configs #6379 and refactor(model): build Ling MoE V2 and EXAONE 4.0 from builder-backed model configs #6401 rebase onto this once it merges; no family-specific change is needed.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[kamran-nvidia]

/ok to test 3eb69a2