Fix DDP crash from CPU-resident rotary inv_freq buffer (#6662)

* Fix DDP crash from CPU-resident rotary inv_freq buffer

DistributedDataParallel broadcasts all named buffers regardless of
persistence or device, but Unsloth's RoPE inv_freq buffer is kept on
CPU on purpose (per-GPU cos/sin caches are precomputed instead). That
mismatch crashed multi-GPU DDP training with "No backend type
associated with device type cpu" during _sync_module_states.

Mark inv_freq/short_inv_freq/long_inv_freq buffers as DDP-ignored
instead of moving them to GPU, so they're skipped during the buffer
broadcast without disabling broadcast_buffers for the rest of the
model.

Fixes #6656

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: harden DDP-ignore against private API drift, re-apply after PEFT wrap

- Wrap the private DistributedDataParallel._set_params_and_buffers_to_ignore_for_model
  call in try/except, falling back to setting _ddp_params_and_buffers_to_ignore
  directly so a future PyTorch API change can't block model loading.
- Move _exclude_rope_inv_freq_from_ddp to loader_utils.py (shared by loader.py,
  llama.py, vision.py without circular imports) and call it again after
  get_peft_model wraps the model in a PeftModel, since the rotary buffers'
  fully qualified names change once nested under "base_model.model...".

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>

unsloth/models/llama.py

@@ -28,7 +28,7 @@ from ._utils import (
     is_bfloat16_supported,
     get_quant_type,
 )
-from .loader_utils import _get_fp8_mode_and_check_settings
+from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
 from ..utils.packing import (
     get_packed_info_from_kwargs,
     mask_packed_sequence_boundaries,
@@ -3049,6 +3049,7 @@ class FastLlamaModel:
                 # Pre-wrapped PEFT model passes through here; still arm the detector so an RL
                 # trainer can reset a compile cache poisoned by a pre-train forward.
                 _unsloth_install_pretrain_detector(model)
+                model = _exclude_rope_inv_freq_from_ddp(model)
                 return model
             else:
                 raise TypeError(
@@ -3404,6 +3405,7 @@ class FastLlamaModel:
         # Detect a stray pre-train forward so train() can drop the torch.compile
         # graph cache it would otherwise poison (see prepare_for_training_mode).
         _unsloth_install_pretrain_detector(model)
+        model = _exclude_rope_inv_freq_from_ddp(model)
         return model
 
     @staticmethod

unsloth/models/loader.py

@@ -33,6 +33,7 @@ from transformers import AutoConfig
 from transformers import __version__ as transformers_version
 from peft import PeftConfig, PeftModel
 from .loader_utils import (
+    _exclude_rope_inv_freq_from_ddp,
     _get_fp8_mode_and_check_settings,
     _offline_quantize_to_fp8,
     _tag_model_with_fp8_torchao_config,
@@ -885,6 +886,7 @@ class FastLanguageModel(FastLlamaModel):
             patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice)
 
         model = _fix_rope_inv_freq(model)
+        model = _exclude_rope_inv_freq_from_ddp(model)
         return model, tokenizer
 
 
@@ -1822,6 +1824,7 @@ class FastModel(FastBaseModel):
             patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice)
 
         model = _fix_rope_inv_freq(model)
+        model = _exclude_rope_inv_freq_from_ddp(model)
         return model, tokenizer
 
 

unsloth/models/loader_utils.py

@@ -499,6 +499,40 @@ def _get_fp8_mode_and_check_settings(
     return fp8_mode
 
 
+# Rotary inv_freq buffers are deliberately kept on CPU - Unsloth pre-builds a
+# cos/sin cache per GPU instead (see LlamaRotaryEmbedding.multi_gpu_cos_cached)
+# so the GPU-resident lookup never needs to move the tiny inv_freq tensor itself.
+# torch.nn.parallel.DistributedDataParallel ignores device entirely when it
+# broadcasts buffers across ranks, so a CPU buffer crashes NCCL's
+# _broadcast_coalesced with "No backend type associated with device type cpu".
+# Telling DDP to skip these specific buffers avoids that crash without moving
+# inv_freq to GPU (which would break the per-GPU cache design) and without
+# disabling buffer broadcast for every other module (the user's workaround).
+# Re-run this after wrapping with PEFT too - the buffers' fully qualified
+# names change once they sit under a PeftModel (eg "base_model.model...").
+# https://github.com/unslothai/unsloth/issues/6656
+_ROTARY_INV_FREQ_BUFFER_NAMES = ("inv_freq", "short_inv_freq", "long_inv_freq")
+
+
+def _exclude_rope_inv_freq_from_ddp(model):
+    ignored = list(getattr(model, "_ddp_params_and_buffers_to_ignore", None) or [])
+    for module_name, module in model.named_modules():
+        for buffer_name, _ in module.named_buffers(recurse = False):
+            if buffer_name in _ROTARY_INV_FREQ_BUFFER_NAMES:
+                fqn = f"{module_name}.{buffer_name}" if module_name else buffer_name
+                if fqn not in ignored:
+                    ignored.append(fqn)
+    if ignored:
+        try:
+            from torch.nn.parallel import DistributedDataParallel
+            DistributedDataParallel._set_params_and_buffers_to_ignore_for_model(model, ignored)
+        except Exception:
+            # Private PyTorch API - fall back to setting the attribute DDP reads
+            # directly if it ever moves or changes signature.
+            model._ddp_params_and_buffers_to_ignore = ignored
+    return model
+
+
 # =============================================================================
 # Offline loading - single source of truth (shared by vision.py, loader.py and
 # the Studio exporter). Decide offline ONCE at the load boundary and force it

unsloth/models/vision.py

@@ -40,7 +40,7 @@ from ._utils import (
     set_task_config_attr,
 )
 from ._utils import *
-from .loader_utils import _get_fp8_mode_and_check_settings
+from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
 from ..save import patch_saving_functions
 from ..models.loader_utils import is_distributed
 from unsloth_zoo.gradient_checkpointing import (
@@ -1741,6 +1741,7 @@ class FastBaseModel:
         # Detect a stray pre-train forward so train() can drop the torch.compile
         # graph cache it would otherwise poison (see prepare_for_training_mode).
         _unsloth_install_pretrain_detector(model)
+        model = _exclude_rope_inv_freq_from_ddp(model)
         return model
 
     @staticmethod