diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 7ce61f8..96ba608 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/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) diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index 8b15093..ff6f4ec 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/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 diff --git a/tests/test_qwen4_exp_ple_fp8_scale.py b/tests/test_qwen4_exp_ple_fp8_scale.py new file mode 100644 index 0000000..fdf877a --- /dev/null +++ b/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