diff --git a/cookbook/mm/fsdp2_gemma4_12b_mm.py b/cookbook/mm/fsdp2_gemma4_12b_mm.py index 62e26d776..c45e5a15d 100644 --- a/cookbook/mm/fsdp2_gemma4_12b_mm.py +++ b/cookbook/mm/fsdp2_gemma4_12b_mm.py @@ -112,7 +112,10 @@ def train(): 'enabled': True, 'router_dtype': 'fp32', 'keep_router_logits': False, - } + }, + # Optional per-field precision overrides (native FSDP2 only), e.g. BF16 compute + # with FP32 gradient reduction; omitted fields follow the `mixed_precision` base. + # 'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}, }, ) diff --git a/cookbook/mm/fsdp2_gemma4_mm.py b/cookbook/mm/fsdp2_gemma4_mm.py index 5c756cfbe..891fdc133 100644 --- a/cookbook/mm/fsdp2_gemma4_mm.py +++ b/cookbook/mm/fsdp2_gemma4_mm.py @@ -102,7 +102,10 @@ def train(): 'enabled': True, 'router_dtype': 'fp32', 'keep_router_logits': False, - } + }, + # Optional per-field precision overrides (native FSDP2 only), e.g. BF16 compute + # with FP32 gradient reduction; omitted fields follow the `mixed_precision` base. + # 'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}, }, ) diff --git a/cookbook/rl/grpo/short_math_grpo_routing_replay.py b/cookbook/rl/grpo/short_math_grpo_routing_replay.py index 18b1cadde..b4ea66075 100644 --- a/cookbook/rl/grpo/short_math_grpo_routing_replay.py +++ b/cookbook/rl/grpo/short_math_grpo_routing_replay.py @@ -177,7 +177,10 @@ def main(): 'enabled': enable_ep, 'router_dtype': 'fp32', 'keep_router_logits': False, - } + }, + # Optional per-field precision overrides (native FSDP2 only), e.g. BF16 compute + # with FP32 gradient reduction; omitted fields follow the `mixed_precision` base. + # 'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}, }, ) diff --git a/cookbook/transformers/ep_fsdp2_lora_deepseek_v4.py b/cookbook/transformers/ep_fsdp2_lora_deepseek_v4.py index 4aecc9a0f..8e54f7dda 100644 --- a/cookbook/transformers/ep_fsdp2_lora_deepseek_v4.py +++ b/cookbook/transformers/ep_fsdp2_lora_deepseek_v4.py @@ -78,7 +78,10 @@ def train(): 'enabled': ENABLE_EP, 'router_dtype': 'fp32', 'keep_router_logits': False, - } + }, + # Optional per-field precision overrides (native FSDP2 only), e.g. BF16 compute + # with FP32 gradient reduction; omitted fields follow the `mixed_precision` base. + # 'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}, }, ) lora_cfg = _build_lora_config() diff --git a/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py b/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py index 12ba8a973..7e3b825a1 100644 --- a/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py +++ b/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py @@ -89,7 +89,10 @@ def train(): 'enabled': ENABLE_EP, 'router_dtype': 'fp32', 'keep_router_logits': False, - } + }, + # Optional per-field precision overrides (native FSDP2 only), e.g. BF16 compute + # with FP32 gradient reduction; omitted fields follow the `mixed_precision` base. + # 'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}, }, ) # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default, diff --git a/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py b/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py index bc06ef5f2..efeeed52a 100644 --- a/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py +++ b/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py @@ -101,7 +101,10 @@ def train(): 'enabled': ENABLE_EP, 'router_dtype': 'fp32', 'keep_router_logits': False, - } + }, + # Optional per-field precision overrides (native FSDP2 only), e.g. BF16 compute + # with FP32 gradient reduction; omitted fields follow the `mixed_precision` base. + # 'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}, }, lora_config=lora_cfg, ) diff --git a/docs/source_en/Components/Model/TransformersModel.md b/docs/source_en/Components/Model/TransformersModel.md index a9005a550..8c10e4b98 100644 --- a/docs/source_en/Components/Model/TransformersModel.md +++ b/docs/source_en/Components/Model/TransformersModel.md @@ -29,7 +29,24 @@ class TransformersModel: - mixed_precision: Mixed precision information, default `bf16`, recommended to keep unchanged if you have GPUs with 30 series or above - strategy: How to encapsulate the model for multi-GPU training, default uses `accelerate` framework. - ddp_config: DDP configuration when strategy is `accelerate`, see: [DDPKwargs](https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L155) -- fsdp_config: FSDP configuration when strategy is `accelerate`, see: [FSDPConfig](https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1566) +- fsdp_config: FSDP configuration when strategy is `accelerate`, see: [FSDPConfig](https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1566). When strategy is `native_fsdp`, it additionally supports the `mixed_precision` sub-config (independent precision settings, mapping to FSDP2's [MixedPrecisionPolicy](https://docs.pytorch.org/docs/stable/distributed.fsdp.html#torch.distributed.fsdp.MixedPrecisionPolicy) fields): + - `param_dtype`: Computation precision. Before forward/backward, FSDP casts the all-gathered full parameters to this dtype for computation; + - `reduce_dtype`: Gradient reduction precision. The cross-rank reduce-scatter accumulation runs in this dtype; `fp32` lowers multi-rank accumulation error at the cost of doubled reduction traffic (i.e. "BF16 compute, FP32 reduce"); + - `output_dtype`: The dtype each FSDP-wrapped unit's forward output is cast to. + + All three fields are optional and accept `'bf16'` / `'fp16'` / `'fp32'` (or the corresponding `torch.dtype`); omitted fields fall back to the base dtype derived from the top-level `mixed_precision` string, so leaving the sub-config out keeps the legacy behavior. Example: + + ```python + TransformersModel( + ..., + strategy='native_fsdp', + fsdp_config={'mixed_precision': {'reduce_dtype': 'fp32'}}, # BF16 compute, FP32 gradient reduction + ) + ``` + + Note: this sub-config only takes effect for `native_fsdp`; with expert parallelism (EP) enabled, NCCL constraints force the expert group's gradient reduction to run in `fp32`. + + Constraint: the effective `output_dtype` must equal `param_dtype` (or be omitted); otherwise construction raises `ValueError` — a mismatch introduces a dtype conversion at FSDP unit boundaries, and gradient checkpointing recompute then fails on mixed-dtype tensors. `reduce_dtype` only affects gradient reduction and is unconstrained. - grad_scaler_config: PyTorch's grad_scaler initialization configuration, see: [PyTorch's GradScaler constructor](https://github.com/pytorch/pytorch/blob/main/torch/cuda/amp/grad_scaler.py#L25) - memory_efficient_init: Whether to enable memory-efficient model initialization for FSDP. When enabled, only rank 0 loads full weights and broadcasts sharded parameters to other ranks, reducing peak memory usage during initialization. Default `False`. Note: The optimization currently only applies to transformers <= 4.57.6; for transformers >= 5.0.0, it may lead to negative performance impact. - kwargs: diff --git "a/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md" "b/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md" index cd0f16ade..85fc81588 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\346\250\241\345\236\213/TransformersModel.md" @@ -29,7 +29,24 @@ class TransformersModel: - mixed_precision: 混合精度信息，默认`bf16`，如果有30系以上显卡建议维持不变 - strategy: 如何封装模型使用多卡训练，默认使用`accelerate`框架。 - ddp_config: strategy为`accelerate`时的DDP配置，参见：[DDPKwargs](https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L155) -- fsdp_config: strategy为`accelerate`时的FSDP配置，参见：[FSDPConfig](https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1566) +- fsdp_config: strategy为`accelerate`时的FSDP配置，参见：[FSDPConfig](https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/dataclasses.py#L1566)。strategy为`native_fsdp`时额外支持`mixed_precision`子配置（独立精度配置，字段对应FSDP2的[MixedPrecisionPolicy](https://docs.pytorch.org/docs/stable/distributed.fsdp.html#torch.distributed.fsdp.MixedPrecisionPolicy)）： + - `param_dtype`: 计算精度。前向/反向前，FSDP将all-gather出的完整参数cast到该精度参与计算； + - `reduce_dtype`: 梯度归约精度。跨卡reduce-scatter累加在该精度下进行，设为`fp32`可降低多卡累加误差，代价是归约通信量翻倍（即"BF16计算、FP32归约"）； + - `output_dtype`: 每个FSDP包装单元forward输出cast到的精度。 + + 三个字段均可缺省，取值支持`'bf16'` / `'fp16'` / `'fp32'`（或对应的`torch.dtype`）；缺省字段回落到顶层`mixed_precision`字符串定出的基准精度，完全不配置时行为与旧版本一致。例如： + + ```python + TransformersModel( + ..., + strategy='native_fsdp', + fsdp_config={'mixed_precision': {'reduce_dtype': 'fp32'}}, # BF16计算、FP32梯度归约 + ) + ``` + + 注意：该子配置仅对`native_fsdp`生效；启用专家并行（EP）时，因NCCL约束专家组的梯度归约恒在`fp32`下进行。 + + 约束：生效的`output_dtype`必须与`param_dtype`一致（或缺省），否则构造时抛出`ValueError`——两者不一致会在FSDP单元边界引入dtype转换，梯度检查点重算时张量dtype不一致会导致训练报错；`reduce_dtype`只作用于梯度归约，不受该约束。 - grad_scaler_config: PyTorch的grad_scaler初始化配置，参见：[PyTorch的GradScaler构造](https://github.com/pytorch/pytorch/blob/main/torch/cuda/amp/grad_scaler.py#L25) - memory_efficient_init: 是否启用FSDP内存高效初始化。启用后仅rank 0加载完整权重，其余rank通过广播获取分片参数，降低初始化阶段的内存和显存峰值。默认`False`。注意：该优化目前仅适用于 transformers <= 4.57.6；对于 transformers >= 5.0.0，可能会导致负面性能影响。 - kwargs: diff --git a/src/twinkle/model/transformers/strategy/mixed_precision.py b/src/twinkle/model/transformers/strategy/mixed_precision.py new file mode 100644 index 000000000..8d431cdf5 --- /dev/null +++ b/src/twinkle/model/transformers/strategy/mixed_precision.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Parsing of the user-facing FSDP mixed-precision overrides. + +``fsdp_config['mixed_precision']`` configures the precision of each FSDP field +independently, e.g. BF16 computation with FP32 gradient reduction:: + + fsdp_config={'mixed_precision': {'param_dtype': 'bf16', 'reduce_dtype': 'fp32'}} + +Every key is optional. Only the native FSDP2 path consumes this sub-config: it +resolves the dict through :func:`parse_mp_overrides` and maps the resolved +dtypes onto torch's ``MixedPrecisionPolicy`` (FSDP1's ``buffer_dtype`` has no +counterpart here because only FSDP2 is supported). The accelerate path is not +adapted and keeps its own plugin-level handling. Combinations whose effective +``output_dtype`` differs from ``param_dtype`` are rejected at strategy +construction (see ``native_fsdp._check_mp_dtypes``): such a unit-boundary dtype +cast breaks gradient-checkpointing recompute. +""" +from typing import Any + +import torch + +# The FSDP2 ``MixedPrecisionPolicy`` fields users may override. ``cast_forward_inputs`` +# is intentionally not exposed: it must stay aligned with ``param_dtype``. +MP_OVERRIDE_FIELDS = ('param_dtype', 'reduce_dtype', 'output_dtype') + +_STR_TO_DTYPE = {'fp32': torch.float32, 'fp16': torch.float16, 'bf16': torch.bfloat16} +_ALLOWED_DTYPES = tuple(_STR_TO_DTYPE.values()) + + +def parse_mp_overrides(config: dict[str, Any] | None) -> dict[str, torch.dtype]: + """Validate ``fsdp_config['mixed_precision']`` and resolve it to torch dtypes. + + Args: + config: ``None``, or a dict mapping any of ``param_dtype`` / ``reduce_dtype`` / + ``output_dtype`` to ``'bf16'`` / ``'fp16'`` / ``'fp32'`` (case-insensitive) + or the equivalent ``torch.dtype``. A key whose value is ``None`` counts as + unset and falls back to the strategy default. + + Returns: + dict[str, torch.dtype]: the explicitly configured fields only; empty when + nothing is configured, so callers can treat it as pure overrides on top of + the ``mixed_precision`` string. + + Raises: + ValueError: on a non-dict config, an unknown key, or an invalid value. + """ + if config is None: + return {} + if not isinstance(config, dict): + raise ValueError(f"fsdp_config['mixed_precision'] must be a dict, got {type(config).__name__}: {config!r}") + unknown = sorted(set(config) - set(MP_OVERRIDE_FIELDS)) + if unknown: + raise ValueError(f"Unknown key(s) {unknown} in fsdp_config['mixed_precision']; " + f'valid keys are {list(MP_OVERRIDE_FIELDS)}.') + resolved: dict[str, torch.dtype] = {} + for field in MP_OVERRIDE_FIELDS: + value = config.get(field) + if value is None: + continue + if isinstance(value, str): + dtype = _STR_TO_DTYPE.get(value.lower()) + elif isinstance(value, torch.dtype) and value in _ALLOWED_DTYPES: + dtype = value + else: + dtype = None + if dtype is None: + raise ValueError(f"Invalid value for fsdp_config['mixed_precision']['{field}']: {value!r}. " + f"Must be one of 'bf16', 'fp16', 'fp32' or the corresponding torch.dtype.") + resolved[field] = dtype + return resolved diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index da49bce79..734d7fe01 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -12,6 +12,7 @@ from twinkle.utils.torch_utils import clone_state_dict_to_cpu from .base import TransformerStrategyBase from .base import fsdp_pretrained_load_context +from .mixed_precision import parse_mp_overrides if TYPE_CHECKING: from torch.distributed.fsdp import MixedPrecisionPolicy @@ -34,6 +35,10 @@ def __init__(self, self.device_mesh = device_mesh self.mixed_precision = mixed_precision self.fsdp_config = fsdp_config or {} + # Per-field precision overrides from fsdp_config['mixed_precision'], resolved (and + # validated) once at construction; empty means the mixed_precision string alone rules. + self._mp_overrides = parse_mp_overrides(self.fsdp_config.get('mixed_precision')) + _check_mp_dtypes(mixed_precision, self._mp_overrides) self._memory_efficient_init = memory_efficient_init self.enable_ep = enable_ep self.ep_fsdp_device_mesh = self._build_ep_fsdp_device_mesh(ep_size) if enable_ep else None @@ -162,7 +167,7 @@ def wrap_model(self, model, optimizer=None, *additional_optimizers): if ep_enabled: _ensure_moe_patched_if_needed(model, self.ep_fsdp_device_mesh) _place_ep_experts_on_local_device(model, self.ep_fsdp_device_mesh) - mp_policy = _build_mp_policy(self.mixed_precision) + mp_policy = _build_mp_policy(self.mixed_precision, self._mp_overrides) logger.info( "FSDP precision: parameter_storage_dtypes=%s, mixed_precision=%s, policy=%s", sorted({str(param.dtype) for param in model.parameters()}), @@ -499,18 +504,60 @@ def _ep_expert_state_dict_gather_dim(name: str) -> int: return 0 -def _build_mp_policy(mixed_precision: str) -> 'MixedPrecisionPolicy': +def _effective_mp_dtypes(mixed_precision: str, + overrides: Optional[Dict[str, torch.dtype]] = None) -> tuple: + """Resolve the base string + overrides into the three effective field dtypes. + + ``mixed_precision`` sets the base dtype shared by all three fields ('no'/'fp8' leave the + base as None, so FSDP follows each parameter's own dtype); ``overrides`` replaces only + the fields the user configured. + """ + base_dtype = {'bf16': torch.bfloat16, 'fp16': torch.float16}.get(mixed_precision) + overrides = overrides or {} + return (overrides.get('param_dtype', base_dtype), + overrides.get('reduce_dtype', base_dtype), + overrides.get('output_dtype', base_dtype)) + + +def _check_mp_dtypes(mixed_precision: str, overrides: Optional[Dict[str, torch.dtype]]) -> None: + """Reject configs whose effective ``output_dtype`` forms a dtype boundary against ``param_dtype``. + + ``output_dtype`` casts every FSDP unit's forward output while ``cast_forward_inputs`` + re-casts the next unit's inputs to ``param_dtype``. When an effective ``output_dtype`` is + set and differs from ``param_dtype``, gradient-checkpointing recompute can run against + tensors captured on the other side of that boundary -- validated on 4x NPU with Qwen3.5: + RoPE-promoted FP32 q/k meet a BF16 v and SDPA raises "Expected query, key, and value to + have the same dtype". An effective ``output_dtype`` of None performs no output cast, so + the tensor stream stays uniform in ``param_dtype`` and remains legal. ``reduce_dtype`` + only affects gradient reduction and is not constrained. + """ + param_dtype, _, output_dtype = _effective_mp_dtypes(mixed_precision, overrides) + if output_dtype is not None and output_dtype != param_dtype: + raise ValueError( + f"fsdp_config['mixed_precision'] must not give output_dtype ({output_dtype}) a different " + f"effective dtype than param_dtype ({param_dtype}) [mixed_precision={mixed_precision!r}]: " + f'the unit-boundary dtype cast breaks gradient-checkpointing recompute. Set output_dtype ' + f'equal to param_dtype or leave both unconfigured; reduce_dtype is unconstrained.') + + +def _build_mp_policy(mixed_precision: str, + overrides: Optional[Dict[str, torch.dtype]] = None) -> 'MixedPrecisionPolicy': + """Build the FSDP2 MixedPrecisionPolicy from the base string plus per-field overrides. + + ``overrides`` is the resolved ``fsdp_config['mixed_precision']`` (see + :func:`parse_mp_overrides`), e.g. BF16 compute with FP32 gradient reduction; param/output + boundary combinations are rejected earlier, at strategy construction (see + :func:`_check_mp_dtypes`). ``cast_forward_inputs`` keeps its legacy value (FSDP2 only acts + on it when ``param_dtype`` is set, so it is inert for the all-None 'no' case). + """ from torch.distributed.fsdp import MixedPrecisionPolicy - if mixed_precision == 'bf16': - dtype = torch.bfloat16 - elif mixed_precision == 'fp16': - dtype = torch.float16 - else: + param_dtype, reduce_dtype, output_dtype = _effective_mp_dtypes(mixed_precision, overrides) + if param_dtype is None and reduce_dtype is None and output_dtype is None: return MixedPrecisionPolicy() return MixedPrecisionPolicy( - param_dtype=dtype, - reduce_dtype=dtype, - output_dtype=dtype, + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, cast_forward_inputs=True, ) diff --git a/tests/transformers/test_fsdp_mixed_precision.py b/tests/transformers/test_fsdp_mixed_precision.py new file mode 100644 index 000000000..0e6f34551 --- /dev/null +++ b/tests/transformers/test_fsdp_mixed_precision.py @@ -0,0 +1,188 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from twinkle.model.transformers.strategy.mixed_precision import parse_mp_overrides +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy, _build_ep_mp_policy, _build_mp_policy + + +def test_none_and_empty_config_resolve_to_no_overrides(): + assert parse_mp_overrides(None) == {} + assert parse_mp_overrides({}) == {} + + +def test_partial_config_resolves_only_set_fields(): + assert parse_mp_overrides({'reduce_dtype': 'fp32'}) == {'reduce_dtype': torch.float32} + + +def test_full_config_string_values(): + config = {'param_dtype': 'bf16', 'reduce_dtype': 'fp32', 'output_dtype': 'fp16'} + assert parse_mp_overrides(config) == { + 'param_dtype': torch.bfloat16, + 'reduce_dtype': torch.float32, + 'output_dtype': torch.float16, + } + + +def test_torch_dtype_values_accepted(): + assert parse_mp_overrides({'param_dtype': torch.bfloat16}) == {'param_dtype': torch.bfloat16} + + +def test_string_values_case_insensitive(): + assert parse_mp_overrides({'reduce_dtype': 'FP32'}) == {'reduce_dtype': torch.float32} + + +def test_none_value_counts_as_unset(): + assert parse_mp_overrides({'param_dtype': 'bf16', 'reduce_dtype': None}) == {'param_dtype': torch.bfloat16} + + +@pytest.mark.parametrize('config', [ + {'param_dtpe': 'bf16'}, # typo'd key + {'buffer_dtype': 'fp32'}, # FSDP1-only field, not valid on the FSDP2 face +]) +def test_unknown_key_raises(config): + with pytest.raises(ValueError, match='Unknown key'): + parse_mp_overrides(config) + + +@pytest.mark.parametrize('value', ['fp64', 'float32', 32, torch.float64]) +def test_invalid_value_raises(value): + with pytest.raises(ValueError, match='reduce_dtype'): + parse_mp_overrides({'reduce_dtype': value}) + + +def test_non_dict_config_raises(): + with pytest.raises(ValueError, match='must be a dict'): + parse_mp_overrides('bf16') + + +# ── native FSDP path: base string + overrides → MixedPrecisionPolicy ────────── + + +def test_native_policy_base_string_only_matches_legacy_behavior(): + policy = _build_mp_policy('bf16') + assert (policy.param_dtype, policy.reduce_dtype, policy.output_dtype) == (torch.bfloat16, ) * 3 + assert policy.cast_forward_inputs + + policy = _build_mp_policy('fp16') + assert (policy.param_dtype, policy.reduce_dtype, policy.output_dtype) == (torch.float16, ) * 3 + assert policy.cast_forward_inputs + + # 'no' must stay byte-identical to the legacy default-constructed policy. + from torch.distributed.fsdp import MixedPrecisionPolicy + assert _build_mp_policy('no') == MixedPrecisionPolicy() + + +def test_native_policy_override_replaces_only_configured_fields(): + policy = _build_mp_policy('bf16', parse_mp_overrides({'reduce_dtype': 'fp32'})) + assert policy.param_dtype is torch.bfloat16 + assert policy.reduce_dtype is torch.float32 + assert policy.output_dtype is torch.bfloat16 + assert policy.cast_forward_inputs + + +def test_native_policy_override_on_no_base_keeps_unset_fields_none(): + policy = _build_mp_policy('no', parse_mp_overrides({'reduce_dtype': 'fp32'})) + assert policy.param_dtype is None + assert policy.reduce_dtype is torch.float32 + assert policy.output_dtype is None + # cast_forward_inputs is inert while param_dtype is None (FSDP2 skips the cast). + assert policy.cast_forward_inputs + + +def test_native_policy_param_override_applies_on_no_base(): + policy = _build_mp_policy('no', parse_mp_overrides({'param_dtype': 'fp16'})) + assert policy.param_dtype is torch.float16 + assert policy.reduce_dtype is None + assert policy.cast_forward_inputs + + +def test_ep_policy_still_forces_fp32_reduce_over_explicit_bf16(): + # NCCL PreMulSum constraint: the expert group must reduce in fp32 even when the + # user explicitly asked for bf16 reduction on the dense side. + base = _build_mp_policy('bf16', parse_mp_overrides({'reduce_dtype': 'bf16'})) + ep_policy = _build_ep_mp_policy(base) + assert ep_policy.reduce_dtype is torch.float32 + assert ep_policy.param_dtype is torch.bfloat16 + assert ep_policy.output_dtype is torch.bfloat16 + + # A user-configured fp32 reduce passes through untouched. + base = _build_mp_policy('bf16', parse_mp_overrides({'reduce_dtype': 'fp32'})) + assert _build_ep_mp_policy(base).reduce_dtype is torch.float32 + + +def test_strategy_parses_overrides_at_construction(): + strategy = NativeFSDPStrategy( + mixed_precision='bf16', + fsdp_config={'mixed_precision': {'reduce_dtype': 'fp32'}}, + enable_ep=False, + ) + assert strategy._mp_overrides == {'reduce_dtype': torch.float32} + + +def test_strategy_rejects_invalid_overrides_at_construction(): + with pytest.raises(ValueError, match='Unknown key'): + NativeFSDPStrategy( + mixed_precision='bf16', + fsdp_config={'mixed_precision': {'buffer_dtype': 'fp32'}}, + enable_ep=False, + ) + + +def test_wrap_model_passes_overridden_policy_to_fully_shard(): + strategy = NativeFSDPStrategy( + device_mesh=MagicMock(), + mixed_precision='bf16', + fsdp_config={'mixed_precision': {'reduce_dtype': 'fp32'}}, + enable_ep=False, + ) + model = torch.nn.Linear(4, 4) + with patch('twinkle.model.transformers.strategy.native_fsdp._build_fsdp_mesh', + return_value=MagicMock()), \ + patch('twinkle.model.transformers.strategy.native_fsdp.fully_shard') as fully_shard: + strategy.wrap_model(model) + + fully_shard.assert_called_once() + policy = fully_shard.call_args.kwargs['mp_policy'] + assert policy.param_dtype is torch.bfloat16 + assert policy.reduce_dtype is torch.float32 + assert policy.output_dtype is torch.bfloat16 + assert policy.cast_forward_inputs + + +# ── param/output dtype boundary guard (validated on 4x NPU: Qwen3.5 SDPA crash) ──── + + +@pytest.mark.parametrize('mixed_precision,overrides', [ + ('bf16', {'param_dtype': 'fp32'}), # output falls back to bf16: the 4x NPU crash config + ('bf16', {'param_dtype': 'fp32', 'output_dtype': 'bf16'}), + ('bf16', {'output_dtype': 'fp32'}), # param stays bf16 + ('fp16', {'param_dtype': 'bf16'}), # output falls back to fp16 + ('no', {'output_dtype': 'bf16'}), # output cast active while param_dtype stays None +]) +def test_strategy_rejects_param_output_dtype_boundary(mixed_precision, overrides): + with pytest.raises(ValueError, match='output_dtype'): + NativeFSDPStrategy( + mixed_precision=mixed_precision, + fsdp_config={'mixed_precision': overrides}, + enable_ep=False, + ) + + +@pytest.mark.parametrize('mixed_precision,overrides', [ + ('bf16', {}), # legacy default + ('bf16', {'reduce_dtype': 'fp32'}), # the production case: BF16 compute, FP32 reduce + ('bf16', {'param_dtype': 'fp32', 'reduce_dtype': 'fp32', 'output_dtype': 'fp32'}), # uniform FP32 + ('no', {'reduce_dtype': 'fp32'}), # effective param/output both None + ('no', {'param_dtype': 'fp16'}), # effective output None performs no cast: no boundary +]) +def test_strategy_accepts_boundary_free_combinations(mixed_precision, overrides): + strategy = NativeFSDPStrategy( + mixed_precision=mixed_precision, + fsdp_config={'mixed_precision': overrides}, + enable_ep=False, + ) + policy = _build_mp_policy(mixed_precision, strategy._mp_overrides) + assert policy.output_dtype is None or policy.output_dtype == policy.param_dtype