Fix two FP8-training issues: PLE ngram dequantization and LoRA fp8 autocast guard

src/mcore_bridge/model/gpts/qwen4_exp.py

@@ -244,6 +244,11 @@ def _get_pp_src_rank(self, has_module: bool) -> int:
             dist.all_reduce(holder, op=dist.ReduceOp.MAX, group=self.pp_group)
         return int(holder.item())
 
+    # The PLE ngram embedding table is stored in the checkpoint as F8_E4M3
+    # shards plus a single scalar `weight_scale` (unlike experts, which use
+    # blockwise `weight_scale_inv`). The true values are `weight * scale`.
+    _PLE_NGRAM_SCALE_KEY = 'ple.ple_embedding.ngram_embedding.weight_scale'
+
     def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_rank: int):
         # The checkpoint shards the (padded) table into `parts` uniform row
         # blocks, so shard boundaries must be derived from the padded size.
@@ -269,6 +274,19 @@ def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_ra
         tp_start = tp_rank * per_partition if emb is not None else 0
         tp_end = min((tp_rank + 1) * per_partition, total) if emb is not None else 0
         if to_mcore:
+            # The fp8 shards must be multiplied by the scalar `weight_scale`
+            # before being written into the bf16 embedding. Checkpoints that
+            # store the table in bf16 directly have no such key; keep the raw
+            # values then (with a warning).
+            scale = None
+            if self._PLE_NGRAM_SCALE_KEY in hf_state_dict:
+                scale = hf_state_dict[self._PLE_NGRAM_SCALE_KEY].load().to(torch.float32)
+                self._ple_ngram_weight_scale = scale
+            else:
+                self._ple_ngram_weight_scale = None
+                logger.warning(
+                    f'`{self._PLE_NGRAM_SCALE_KEY}` not found in the checkpoint; assuming the PLE ngram '
+                    'embedding is already dequantized and loading it as-is.')
             for i in range(parts):
                 key = f'ple.ple_embedding.ngram_embedding.shard_{i}.weight'
                 if key not in hf_state_dict:
@@ -277,8 +295,20 @@ def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_ra
                 s, e = max(cs, tp_start), min(ce, tp_end)
                 if s < e:
                     weight = hf_state_dict[key].load()
+                    if scale is not None:
+                        weight = weight.to(torch.float32) * scale
                     emb.weight.data[s - tp_start:e - tp_start] = weight[s - cs:e - cs].to(emb.weight.dtype)
         else:
+            # Inverse of the to_mcore path: the checkpoint format is fp8
+            # shards + scalar `weight_scale`, so divide by the scale (stashed
+            # during loading) and cast back to fp8. Without a known scale the
+            # values cannot be represented as fp8 + scale; keep them in the
+            # current dtype and warn.
+            scale = getattr(self, '_ple_ngram_weight_scale', None)
+            if scale is None:
+                logger.warning(
+                    f'`{self._PLE_NGRAM_SCALE_KEY}` was not seen during loading; exporting the PLE ngram '
+                    'embedding without re-quantizing to fp8.')
             for i in range(parts):
                 cs, ce = i * shard_size, min((i + 1) * shard_size, total)
                 pieces = []
@@ -297,11 +327,15 @@ def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_ra
                 shard = torch.cat(pieces, dim=0)
                 if self.pp_size > 1:
                     dist.broadcast(shard, src=pp_src_rank, group=self.pp_group)
+                if scale is not None:
+                    shard = (shard.to(torch.float32) / scale.to(shard.device)).to(torch.float8_e4m3fn)
                 # Written directly into the state dict (bypasses _get_weight,
                 # which normally applies _target_device).
                 if self._target_device is not None:
                     shard = shard.to(self._target_device)
                 hf_state_dict[f'ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = shard
+            if scale is not None and self._PLE_NGRAM_SCALE_KEY not in hf_state_dict:
+                hf_state_dict[self._PLE_NGRAM_SCALE_KEY] = scale.reshape(())
 
     def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool):
         ple = None if mg_layer is None else getattr(mg_layer, 'ple', None)

