From 178d836dd90a43f999672aae80dfd089b7e9f92d Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 13 Apr 2026 03:47:20 +0000 Subject: [PATCH 1/4] fix qwen35 and support general fp32 --- .../megatron_utils/fp32_param_utils.py | 52 ++++ miles/backends/megatron_utils/model.py | 5 + miles_plugins/mbridge/qwen3_5.py | 6 + miles_plugins/models/qwen3_5.py | 4 +- .../megatron_utils/test_fp32_param_utils.py | 262 ++++++++++++++++++ tools/convert_hf_to_torch_dist.py | 25 +- 6 files changed, 330 insertions(+), 24 deletions(-) create mode 100644 miles/backends/megatron_utils/fp32_param_utils.py create mode 100644 tests/fast/backends/megatron_utils/test_fp32_param_utils.py diff --git a/miles/backends/megatron_utils/fp32_param_utils.py b/miles/backends/megatron_utils/fp32_param_utils.py new file mode 100644 index 00000000000..afd6bde7f0f --- /dev/null +++ b/miles/backends/megatron_utils/fp32_param_utils.py @@ -0,0 +1,52 @@ +import logging +from collections.abc import Sequence + +import torch +import torch.distributed as dist + +logger = logging.getLogger(__name__) + + +# Parameter attribute used by model definitions to pin parameter dtype. +FORCED_PARAM_DTYPE_ATTR = "_miles_forced_param_dtype" + + +def mark_param_dtype(param: torch.nn.Parameter, dtype: torch.dtype) -> None: + """Mark a parameter with its required runtime dtype.""" + setattr(param, FORCED_PARAM_DTYPE_ATTR, dtype) + + +def enforce_marked_param_dtypes(model_chunks: Sequence[torch.nn.Module]) -> list[str]: + """Apply dtype overrides declared on parameters via ``mark_param_dtype``. + + This keeps the policy in model definitions and avoids model-name checks in + the training/conversion mainline. + + Motivation: Megatron's ``Float16Module`` unconditionally casts every + floating-point parameter to bf16/fp16 at wrap time, and there is no + declarative opt-out in nn.Module or Megatron. Megatron's MoE router hits the + same problem and solves it with ``_maintain_float32_expert_bias`` (see + ``megatron/core/transformer/moe/router.py``), which post-hoc casts the + expert_bias back to fp32. This function generalizes that pattern: callers + mark params with their required dtype at the model-definition site, and we + re-cast after ``get_model`` so the rest of the stack (optimizer, DDP, mbridge + load path) sees the intended dtype. + """ + updated_names: list[str] = [] + for chunk in model_chunks: + for name, param in chunk.named_parameters(): + target_dtype = getattr(param, FORCED_PARAM_DTYPE_ATTR, None) + if target_dtype is None: + continue + + if param.dtype != target_dtype: + # Keep Parameter identity to avoid breaking optimizer/DDP maps. + param.data = param.data.to(dtype=target_dtype) + updated_names.append(name) + + rank = 0 + if dist.is_available() and dist.is_initialized(): + rank = dist.get_rank() + if rank == 0 and updated_names: + logger.info("Enforced marked parameter dtypes for %d tensors.", len(updated_names)) + return updated_names diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index e95158bbfcb..c7ca14b486e 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -36,6 +36,7 @@ compute_model_hashes_by_layer, save_model_hashes, ) +from .fp32_param_utils import enforce_marked_param_dtypes from .initialize import is_megatron_main_rank from .lora_utils import is_lora_enabled, is_lora_model from .model_provider import get_model_provider_func @@ -125,6 +126,9 @@ def setup_model_and_optimizer( else: model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder) + # Apply parameter-level dtype overrides declared in model definitions. + enforce_marked_param_dtypes(model) + # Optimizer kwargs = {} for f in dataclasses.fields(OptimizerConfig): @@ -132,6 +136,7 @@ def setup_model_and_optimizer( kwargs[f.name] = getattr(args, f.name) config = OptimizerConfig(**kwargs) config.timers = None + optimizer = get_megatron_optimizer( config=config, model_chunks=model, diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index ee629d009f2..8da5b7204b4 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -254,6 +254,12 @@ def _convert_mtp_param(self, name: str) -> list[str]: def _weight_to_mcore_format( self, mcore_weights_name: str, hf_weights: list[torch.Tensor] ) -> tuple[list[str], list[torch.Tensor]]: + if mcore_weights_name.endswith("self_attention.linear_attn.A_log"): + assert len(hf_weights) == 1 + # Keep A_log in fp32 before TP scatter; this avoids precision loss + # from Bridge's global pre-cast to self.dtype. + return hf_weights[0].to(dtype=torch.float32).contiguous() + if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name: # merge qkv assert len(hf_weights) == 3 diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index 794cf738081..f8457dc0336 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -15,6 +15,7 @@ except ImportError: pass +from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype from miles.backends.training_utils.cp_utils import build_gdn_cp_context from .hf_attention import HuggingfaceAttention, _load_hf_config @@ -71,7 +72,8 @@ def __init__(self, config, layer_idx: int): self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) A = torch.empty(self.num_v_heads).uniform_(0, 16) - self.A_log = nn.Parameter(torch.log(A)) + self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) + mark_param_dtype(self.A_log, torch.float32) self.norm = FusedRMSNormGated( self.head_v_dim, diff --git a/tests/fast/backends/megatron_utils/test_fp32_param_utils.py b/tests/fast/backends/megatron_utils/test_fp32_param_utils.py new file mode 100644 index 00000000000..c98b2c3e5d9 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_fp32_param_utils.py @@ -0,0 +1,262 @@ +"""Tests for the A_log fp32 preservation chain. + +Feature: Qwen3.5's ``A_log`` must end up as fp32 in the Megatron parameter +after hf->mcore conversion, because the chunk-gated-delta-rule kernel relies +on that precision. Two complementary pieces keep this invariant: + +- Downstream — ``enforce_marked_param_dtypes`` (this module): + Megatron's ``Float16Module`` unconditionally casts every floating-point + parameter to bf16/fp16 at wrap time. There is no declarative opt-out in + nn.Module or Megatron; even Megatron's own MoE router uses the same + post-hoc ``.data = ...to(float32)`` pattern in + ``_maintain_float32_expert_bias``. We generalize that by letting model + definitions declare intent via ``mark_param_dtype`` and re-casting after + ``get_model`` returns. +- Upstream — ``Qwen3_5Bridge._weight_to_mcore_format``: + mbridge's base ``_weight_to_mcore_format`` pre-casts every HF tensor to + ``self.dtype`` (bf16) before TP scatter. For A_log that pre-cast rounds + the fp32 HF value. The override returns A_log as fp32 early, bypassing + that pre-cast entirely. + +The end-to-end test ties both halves together and checks bit-exact equality +with the HF fp32 source — this is the regression guard against the original +``patch_weight_to_mcore_format_preserve_fp32`` failure mode, where only the +upstream cast was intercepted and the downstream ``t.to(param.dtype)`` in +``Bridge.load_weights`` still demoted A_log back to bf16. +""" + +import pytest +import torch +import torch.nn as nn + +from miles.backends.megatron_utils.fp32_param_utils import ( + FORCED_PARAM_DTYPE_ATTR, + enforce_marked_param_dtypes, + mark_param_dtype, +) + + +# --------------------------------------------------------------------------- +# Downstream: mark_param_dtype + enforce_marked_param_dtypes +# --------------------------------------------------------------------------- + + +class _ToyModule(nn.Module): + """Minimal stand-in for Qwen3_5GatedDeltaNet: one marked fp32 param plus + one regular bf16-target param, so we can check the collateral damage + boundary of ``enforce_marked_param_dtypes``.""" + + def __init__(self, num_heads: int = 8): + super().__init__() + A = torch.empty(num_heads).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) + mark_param_dtype(self.A_log, torch.float32) + self.in_proj = nn.Linear(16, num_heads, bias=False) + + +class TestMarkParamDtype: + def test_attaches_expected_attribute(self): + p = nn.Parameter(torch.zeros(4)) + mark_param_dtype(p, torch.float32) + assert getattr(p, FORCED_PARAM_DTYPE_ATTR) is torch.float32 + + def test_overwrites_previous_mark(self): + p = nn.Parameter(torch.zeros(4)) + mark_param_dtype(p, torch.float32) + mark_param_dtype(p, torch.float64) + assert getattr(p, FORCED_PARAM_DTYPE_ATTR) is torch.float64 + + +class TestEnforceMarkedParamDtypes: + def test_recasts_marked_param_back_to_fp32_after_float16_wrap(self): + """Simulates the full Megatron path: construct -> bfloat16() (what + ``Float16Module(...)`` does) -> enforce. A_log must come out fp32.""" + m = _ToyModule() + assert m.A_log.dtype == torch.float32 + + # Simulate Float16Module(config, m) — module.bfloat16() in the ctor + # demotes every floating param including the marked one. + m.bfloat16() + assert m.A_log.dtype == torch.bfloat16 + + enforce_marked_param_dtypes([m]) + assert m.A_log.dtype == torch.float32 + + def test_preserves_parameter_identity(self): + """Optimizer and DDP bucket parameters by Python identity, set up + AFTER ``enforce_marked_param_dtypes`` runs. If we re-bind via + ``self.A_log = nn.Parameter(...)`` the id changes and the optimizer + map breaks. We must only mutate ``.data``.""" + m = _ToyModule() + m.bfloat16() + before_id = id(m.A_log) + before_param_obj = m.A_log + + enforce_marked_param_dtypes([m]) + + assert id(m.A_log) == before_id + assert m.A_log is before_param_obj + + def test_leaves_unmarked_params_alone(self): + m = _ToyModule() + m.bfloat16() + assert m.in_proj.weight.dtype == torch.bfloat16 + + enforce_marked_param_dtypes([m]) + assert m.in_proj.weight.dtype == torch.bfloat16 + + def test_is_noop_when_already_target_dtype(self): + """Idempotency — second call must not re-allocate or change anything. + Guards against accidental double-work when the hook is called on + both the training and conversion entrypoints in the same process.""" + m = _ToyModule() + m.bfloat16() + enforce_marked_param_dtypes([m]) + + data_before = m.A_log.data + updated = enforce_marked_param_dtypes([m]) + assert m.A_log.dtype == torch.float32 + # ``.data`` should be the same tensor object (no unnecessary realloc). + assert m.A_log.data.data_ptr() == data_before.data_ptr() + # Name is still reported even on the no-realloc path — this is by + # design so the rank-0 log line reflects policy coverage, not churn. + assert any(n.endswith("A_log") for n in updated) + + def test_walks_multiple_model_chunks(self): + """``setup_model_and_optimizer`` passes a list of model chunks (for + virtual pipeline parallelism). The helper must iterate all of them.""" + chunks = [_ToyModule(), _ToyModule()] + for c in chunks: + c.bfloat16() + + enforce_marked_param_dtypes(chunks) + for c in chunks: + assert c.A_log.dtype == torch.float32 + + def test_returns_empty_when_no_marks(self): + m = nn.Linear(4, 4) + m.bfloat16() + assert enforce_marked_param_dtypes([m]) == [] + + +# --------------------------------------------------------------------------- +# Upstream: Qwen3_5Bridge._weight_to_mcore_format +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def bridge_stub(): + """Build a ``Qwen3_5Bridge`` without invoking ``__init__`` — ``__init__`` + needs a real HF config. The A_log branch only reads ``self.dtype``, which + we set directly, so skipping init is safe and lets this test stay + CPU-only and dep-free.""" + pytest.importorskip("mbridge") + from miles_plugins.mbridge.qwen3_5 import Qwen3_5Bridge + + bridge = Qwen3_5Bridge.__new__(Qwen3_5Bridge) + return bridge + + +class TestQwen3_5BridgeALogOverride: + A_LOG_NAME = "decoder.layers.0.self_attention.linear_attn.A_log" + + def test_returns_fp32_when_bridge_dtype_is_bf16(self, bridge_stub): + """The override must bypass mbridge's ``w.to(self.dtype)`` pre-cast + that would otherwise round HF fp32 to bf16 here.""" + bridge_stub.dtype = torch.bfloat16 + hf_tensor = torch.randn(32, dtype=torch.float32) + + out = bridge_stub._weight_to_mcore_format(self.A_LOG_NAME, [hf_tensor]) + + assert out.dtype == torch.float32 + assert torch.equal(out, hf_tensor) + assert out.is_contiguous() + + def test_upcasts_when_hf_input_is_bf16(self, bridge_stub): + """A_log arriving as bf16 (non-canonical ckpt) is still forced to + fp32 — the invariant is the output dtype, not the input's.""" + bridge_stub.dtype = torch.bfloat16 + hf_tensor = torch.randn(32, dtype=torch.bfloat16) + + out = bridge_stub._weight_to_mcore_format(self.A_LOG_NAME, [hf_tensor]) + + assert out.dtype == torch.float32 + + def test_mtp_layer_a_log_also_matches(self, bridge_stub): + """The override uses ``endswith`` so MTP-layer A_log + (``mtp.layers.{idx}...``) also matches — MTP is a real Qwen3.5 + variant and must not silently skip the override.""" + bridge_stub.dtype = torch.bfloat16 + hf_tensor = torch.randn(32, dtype=torch.float32) + + out = bridge_stub._weight_to_mcore_format("mtp.layers.0.self_attention.linear_attn.A_log", [hf_tensor]) + assert out.dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# End-to-end: the two halves together, matching ``Bridge.load_weights``. +# --------------------------------------------------------------------------- + + +class TestALogLoadPathEndToEnd: + """Replays the dtype-relevant subset of ``Bridge.load_weights`` on a toy + model, as documented in ``tools/debug_a_log_old_flow.py``. No distributed + or real safetensor IO — only the two cast points we care about. + + Expected outcome: HF fp32 value lands in the Megatron A_log param + bit-exactly. Regression target: the OLD ``patch_weight_to_mcore_format_preserve_fp32`` + failed here because ``bridge.py:246`` still cast down to ``param.dtype == bf16``. + """ + + def test_lossless_roundtrip(self, bridge_stub): + a_log_name = "decoder.layers.0.self_attention.linear_attn.A_log" + hf_tensor = torch.randn(32, dtype=torch.float32) + + # 1. Build model (A_log marked fp32 at definition site). + model = _ToyModule(num_heads=32) + + # 2. Megatron wraps with Float16Module → .bfloat16(). + model.bfloat16() + + # 3. enforce_marked_param_dtypes restores A_log to fp32 BEFORE + # load_weights runs, so ``param.dtype`` at bridge.py:246 is fp32. + enforce_marked_param_dtypes([model]) + assert model.A_log.dtype == torch.float32 + + # 4. mbridge: _weight_to_mcore_format (with override → fp32). + bridge_stub.dtype = torch.bfloat16 # would demote without override + mcore_weight = bridge_stub._weight_to_mcore_format(a_log_name, [hf_tensor]) + assert mcore_weight.dtype == torch.float32 + + # 5. mbridge bridge.py:246 — ``t.to(param.device, dtype=param.dtype)``. + param = model.A_log + staged = mcore_weight.to(param.device, dtype=param.dtype).contiguous() + assert staged.dtype == torch.float32 # no-op cast + + # 6. mbridge bridge.py:258 — ``param.copy_(param_to_load)``. + param.data.copy_(staged) + + # Bit-exact round-trip: both halves were required to get here. + assert model.A_log.dtype == torch.float32 + assert torch.equal(model.A_log.data, hf_tensor) + + def test_old_patch_only_regresses_without_enforce(self, bridge_stub): + """Negative control: if we DROP ``enforce_marked_param_dtypes`` and + only keep the upstream override (the shape of the old patch), the + downstream ``t.to(param.dtype)`` still rounds to bf16. This pins the + old failure mode so it cannot be re-introduced by accident.""" + a_log_name = "decoder.layers.0.self_attention.linear_attn.A_log" + # Use a value where bf16 rounding is observable. + hf_tensor = torch.tensor([0.970378123] * 8, dtype=torch.float32) + + model = _ToyModule(num_heads=8) + model.bfloat16() # A_log is bf16; no enforce call here on purpose. + + bridge_stub.dtype = torch.bfloat16 + mcore_weight = bridge_stub._weight_to_mcore_format(a_log_name, [hf_tensor]) + assert mcore_weight.dtype == torch.float32 + + staged = mcore_weight.to(model.A_log.device, dtype=model.A_log.dtype).contiguous() + # Regression check: demoted to bf16 because param.dtype is bf16. + assert staged.dtype == torch.bfloat16 + assert not torch.equal(staged.to(torch.float32), hf_tensor) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 354c216a651..a0682cf665f 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -1,7 +1,6 @@ import gc import os import shutil -from functools import wraps import torch import torch.distributed as dist @@ -12,8 +11,8 @@ import miles_plugins.mbridge # noqa: F401 from mbridge import AutoBridge -from mbridge.core.bridge import Bridge from miles.backends.megatron_utils.arguments import set_default_megatron_args +from miles.backends.megatron_utils.fp32_param_utils import enforce_marked_param_dtypes from miles.backends.megatron_utils.initialize import init from miles.backends.megatron_utils.model_provider import get_model_provider_func from miles.utils.logging_utils import configure_logger @@ -21,24 +20,6 @@ from miles_plugins.models.hf_attention import _load_hf_config -def patch_weight_to_mcore_format_preserve_fp32(): - - original_method = Bridge._weight_to_mcore_format - - @wraps(original_method) - def patched_method(self, mcore_weights_name, hf_weights): - original_dtype = getattr(self, "dtype", None) - self.dtype = None - try: - result = original_method(self, mcore_weights_name, hf_weights) - finally: - self.dtype = original_dtype - return result - - Bridge._weight_to_mcore_format = patched_method - print("[Patch] Applied patch to preserve FP32 precision in _weight_to_mcore_format") - - def add_convertion_args(parser): """Add conversion arguments to the parser""" parser.add_argument("--hf-checkpoint", type=str, required=True, help="HuggingFace model path") @@ -129,6 +110,7 @@ def main(): args = get_args() init(args) model = get_model(get_model_provider_func(args), ModelType.encoder_or_decoder, wrap_with_ddp=False) + enforce_marked_param_dtypes(model) # Load model hf_model_path = args.hf_checkpoint @@ -138,9 +120,6 @@ def main(): # Fallback for configs with model_type unknown to installed transformers. bridge = AutoBridge.from_config(_load_hf_config(hf_model_path)) - # Patch to preserve FP32 precision for _keep_fp32 params - patch_weight_to_mcore_format_preserve_fp32() - bridge.load_weights(model, hf_model_path, memory_efficient=True) print(f"Model loaded: {hf_model_path}") From 22691274eb360762e94bed7c3e63159d5aa5903e Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 13 Apr 2026 03:56:44 +0000 Subject: [PATCH 2/4] cmt --- miles_plugins/models/qwen3_5.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index f8457dc0336..5c43d732dcd 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -75,6 +75,9 @@ def __init__(self, config, layer_idx: int): self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) mark_param_dtype(self.A_log, torch.float32) + # HF stores this norm in fp32, but unlike A_log its precision impact is + # negligible and sglang runs it in bf16 on the rollout side — follow + # config.dtype (bf16) to stay equivalent to rollout. self.norm = FusedRMSNormGated( self.head_v_dim, eps=self.layer_norm_epsilon, From a6b175e956c4f64abc8ce755e3dc4fe468c9363d Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 14 Apr 2026 19:13:15 +0000 Subject: [PATCH 3/4] docs: add mixed-precision fp32 parameter preservation guide Rewrite fp32_param_preservation.md as a practical how-to guide for supporting fp32 parameters in bf16 models. Leads with a 2-step quick start, uses Qwen3.5 A_log as a complete example, and keeps the cast chain analysis as background. Register the doc in index.rst. Made-with: Cursor --- .../fp32_param_preservation.md | 175 ++++++++++++++++++ docs/en/index.rst | 1 + 2 files changed, 176 insertions(+) create mode 100644 docs/en/developer_guide/fp32_param_preservation.md diff --git a/docs/en/developer_guide/fp32_param_preservation.md b/docs/en/developer_guide/fp32_param_preservation.md new file mode 100644 index 00000000000..0da80cafb7b --- /dev/null +++ b/docs/en/developer_guide/fp32_param_preservation.md @@ -0,0 +1,175 @@ +# Supporting Mixed-Precision Parameters (fp32 in bf16 models) + +Some model architectures require specific parameters to remain in fp32 even when the rest of the model runs in bf16/fp16. For example, Qwen3.5's `A_log` parameter must stay fp32 — if it gets rounded to bf16, Megatron-side activations no longer match sglang's fp32 `A_log` on the rollout side, causing precision drift. + +miles provides a lightweight, model-agnostic utility in `miles/backends/megatron_utils/fp32_param_utils.py` to handle this. The utility requires **no changes** to Megatron or mbridge base code. + +## Quick start — 2 steps to support a new fp32 parameter + +### Step 1: Mark the parameter in your model definition + +At the model definition site (`miles_plugins/models/your_model.py`), call `mark_param_dtype` right after creating the parameter: + +```python +from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype + +# In your model's __init__: +self.X = nn.Parameter(some_init_tensor.to(torch.float32)) +mark_param_dtype(self.X, torch.float32) +``` + +This tags the parameter so that `enforce_marked_param_dtypes` (already wired into the training and conversion entry points) will restore it to fp32 after Megatron's `Float16Module` casts the entire model to bf16. + +### Step 2: Override the Bridge to preserve fp32 during weight loading + +In your Bridge subclass (`miles_plugins/mbridge/your_model.py`), add a name-matched early-return in `_weight_to_mcore_format` so the HF checkpoint value is not pre-cast to bf16: + +```python +def _weight_to_mcore_format( + self, mcore_weights_name: str, hf_weights: list[torch.Tensor] +) -> tuple[list[str], list[torch.Tensor]]: + if mcore_weights_name.endswith("your_module.X"): + assert len(hf_weights) == 1 + return hf_weights[0].to(dtype=torch.float32).contiguous() + + return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) +``` + +That's it. No changes to `enforce_marked_param_dtypes`, `model.py`, `convert_hf_to_torch_dist.py`, Megatron, or mbridge base are needed. + +## Complete example: Qwen3.5 `A_log` + +### Model definition (`miles_plugins/models/qwen3_5.py`) + +```python +from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype + +class Qwen3_5GatedDeltaNet(nn.Module): + def __init__(self, config, layer_idx: int): + ... + A = torch.empty(self.num_v_heads).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) + mark_param_dtype(self.A_log, torch.float32) +``` + +### Bridge override (`miles_plugins/mbridge/qwen3_5.py`) + +```python +class Qwen3_5Bridge(Qwen2MoEBridge): + def _weight_to_mcore_format( + self, mcore_weights_name: str, hf_weights: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: + if mcore_weights_name.endswith("self_attention.linear_attn.A_log"): + assert len(hf_weights) == 1 + # Keep A_log in fp32 before TP scatter; this avoids precision loss + # from Bridge's global pre-cast to self.dtype. + return hf_weights[0].to(dtype=torch.float32).contiguous() + + # ... other weight conversions ... + return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) +``` + +### Integration points (already wired — no action needed) + +`enforce_marked_param_dtypes(model)` is called right after `get_model` in both: + +- `miles/backends/megatron_utils/model.py` — the training entry point +- `tools/convert_hf_to_torch_dist.py` — the HF → Megatron checkpoint conversion tool + +## Code path: from parameter definition to loaded weight + +The following traces the real execution path of a parameter from model definition through to the final loaded weight. Understanding this path makes it clear where each cast happens and why both steps above are needed. + +### Phase 1 — Model construction and Float16Module wrap + +``` +model_provider_func() # builds your nn.Module (fp32 params) + └─ Qwen3_5GatedDeltaNet.__init__() + └─ self.A_log = nn.Parameter(...) # fp32 at creation + └─ mark_param_dtype(self.A_log, fp32) # tags _miles_forced_param_dtype attr + │ + ▼ +get_model() # megatron/training/training.py:1162 + ├─ build_model() # calls model_provider_func + ├─ model_module.cuda() # move to GPU, still fp32 + └─ Float16Module(config, model_module) # :1264 + └─ module.bfloat16() # ← CAST 1: every param.data → bf16 + │ + ▼ +enforce_marked_param_dtypes(model) # miles model.py / convert tool + └─ for each param with _miles_forced_param_dtype: + param.data = param.data.to(fp32) # ← UNDO CAST 1: restore tagged params +``` + +After this phase, tagged parameters are fp32 in the live model. Untagged parameters stay bf16. The `Parameter` object identity is preserved so optimizer and DDP registered afterwards see stable references. + +### Phase 2 — HF weight loading via mbridge + +``` +bridge.load_weights(model, hf_path) # mbridge/core/bridge.py:152 + │ + │ for each (local_name, hf_names): + │ ┌─ load HF tensors from safetensors (original dtype, e.g. fp32) + │ │ + │ ├─ _weight_to_mcore_format(name, hf_weights) # bridge.py:816 + │ │ │ + │ │ ├─ Bridge base: w.to(self.dtype) # ← CAST 2: hf tensor → bf16 + │ │ │ (self.dtype is typically bf16) + │ │ │ + │ │ └─ Subclass override (Step 2): + │ │ if name matches "...A_log": + │ │ return hf_weights[0].to(fp32) # ← BYPASS CAST 2 + │ │ + │ ├─ _weight_split_across_tp(name, mcore_weight, param, tp_size) + │ │ + │ ├─ t.to(param.device, dtype=param.dtype) # ← CAST 3: align to param dtype + │ │ param.dtype is fp32 (thanks to enforce above), so this is a no-op for + │ │ tagged params; for untagged params it casts to bf16 as expected + │ │ + │ └─ scatter across TP ranks → param.copy_(result) +``` + +### End-to-end dtype flow for a tagged fp32 parameter + +| Stage | What happens | Resulting dtype | +|---|---|---| +| `nn.Parameter(...)` | Created in model definition | fp32 | +| `mark_param_dtype(...)` | Tags the param, no dtype change | fp32 | +| `module.bfloat16()` (cast 1) | `Float16Module` wraps the model | **bf16** | +| `enforce_marked_param_dtypes` | Restores tagged params | fp32 | +| `_weight_to_mcore_format` (cast 2) | Bridge subclass early-returns fp32 | fp32 | +| `t.to(dtype=param.dtype)` (cast 3) | `param.dtype` is fp32, no-op | fp32 | +| `param.copy_(...)` | Final loaded value | fp32 | + +Without either step, one of the intermediate stages silently rounds the value to bf16 — see the table below. + +## Background: why both steps are required + +Megatron's bf16/fp16 training stack introduces **three implicit cast points** that can silently round fp32 parameters: + +| # | Location | What it casts | +|---|---|---| +| 1 | `Float16Module` ctor — `module.bfloat16()` | Every `nn.Parameter.data` → bf16 at wrap time | +| 2 | `Bridge._weight_to_mcore_format` — `w.to(self.dtype)` | HF tensor → Bridge's `self.dtype` (bf16) | +| 3 | `Bridge.load_weights` — `t.to(param.device, dtype=param.dtype)` | mcore tensor → Megatron `param.dtype` (bf16, due to cast 1) | + +Cast 1 has no declarative opt-out — even Megatron's own `_maintain_float32_expert_bias` uses a post-hoc `.data.to(float32)` workaround. `enforce_marked_param_dtypes` generalizes this pattern. + +**Step 1** (`mark_param_dtype` + `enforce_marked_param_dtypes`) closes cast 1 and 3: once `param.dtype == fp32`, the `load_weights` in-place cast is a no-op for tagged params. + +**Step 2** (Bridge override) closes cast 2: the HF tensor is kept fp32 before being scattered across TP ranks. + +Both steps are necessary. Doing only one leaves a silent precision trap: + +| Config | mcore_weight | After load_weights cast | Final dtype | Value-accurate | +|---|---|---|---|---| +| Nothing | bf16 | bf16 | bf16 | no | +| Step 1 only (no bridge override) | **bf16** | fp32 up-cast | fp32 | **no** — bits already rounded at cast 2 | +| Step 2 only (no mark/enforce) | fp32 | **bf16** | **bf16** | no | +| Both steps | fp32 | fp32 | fp32 | yes | + +The "Step 1 only" row is the subtle trap: the final dtype *looks* correct (fp32), but the values were already rounded to bf16 precision at cast 2 and then up-cast back into an fp32 container. + +## Tests + +`tests/fast/backends/megatron_utils/test_fp32_param_utils.py` — 13 CPU-only tests covering the downstream helper, the upstream bridge override, a bit-exact end-to-end round-trip, and a negative regression guard against the Step-1-only failure mode. diff --git a/docs/en/index.rst b/docs/en/index.rst index 7c48585b650..b99aeba171e 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -70,6 +70,7 @@ miles is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models :caption: Developer Guide developer_guide/debug.md + developer_guide/fp32_param_preservation.md .. toctree:: :maxdepth: 1 From f37cfee525ac551c341ca46af898fd34230f0e3d Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Tue, 14 Apr 2026 12:42:03 -0700 Subject: [PATCH 4/4] =?UTF-8?q?change=20docs=E2=80=9C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../advanced/arch-support-beyond-megatron.md | 29 +++ .../fp32_param_preservation.md | 175 ------------------ docs/en/index.rst | 1 - 3 files changed, 29 insertions(+), 176 deletions(-) delete mode 100644 docs/en/developer_guide/fp32_param_preservation.md diff --git a/docs/en/advanced/arch-support-beyond-megatron.md b/docs/en/advanced/arch-support-beyond-megatron.md index 0db8c8a40a3..4b3e2a02ca6 100644 --- a/docs/en/advanced/arch-support-beyond-megatron.md +++ b/docs/en/advanced/arch-support-beyond-megatron.md @@ -27,6 +27,35 @@ miles leverages this mechanism by **hijacking the spec generation stage to repla Through the coordination of these three components, we can successfully run a complex model architecture not natively supported by Megatron—using its HuggingFace implementation as the vehicle—on top of Megatron's parallel framework. This is achieved while fully retaining all key capabilities like model parallelism, MoE acceleration, and pipeline scheduling. +## Mixed-Precision: Preserving fp32 Parameters in bf16 Models + +Some model architectures require specific parameters to remain in fp32 even when the rest of the model runs in bf16. For example, Qwen3.5's `A_log` parameter must stay fp32 — if rounded to bf16, Megatron-side activations diverge from sglang's fp32 `A_log` on the rollout side, causing precision drift. + +Megatron's training stack has **three implicit cast points** that silently round fp32 parameters to bf16: `Float16Module` construction, `Bridge._weight_to_mcore_format`, and `Bridge.load_weights`. Both steps below are required — doing only one leaves a silent precision trap where the final dtype *looks* correct (fp32) but values were already rounded to bf16 precision. + +### Step 1: Mark the parameter in your model definition + +```python +from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype + +# In your model's __init__: +self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) +mark_param_dtype(self.A_log, torch.float32) +``` + +`enforce_marked_param_dtypes(model)` — already wired into training and checkpoint conversion entry points — restores tagged params to fp32 after `Float16Module` casts the entire model to bf16. + +### Step 2: Override the Bridge to bypass bf16 pre-cast during weight loading + +```python +class Qwen3_5Bridge(Qwen2MoEBridge): + def _weight_to_mcore_format(self, mcore_weights_name, hf_weights): + if mcore_weights_name.endswith("self_attention.linear_attn.A_log"): + assert len(hf_weights) == 1 + return hf_weights[0].to(dtype=torch.float32).contiguous() + return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) +``` + ## Current Limitations * This approach does not currently support Tensor Parallelism (TP) within the replaced module itself (e.g., the Attention layer in this case). diff --git a/docs/en/developer_guide/fp32_param_preservation.md b/docs/en/developer_guide/fp32_param_preservation.md deleted file mode 100644 index 0da80cafb7b..00000000000 --- a/docs/en/developer_guide/fp32_param_preservation.md +++ /dev/null @@ -1,175 +0,0 @@ -# Supporting Mixed-Precision Parameters (fp32 in bf16 models) - -Some model architectures require specific parameters to remain in fp32 even when the rest of the model runs in bf16/fp16. For example, Qwen3.5's `A_log` parameter must stay fp32 — if it gets rounded to bf16, Megatron-side activations no longer match sglang's fp32 `A_log` on the rollout side, causing precision drift. - -miles provides a lightweight, model-agnostic utility in `miles/backends/megatron_utils/fp32_param_utils.py` to handle this. The utility requires **no changes** to Megatron or mbridge base code. - -## Quick start — 2 steps to support a new fp32 parameter - -### Step 1: Mark the parameter in your model definition - -At the model definition site (`miles_plugins/models/your_model.py`), call `mark_param_dtype` right after creating the parameter: - -```python -from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype - -# In your model's __init__: -self.X = nn.Parameter(some_init_tensor.to(torch.float32)) -mark_param_dtype(self.X, torch.float32) -``` - -This tags the parameter so that `enforce_marked_param_dtypes` (already wired into the training and conversion entry points) will restore it to fp32 after Megatron's `Float16Module` casts the entire model to bf16. - -### Step 2: Override the Bridge to preserve fp32 during weight loading - -In your Bridge subclass (`miles_plugins/mbridge/your_model.py`), add a name-matched early-return in `_weight_to_mcore_format` so the HF checkpoint value is not pre-cast to bf16: - -```python -def _weight_to_mcore_format( - self, mcore_weights_name: str, hf_weights: list[torch.Tensor] -) -> tuple[list[str], list[torch.Tensor]]: - if mcore_weights_name.endswith("your_module.X"): - assert len(hf_weights) == 1 - return hf_weights[0].to(dtype=torch.float32).contiguous() - - return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) -``` - -That's it. No changes to `enforce_marked_param_dtypes`, `model.py`, `convert_hf_to_torch_dist.py`, Megatron, or mbridge base are needed. - -## Complete example: Qwen3.5 `A_log` - -### Model definition (`miles_plugins/models/qwen3_5.py`) - -```python -from miles.backends.megatron_utils.fp32_param_utils import mark_param_dtype - -class Qwen3_5GatedDeltaNet(nn.Module): - def __init__(self, config, layer_idx: int): - ... - A = torch.empty(self.num_v_heads).uniform_(0, 16) - self.A_log = nn.Parameter(torch.log(A).to(torch.float32)) - mark_param_dtype(self.A_log, torch.float32) -``` - -### Bridge override (`miles_plugins/mbridge/qwen3_5.py`) - -```python -class Qwen3_5Bridge(Qwen2MoEBridge): - def _weight_to_mcore_format( - self, mcore_weights_name: str, hf_weights: list[torch.Tensor] - ) -> tuple[list[str], list[torch.Tensor]]: - if mcore_weights_name.endswith("self_attention.linear_attn.A_log"): - assert len(hf_weights) == 1 - # Keep A_log in fp32 before TP scatter; this avoids precision loss - # from Bridge's global pre-cast to self.dtype. - return hf_weights[0].to(dtype=torch.float32).contiguous() - - # ... other weight conversions ... - return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) -``` - -### Integration points (already wired — no action needed) - -`enforce_marked_param_dtypes(model)` is called right after `get_model` in both: - -- `miles/backends/megatron_utils/model.py` — the training entry point -- `tools/convert_hf_to_torch_dist.py` — the HF → Megatron checkpoint conversion tool - -## Code path: from parameter definition to loaded weight - -The following traces the real execution path of a parameter from model definition through to the final loaded weight. Understanding this path makes it clear where each cast happens and why both steps above are needed. - -### Phase 1 — Model construction and Float16Module wrap - -``` -model_provider_func() # builds your nn.Module (fp32 params) - └─ Qwen3_5GatedDeltaNet.__init__() - └─ self.A_log = nn.Parameter(...) # fp32 at creation - └─ mark_param_dtype(self.A_log, fp32) # tags _miles_forced_param_dtype attr - │ - ▼ -get_model() # megatron/training/training.py:1162 - ├─ build_model() # calls model_provider_func - ├─ model_module.cuda() # move to GPU, still fp32 - └─ Float16Module(config, model_module) # :1264 - └─ module.bfloat16() # ← CAST 1: every param.data → bf16 - │ - ▼ -enforce_marked_param_dtypes(model) # miles model.py / convert tool - └─ for each param with _miles_forced_param_dtype: - param.data = param.data.to(fp32) # ← UNDO CAST 1: restore tagged params -``` - -After this phase, tagged parameters are fp32 in the live model. Untagged parameters stay bf16. The `Parameter` object identity is preserved so optimizer and DDP registered afterwards see stable references. - -### Phase 2 — HF weight loading via mbridge - -``` -bridge.load_weights(model, hf_path) # mbridge/core/bridge.py:152 - │ - │ for each (local_name, hf_names): - │ ┌─ load HF tensors from safetensors (original dtype, e.g. fp32) - │ │ - │ ├─ _weight_to_mcore_format(name, hf_weights) # bridge.py:816 - │ │ │ - │ │ ├─ Bridge base: w.to(self.dtype) # ← CAST 2: hf tensor → bf16 - │ │ │ (self.dtype is typically bf16) - │ │ │ - │ │ └─ Subclass override (Step 2): - │ │ if name matches "...A_log": - │ │ return hf_weights[0].to(fp32) # ← BYPASS CAST 2 - │ │ - │ ├─ _weight_split_across_tp(name, mcore_weight, param, tp_size) - │ │ - │ ├─ t.to(param.device, dtype=param.dtype) # ← CAST 3: align to param dtype - │ │ param.dtype is fp32 (thanks to enforce above), so this is a no-op for - │ │ tagged params; for untagged params it casts to bf16 as expected - │ │ - │ └─ scatter across TP ranks → param.copy_(result) -``` - -### End-to-end dtype flow for a tagged fp32 parameter - -| Stage | What happens | Resulting dtype | -|---|---|---| -| `nn.Parameter(...)` | Created in model definition | fp32 | -| `mark_param_dtype(...)` | Tags the param, no dtype change | fp32 | -| `module.bfloat16()` (cast 1) | `Float16Module` wraps the model | **bf16** | -| `enforce_marked_param_dtypes` | Restores tagged params | fp32 | -| `_weight_to_mcore_format` (cast 2) | Bridge subclass early-returns fp32 | fp32 | -| `t.to(dtype=param.dtype)` (cast 3) | `param.dtype` is fp32, no-op | fp32 | -| `param.copy_(...)` | Final loaded value | fp32 | - -Without either step, one of the intermediate stages silently rounds the value to bf16 — see the table below. - -## Background: why both steps are required - -Megatron's bf16/fp16 training stack introduces **three implicit cast points** that can silently round fp32 parameters: - -| # | Location | What it casts | -|---|---|---| -| 1 | `Float16Module` ctor — `module.bfloat16()` | Every `nn.Parameter.data` → bf16 at wrap time | -| 2 | `Bridge._weight_to_mcore_format` — `w.to(self.dtype)` | HF tensor → Bridge's `self.dtype` (bf16) | -| 3 | `Bridge.load_weights` — `t.to(param.device, dtype=param.dtype)` | mcore tensor → Megatron `param.dtype` (bf16, due to cast 1) | - -Cast 1 has no declarative opt-out — even Megatron's own `_maintain_float32_expert_bias` uses a post-hoc `.data.to(float32)` workaround. `enforce_marked_param_dtypes` generalizes this pattern. - -**Step 1** (`mark_param_dtype` + `enforce_marked_param_dtypes`) closes cast 1 and 3: once `param.dtype == fp32`, the `load_weights` in-place cast is a no-op for tagged params. - -**Step 2** (Bridge override) closes cast 2: the HF tensor is kept fp32 before being scattered across TP ranks. - -Both steps are necessary. Doing only one leaves a silent precision trap: - -| Config | mcore_weight | After load_weights cast | Final dtype | Value-accurate | -|---|---|---|---|---| -| Nothing | bf16 | bf16 | bf16 | no | -| Step 1 only (no bridge override) | **bf16** | fp32 up-cast | fp32 | **no** — bits already rounded at cast 2 | -| Step 2 only (no mark/enforce) | fp32 | **bf16** | **bf16** | no | -| Both steps | fp32 | fp32 | fp32 | yes | - -The "Step 1 only" row is the subtle trap: the final dtype *looks* correct (fp32), but the values were already rounded to bf16 precision at cast 2 and then up-cast back into an fp32 container. - -## Tests - -`tests/fast/backends/megatron_utils/test_fp32_param_utils.py` — 13 CPU-only tests covering the downstream helper, the upstream bridge override, a bit-exact end-to-end round-trip, and a negative regression guard against the Step-1-only failure mode. diff --git a/docs/en/index.rst b/docs/en/index.rst index b99aeba171e..7c48585b650 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -70,7 +70,6 @@ miles is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models :caption: Developer Guide developer_guide/debug.md - developer_guide/fp32_param_preservation.md .. toctree:: :maxdepth: 1