Model-Optimizer/examples/speculative_decoding/fsdp2_buffer_patch.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Monkey-patch for accelerate's fsdp2_load_full_state_dict buffer handling.

Applicability (scope of this patch)
-----------------------------------
This is **not** needed for FSDP2 in general. It is required only for the narrow
combination of: **FSDP2 configured via an accelerate YAML config** (not torch-native
``ParallelismConfig``) **with** ``cpu_ram_efficient_loading=True``. Today that path is
forced by **models that require transformers 4.57.x** (their ``trust_remote_code`` code
predates transformers 5.x ``ParallelismConfig`` support) **and** are too large to load on
every rank — currently only **MiniMax-M2.7** (229B MoE). Models that run on transformers
5.x (Qwen, Llama, Nemotron, ...) use native ``ParallelismConfig`` (``dp_shard_size > 1``),
which handles buffers/dtypes correctly and never enters ``fsdp2_load_full_state_dict`` —
they need none of this. Gated off by default; activate with ``PATCH_FSDP2_BUFFERS_TF457=1``.

Problem
-------
accelerate's ``fsdp2_load_full_state_dict`` (called during model preparation
when ``cpu_ram_efficient_loading=True``) iterates ``model.state_dict()`` and
unconditionally accesses ``.device_mesh`` on every entry, assuming they are all
DTensors.  After ``fully_shard()``, **parameters** become DTensors but
**persistent buffers** (e.g., rotary-embedding ``inv_freq``) remain plain
``torch.Tensor``.  This crashes with::

    AttributeError: 'Tensor' object has no attribute 'device_mesh'

