From 43adc88c02abf199079dc7d178a8c42d7fc6710e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 04:20:11 +0000 Subject: [PATCH 1/5] Load Qwen3.5 adapters with the pinned model's attention shape Adapter conversion looked the base model's config up by name, ignoring the adapter's pinned revision. An offline cache holding only the pinned snapshot failed, and a different refs/main silently gave the wrong shape. Checkpoint loading now fills missing attention dimensions from the running model, and the by-name fallback resolves the adapter's revision. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/prek.yml | 2 + .../model_support/handlers/qwen3_5.py | 10 +- src/art/megatron/model_support/lora_disk.py | 5 +- src/art/trainer_rank/_checkpoint.py | 14 +- src/art/trainer_rank/_impl.py | 25 ++- tests/unit/test_qwen35_adapter_config.py | 157 ++++++++++++++++++ 6 files changed, 196 insertions(+), 17 deletions(-) create mode 100644 tests/unit/test_qwen35_adapter_config.py diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 17e6913da..07403f7cf 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -224,6 +224,7 @@ jobs: tests/unit/test_prefix_tree_grad_parity.py \ tests/unit/test_prefix_tree_packing.py \ tests/unit/test_trainer_rank_handoff_budget.py \ + tests/unit/test_qwen35_adapter_config.py \ tests/unit/test_trainer_rank_physical_reserve.py \ tests/unit/test_trainer_rank_validation.py \ tests/unit/test_trainer_rank_weird_shapes.py \ @@ -265,6 +266,7 @@ jobs: --ignore=tests/unit/test_prefix_tree_grad_parity.py \ --ignore=tests/unit/test_prefix_tree_packing.py \ --ignore=tests/unit/test_trainer_rank_handoff_budget.py \ + --ignore=tests/unit/test_qwen35_adapter_config.py \ --ignore=tests/unit/test_trainer_rank_physical_reserve.py \ --ignore=tests/unit/test_trainer_rank_validation.py \ --ignore=tests/unit/test_trainer_rank_weird_shapes.py \ diff --git a/src/art/megatron/model_support/handlers/qwen3_5.py b/src/art/megatron/model_support/handlers/qwen3_5.py index 309046311..310020249 100644 --- a/src/art/megatron/model_support/handlers/qwen3_5.py +++ b/src/art/megatron/model_support/handlers/qwen3_5.py @@ -509,11 +509,12 @@ def _is_self_attn_q_proj_lora_b(key: str) -> bool: @lru_cache(maxsize=8) -def _qwen35_text_config(base_model_name_or_path: str) -> Any: +def _qwen35_text_config(base_model_name_or_path: str, revision: str | None) -> Any: from transformers import AutoConfig config = AutoConfig.from_pretrained( base_model_name_or_path, + revision=revision, local_files_only=True, trust_remote_code=True, ) @@ -529,7 +530,12 @@ def _qwen35_attention_dims(adapter_config: dict[str, Any]) -> tuple[int, int, in base_model = adapter_config.get("base_model_name_or_path") if not base_model: raise RuntimeError("Qwen3.5 LoRA adapter config is missing base model path") - config = _qwen35_text_config(str(base_model)) + # Resolve the adapter's pinned snapshot, not whatever the name + # currently points at (an offline cache may hold only the pin). + revision = adapter_config.get("revision") + config = _qwen35_text_config( + str(base_model), None if revision is None else str(revision) + ) num_heads = getattr(config, "num_attention_heads") num_groups = getattr(config, "num_key_value_heads", num_heads) head_dim = getattr(config, "head_dim", None) diff --git a/src/art/megatron/model_support/lora_disk.py b/src/art/megatron/model_support/lora_disk.py index f0b01183b..970d0882f 100644 --- a/src/art/megatron/model_support/lora_disk.py +++ b/src/art/megatron/model_support/lora_disk.py @@ -127,6 +127,7 @@ def load_lora_tensors_for_megatron( lora_path: str | Path, *, handler: ModelSupportHandler | None = None, + adapter_config: dict[str, Any] | None = None, allow_unvalidated_arch: bool = False, ) -> dict[str, torch.Tensor]: resolved_handler = resolve_lora_handler( @@ -136,5 +137,7 @@ def load_lora_tensors_for_megatron( ) return resolved_handler.from_vllm_lora_tensors( load_vllm_lora_tensors(lora_path), - adapter_config=load_adapter_config(lora_path), + adapter_config=( + load_adapter_config(lora_path) if adapter_config is None else adapter_config + ), ) diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py index 1c3b33d97..cc26179c5 100644 --- a/src/art/trainer_rank/_checkpoint.py +++ b/src/art/trainer_rank/_checkpoint.py @@ -1488,11 +1488,23 @@ def _load_adapter( ) -> dict[str, torch.Tensor]: if source.manifest is None: from art.megatron.model_support.lora_disk import ( + load_adapter_config, load_lora_tensors_for_megatron, ) + from art.trainer_rank._impl import _provider_attention_dimensions + # Convert with the running model's attention shape where the adapter + # omits it, instead of looking the base model up again by name. + adapter_config = { + **_provider_attention_dimensions( + getattr(trainer.runtime, "provider", None) + ), + **load_adapter_config(source.path), + } loaded = load_lora_tensors_for_megatron( - source.path, handler=trainer.runtime.model_support_handler + source.path, + handler=trainer.runtime.model_support_handler, + adapter_config=adapter_config, ) return {key: value for key, value in loaded.items() if key in set(keys)} safe_open = importlib.import_module("safetensors").safe_open diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index d4718cab0..de0807644 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -162,6 +162,17 @@ def _checkpoint_prefetch_executor() -> ThreadPoolExecutor: return _CHECKPOINT_PREFETCH_EXECUTOR[1] +def _provider_attention_dimensions(provider: Any) -> dict[str, int]: + """The running model's attention shape, in adapter-config keys.""" + dimensions = { + "num_attention_heads": getattr(provider, "num_attention_heads", None), + "num_key_value_heads": getattr(provider, "num_query_groups", None), + "head_dim": getattr(provider, "kv_channels", None), + "hidden_size": getattr(provider, "hidden_size", None), + } + return {key: int(value) for key, value in dimensions.items() if value is not None} + + class _AdapterConfig(TypedDict): base_model_name_or_path: str revision: NotRequired[str | None] @@ -2626,19 +2637,7 @@ def _validate_checkpoint_adapter_config( "adapter_config['base_model_name_or_path'] must be a string" ) if base_model.startswith(("Qwen/Qwen3.5-", "Qwen/Qwen3.6-", "Qwen/Qwen3.8-")): - dimensions = { - "num_attention_heads": getattr( - self.runtime.provider, "num_attention_heads", None - ), - "num_key_value_heads": getattr( - self.runtime.provider, "num_query_groups", None - ), - "head_dim": getattr(self.runtime.provider, "kv_channels", None), - "hidden_size": getattr(self.runtime.provider, "hidden_size", None), - } - for key, value in dimensions.items(): - if value is not None: - config[key] = int(value) + config.update(_provider_attention_dimensions(self.runtime.provider)) if not isinstance(rank, int) or isinstance(rank, bool): raise TypeError("adapter_config['r'] must be an integer") if not isinstance(config_alpha_value, int | float) or isinstance( diff --git a/tests/unit/test_qwen35_adapter_config.py b/tests/unit/test_qwen35_adapter_config.py new file mode 100644 index 000000000..79db0a3f3 --- /dev/null +++ b/tests/unit/test_qwen35_adapter_config.py @@ -0,0 +1,157 @@ +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +qwen35 = pytest.importorskip("art.megatron.model_support.handlers.qwen3_5") + +REPO = "Qwen/Qwen3.6-35B-A3B" +PIN = "995ad96eacd98c81ed38be0c5b274b04031597b0" +OTHER = "1" * 40 +LAYER = "base_model.model.model.language_model.layers.0.self_attn.q_proj" + + +@pytest.fixture +def hub_cache(tmp_path, monkeypatch): + """An offline cache that holds only the snapshots a test adds.""" + import huggingface_hub.constants + + repo = tmp_path / "models--Qwen--Qwen3.6-35B-A3B" + + def add_snapshot(revision: str, groups: int, *, main: bool = False) -> None: + snapshot = repo / "snapshots" / revision + snapshot.mkdir(parents=True) + (snapshot / "config.json").write_text( + json.dumps( + { + "model_type": "llama", + "num_attention_heads": 16, + "num_key_value_heads": groups, + "head_dim": 256, + "hidden_size": 2048, + } + ) + ) + if main: + (repo / "refs").mkdir(exist_ok=True) + (repo / "refs" / "main").write_text(revision) + + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + qwen35._qwen35_text_config.cache_clear() + yield add_snapshot + qwen35._qwen35_text_config.cache_clear() + + +def _dims(revision: str | None = None) -> tuple[int, int, int]: + config: dict[str, object] = {"base_model_name_or_path": REPO} + if revision is not None: + config["revision"] = revision + return qwen35._qwen35_attention_dims(config) + + +def test_attention_dims_resolve_the_pinned_snapshot_offline(hub_cache): + hub_cache(PIN, groups=2) + assert _dims(PIN) == (16, 2, 256) + # By name alone, the lookup needs refs/main, which this cache lacks. + with pytest.raises(OSError): + _dims() + + +def test_a_different_main_revision_does_not_change_a_pinned_adapter(hub_cache): + hub_cache(PIN, groups=2) + hub_cache(OTHER, groups=4, main=True) + assert [_dims(PIN), _dims(OTHER), _dims(PIN)] == [ + (16, 2, 256), + (16, 4, 256), + (16, 2, 256), + ] + assert _dims() == (16, 4, 256) # unpinned adapters still follow main + + +def _write_adapter(path: Path, rows: int) -> dict[str, torch.Tensor]: + from art.megatron.model_support.lora_disk import save_vllm_lora_tensors + + tensors = { + f"{LAYER}.lora_A.weight": torch.randn(2, 32), + f"{LAYER}.lora_B.weight": torch.randn(rows, 2), + } + # As published for a pinned base model, without attention dimensions. + config = { + "base_model_name_or_path": REPO, + "revision": PIN, + "r": 2, + "lora_alpha": 32, + "target_modules": ["q_proj"], + "art_lora_format": "vllm", + } + save_vllm_lora_tensors(path, tensors, config) + return tensors + + +@pytest.mark.parametrize("handler", ["QWEN3_5_DENSE_HANDLER", "QWEN3_5_MOE_HANDLER"]) +def test_checkpoint_load_converts_with_the_running_models_attention_shape( + tmp_path, monkeypatch, handler +): + from art.trainer_rank import _checkpoint + + def lookup(*_args): + raise AssertionError("adapter conversion looked the base model up again") + + monkeypatch.setattr(qwen35, "_qwen35_text_config", lookup) + provider = SimpleNamespace( + num_attention_heads=4, num_query_groups=2, kv_channels=8, hidden_size=32 + ) + trainer = SimpleNamespace( + runtime=SimpleNamespace( + provider=provider, model_support_handler=getattr(qwen35, handler) + ) + ) + # Two query groups of two heads, each with query and gate rows of size 8. + tensors = _write_adapter(tmp_path, rows=2 * 2 * 2 * 8) + art_key = f"{LAYER}.lora_B.weight".replace(".language_model.layers.", ".layers.") + loaded = _checkpoint._load_adapter( + cast(Any, trainer), + cast(Any, SimpleNamespace(manifest=None, path=tmp_path)), + [art_key], + ) + expected = qwen35._qwen35_q_proj_lora_b_from_vllm( + tensors[f"{LAYER}.lora_B.weight"], + {"num_attention_heads": 4, "num_key_value_heads": 2, "head_dim": 8}, + ) + torch.testing.assert_close(loaded[art_key], expected) + + +def test_adapter_dimensions_take_precedence_over_the_running_model( + tmp_path, monkeypatch +): + from art.megatron.model_support import lora_disk + from art.trainer_rank import _checkpoint + + seen = [] + monkeypatch.setattr( + lora_disk, + "load_lora_tensors_for_megatron", + lambda path, **kwargs: seen.append(kwargs["adapter_config"]) or {}, + ) + _write_adapter(tmp_path, rows=8) + config = json.loads((tmp_path / "adapter_config.json").read_text()) + config["num_attention_heads"] = 1 + (tmp_path / "adapter_config.json").write_text(json.dumps(config)) + trainer = SimpleNamespace( + runtime=SimpleNamespace( + provider=SimpleNamespace(num_attention_heads=4, num_query_groups=2), + model_support_handler=qwen35.QWEN3_5_MOE_HANDLER, + ) + ) + _checkpoint._load_adapter( + cast(Any, trainer), + cast(Any, SimpleNamespace(manifest=None, path=tmp_path)), + [], + ) + assert seen[0]["num_attention_heads"] == 1 + assert seen[0]["num_key_value_heads"] == 2 + assert seen[0]["revision"] == PIN From 4c4b65d7ff13cbea2c3dbc36e54718162cf42a36 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 04:49:15 +0000 Subject: [PATCH 2/5] Fill adapter dimensions only from a complete model shape, on both load paths Treat null adapter dimensions as missing, use the running model's shape only when it has heads, query groups and head size, and apply the same fill on the Megatron service's adapter load. An empty revision means unpinned. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../model_support/handlers/qwen3_5.py | 2 +- src/art/megatron/model_support/lora_disk.py | 46 +++++- src/art/megatron/train.py | 7 +- src/art/trainer_rank/_checkpoint.py | 12 +- src/art/trainer_rank/_impl.py | 17 +-- tests/unit/test_qwen35_adapter_config.py | 140 +++++++++++++----- 6 files changed, 155 insertions(+), 69 deletions(-) diff --git a/src/art/megatron/model_support/handlers/qwen3_5.py b/src/art/megatron/model_support/handlers/qwen3_5.py index 310020249..e7ffebf58 100644 --- a/src/art/megatron/model_support/handlers/qwen3_5.py +++ b/src/art/megatron/model_support/handlers/qwen3_5.py @@ -532,7 +532,7 @@ def _qwen35_attention_dims(adapter_config: dict[str, Any]) -> tuple[int, int, in raise RuntimeError("Qwen3.5 LoRA adapter config is missing base model path") # Resolve the adapter's pinned snapshot, not whatever the name # currently points at (an offline cache may hold only the pin). - revision = adapter_config.get("revision") + revision = adapter_config.get("revision") or None config = _qwen35_text_config( str(base_model), None if revision is None else str(revision) ) diff --git a/src/art/megatron/model_support/lora_disk.py b/src/art/megatron/model_support/lora_disk.py index 970d0882f..9e16d42d8 100644 --- a/src/art/megatron/model_support/lora_disk.py +++ b/src/art/megatron/model_support/lora_disk.py @@ -19,6 +19,37 @@ safe_open = safetensors.safe_open +def model_attention_dimensions(provider: Any) -> dict[str, int]: + """The running model's attention shape, in adapter-config keys.""" + dimensions = { + "num_attention_heads": getattr(provider, "num_attention_heads", None), + "num_key_value_heads": getattr(provider, "num_query_groups", None), + "head_dim": getattr(provider, "kv_channels", None), + "hidden_size": getattr(provider, "hidden_size", None), + } + return {key: int(value) for key, value in dimensions.items() if value is not None} + + +def with_model_attention_dimensions( + adapter_config: dict[str, Any], provider: Any +) -> dict[str, Any]: + """Fill attention dimensions the adapter config omits or nulls from the model. + + Only a complete shape is used: a partial one could silently turn grouped + query attention into full multi-head attention. + """ + dimensions = model_attention_dimensions(provider) + if not {"num_attention_heads", "num_key_value_heads", "head_dim"} <= set( + dimensions + ): + return adapter_config + config = dict(adapter_config) + for key, value in dimensions.items(): + if config.get(key) is None: + config[key] = value + return config + + def _jsonable_config(value: Any) -> Any: if isinstance(value, dict): return {key: _jsonable_config(item) for key, item in value.items()} @@ -127,17 +158,24 @@ def load_lora_tensors_for_megatron( lora_path: str | Path, *, handler: ModelSupportHandler | None = None, - adapter_config: dict[str, Any] | None = None, + provider: Any = None, allow_unvalidated_arch: bool = False, ) -> dict[str, torch.Tensor]: + """Load an adapter in Megatron layout. + + With the running model's ``provider``, conversion uses its attention shape + wherever the adapter config lacks one, instead of looking the base model + up again by name. + """ resolved_handler = resolve_lora_handler( lora_path, handler, allow_unvalidated_arch=allow_unvalidated_arch, ) + adapter_config = load_adapter_config(lora_path) + if provider is not None: + adapter_config = with_model_attention_dimensions(adapter_config, provider) return resolved_handler.from_vllm_lora_tensors( load_vllm_lora_tensors(lora_path), - adapter_config=( - load_adapter_config(lora_path) if adapter_config is None else adapter_config - ), + adapter_config=adapter_config, ) diff --git a/src/art/megatron/train.py b/src/art/megatron/train.py index cf7d9452a..e203bcdd5 100644 --- a/src/art/megatron/train.py +++ b/src/art/megatron/train.py @@ -1006,6 +1006,7 @@ def _prepare_rl_training_state( job.source_adapter_path, runtime.rank, handler=runtime.model_support_handler, + provider=runtime.provider, optimizer=runtime.optimizer, ) if runtime.optimizer is None: @@ -1069,10 +1070,13 @@ def _load_adapter_into_model( rank: int, *, handler: Any | None = None, + provider: Any = None, optimizer: Any | None = None, ) -> dict[str, torch.Tensor]: print0(rank, "Loading adapter model from", lora_path) - adapter_model = load_lora_tensors_for_megatron(lora_path, handler=handler) + adapter_model = load_lora_tensors_for_megatron( + lora_path, handler=handler, provider=provider + ) load_adapter_into_model( model_chunks, adapter_model, @@ -2319,6 +2323,7 @@ def _prepare_kl_reference_logprobs( ref_adapter_path, runtime.rank, handler=runtime.model_support_handler, + provider=runtime.provider, ) loaded_ref_adapter = True return _precompute_reference_logprobs( diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py index cc26179c5..530a7f35e 100644 --- a/src/art/trainer_rank/_checkpoint.py +++ b/src/art/trainer_rank/_checkpoint.py @@ -1488,23 +1488,13 @@ def _load_adapter( ) -> dict[str, torch.Tensor]: if source.manifest is None: from art.megatron.model_support.lora_disk import ( - load_adapter_config, load_lora_tensors_for_megatron, ) - from art.trainer_rank._impl import _provider_attention_dimensions - # Convert with the running model's attention shape where the adapter - # omits it, instead of looking the base model up again by name. - adapter_config = { - **_provider_attention_dimensions( - getattr(trainer.runtime, "provider", None) - ), - **load_adapter_config(source.path), - } loaded = load_lora_tensors_for_megatron( source.path, handler=trainer.runtime.model_support_handler, - adapter_config=adapter_config, + provider=getattr(trainer.runtime, "provider", None), ) return {key: value for key, value in loaded.items() if key in set(keys)} safe_open = importlib.import_module("safetensors").safe_open diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index de0807644..b51d21605 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -162,17 +162,6 @@ def _checkpoint_prefetch_executor() -> ThreadPoolExecutor: return _CHECKPOINT_PREFETCH_EXECUTOR[1] -def _provider_attention_dimensions(provider: Any) -> dict[str, int]: - """The running model's attention shape, in adapter-config keys.""" - dimensions = { - "num_attention_heads": getattr(provider, "num_attention_heads", None), - "num_key_value_heads": getattr(provider, "num_query_groups", None), - "head_dim": getattr(provider, "kv_channels", None), - "hidden_size": getattr(provider, "hidden_size", None), - } - return {key: int(value) for key, value in dimensions.items() if value is not None} - - class _AdapterConfig(TypedDict): base_model_name_or_path: str revision: NotRequired[str | None] @@ -2637,7 +2626,11 @@ def _validate_checkpoint_adapter_config( "adapter_config['base_model_name_or_path'] must be a string" ) if base_model.startswith(("Qwen/Qwen3.5-", "Qwen/Qwen3.6-", "Qwen/Qwen3.8-")): - config.update(_provider_attention_dimensions(self.runtime.provider)) + from art.megatron.model_support.lora_disk import ( + model_attention_dimensions, + ) + + config.update(model_attention_dimensions(self.runtime.provider)) if not isinstance(rank, int) or isinstance(rank, bool): raise TypeError("adapter_config['r'] must be an integer") if not isinstance(config_alpha_value, int | float) or isinstance( diff --git a/tests/unit/test_qwen35_adapter_config.py b/tests/unit/test_qwen35_adapter_config.py index 79db0a3f3..4b4617547 100644 --- a/tests/unit/test_qwen35_adapter_config.py +++ b/tests/unit/test_qwen35_adapter_config.py @@ -72,7 +72,9 @@ def test_a_different_main_revision_does_not_change_a_pinned_adapter(hub_cache): assert _dims() == (16, 4, 256) # unpinned adapters still follow main -def _write_adapter(path: Path, rows: int) -> dict[str, torch.Tensor]: +def _write_adapter( + path: Path, rows: int, **dimensions: object +) -> dict[str, torch.Tensor]: from art.megatron.model_support.lora_disk import save_vllm_lora_tensors tensors = { @@ -87,71 +89,129 @@ def _write_adapter(path: Path, rows: int) -> dict[str, torch.Tensor]: "lora_alpha": 32, "target_modules": ["q_proj"], "art_lora_format": "vllm", + **dimensions, } save_vllm_lora_tensors(path, tensors, config) return tensors +# Four heads in two query groups, with head_dim 8: the runtime's shape. +PROVIDER = SimpleNamespace( + num_attention_heads=4, num_query_groups=2, kv_channels=8, hidden_size=32 +) +ROWS = 2 * 2 * 2 * 8 # groups x (query + gate) x heads per group x head_dim +ART_KEY = f"{LAYER}.lora_B.weight".replace(".language_model.layers.", ".layers.") + + +def _expected(tensors: dict[str, torch.Tensor]) -> torch.Tensor: + return qwen35._qwen35_q_proj_lora_b_from_vllm( + tensors[f"{LAYER}.lora_B.weight"], + {"num_attention_heads": 4, "num_key_value_heads": 2, "head_dim": 8}, + ) + + +def _forbid_lookup(monkeypatch) -> None: + def lookup(*_args): + raise AssertionError("adapter conversion looked the base model up again") + + monkeypatch.setattr(qwen35, "_qwen35_text_config", lookup) + + @pytest.mark.parametrize("handler", ["QWEN3_5_DENSE_HANDLER", "QWEN3_5_MOE_HANDLER"]) def test_checkpoint_load_converts_with_the_running_models_attention_shape( tmp_path, monkeypatch, handler ): from art.trainer_rank import _checkpoint - def lookup(*_args): - raise AssertionError("adapter conversion looked the base model up again") - - monkeypatch.setattr(qwen35, "_qwen35_text_config", lookup) - provider = SimpleNamespace( - num_attention_heads=4, num_query_groups=2, kv_channels=8, hidden_size=32 - ) + _forbid_lookup(monkeypatch) trainer = SimpleNamespace( runtime=SimpleNamespace( - provider=provider, model_support_handler=getattr(qwen35, handler) + provider=PROVIDER, model_support_handler=getattr(qwen35, handler) ) ) - # Two query groups of two heads, each with query and gate rows of size 8. - tensors = _write_adapter(tmp_path, rows=2 * 2 * 2 * 8) - art_key = f"{LAYER}.lora_B.weight".replace(".language_model.layers.", ".layers.") + tensors = _write_adapter(tmp_path, ROWS) loaded = _checkpoint._load_adapter( cast(Any, trainer), cast(Any, SimpleNamespace(manifest=None, path=tmp_path)), - [art_key], + [ART_KEY], ) - expected = qwen35._qwen35_q_proj_lora_b_from_vllm( - tensors[f"{LAYER}.lora_B.weight"], - {"num_attention_heads": 4, "num_key_value_heads": 2, "head_dim": 8}, + torch.testing.assert_close(loaded[ART_KEY], _expected(tensors)) + + +def test_null_adapter_dimensions_are_filled_from_the_running_model( + tmp_path, monkeypatch +): + from art.megatron.model_support.lora_disk import load_lora_tensors_for_megatron + + _forbid_lookup(monkeypatch) + # A null group count must not fall back to one group per head. + tensors = _write_adapter(tmp_path, ROWS, num_key_value_heads=None, head_dim=None) + loaded = load_lora_tensors_for_megatron( + tmp_path, handler=qwen35.QWEN3_5_MOE_HANDLER, provider=PROVIDER ) - torch.testing.assert_close(loaded[art_key], expected) + torch.testing.assert_close(loaded[ART_KEY], _expected(tensors)) -def test_adapter_dimensions_take_precedence_over_the_running_model( +def test_adapter_dimensions_take_precedence_over_the_running_model(): + from art.megatron.model_support.lora_disk import with_model_attention_dimensions + + config = with_model_attention_dimensions( + {"revision": PIN, "num_attention_heads": 1, "num_key_value_heads": None}, + PROVIDER, + ) + assert config == { + "revision": PIN, + "num_attention_heads": 1, + "num_key_value_heads": 2, + "head_dim": 8, + "hidden_size": 32, + } + + +def test_an_incomplete_model_shape_falls_back_to_the_pinned_lookup( tmp_path, monkeypatch ): - from art.megatron.model_support import lora_disk - from art.trainer_rank import _checkpoint + from art.megatron.model_support.lora_disk import load_lora_tensors_for_megatron + + looked_up = [] + + def lookup(name, revision): + looked_up.append((name, revision)) + return SimpleNamespace(num_attention_heads=4, num_key_value_heads=2, head_dim=8) + + monkeypatch.setattr(qwen35, "_qwen35_text_config", lookup) + tensors = _write_adapter(tmp_path, ROWS) + # Without query groups, filling heads alone would imply one group per head. + provider = SimpleNamespace(num_attention_heads=4, kv_channels=8) + loaded = load_lora_tensors_for_megatron( + tmp_path, handler=qwen35.QWEN3_5_MOE_HANDLER, provider=provider + ) + assert looked_up == [(REPO, PIN)] + torch.testing.assert_close(loaded[ART_KEY], _expected(tensors)) + +def test_export_of_a_pinned_adapter_resolves_its_revision(hub_cache): + hub_cache(PIN, groups=2) + tensor = torch.randn(2 * 2 * 8 * 256, 2) # 16 heads in 2 groups, head_dim 256 + config = {"base_model_name_or_path": REPO, "revision": PIN} + art_key = "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight" + exported, _ = qwen35.QWEN3_5_MOE_HANDLER.to_vllm_lora_tensors( + {art_key: tensor}, adapter_config=config + ) + expected = qwen35._qwen35_q_proj_lora_b_to_vllm( + tensor, {"num_attention_heads": 16, "num_key_value_heads": 2, "head_dim": 256} + ) + torch.testing.assert_close(exported[f"{LAYER}.lora_B.weight"], expected) + + +def test_megatron_service_adapter_load_passes_the_running_model(monkeypatch): + train = pytest.importorskip("art.megatron.train") seen = [] monkeypatch.setattr( - lora_disk, + train, "load_lora_tensors_for_megatron", - lambda path, **kwargs: seen.append(kwargs["adapter_config"]) or {}, - ) - _write_adapter(tmp_path, rows=8) - config = json.loads((tmp_path / "adapter_config.json").read_text()) - config["num_attention_heads"] = 1 - (tmp_path / "adapter_config.json").write_text(json.dumps(config)) - trainer = SimpleNamespace( - runtime=SimpleNamespace( - provider=SimpleNamespace(num_attention_heads=4, num_query_groups=2), - model_support_handler=qwen35.QWEN3_5_MOE_HANDLER, - ) - ) - _checkpoint._load_adapter( - cast(Any, trainer), - cast(Any, SimpleNamespace(manifest=None, path=tmp_path)), - [], + lambda path, **kwargs: seen.append(kwargs) or {}, ) - assert seen[0]["num_attention_heads"] == 1 - assert seen[0]["num_key_value_heads"] == 2 - assert seen[0]["revision"] == PIN + monkeypatch.setattr(train, "load_adapter_into_model", lambda *a, **k: None) + train._load_adapter_into_model([], "adapter", 0, handler=None, provider=PROVIDER) + assert seen[0]["provider"] is PROVIDER From e2264fb3adaff7eb0d8365db0af1ed66fb99d7ed Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 04:56:42 +0000 Subject: [PATCH 3/5] Cover missing head size and empty revisions; pass the provider in the harness Co-Authored-By: Claude Opus 5.5 (1M context) --- .../megatron/train_inf_mismatch/real_path.py | 1 + tests/unit/test_qwen35_adapter_config.py | 23 ++++++++++++++----- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py index e56ea0456..7eb69063b 100644 --- a/tests/integration/megatron/train_inf_mismatch/real_path.py +++ b/tests/integration/megatron/train_inf_mismatch/real_path.py @@ -1573,6 +1573,7 @@ def _configure_worker_bundle(bundle: Any) -> None: adapter_model = load_lora_tensors_for_megatron( str(adapter_path), handler=runtime.model_support_handler, + provider=runtime.provider, allow_unvalidated_arch=request.config.allow_unvalidated_arch, ) megatron_train.load_adapter_into_model( diff --git a/tests/unit/test_qwen35_adapter_config.py b/tests/unit/test_qwen35_adapter_config.py index 4b4617547..2df34112f 100644 --- a/tests/unit/test_qwen35_adapter_config.py +++ b/tests/unit/test_qwen35_adapter_config.py @@ -70,6 +70,7 @@ def test_a_different_main_revision_does_not_change_a_pinned_adapter(hub_cache): (16, 2, 256), ] assert _dims() == (16, 4, 256) # unpinned adapters still follow main + assert _dims("") == (16, 4, 256) # an empty revision is unpinned def _write_adapter( @@ -95,9 +96,11 @@ def _write_adapter( return tensors -# Four heads in two query groups, with head_dim 8: the runtime's shape. +# Four heads in two query groups, with head_dim 8: the runtime's shape. The +# hidden size is not heads x head_dim, as in Qwen3.5, so deriving head_dim from +# it would be caught. PROVIDER = SimpleNamespace( - num_attention_heads=4, num_query_groups=2, kv_channels=8, hidden_size=32 + num_attention_heads=4, num_query_groups=2, kv_channels=8, hidden_size=48 ) ROWS = 2 * 2 * 2 * 8 # groups x (query + gate) x heads per group x head_dim ART_KEY = f"{LAYER}.lora_B.weight".replace(".language_model.layers.", ".layers.") @@ -164,12 +167,22 @@ def test_adapter_dimensions_take_precedence_over_the_running_model(): "num_attention_heads": 1, "num_key_value_heads": 2, "head_dim": 8, - "hidden_size": 32, + "hidden_size": 48, } +@pytest.mark.parametrize( + "provider", + [ + # Without query groups, filling heads alone would imply MHA. + SimpleNamespace(num_attention_heads=4, kv_channels=8, hidden_size=48), + # Without a head size, it would be derived from the hidden size. + SimpleNamespace(num_attention_heads=4, num_query_groups=2, hidden_size=48), + ], + ids=["no-query-groups", "no-head-size"], +) def test_an_incomplete_model_shape_falls_back_to_the_pinned_lookup( - tmp_path, monkeypatch + tmp_path, monkeypatch, provider ): from art.megatron.model_support.lora_disk import load_lora_tensors_for_megatron @@ -181,8 +194,6 @@ def lookup(name, revision): monkeypatch.setattr(qwen35, "_qwen35_text_config", lookup) tensors = _write_adapter(tmp_path, ROWS) - # Without query groups, filling heads alone would imply one group per head. - provider = SimpleNamespace(num_attention_heads=4, kv_channels=8) loaded = load_lora_tensors_for_megatron( tmp_path, handler=qwen35.QWEN3_5_MOE_HANDLER, provider=provider ) From b7c3915821ef09978da3542d8eb9c1a2d3049cbb Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 05:21:26 +0000 Subject: [PATCH 4/5] Resolve each missing Qwen3.5 attention dimension from the pinned config Take every dimension from the adapter when set, else the running model, else the base model's config at the adapter's revision. Previously a model missing its head size dropped all of its dimensions, and the handler then defaulted query groups to one per head without a lookup. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../model_support/handlers/qwen3_5.py | 26 ++++++++++----- src/art/megatron/model_support/lora_disk.py | 11 ++----- tests/unit/test_qwen35_adapter_config.py | 32 +++++++++++++++++-- 3 files changed, 50 insertions(+), 19 deletions(-) diff --git a/src/art/megatron/model_support/handlers/qwen3_5.py b/src/art/megatron/model_support/handlers/qwen3_5.py index e7ffebf58..8015da2e1 100644 --- a/src/art/megatron/model_support/handlers/qwen3_5.py +++ b/src/art/megatron/model_support/handlers/qwen3_5.py @@ -522,11 +522,15 @@ def _qwen35_text_config(base_model_name_or_path: str, revision: str | None) -> A def _qwen35_attention_dims(adapter_config: dict[str, Any]) -> tuple[int, int, int]: - num_heads = adapter_config.get("num_attention_heads") - num_groups = adapter_config.get("num_key_value_heads") - head_dim = adapter_config.get("head_dim") + dims = { + key: adapter_config.get(key) + for key in ("num_attention_heads", "num_key_value_heads", "head_dim") + } hidden_size = adapter_config.get("hidden_size") - if num_heads is None: + if None in dims.values(): + # Take each missing dimension from the base model's config rather than + # defaulting it: Qwen3.5 uses grouped queries and a head size that is + # not hidden_size / heads. base_model = adapter_config.get("base_model_name_or_path") if not base_model: raise RuntimeError("Qwen3.5 LoRA adapter config is missing base model path") @@ -536,10 +540,16 @@ def _qwen35_attention_dims(adapter_config: dict[str, Any]) -> tuple[int, int, in config = _qwen35_text_config( str(base_model), None if revision is None else str(revision) ) - num_heads = getattr(config, "num_attention_heads") - num_groups = getattr(config, "num_key_value_heads", num_heads) - head_dim = getattr(config, "head_dim", None) - hidden_size = getattr(config, "hidden_size", None) + for key, value in dims.items(): + if value is None: + dims[key] = getattr(config, key, None) + if hidden_size is None: + hidden_size = getattr(config, "hidden_size", None) + num_heads = dims["num_attention_heads"] + num_groups = dims["num_key_value_heads"] + head_dim = dims["head_dim"] + if num_heads is None: + raise RuntimeError("Qwen3.5 config is missing num_attention_heads") num_heads = int(num_heads) num_groups = int(num_groups if num_groups is not None else num_heads) if head_dim is None: diff --git a/src/art/megatron/model_support/lora_disk.py b/src/art/megatron/model_support/lora_disk.py index 9e16d42d8..dbf3f9846 100644 --- a/src/art/megatron/model_support/lora_disk.py +++ b/src/art/megatron/model_support/lora_disk.py @@ -35,16 +35,11 @@ def with_model_attention_dimensions( ) -> dict[str, Any]: """Fill attention dimensions the adapter config omits or nulls from the model. - Only a complete shape is used: a partial one could silently turn grouped - query attention into full multi-head attention. + Values the adapter sets win. The handler resolves anything still missing + from the base model's config at the adapter's revision. """ - dimensions = model_attention_dimensions(provider) - if not {"num_attention_heads", "num_key_value_heads", "head_dim"} <= set( - dimensions - ): - return adapter_config config = dict(adapter_config) - for key, value in dimensions.items(): + for key, value in model_attention_dimensions(provider).items(): if config.get(key) is None: config[key] = value return config diff --git a/tests/unit/test_qwen35_adapter_config.py b/tests/unit/test_qwen35_adapter_config.py index 2df34112f..c65b8d888 100644 --- a/tests/unit/test_qwen35_adapter_config.py +++ b/tests/unit/test_qwen35_adapter_config.py @@ -73,6 +73,17 @@ def test_a_different_main_revision_does_not_change_a_pinned_adapter(hub_cache): assert _dims("") == (16, 4, 256) # an empty revision is unpinned +def test_missing_dimensions_come_from_the_pinned_config_not_defaults(hub_cache): + hub_cache(PIN, groups=2) + config = { + "base_model_name_or_path": REPO, + "revision": PIN, + "num_attention_heads": 16, + } + # Not one group per head, nor head_dim = hidden_size / heads = 128. + assert qwen35._qwen35_attention_dims(config) == (16, 2, 256) + + def _write_adapter( path: Path, rows: int, **dimensions: object ) -> dict[str, torch.Tensor]: @@ -155,6 +166,21 @@ def test_null_adapter_dimensions_are_filled_from_the_running_model( torch.testing.assert_close(loaded[ART_KEY], _expected(tensors)) +@pytest.mark.parametrize("handler", ["QWEN3_5_DENSE_HANDLER", "QWEN3_5_MOE_HANDLER"]) +def test_adapter_and_model_dimensions_combine(tmp_path, monkeypatch, handler): + from art.megatron.model_support.lora_disk import load_lora_tensors_for_megatron + + _forbid_lookup(monkeypatch) + # The adapter has heads and head size, the model has query groups but no + # head size: together they are complete. + tensors = _write_adapter(tmp_path, ROWS, num_attention_heads=4, head_dim=8) + provider = SimpleNamespace(num_attention_heads=4, num_query_groups=2) + loaded = load_lora_tensors_for_megatron( + tmp_path, handler=getattr(qwen35, handler), provider=provider + ) + torch.testing.assert_close(loaded[ART_KEY], _expected(tensors)) + + def test_adapter_dimensions_take_precedence_over_the_running_model(): from art.megatron.model_support.lora_disk import with_model_attention_dimensions @@ -174,14 +200,14 @@ def test_adapter_dimensions_take_precedence_over_the_running_model(): @pytest.mark.parametrize( "provider", [ - # Without query groups, filling heads alone would imply MHA. + # Query groups must not default to one per head. SimpleNamespace(num_attention_heads=4, kv_channels=8, hidden_size=48), - # Without a head size, it would be derived from the hidden size. + # The head size must not be derived from the hidden size. SimpleNamespace(num_attention_heads=4, num_query_groups=2, hidden_size=48), ], ids=["no-query-groups", "no-head-size"], ) -def test_an_incomplete_model_shape_falls_back_to_the_pinned_lookup( +def test_dimensions_the_model_lacks_come_from_the_pinned_lookup( tmp_path, monkeypatch, provider ): from art.megatron.model_support.lora_disk import load_lora_tensors_for_megatron From a23028de551b49d36ab48d0b27c365c1db09f0ab Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 05:28:11 +0000 Subject: [PATCH 5/5] Exercise the old wrong head-size fallback in the heads-only test Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_qwen35_adapter_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_qwen35_adapter_config.py b/tests/unit/test_qwen35_adapter_config.py index c65b8d888..2880d073a 100644 --- a/tests/unit/test_qwen35_adapter_config.py +++ b/tests/unit/test_qwen35_adapter_config.py @@ -79,6 +79,7 @@ def test_missing_dimensions_come_from_the_pinned_config_not_defaults(hub_cache): "base_model_name_or_path": REPO, "revision": PIN, "num_attention_heads": 16, + "hidden_size": 2048, } # Not one group per head, nor head_dim = hidden_size / heads = 128. assert qwen35._qwen35_attention_dims(config) == (16, 2, 256)