src/mcore_bridge/tuners/lora.py

@@ -5,6 +5,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import transformer_engine
 import warnings
 from contextlib import contextmanager, nullcontext
 from importlib import metadata
@@ -431,12 +432,19 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
                                                                     NpuGroupedLoraLinear)) else lora_A.weight.dtype
                 x = x.to(dtype)
 
-                lora_result = lora_A(dropout(x), *args, **kwargs) if isinstance(
-                    lora_A, (TEGroupedLinear, NpuGroupedLoraLinear)) else lora_A(dropout(x))
-                if isinstance(lora_result, tuple):
-                    lora_result = lora_result[0]
-                lora_result = lora_B(lora_result, *args, **kwargs) if isinstance(
-                    lora_B, (TEGroupedLinear, NpuGroupedLoraLinear)) else lora_B(lora_result)
+                # LoRA A/B weights are rank-sized (e.g. [4, 2560]), which
+                # violates TE's FP8 GEMM divisibility rules; never run them
+                # under fp8 autocast (mirrors the in_proj_ba guard in
+                # gpts/qwen4_exp.py's GDN path).
+                fp8_context = (transformer_engine.pytorch.fp8_autocast(enabled=False)
+                               if getattr(self.config, 'fp8_param', False) else nullcontext())
+                with fp8_context:
+                    lora_result = lora_A(dropout(x), *args, **kwargs) if isinstance(
+                        lora_A, (TEGroupedLinear, NpuGroupedLoraLinear)) else lora_A(dropout(x))
+                    if isinstance(lora_result, tuple):
+                        lora_result = lora_result[0]
+                    lora_result = lora_B(lora_result, *args, **kwargs) if isinstance(
+                        lora_B, (TEGroupedLinear, NpuGroupedLoraLinear)) else lora_B(lora_result)
                 if isinstance(lora_result, tuple):
                     lora_result = lora_result[0]
                 lora_result = lora_result * scaling

tests/test_qwen4_exp_ple_fp8_scale.py