Additionally, ``cpu_ram_efficient_loading`` causes a dtype divergence: rank 0
loads the model on CPU (inheriting the checkpoint's dtype, e.g. bfloat16) while
other ranks use ``meta`` device (defaulting to float32 for newly-added modules
like the DFlash head).  After ``fully_shard()``, the DTensor dtypes differ
across ranks for these modules.  Since ``dist.broadcast()`` requires matching
dtypes and element sizes on all ranks, broadcasting a bfloat16 tensor (2
bytes/elem) to a float32 receive buffer (4 bytes/elem) causes an NCCL deadlock.

Why we need FSDP2 via accelerate config (not ParallelismConfig)
---------------------------------------------------------------
MiniMax-M2.7's ``trust_remote_code`` model code requires **transformers 4.57.x**.
Transformers' native FSDP2 support via ``ParallelismConfig`` requires
**transformers 5.x**.  So we fall back to configuring FSDP2 through an
``accelerate`` YAML config file (``accelerate_fsdp2.yaml``), which works with
transformers 4.57.x.  We set ``dp_shard_size=1`` to prevent ``main.py`` from
creating a ``ParallelismConfig``, letting the accelerate config handle sharding.

Why we need cpu_ram_efficient_loading
-------------------------------------
MiniMax-M2.7 is a 229B MoE model (~230 GB in FP8).  Each GB200 node has 4 GPUs
and ~800 GB system RAM.  Without ``cpu_ram_efficient_loading``, all 4 ranks per
node load the model to CPU simultaneously (4 × 230 GB ≈ 920 GB), exceeding
system RAM and triggering OOM kills.  With ``cpu_ram_efficient_loading``, only
rank 0 loads the model; other ranks initialize on ``meta`` device.  The weights
are then broadcast via ``fsdp2_load_full_state_dict`` — which is where the bug
hits.

What this patch does
--------------------
1. Before accessing ``.device_mesh``, checks ``isinstance(entry, DTensor)``.
   For non-DTensor entries (persistent buffers), broadcasts the raw tensor from
   rank 0 without calling ``distribute_tensor()``.

2. All ranks iterate ``model.state_dict()`` (post-shard) in the same order so
   broadcast calls match 1-to-1.  Rank 0 looks up the full parameter **by key
   name** from the pre-shard state dict — never by positional ``zip``, because
   ``model.to("meta")`` + ``fully_shard()`` can reorder keys.

3. **Dtype synchronization**: rank 0 broadcasts a dtype code for each entry
   before the main loop.  All ranks then use the same dtype for their broadcast
   tensors.  This fixes the dtype divergence caused by rank 0 loading in
   bfloat16 while other ranks default to float32 for newly-added modules.

Accelerate config constraints (for reference)
----------------------------------------------
``accelerate_fsdp2.yaml`` also requires:

- ``fsdp_use_orig_params: true`` — without this, FSDP flattens all params into
  FlatParameter, losing per-parameter ``requires_grad`` flags.  The DFlash head
  can't train because its gradients mix with frozen base model zeros.
- ``fsdp_transformer_layer_cls_to_wrap: MiniMaxM2DecoderLayer,DFlashModule`` —
  DFlash head params at the model root must be in the wrap policy so they become
  DTensors.  Without this, ``fsdp2_load_full_state_dict`` also crashes.
- ``fsdp_sync_module_states: true`` — accelerate's launch validator requires it
  when ``cpu_ram_efficient_loading`` is enabled, even though FSDP2 ignores it at
  runtime (sets it to None with a warning).

Does NOT affect models on transformers 5.x
-------------------------------------------
This entire workaround exists ONLY because MiniMax-M2.7 requires
transformers 4.57.x.  Models that support transformers 5.x (Qwen, Llama,
Nemotron, etc.) use ``ParallelismConfig`` natively by setting
``dp_shard_size > 1`` in the training args.  That code path handles buffers
correctly and does not go through ``fsdp2_load_full_state_dict`` at all.
No accelerate config file, no ``PATCH_FSDP2_BUFFERS_TF457``, no
``OVERRIDE_TRANSFORMERS`` needed.

When to remove
--------------
This patch can be removed when EITHER of these happens:

1. MiniMax updates ``trust_remote_code`` for transformers 5.x, allowing native
   ``ParallelismConfig`` (which handles this correctly).
2. Upstream accelerate fixes ``fsdp2_load_full_state_dict`` to skip non-DTensor
   entries.  Track: https://github.com/huggingface/accelerate

Activation
----------
Set ``PATCH_FSDP2_BUFFERS_TF457=1`` in the environment to activate.  Off by default.
Only needed in MiniMax-M2.7 pipeline YAMLs.
"""

import torch

# Dtype encoding for the broadcast dtype-sync step.
_DTYPE_TO_CODE = {
    torch.float32: 0,
    torch.bfloat16: 1,
    torch.float16: 2,
    torch.float8_e4m3fn: 3 if hasattr(torch, "float8_e4m3fn") else -1,
}
_CODE_TO_DTYPE = {v: k for k, v in _DTYPE_TO_CODE.items() if v >= 0}


def apply():
    """Patch fsdp2_load_full_state_dict if the buffer bug is present."""
    try:
        import accelerate.utils.fsdp_utils as fsdp_utils
        from torch.distributed.tensor import DTensor

        def _dtype_code(dt):
            """Map a dtype to its broadcast sync code, raising on anything unmapped.

            Silently coercing an unknown dtype to fp32 would cast data on the wire (or
            make NCCL refuse on an element-size mismatch), so fail loudly instead.
            """
            code = _DTYPE_TO_CODE.get(dt)
            if code is None or code < 0:
                raise ValueError(
                    f"fsdp2_buffer_patch: unsupported dtype {dt} in the broadcast "
                    f"dtype-sync; add it to _DTYPE_TO_CODE."
                )
            return code

        def _patched(accelerator, model, full_sd, cpu_offload=False):
            import time

            import torch.distributed as dist
            from torch.distributed.tensor import distribute_tensor

            meta_sharded_sd = model.state_dict()
            sharded_sd = {}
            n_total = len(meta_sharded_sd)
            n_dtensor = sum(1 for v in meta_sharded_sd.values() if isinstance(v, DTensor))
            n_buffer = n_total - n_dtensor

            if accelerator.is_main_process:
                print(
                    f"[fsdp2_buffer_patch] State dict: {n_total} entries "
                    f"({n_dtensor} DTensor, {n_buffer} buffer), full_sd: {len(full_sd)}"
                )
            else:
                print(
                    f"[fsdp2_buffer_patch] State dict: {n_total} entries "
                    f"({n_dtensor} DTensor, {n_buffer} buffer)"
                )
            t0 = time.time()

            # --- Step 0: broadcast dtype codes from rank 0 ---
            # cpu_ram_efficient_loading causes rank 0 to load in bfloat16 while
            # other ranks default to float32 for newly-added modules (DFlash).
            # After fully_shard(), DTensor dtypes diverge across ranks.
            # Broadcast rank 0's dtypes so all ranks use the same dtype for
            # each broadcast tensor.
            if accelerator.is_main_process:
                dtype_codes = torch.tensor(
                    [_dtype_code(full_sd[name].dtype) for name in meta_sharded_sd],
                    dtype=torch.int32,
                    device=accelerator.device,
                )
            else:
                dtype_codes = torch.empty(
                    n_total,
                    dtype=torch.int32,
                    device=accelerator.device,
                )
            dist.broadcast(dtype_codes, src=0, group=dist.group.WORLD)
            broadcast_dtypes = [_CODE_TO_DTYPE[c.item()] for c in dtype_codes]
            del dtype_codes

            # Infer dtype/contiguity for cast — copied from upstream
            def _infer_parameter_dtype(mdl, param_name, empty_param):
                try:
                    old_param = mdl.get_parameter_or_buffer(param_name)
                except AttributeError:
                    base, local = param_name.rsplit(".", 1)
                    old_param = getattr(mdl.get_submodule(base), local)
                is_f8 = hasattr(torch, "float8_e4m3fn") and empty_param.dtype == torch.float8_e4m3fn
                casting_dtype = (
                    old_param.dtype if (empty_param.dtype.is_floating_point and not is_f8) else None
                )
                return old_param is not None and old_param.is_contiguous(), casting_dtype

            def _finish(st, contig, dtype, offload):
                if dtype is not None:
                    st = st.to(dtype=dtype)
                if contig:
                    st = st.contiguous()
                if offload:
                    st = st.to("cpu")
                return st

            # --- Step 1: broadcast all entries ---
            # All ranks iterate meta_sharded_sd in the same order to ensure
            # identical broadcast sequences.  Rank 0 looks up the full parameter
            # by name — never positional zip (model.to("meta") + fully_shard()
            # can reorder keys).  broadcast_dtypes[idx] is used for the tensor
            # dtype on ALL ranks to prevent NCCL deadlocks from dtype divergence.
            for idx, (param_name, sharded_param) in enumerate(meta_sharded_sd.items()):
                is_dtensor = isinstance(sharded_param, DTensor)
                bcast_dtype = broadcast_dtypes[idx]

                if not is_dtensor:
                    # Persistent buffer — broadcast raw, no distribute_tensor
                    if accelerator.is_main_process:
                        t = (
                            full_sd[param_name]
                            .detach()
                            .to(device=accelerator.device, dtype=bcast_dtype)
                        )
                    else:
                        t = torch.empty(
                            sharded_param.size(),
                            device=accelerator.device,
                            dtype=bcast_dtype,
                        )
                    dist.broadcast(t, src=0, group=dist.group.WORLD)
                    sharded_sd[param_name] = t
                    continue

                device_mesh = sharded_param.device_mesh
                if accelerator.is_main_process:
                    ft = (
                        full_sd[param_name]
                        .detach()
                        .to(device=device_mesh.device_type, dtype=bcast_dtype)
                    )
                    if isinstance(ft, DTensor):
                        ft = ft.to_local()
                else:
                    ft = torch.empty(
                        sharded_param.size(),
                        device=device_mesh.device_type,
                        dtype=bcast_dtype,
                    )
                dist.broadcast(ft, src=0, group=dist.group.WORLD)
                st = distribute_tensor(ft, device_mesh, sharded_param.placements)
                contig, _ = _infer_parameter_dtype(model, param_name, ft)
                # Use bcast_dtype (from rank 0) instead of the model's local
                # param dtype — with cpu_ram_efficient_loading, non-rank-0
                # processes have fp32 meta-device params for DFlash, and
                # _infer_parameter_dtype would incorrectly cast bf16 back to fp32.
                is_f8 = hasattr(torch, "float8_e4m3fn") and bcast_dtype == torch.float8_e4m3fn
                final_dtype = None if is_f8 else bcast_dtype
                sharded_sd[param_name] = _finish(st, contig, final_dtype, cpu_offload)

            elapsed = time.time() - t0
            print(
                f"[fsdp2_buffer_patch] Broadcast done in {elapsed:.1f}s, "
                f"loading {len(sharded_sd)} entries into model..."
            )
            model.load_state_dict(sharded_sd, assign=True)
            print(
                f"[fsdp2_buffer_patch] State dict loaded successfully "
                f"({time.time() - t0:.1f}s total)"
            )
            return model

        fsdp_utils.fsdp2_load_full_state_dict = _patched
        print("[fsdp2_buffer_patch] Patched fsdp2_load_full_state_dict for buffer compatibility")
    except Exception as e:
        print(f"[fsdp2_buffer_patch] Patch skipped: {e}")


_clip_grad_norm_call_count = 0


def _clip_grad_norm(parameters, max_norm, norm_type=2):
    """Clip gradient norms for FSDP2 DTensor parameters.

    Bypasses DTensor dispatch (which deadlocks with partially-frozen models
    on the accelerate FSDP2 path) by extracting local tensor shards and
    doing an explicit all_reduce for the global norm.

    Handles Shard (need all_reduce) and Replicate/regular (already global)
    placements.  Safe for DFlash-only and LoRA co-training.
    """
    global _clip_grad_norm_call_count
    import torch.distributed as dist
    from torch.distributed.tensor import DTensor
    from torch.distributed.tensor.placement_types import Shard

    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    parameters = list(parameters)  # materialize generator
    max_norm = float(max_norm)
    norm_type = float(norm_type)

    grads = [p.grad for p in parameters if p.grad is not None]
    # Every rank MUST reach the all_reduce below: under FSDP2 sharding (and especially
    # MoE + LoRA co-training) a rank can legitimately have no grads on a step — e.g. an
    # expert that received no tokens, so the shard it owns gets no gradient. Early-
    # returning here while other ranks call all_reduce would deadlock the job. So we
    # never short-circuit before the collective; an empty-grad rank simply contributes a
    # zero local norm and clips nothing.
    if grads:
        device = grads[0]._local_tensor.device if isinstance(grads[0], DTensor) else grads[0].device
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Shard DTensors hold partial data — need all_reduce for global norm.
    # Replicate DTensors and regular tensors already hold full data.
    sharded_norm_p = torch.tensor(0.0, device=device)
    local_norm_p = torch.tensor(0.0, device=device)

    n_sharded = 0
    n_replicate = 0
    n_regular = 0

    for g in grads:
        if isinstance(g, DTensor):
            is_sharded = any(isinstance(p, Shard) for p in g.placements)
            t = g._local_tensor.detach().to(torch.float32)
            n = torch.linalg.vector_norm(t, norm_type)
            if is_sharded:
                sharded_norm_p += n.pow(norm_type)
                n_sharded += 1
            else:
                local_norm_p += n.pow(norm_type)
                n_replicate += 1
        else:
            n = torch.linalg.vector_norm(g.detach().to(torch.float32), norm_type)
            local_norm_p += n.pow(norm_type)
            n_regular += 1

    # Symmetric across ranks: reached on every rank regardless of whether this rank had
    # grads (see note above). Guard for the non-distributed case where local == global.
    if dist.is_available() and dist.is_initialized():
        dist.all_reduce(sharded_norm_p, op=dist.ReduceOp.SUM)
    total_norm = (sharded_norm_p + local_norm_p).pow(1.0 / norm_type)

    clip_coef = torch.clamp(max_norm / (total_norm + 1e-6), max=1.0)

    # Debug: log computation breakdown on first 5 calls (no collectives — safe).
    _clip_grad_norm_call_count += 1
    _rank0 = not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0
    if _clip_grad_norm_call_count <= 5 and _rank0:
        print(
            f"[clip_grad_norm debug] call={_clip_grad_norm_call_count} "
            f"total_norm={total_norm.item():.6f} "
            f"sharded_norm_p={sharded_norm_p.item():.6f} local_norm_p={local_norm_p.item():.6f} "
            f"grads={len(grads)} (sharded={n_sharded} replicate={n_replicate} regular={n_regular}) "
            f"max_norm={max_norm} clip_coef={clip_coef.item():.6f}"
        )
    for g in grads:
        if isinstance(g, DTensor):
            g._local_tensor.mul_(clip_coef)
        else:
            g.mul_(clip_coef)

    return total_norm


def patch_accelerator(accelerator):
    """Replace accelerator's clip_grad_norm_ with FSDP2-safe version."""
    accelerator.clip_grad_norm_ = _clip_grad_norm
    print(
        "[fsdp2_buffer_patch] Patched accelerator.clip_grad_norm_ for FSDP2 DTensor compatibility"
    )


def log_param_dtypes(model):
    """Debug aid: log per-rank parameter dtype counts (gated by DFLASH_LOG_PARAM_DTYPES=1).

    Used to verify the FSDP2 dtype synchronization above — after ``fully_shard()`` params
    are DTensors whose dtype lives on ``_local_tensor``. Off by default; this is purely
    diagnostic and has no effect on training.
    """
    import os

    if os.environ.get("DFLASH_LOG_PARAM_DTYPES") != "1":
        return
    rank = int(os.environ.get("RANK", 0))
    dtypes = {}
    for name, p in model.named_parameters():
        dt_key = str(p.dtype) if not hasattr(p, "_local_tensor") else str(p._local_tensor.dtype)
        dtypes.setdefault(dt_key, []).append(name)
    for dt, names in dtypes.items():
        print(f"[dtype_check rank={rank}] {dt}: {len(names)} params (e.g. {names[0]})")