[fix] support general logic to bypass fp32 downcast and fix qwen35 A_log dtype

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).

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

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,13 +126,17 @@ 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):
         if hasattr(args, f.name):
             kwargs[f.name] = getattr(args, f.name)
     config = OptimizerConfig(**kwargs)
     config.timers = None
+
     optimizer = get_megatron_optimizer(
         config=config,
         model_chunks=model,

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

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,8 +72,12 @@ 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)
 
+        # 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,