Clean up PTQ buffer naming.

Signed-off-by: Cory Ye <cye@nvidia.com>

Author: cspades

megatron/core/extensions/transformer_engine.py

@@ -310,6 +310,19 @@ def _get_global_layer_fqn(module_name: str, named_modules: Dict[str, torch.nn.Mo
     return module_name
 
 
+def _parse_ptq_calibration_buffer_name(buffer_name: str) -> Tuple[str, str, str]:
+    """Parse a TransformerEngine PTQ calibration buffer name."""
+    match = re.fullmatch(
+        r"(?P<tensor_name>.+)_"
+        r"(?P<metadata_name>amax(?:_rowwise)?|scale_inv)_"
+        r"(?P<recipe_name>.+)_te_ptq_calibrated",
+        buffer_name,
+    )
+    if match is None:
+        raise ValueError(f"Malformed TransformerEngine PTQ calibration buffer {buffer_name!r}")
+    return (match.group("tensor_name"), match.group("metadata_name"), match.group("recipe_name"))
+
+
 def add_ptq_calibration_metadata_to_state_dict(
     state_dict: Dict[str, Any], model: List[torch.nn.Module]
 ) -> None:
@@ -341,22 +354,21 @@ def add_ptq_calibration_metadata_to_state_dict(
                 continue
             processed_buffers = []
             for key, data in sorted(calibration_buffers):
+                tensor_name, metadata_name, recipe_name = _parse_ptq_calibration_buffer_name(key)
                 metadata = data.detach()
                 if metadata.numel() > 1:
                     # Only global FP32 per-tensor scales are checkpointed.
                     metadata = metadata.amax().reshape(1)
 
-                if "_scale_inv_" in key:
+                if metadata_name == "scale_inv":
                     # Inverted scaling factor, i.e. absmax / quantmax.
-                    scale_inv = metadata
+                    scale_inv = metadata.clone()
                 else:
                     # Convert absmax to a dequantization scaling factor.
-                    if "nvfp4" in key:
+                    if recipe_name == "nvfp4":
                         # E2M1_MAX * E4M3_MAX
                         quant_max = 6.0 * 448.0
-                    elif any(
-                        recipe in key for recipe in ("fp8_delayed_scaling", "fp8_current_scaling")
-                    ):
+                    elif recipe_name == "fp8_delayed_scaling":
                         quant_max = 448.0
                     else:
                         raise ValueError(f"Unsupported amax calibration buffer {key!r}")
@@ -370,25 +382,22 @@ def add_ptq_calibration_metadata_to_state_dict(
                         scale_inv, op=torch.distributed.ReduceOp.MAX, group=reduction_group
                     )
 
-                tensor_name = re.split(
-                    r"_tensor_(?:amax(?:_rowwise)?|scale_inv)_", key, maxsplit=1
-                )[0]
                 gemm_match = re.fullmatch(r"(input|weight)_gemm(\d+)", tensor_name)
-                processed_buffers.append((key, scale_inv, gemm_match))
+                processed_buffers.append((tensor_name, recipe_name, scale_inv, gemm_match))
 
             expert_buffers = [
                 entry
                 for entry in processed_buffers
-                if hasattr(module, "_pg_collection") and entry[2] is not None
+                if hasattr(module, "_pg_collection") and entry[3] is not None
             ]
             non_expert_buffers = [
                 entry
                 for entry in processed_buffers
-                if not (hasattr(module, "_pg_collection") and entry[2] is not None)
+                if not (hasattr(module, "_pg_collection") and entry[3] is not None)
             ]
             if expert_buffers:
                 # Gather every local GEMM's input and weight scales in one EP collective.
-                local_scales = torch.cat([scale for _, scale, _ in expert_buffers])
+                local_scales = torch.cat([scale for _, _, scale, _ in expert_buffers])
                 ep_group = module._pg_collection.ep
                 ep_size = get_pg_size(ep_group)
                 num_local_scales = local_scales.numel()
@@ -402,24 +411,27 @@ def add_ptq_calibration_metadata_to_state_dict(
                 gathered_scales = gathered_scales.cpu()
 
                 for ep_rank in range(ep_size):
-                    for scale_idx, (key, _, gemm_match) in enumerate(expert_buffers):
+                    for scale_idx, (_, recipe_name, _, gemm_match) in enumerate(expert_buffers):
                         tensor_kind = gemm_match.group(1)
                         local_gemm_idx = int(gemm_match.group(2))
                         global_expert_idx = ep_rank * module.num_gemms + local_gemm_idx
                         expert_fqn = _get_grouped_linear_expert_fqn(export_name, global_expert_idx)
                         # Rename dense keys to match vLLM, TensorRT-LLM, and other PTQ formats.
                         # NOTE(@cspades): I don't like this but these are the conventional names.
-                        scale_suffix = "_2" if tensor_kind == "weight" and "nvfp4" in key else ""
+                        scale_suffix = (
+                            "_2" if tensor_kind == "weight" and recipe_name == "nvfp4" else ""
+                        )
                         local_calibration_state[
                             f"{expert_fqn}.{tensor_kind}_scale{scale_suffix}"
                         ] = gathered_scales[ep_rank * num_local_scales + scale_idx]
 
-            for key, scale_inv, _ in non_expert_buffers:
+            for tensor_name, recipe_name, scale_inv, _ in non_expert_buffers:
                 # Rename dense keys to match vLLM, TensorRT-LLM, and other PTQ formats.
                 # NOTE(@cspades): I don't like this but these are the conventional names.
-                scale_prefix = "input_" if "input" in key else "weight_"
-                scale_suffix = "_2" if "weight" in key and "nvfp4" in key else ""
-                local_calibration_state[f"{export_name}.{scale_prefix}scale{scale_suffix}"] = (
+                scale_suffix = (
+                    "_2" if tensor_name.endswith("weight") and recipe_name == "nvfp4" else ""
+                )
+                local_calibration_state[f"{export_name}.{tensor_name}_scale{scale_suffix}"] = (
                     scale_inv.cpu()
                 )
 

megatron/core/fp8_utils.py

@@ -107,8 +107,7 @@ def _make_te_calibration_config(transformer_engine_calibration_decay: float = 0.
     """Construct an immutable TE calibration config once per decay value."""
     if not HAVE_TE:
         raise RuntimeError(
-            "ckpt_transformer_engine_calibration requires Transformer Engine "
-            "to be installed"
+            "ckpt_transformer_engine_calibration requires Transformer Engine to be installed"
         )
     calibration_config_cls = getattr(
         transformer_engine.pytorch, "QuantizationCalibrationConfig", None

tests/unit_tests/test_checkpointing.py

@@ -683,12 +683,15 @@ def test_save_checkpoint(
 
         # Text-only checkpoint saves still bypass dataloader state, while passing CP through.
         cp_group = mock.sentinel.cp_group
-        with mock.patch(
-            "megatron.training.checkpointing.maybe_save_dataloader_state",
-            wraps=maybe_save_dataloader_state,
-        ) as save_dataloader_state, mock.patch(
-            "megatron.training.checkpointing.add_ptq_calibration_metadata_to_state_dict"
-        ) as export_calibration_metadata:
+        with (
+            mock.patch(
+                "megatron.training.checkpointing.maybe_save_dataloader_state",
+                wraps=maybe_save_dataloader_state,
+            ) as save_dataloader_state,
+            mock.patch(
+                "megatron.training.checkpointing.add_ptq_calibration_metadata_to_state_dict"
+            ) as export_calibration_metadata,
+        ):
             save_checkpoint(
                 iteration,
                 [model],

tests/unit_tests/transformer/test_ptq_calibration_metadata_buffering.py

@@ -30,17 +30,15 @@ def test_calibration_options_default_to_disabled():
 
 def test_disabled_megatron_option_returns_no_te_calibration_config():
     config = SimpleNamespace(
-        ckpt_transformer_engine_calibration=False,
-        transformer_engine_calibration_decay=0.5,
+        ckpt_transformer_engine_calibration=False, transformer_engine_calibration_decay=0.5
     )
 
     assert fp8_utils.get_te_calibration_config(config) is None
 
 
 def test_calibration_option_requires_te_config_support(monkeypatch):
     config = SimpleNamespace(
-        ckpt_transformer_engine_calibration=True,
-        transformer_engine_calibration_decay=0.123,
+        ckpt_transformer_engine_calibration=True, transformer_engine_calibration_decay=0.123
     )
     monkeypatch.setattr(fp8_utils.transformer_engine.pytorch, "QuantizationCalibrationConfig", None)
     fp8_utils._make_te_calibration_config.cache_clear()
@@ -51,8 +49,7 @@ def test_calibration_option_requires_te_config_support(monkeypatch):
 
 def test_megatron_options_create_te_calibration_config():
     config = SimpleNamespace(
-        ckpt_transformer_engine_calibration=True,
-        transformer_engine_calibration_decay=0.5,
+        ckpt_transformer_engine_calibration=True, transformer_engine_calibration_decay=0.5
     )
 
     calibration_config = fp8_utils.get_te_calibration_config(config)
@@ -109,8 +106,7 @@ def legacy_fp8_autocast(*, enabled, fp8_recipe, fp8_group):
     monkeypatch.setattr(fp8_utils.transformer_engine.pytorch, "fp8_autocast", legacy_fp8_autocast)
 
     with pytest.raises(
-        RuntimeError,
-        match="--ckpt-transformer-engine-calibration.*accepts calibration_config",
+        RuntimeError, match="--ckpt-transformer-engine-calibration.*accepts calibration_config"
     ):
         fp8_utils.get_fp8_context(config)
 
@@ -233,10 +229,10 @@ def test_calibration_amax_conversion_and_scale_inv_passthrough(monkeypatch):
     _disable_scale_collectives(monkeypatch)
     linear = torch.nn.Module()
     linear.register_buffer(
-        "input_tensor_amax_fp8_delayed_scaling_te_ptq_calibrated", torch.tensor([896.0])
+        "input_amax_fp8_delayed_scaling_te_ptq_calibrated", torch.tensor([896.0])
     )
     linear.register_buffer(
-        "weight_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
+        "weight_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
     )
     model = torch.nn.Module()
     model.add_module("linear", linear)
@@ -253,7 +249,7 @@ def test_calibration_amax_conversion_and_scale_inv_passthrough(monkeypatch):
 def test_nvfp4_weight_amax_exports_secondary_weight_scale(monkeypatch):
     _disable_scale_collectives(monkeypatch)
     linear = torch.nn.Module()
-    linear.register_buffer("weight_tensor_amax_nvfp4_te_ptq_calibrated", torch.tensor([2688.0]))
+    linear.register_buffer("weight_amax_nvfp4_te_ptq_calibrated", torch.tensor([2688.0]))
     model = torch.nn.Module()
     model.add_module("linear", linear)
     state_dict = {"model": {}}
@@ -263,6 +259,25 @@ def test_nvfp4_weight_amax_exports_secondary_weight_scale(monkeypatch):
     torch.testing.assert_close(state_dict["model"]["linear.weight_scale_2"], torch.ones(1))
 
 
+def test_fused_module_tensor_names_export_distinct_scales(monkeypatch):
+    _disable_scale_collectives(monkeypatch)
+    mlp = torch.nn.Module()
+    mlp.register_buffer(
+        "fc1_input_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
+    )
+    mlp.register_buffer(
+        "fc2_input_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.5])
+    )
+    model = torch.nn.Module()
+    model.add_module("mlp", mlp)
+    state_dict = {"model": {}}
+
+    te_ext.add_ptq_calibration_metadata_to_state_dict(state_dict, [model])
+
+    torch.testing.assert_close(state_dict["model"]["mlp.fc1_input_scale"], torch.tensor([0.25]))
+    torch.testing.assert_close(state_dict["model"]["mlp.fc2_input_scale"], torch.tensor([0.5]))
+
+
 def test_global_layer_fqn_uses_layer_number():
     model = torch.nn.Module()
     model.layers = torch.nn.ModuleList([torch.nn.Module()])
@@ -284,10 +299,10 @@ def test_expert_parallel_calibration_scales_are_gathered(monkeypatch):
     grouped_linear.num_gemms = 2
     grouped_linear._pg_collection = SimpleNamespace(ep=ep_group)
     grouped_linear.register_buffer(
-        "input_gemm0_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
+        "input_gemm0_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
     )
     grouped_linear.register_buffer(
-        "input_gemm1_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.5])
+        "input_gemm1_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.5])
     )
     mlp = torch.nn.Module()
     mlp.add_module("linear_fc1", grouped_linear)
@@ -321,7 +336,7 @@ def test_pipeline_parallel_calibration_states_are_merged(monkeypatch):
     model.pg_collection = SimpleNamespace(pp=pp_group)
     linear = torch.nn.Module()
     linear.register_buffer(
-        "input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
+        "input_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
     )
     model.add_module("linear", linear)
 
@@ -352,7 +367,7 @@ def test_pipeline_parallel_conflicting_calibration_states_are_rejected(monkeypat
     model.pg_collection = SimpleNamespace(pp=pp_group)
     linear = torch.nn.Module()
     linear.register_buffer(
-        "input_tensor_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
+        "input_scale_inv_fp8_current_scaling_te_ptq_calibrated", torch.tensor([0.25])
     )
     model.add_module("linear", linear)