diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 239933498..d5eca8d6d 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 75eba2d9f..562afdd64 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 7000c1587..d6d6ce877 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 c72c1af4b..39004b4d4 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