fix(model): reject unaligned FP8 QKV scale export

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/models/conversion/model_bridge.py

@@ -54,6 +54,7 @@
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.param_mapping import (
     MegatronParamMapping,
+    QKVMapping,
 )
 from megatron.bridge.models.conversion.peft_bridge import (
     AdapterWeight,
@@ -2118,6 +2119,21 @@ def build_export_fp8_tasks(
             fp8_scale_inv_attr,
         )
 
+        for global_name, fp8_flag in global_fp8_flags.items():
+            block_size = self._fp8_scale_block_size(fp8_flag)
+            if block_size is None:
+                continue
+            mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
+            if not isinstance(mapping, QKVMapping):
+                continue
+            head_size = model_config.kv_channels or (model_config.hidden_size // model_config.num_attention_heads)
+            if head_size % block_size != 0:
+                raise ValueError(
+                    f"Cannot export blockwise FP8 QKV scales: QKV head size {head_size} "
+                    f"is not divisible by FP8 block size {block_size}. Scale blocks that cross "
+                    "QKV head boundaries cannot be split into separate HF projections."
+                )
+
         # 2) Expand the global name list with `*.scale_inv` entries.
         #    This defines the final deterministic task ordering.
         expanded_global_names: list[str] = []

tests/unit_tests/models/test_fp8_param_export.py

@@ -30,7 +30,7 @@
     WeightConversionTask,
     _HFNameSuffixMapping,
 )
-from megatron.bridge.models.conversion.param_mapping import split_qkv_weights
+from megatron.bridge.models.conversion.param_mapping import QKVMapping, split_qkv_weights
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 
 
@@ -160,6 +160,31 @@ def megatron_to_hf(self, megatron_weights, megatron_module):
 
 
 class TestFp8ParamExport:
+    def test_build_export_fp8_tasks_rejects_unaligned_qkv_scale_blocks(self, monkeypatch):
+        bridge = DummyBridge()
+        global_name = "decoder.layers.0.self_attention.linear_qkv.weight"
+        mapping = QKVMapping(global_name, "hf.q.weight", "hf.k.weight", "hf.v.weight")
+        _patch_export_task_context(
+            monkeypatch,
+            bridge,
+            global_name,
+            registry_factory=lambda: MegatronMappingRegistry(mapping),
+            detect_fp8=lambda *_a, **_k: {global_name: 128},
+        )
+        model = SimpleNamespace(
+            config=SimpleNamespace(
+                num_attention_heads=64,
+                num_query_groups=8,
+                hidden_size=2880,
+                kv_channels=64,
+                share_embeddings_and_output_weights=False,
+            ),
+            named_parameters=lambda: [],
+        )
+
+        with pytest.raises(ValueError, match="QKV head size 64.*FP8 block size 128"):
+            bridge.build_export_fp8_tasks(SimpleNamespace(state=SimpleNamespace(source=SimpleNamespace())), [model])
+
     @pytest.mark.parametrize(
         "export_weight_dtype, expect_unquantized",
         [("fp8", True), ("bf16", False)],