@@ -0,0 +1,124 @@
+"""Verify the PLE ngram embedding FP8 <-> bf16 `weight_scale` handling.
+
+FP8 checkpoints store the PLE ngram embedding table as F8_E4M3 shards plus a
+single scalar `ple.ple_embedding.ngram_embedding.weight_scale` (true value =
+weight * scale). This only matters when training starts from an FP8-format
+checkpoint; bf16 checkpoints have no such key.
+
+No model is downloaded: the checkpoint and the Megatron-side embedding are
+synthetic tensors, exercised through `_set_ple_ngram_embedding` on a single
+gloo rank.
+"""
+import os
+
+os.environ.setdefault('RANK', '0')
+os.environ.setdefault('LOCAL_RANK', '0')
+os.environ.setdefault('WORLD_SIZE', '1')
+os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
+os.environ.setdefault('MASTER_PORT', '29517')
+
+from types import SimpleNamespace  # noqa: E402
+
+import torch  # noqa: E402
+import torch.distributed as dist  # noqa: E402
+from megatron.core import mpu  # noqa: E402
+
+from mcore_bridge.model.gpts.qwen4_exp import Qwen4ExpBridge  # noqa: E402
+
+SCALE_KEY = Qwen4ExpBridge._PLE_NGRAM_SCALE_KEY
+
+
+class _LazyTensor:
+    """Stand-in for the lazy checkpoint tensors: only `.load()` is used."""
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor = tensor
+
+    def load(self) -> torch.Tensor:
+        return self.tensor
+
+
+def _init_single_rank():
+    if not dist.is_initialized():
+        dist.init_process_group('gloo')
+    if not mpu.model_parallel_is_initialized():
+        mpu.initialize_model_parallel(1)
+
+
+def _shard_keys(parts: int):
+    return [f'ple.ple_embedding.ngram_embedding.shard_{i}.weight' for i in range(parts)]
+
+
+def _make_fake_ple(total: int, dim: int, parts: int):
+    ngram_embedding = SimpleNamespace(
+        weight=torch.zeros(total, dim, dtype=torch.bfloat16),
+        num_embeddings=total,
+        num_embeddings_per_partition=total)
+    return SimpleNamespace(
+        ple_embedding=SimpleNamespace(ngram_embedding=ngram_embedding, split_ngram_parts=parts, head_dim=dim))
+
+
+def _make_fake_bridge():
+    return SimpleNamespace(
+        pp_size=1,
+        tp_group=dist.group.WORLD,
+        pp_group=dist.group.WORLD,
+        _target_device=None,
+        config=SimpleNamespace(params_dtype=torch.bfloat16),
+        _PLE_NGRAM_SCALE_KEY=SCALE_KEY)
+
+
+def test_ple_ngram_fp8_scale_roundtrip():
+    _init_single_rank()
+    torch.manual_seed(0)
+    total, dim, parts = 16, 8, 4
+    rows_per_shard = total // parts
+    scale = torch.tensor(0.05, dtype=torch.float32)
+    fp8_table = torch.randn(total, dim).to(torch.float8_e4m3fn)
+    hf_state_dict = {
+        key: _LazyTensor(fp8_table[i * rows_per_shard:(i + 1) * rows_per_shard].clone())
+        for i, key in enumerate(_shard_keys(parts))
+    }
+    hf_state_dict[SCALE_KEY] = _LazyTensor(scale.clone())
+
+    # to_mcore: fp8 shards must be dequantized with the scalar weight_scale.
+    ple = _make_fake_ple(total, dim, parts)
+    bridge = _make_fake_bridge()
+    Qwen4ExpBridge._set_ple_ngram_embedding(bridge, ple, hf_state_dict, True, 0)
+    expected = (fp8_table.float() * scale).to(torch.bfloat16)
+    assert torch.equal(ple.ple_embedding.ngram_embedding.weight.data, expected)
+
+    # to_hf: symmetric re-quantization, the scale key is written back.
+    exported = {}
+    Qwen4ExpBridge._set_ple_ngram_embedding(bridge, ple, exported, False, 0)
+    assert torch.equal(exported[SCALE_KEY], scale)
+    # bf16(fp8 * scale) / scale stays within half an FP8 ulp of the original
+    # grid point, so re-quantization must recover the exact FP8 values.
+    for i, key in enumerate(_shard_keys(parts)):
+        shard = exported[key]
+        assert shard.dtype == torch.float8_e4m3fn
+        assert torch.equal(shard, fp8_table[i * rows_per_shard:(i + 1) * rows_per_shard])
+
+
+def test_ple_ngram_bf16_checkpoint_loads_as_is():
+    _init_single_rank()
+    torch.manual_seed(0)
+    total, dim, parts = 8, 4, 2
+    table = torch.randn(total, dim, dtype=torch.bfloat16)
+    hf_state_dict = {
+        key: _LazyTensor(table[i * (total // parts):(i + 1) * (total // parts)].clone())
+        for i, key in enumerate(_shard_keys(parts))
+    }  # no weight_scale key: a plain bf16 checkpoint
+
+    ple = _make_fake_ple(total, dim, parts)
+    bridge = _make_fake_bridge()
+    Qwen4ExpBridge._set_ple_ngram_embedding(bridge, ple, hf_state_dict, True, 0)
+    assert torch.equal(ple.ple_embedding.ngram_embedding.weight.data, table)
+    assert bridge._ple_ngram_weight_scale is None
+
+    # Without a known scale the export keeps the values as-is (no fp8 cast).
+    exported = {}
+    Qwen4ExpBridge._set_ple_ngram_embedding(bridge, ple, exported, False, 0)
+    assert SCALE_KEY not in exported
+    for key in _shard_keys(parts):
+        assert exported[key].dtype == torch.bfloat16