From a9c8bcf0e1b517071060b456e2cb5d472cf96fd0 Mon Sep 17 00:00:00 2001 From: Abdul Moiz Date: Fri, 26 Jun 2026 13:44:41 +0500 Subject: [PATCH] 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 --- unsloth/models/llama.py | 4 +++- unsloth/models/loader.py | 3 +++ unsloth/models/loader_utils.py | 34 ++++++++++++++++++++++++++++++++++ unsloth/models/vision.py | 3 ++- 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 2399334983..d5eca8d6df 100644 --- a/unsloth/models/llama.py +++ b/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 diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 75eba2d9f8..562afdd645 100644 --- a/unsloth/models/loader.py +++ b/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 diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 7000c1587d..d6d6ce877e 100644 --- a/unsloth/models/loader_utils.py +++ b/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 diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index c72c1af4b1..39004b4d45 100644 --- a/unsloth/models/vision.py +++ b/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