# Copyright 2025 the LlamaFactory team. # # 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. import copy import gc import os import torch import torch.distributed.checkpoint as dcp import torch.nn as nn from peft.tuners.lora import LoraLayer from torch.distributed.checkpoint.state_dict import ( StateDictOptions, get_model_state_dict, get_optimizer_state_dict, set_model_state_dict, set_optimizer_state_dict, ) from torch.distributed.fsdp import ( CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard, ) from ....accelerator.helper import get_current_accelerator from ....accelerator.interface import DistributedInterface from ....utils.logging import get_logger from ....utils.types import HFModel, Processor logger = get_logger(__name__) def _fallback_dot_natural_key(name: str): parts = [] for part in name.split("."): if part.isdigit(): parts.append((0, int(part))) else: parts.append((1, part)) return parts def _get_checkpoint_sort_key(): try: from transformers.core_model_loading import dot_natural_key return dot_natural_key except ImportError: return _fallback_dot_natural_key def _make_safetensor_loader(checkpoint_file: str, tensor_key: str): # Delay tensor materialization until converter.convert() to reduce peak CPU memory. # This works because HF WeightConverter accepts callables and materializes them later. def _load_tensor(): from safetensors import safe_open with safe_open(checkpoint_file, framework="pt", device="cpu") as f: return f.get_tensor(tensor_key) return _load_tensor def get_transformer_layer_cls(model: HFModel) -> type[nn.Module] | None: no_split_modules = getattr(model, "_no_split_modules", None) if no_split_modules: if isinstance(no_split_modules, (list, tuple)): for name, module in model.named_modules(): for cls_name in no_split_modules: if module.__class__.__name__ == cls_name: return module.__class__ if hasattr(model, "model") and hasattr(model.model, "layers"): return type(model.model.layers[0]) if hasattr(model, "layers"): return type(model.layers[0]) return None def save_model(model: HFModel, output_dir: str, processor: Processor) -> None: if DistributedInterface().get_rank() == 0: logger.info("Gathering state dict for saving...") options = StateDictOptions(full_state_dict=True, cpu_offload=True) state_dict = get_model_state_dict(model, options=options) if DistributedInterface().get_rank() == 0: model_to_save = model.module if hasattr(model, "module") else model model_to_save.save_pretrained(output_dir, state_dict=state_dict, max_shard_size="4GB") processor.save_pretrained(output_dir, max_shard_size="4GB") logger.info(f"Model saved to {output_dir}") def save_checkpoint(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None: save_ckpt_as_hf = kwargs.get("save_ckpt_as_hf", False) processor = kwargs.get("processor", None) # Always save DCP format for resume capability options = StateDictOptions(full_state_dict=False, cpu_offload=True) model_state = get_model_state_dict(model, options=options) dcp.save(state_dict=model_state, checkpoint_id=os.path.join(ckpt_dir, "model")) optim_state = get_optimizer_state_dict(model, optimizer, options=options) dcp.save(state_dict=optim_state, checkpoint_id=os.path.join(ckpt_dir, "optimizer")) # Additionally save HF format if requested if save_ckpt_as_hf: if DistributedInterface().get_rank() == 0: logger.info("Gathering state dict for saving additional HF format checkpoint...") hf_options = StateDictOptions(full_state_dict=True, cpu_offload=True) hf_state_dict = get_model_state_dict(model, options=hf_options) if DistributedInterface().get_rank() == 0: model_to_save = model.module if hasattr(model, "module") else model hf_dir = os.path.join(ckpt_dir, "hf_model") model_to_save.save_pretrained(hf_dir, state_dict=hf_state_dict, max_shard_size="4GB") if processor is not None: processor.save_pretrained(hf_dir, max_shard_size="4GB") logger.info(f"Additional HF format checkpoint saved to {hf_dir}") def load_checkpoint(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None: options = StateDictOptions(full_state_dict=False, cpu_offload=True) ckpt_model_dir = os.path.join(ckpt_dir, "model") model_state = get_model_state_dict(model, options=options) dcp.load(state_dict=model_state, checkpoint_id=ckpt_model_dir) set_model_state_dict(model, model_state, options=options) ckpt_optim_dir = os.path.join(ckpt_dir, "optimizer") optim_state = get_optimizer_state_dict(model, optimizer, options=options) dcp.load(state_dict=optim_state, checkpoint_id=ckpt_optim_dir) set_optimizer_state_dict(model, optimizer, optim_state, options=options) class FSDP2Engine: def __init__(self, dist_config: dict, bf16: bool = False): self.dist_interface = DistributedInterface() self.rank = self.dist_interface.get_rank() self.local_rank = self.dist_interface.get_local_rank() self.world_size = self.dist_interface.get_world_size() self.mixed_precision = "bf16" if bf16 else "fp32" self.reshard_after_forward = dist_config.get("reshard_after_forward", True) self.offload_params = dist_config.get("offload_params", False) self.pin_memory = dist_config.get("pin_memory", True) self.dcp_path = dist_config.get("dcp_path", None) self.device_mesh = self.dist_interface.model_device_mesh if self.device_mesh is None: logger.warning( "Device Mesh not found in DistributedInterface. FSDP2 might fail if not running in distributed mode." ) if self.device_mesh is not None: self.fsdp_mesh = self.device_mesh logger.info(f"Using Device Mesh: {self.fsdp_mesh}") else: self.fsdp_mesh = None def get_mp_policy(self) -> MixedPrecisionPolicy: if self.mixed_precision == "bf16": param_dtype = torch.bfloat16 reduce_dtype = torch.float32 elif self.mixed_precision == "fp32": param_dtype = torch.float32 reduce_dtype = torch.float32 return MixedPrecisionPolicy( param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True, ) def is_lora_module_wrap(self, model) -> bool: return any(isinstance(module, LoraLayer) for module in model.modules()) def prepare_model(self, model: HFModel) -> HFModel: if self.fsdp_mesh is None: logger.warning("No FSDP Mesh available, skipping FSDP wrapping.") return model mp_policy = self.get_mp_policy() layer_cls = get_transformer_layer_cls(model) if layer_cls is None: logger.warning( "Could not identify Transformer Layer class, applying FSDP to the whole model structure only." ) transformer_layer_cls_to_wrap = set() else: logger.info(f"Applying per-layer FSDP to {layer_cls.__name__}") transformer_layer_cls_to_wrap = {layer_cls} if self.is_lora_module_wrap(model): lora_modules = [] for module in model.modules(): if len(list(module.children())) != 0: continue if any(param.requires_grad for param in module.parameters(recurse=False)): lora_modules.append(module) for module in lora_modules: fully_shard( module, mesh=self.fsdp_mesh, reshard_after_forward=self.reshard_after_forward, mp_policy=mp_policy, offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None, ) logger.info("Applying FSDP wrap for LoRA layer separately.") for name, module in model.named_modules(): should_wrap = False if type(module) in transformer_layer_cls_to_wrap: should_wrap = True elif isinstance(module, nn.Embedding): if not getattr(model.config, "tie_word_embeddings", True): should_wrap = True if should_wrap: fully_shard( module, mesh=self.fsdp_mesh, reshard_after_forward=self.reshard_after_forward, mp_policy=mp_policy, offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None, ) # BaseTrainer is the single source of truth for gradient checkpointing. # FSDP2 only applies the input-grad compatibility hook when checkpointing is already enabled. if getattr(model, "is_gradient_checkpointing", False): if self.rank == 0: logger.info("Gradient checkpointing is enabled. Applying FSDP2 input grad preparation.") if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() else: def make_inputs_require_grad(module, input, output): output.requires_grad_(True) model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) fully_shard( model, mesh=self.fsdp_mesh, reshard_after_forward=self.reshard_after_forward, mp_policy=mp_policy, offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None, ) return model @torch.no_grad() def materialize_and_load(self, model: HFModel, hf_model_path: str, dcp_path: str = None): if self.rank == 0: logger.info("Materializing sharded model params...") device = get_current_accelerator() model.to_empty(device=device) if dcp_path and os.path.exists(dcp_path): if self.rank == 0: logger.info(f"DCP path found at {dcp_path}. Using efficient Sharded Loading (DCP Load).") self._load_from_dcp(model, dcp_path) else: if self.rank == 0: if dcp_path: logger.warning(f"DCP path {dcp_path} not found.") logger.info("Using HF Meta Loading (Chunk Load).") self._load_weights_from_hf_checkpoint(model, hf_model_path) return model def _save_non_persistent_buffers(self, model: HFModel) -> dict: """Save non-persistent buffers, such as inv_freq.""" saved = {} for mod_name, module in model.named_modules(): for buf_name in module._non_persistent_buffers_set: fqn = f"{mod_name}.{buf_name}" if mod_name else buf_name buf = getattr(module, buf_name, None) if buf is not None: saved[fqn] = copy.deepcopy(buf) if self.rank == 0 and saved: logger.info(f"Saved {len(saved)} non-persistent buffers") return saved def _restore_non_persistent_buffers(self, model: HFModel, saved_buffers: dict): """Register saved non-persistent buffers to model.""" if not saved_buffers: return device = get_current_accelerator() for fqn, buf in saved_buffers.items(): buf = buf.to(device) if "." in fqn: parent_fqn, buf_name = fqn.rsplit(".", 1) parent_module = model.get_submodule(parent_fqn) else: buf_name = fqn parent_module = model parent_module.register_buffer(buf_name, buf, persistent=False) if self.rank == 0: logger.info(f"Restored {len(saved_buffers)} non-persistent buffers") def shard_model(self, model: HFModel) -> HFModel: init_mode = getattr(model, "_init_mode", "init_on_default") if init_mode == "init_on_rank0": if getattr(model.config, "tie_word_embeddings", False): model.tie_weights() if self.rank == 0: logger.info("init_on_rank0 detected: sharding then scattering Rank 0 CPU weights.") full_sd = {k: v.clone() for k, v in model.state_dict().items()} else: full_sd = {} model = self.prepare_model(model) device = get_current_accelerator() model.to_empty(device=device) # Scatter params from Rank 0 into all DTensor shards # Broadcast the full state dict from the global rank-0 process to all ranks in this group. options = StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True) set_model_state_dict(model, full_sd, options=options) if self.rank == 0: logger.info("init_on_rank0 sync complete.") elif init_mode == "init_on_meta": non_persistent_buffers = self._save_non_persistent_buffers(model) if getattr(model.config, "tie_word_embeddings", False): model.tie_weights() model = self.prepare_model(model) model = self.materialize_and_load(model, hf_model_path=model.config.name_or_path, dcp_path=self.dcp_path) # fix tied broken for no-fsdp-wrap case if getattr(model.config, "tie_word_embeddings", False): model.tie_weights() self._restore_non_persistent_buffers(model, non_persistent_buffers) else: model = self.prepare_model(model) self._warmup_grad_norm(model) return model def _warmup_grad_norm(self, model: HFModel) -> None: """Warmup grad norm computation to initialize NCCL communication groups.""" if self.fsdp_mesh is None: return logger.info_rank0("Warming up grad norm computation...") for param in model.parameters(): if param.requires_grad: param.grad = torch.zeros_like(param) with torch.no_grad(): grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) if isinstance(grad_norm, torch.distributed.tensor.DTensor): grad_norm = grad_norm.full_tensor() for param in model.parameters(): if param.requires_grad: param.grad = None logger.info_rank0("Grad norm warmup completed.") def _load_from_dcp(self, model: HFModel, dcp_path: str): import torch.distributed.checkpoint as dcp try: if self.rank == 0: logger.info(f"Loading distributed checkpoint from {dcp_path} ...") options = StateDictOptions(full_state_dict=False, cpu_offload=True) local_state_dict = get_model_state_dict(model, options=options) dcp.load(state_dict=local_state_dict, checkpoint_id=dcp_path) set_model_state_dict(model, local_state_dict, options=options) if self.rank == 0: logger.info("DCP weights loaded successfully.") except Exception as e: logger.error(f"Failed to load from DCP: {e}") raise e def _try_build_hf_weight_conversion_context(self, model: HFModel) -> dict | None: try: from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key except ImportError: return None weight_mapping = get_model_conversion_mapping(model) if not weight_mapping: return None renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)] converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)] return { "prefix": getattr(model, "base_model_prefix", ""), "meta_state_dict": model.state_dict(), "rename_source_key": rename_source_key, "renamings": renamings, "converters": converters, "converter_templates": { pattern: converter for converter in converters for pattern in converter.source_patterns }, "pending_converters": {}, } def _load_weights_from_hf_checkpoint(self, model: HFModel, hf_model_path: str): import glob import json hf_model_path = self._resolve_hf_checkpoint_dir(hf_model_path) sort_key = _get_checkpoint_sort_key() if self.rank == 0: logger.info(f"Loading weights from {hf_model_path} ...") index_file = os.path.join(hf_model_path, "model.safetensors.index.json") is_safetensors = True checkpoint_files = [] if os.path.exists(index_file): with open(index_file) as f: index = json.load(f) checkpoint_files = sorted(set(index["weight_map"].values())) checkpoint_files = [os.path.join(hf_model_path, f) for f in checkpoint_files] elif os.path.exists(os.path.join(hf_model_path, "model.safetensors")): checkpoint_files = [os.path.join(hf_model_path, "model.safetensors")] else: is_safetensors = False index_file = os.path.join(hf_model_path, "pytorch_model.bin.index.json") if os.path.exists(index_file): with open(index_file) as f: index = json.load(f) checkpoint_files = sorted(set(index["weight_map"].values())) checkpoint_files = [os.path.join(hf_model_path, f) for f in checkpoint_files] elif os.path.exists(os.path.join(hf_model_path, "pytorch_model.bin")): checkpoint_files = [os.path.join(hf_model_path, "pytorch_model.bin")] else: checkpoint_files = sorted(glob.glob(os.path.join(hf_model_path, "*.safetensors"))) if checkpoint_files: is_safetensors = True else: checkpoint_files = sorted(glob.glob(os.path.join(hf_model_path, "*.bin"))) if not checkpoint_files: raise ValueError(f"No checkpoint files found in {hf_model_path}") param_map = dict(model.named_parameters()) conversion_ctx = self._try_build_hf_weight_conversion_context(model) total_files = len(checkpoint_files) for i, ckpt_file in enumerate(checkpoint_files): if self.rank == 0: logger.info(f"[{i + 1}/{total_files}] Loading {os.path.basename(ckpt_file)} ...") if is_safetensors: from safetensors import safe_open with safe_open(ckpt_file, framework="pt", device="cpu") as f: for key in sorted(f.keys(), key=sort_key): renamed_key = key source_pattern = None if conversion_ctx is not None: renamed_key, source_pattern = conversion_ctx["rename_source_key"]( key, conversion_ctx["renamings"], conversion_ctx["converters"], prefix=conversion_ctx["prefix"], meta_state_dict=conversion_ctx["meta_state_dict"], ) if source_pattern is not None: template = conversion_ctx["converter_templates"][source_pattern] converter = conversion_ctx["pending_converters"].setdefault( renamed_key, copy.deepcopy(template) ) converter.add_tensor( renamed_key, key, source_pattern, _make_safetensor_loader(ckpt_file, key), ) elif renamed_key in param_map: tensor = f.get_tensor(key) self._copy_weights(param_map[renamed_key], tensor) else: state_dict = torch.load(ckpt_file, map_location="cpu") for key, tensor in sorted(state_dict.items(), key=lambda item: sort_key(item[0])): renamed_key = key source_pattern = None if conversion_ctx is not None: renamed_key, source_pattern = conversion_ctx["rename_source_key"]( key, conversion_ctx["renamings"], conversion_ctx["converters"], prefix=conversion_ctx["prefix"], meta_state_dict=conversion_ctx["meta_state_dict"], ) if source_pattern is not None: template = conversion_ctx["converter_templates"][source_pattern] converter = conversion_ctx["pending_converters"].setdefault( renamed_key, copy.deepcopy(template) ) converter.add_tensor(renamed_key, key, source_pattern, tensor) elif renamed_key in param_map: self._copy_weights(param_map[renamed_key], tensor) del state_dict gc.collect() if conversion_ctx is not None: pending_count = len(conversion_ctx["pending_converters"]) log_fn = getattr(logger, "info_rank0", logger.info) log_fn(f"Applying {pending_count} deferred HF weight conversions.") for layer_name, converter in sorted(conversion_ctx["pending_converters"].items()): realized_tensors = converter.convert(layer_name, model=model, config=model.config) for target_name, tensor in realized_tensors.items(): if isinstance(tensor, list): tensor = tensor[0] if target_name in param_map: self._copy_weights(param_map[target_name], tensor) del realized_tensors gc.collect() def _resolve_hf_checkpoint_dir(self, hf_model_path: str) -> str: """Resolve a HF model identifier or local path to a local directory containing checkpoint files. - If `hf_model_path` is an existing directory, return it. - If it's a file path, return its parent directory. - Otherwise treat it as a Hugging Face Hub repo id and download/resolve to the local cache dir. """ if not hf_model_path: return hf_model_path # Local directory or file path. if os.path.isdir(hf_model_path): return hf_model_path if os.path.isfile(hf_model_path): return os.path.dirname(hf_model_path) # HuggingFace Hub repo id: snapshot to local cache so we can glob/index files. try: from huggingface_hub import snapshot_download except ImportError as e: raise ValueError( f"hf_model_path='{hf_model_path}' does not exist locally and huggingface_hub is not available " f"to download it. Please provide a local model directory or install huggingface_hub. Error: {e}" ) from e revision = os.getenv("HF_REVISION") offline = os.getenv("HF_HUB_OFFLINE") == "1" or os.getenv("TRANSFORMERS_OFFLINE") == "1" # In distributed runs, let rank0 download first to avoid N-way concurrent downloads. if torch.distributed.is_available() and torch.distributed.is_initialized(): if self.rank == 0: local_dir = snapshot_download( repo_id=hf_model_path, revision=revision, local_files_only=offline, allow_patterns=[ "*.safetensors", "*.bin", "*.index.json", "model.safetensors", "model.safetensors.index.json", "pytorch_model.bin", "pytorch_model.bin.index.json", "config.json", ], ) logger.info(f"Resolved HF repo id '{hf_model_path}' to local dir: {local_dir}") torch.distributed.barrier() if self.rank != 0: local_dir = snapshot_download( repo_id=hf_model_path, revision=revision, local_files_only=True, allow_patterns=[ "*.safetensors", "*.bin", "*.index.json", "model.safetensors", "model.safetensors.index.json", "pytorch_model.bin", "pytorch_model.bin.index.json", "config.json", ], ) return local_dir local_dir = snapshot_download( repo_id=hf_model_path, revision=revision, local_files_only=offline, allow_patterns=[ "*.safetensors", "*.bin", "*.index.json", "model.safetensors", "model.safetensors.index.json", "pytorch_model.bin", "pytorch_model.bin.index.json", "config.json", ], ) if self.rank == 0: logger.info(f"Resolved HF repo id '{hf_model_path}' to local dir: {local_dir}") return local_dir def _copy_weights(self, param, loaded_tensor): from torch.distributed._tensor import DTensor, Shard if loaded_tensor.dtype != param.dtype: loaded_tensor = loaded_tensor.to(param.dtype) if isinstance(param, DTensor): shard_placement = None mesh_dim = -1 for i, placement in enumerate(param.placements): if isinstance(placement, Shard): shard_placement = placement mesh_dim = i break local_tensor = param.to_local() if shard_placement is None: local_tensor.copy_(loaded_tensor) else: dim = shard_placement.dim mesh = param.device_mesh my_coordinate = mesh.get_coordinate() if my_coordinate is None: return rank_in_dim = my_coordinate[mesh_dim] world_size_in_dim = mesh.size(mesh_dim) full_size = param.shape[dim] chunk_size = (full_size + world_size_in_dim - 1) // world_size_in_dim start = rank_in_dim * chunk_size end = min(start + chunk_size, full_size) if start >= full_size: return sliced_tensor = loaded_tensor.narrow(dim, start, end - start) slices = [slice(None)] * local_tensor.ndim slices[dim] = slice(0, sliced_tensor.shape[dim]) local_tensor[tuple(slices)].copy_(sliced_tensor) else: param.data.copy_(loaded_tensor)