Load Qwen3.5 adapters with the pinned model's attention shape

[bradhilton]

Loading a Qwen3.5/3.6 LoRA adapter can fail offline, or silently convert with the wrong attention shape, even though the base model loaded from its pinned snapshot.

Cause

This affects an adapter stored as a vLLM/PEFT checkpoint (no ART manifest). It is converted by load_lora_tensors_for_megatron, which reads the raw adapter_config.json. Such configs carry base_model_name_or_path and revision, but no attention dimensions. The Qwen3.5 handler needs those dimensions to reorder q_proj LoRA-B rows, so it fell back to AutoConfig.from_pretrained(base_model_name_or_path, local_files_only=True):

• No revision passed. The adapter's revision was never passed on, and the helper's cache key was the name alone. The lookup went through the cache's refs/main, not the pinned snapshot.
• Offline failure. A cache holding only the pinned snapshot, with no refs/main, raised OSError during adapter load, after the pinned model had initialized. This is the failure Krennic hit in 062's MoE qualification, through TrainerRank.load_checkpoint.
• Wrong shape. If refs/main pointed at a different revision, the lookup silently returned that revision's shape for the pinned adapter.
• Save path. TrainerRank already writes the running model's attention dimensions into adapter configs it saves. The disk load paths never did.

Change

Each attention dimension now comes from the first of these that has it:

1. the adapter config, when set and not null, because it describes the tensors on disk;
2. the running model's provider: num_attention_heads, num_key_value_heads (from num_query_groups), head_dim (from kv_channels), hidden_size;
3. the base model's config at the adapter's revision.

• Loads pass the running model. load_lora_tensors_for_megatron(..., provider=...) fills dimensions the adapter config omits or nulls. The following all pass their provider:
  • TrainerRank's _checkpoint._load_adapter;
  • the Megatron service's _load_adapter_into_model in megatron/train.py;
  • the train/inference-mismatch integration harness.
• Anything still missing is looked up, not defaulted. Before, the handler looked up the base config only when heads were missing. A config with heads but no query groups defaulted to one group per head, and one without a head size used hidden_size // heads. Both are wrong for Qwen3.5/3.6, which use grouped queries and head_dim 256. Now any missing dimension comes from the pinned config. hidden_size // heads remains only as a last resort, when no source has a head size. Likewise, one group per head remains only when even the fetched config has no query-group count. One behavior change follows: if the adapter and model still leave a dimension missing, that lookup must succeed. Offline, it now fails loudly where the old code silently used those defaults, which are wrong for this model family.
• Save path unchanged. It still writes the provider's values into the configs it stores. Both paths share model_attention_dimensions.
• The lookup honors the pin. _qwen35_text_config takes the adapter's revision, passes it to AutoConfig.from_pretrained and caches by (name, revision).
  • A full commit hash resolves offline from its snapshot. Branches and tags still need their refs cached.
  • Unpinned adapters still resolve through refs/main, as before. An empty revision counts as unpinned.

Tests

tests/unit/test_qwen35_adapter_config.py runs on CPU only, with an isolated offline HF cache. There is no network access and no model weights are loaded.

• An offline cache with only the pinned snapshot resolves the pinned adapter's dimensions. The unpinned lookup still fails there, which documents the condition.
• refs/main points at a second revision with a different query-group count. Dimension lookups pinned to A, then B, then A return A's, B's and A's shape, so there is no cache collision. Unpinned and empty-revision lookups still follow refs/main.
• TrainerRank's _load_adapter converts a q_proj LoRA-B tensor with the running model's shape, with no base-model lookup, for both the dense and MoE Qwen3.5 handlers.
• null adapter dimensions are filled from the running model, and the conversion matches the model's grouped layout. The model's hidden size is not heads × head_dim, so deriving head_dim from it would also fail the test.
• Non-null adapter dimensions take precedence, and the revision is preserved.
• A model missing query groups, or missing a head size, gets that dimension from the lookup pinned to the adapter's revision.
• The adapter has heads and a head size, and the model has query groups but no head size. Together they are complete: both handlers convert with no lookup.
• An adapter with only heads gets query groups and head size from the pinned config, not one group per head or hidden_size // heads.
• Exporting a pinned adapter (to_vllm_lora_tensors) resolves its revision in the snapshot-only cache.
• The Megatron service's adapter load forwards the running model's provider.

All 13 test cases fail on main. Existing unit tests that touch adapter loading, checkpoints and the Qwen handlers pass locally (341). The module runs in CI's Megatron lightweight lane with the other trainer-rank tests, and the root step ignores it.

Krennic's independent CPU reproducer, from the 062 handoff, exercises the real snapshot-only failure and the drift case. It passed against this branch's first commit. His review of e2264fb3a then found the partial-model corner that this head fixes.

Not covered

• Other handlers read config.json by name without a revision: gpt_oss, gemma4 and nemotron_h.
• create_identity_lora passes normalization metadata without dimensions or revision.

Neither is on the failing path; each could follow the same pattern separately.

🤖 Generated with Claude Code