diff --git a/finetune_overfit.py b/finetune_overfit.py new file mode 100644 index 000000000000..4e1253d61321 --- /dev/null +++ b/finetune_overfit.py @@ -0,0 +1,221 @@ +import os +import shutil +import tempfile +from dataclasses import dataclass + +import torch +from accelerate import ParallelismConfig +from datasets import load_dataset + +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoTokenizer, + DataCollatorForLanguageModeling, + HfArgumentParser, + Trainer, + TrainingArguments, +) +from transformers.distributed import DistributedConfig +from transformers.distributed.utils import _distributed_barrier, _ensure_torch_distributed +from transformers.utils import is_torch_neuron_available + + +@dataclass +class ScriptArguments: + dataset_name: str + dataset_config: str | None = None + dataset_split: str = "train" + num_examples: int = 16 # tiny fixed subset to overfit on + max_length: int = 1024 + + +@dataclass +class ModelArguments: + model_name_or_path: str + model_revision: str = "main" + trust_remote_code: bool = False + use_lora: bool = False + + +def main(script_args, training_args, model_args): + if not torch.cuda.is_available() and is_torch_neuron_available(check_device=True): + import torch_neuronx # noqa: F401 + + tp_size = int(os.environ.get("TP_SIZE", "1")) + fsdp_size = int(os.environ.get("FSDP_SIZE", "1")) + if tp_size > 1 and fsdp_size > 1: + raise ValueError( + f"TP_SIZE ({tp_size}) > 1 together with FSDP_SIZE ({fsdp_size}) > 1 is not supported: " + "1D `DistributedConfig` only. Use one or the other." + ) + if (tp_size > 1 or fsdp_size > 1) and training_args.fsdp: + raise ValueError( + f"TP_SIZE ({tp_size}) / FSDP_SIZE ({fsdp_size}) together with --fsdp is not supported: " + "that flag configures Accelerate's own FSDP plugin, a separate mechanism from " + "`distributed_config`. Use one or the other." + ) + + kwargs = {} + if tp_size > 1: + training_args.parallelism_config = ParallelismConfig(tp_size=tp_size) + kwargs["distributed_config"] = DistributedConfig(tp_size=tp_size) + elif fsdp_size > 1: + kwargs["distributed_config"] = DistributedConfig(fsdp_size=fsdp_size) + + if tp_size > 1 or fsdp_size > 1: + _ensure_torch_distributed() + + config = AutoConfig.from_pretrained( + model_args.model_name_or_path, + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + ) + dtype = torch.bfloat16 if training_args.bf16 else torch.float32 + + from peft import LoraConfig, get_peft_model, PeftModel + + lora_config = LoraConfig( + r=64, + lora_alpha=128, + target_modules=[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "embed_tokens", + ], + ) + + # Loading the model once to make a copy of the original state dict for later comparison. + # No sharding is performed. + if model_args.use_lora: + original_model_path = os.path.join(training_args.output_dir, "original_model") + if int(os.environ.get("RANK", "0")) == 0: + original_model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, + config=config, + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + torch_dtype=dtype, + ) + + original_model = get_peft_model(original_model, lora_config) + + shutil.rmtree(original_model_path, ignore_errors=True) + original_model.save_pretrained(original_model_path) + + if tp_size > 1 or fsdp_size > 1: + _distributed_barrier() + + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, + config=config, + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + torch_dtype=dtype, + low_cpu_mem_usage=True, + **kwargs, + ) + + torch_device = f"cuda:{os.environ.get('LOCAL_RANK')}" if torch.cuda.is_available() else None + model = PeftModel.from_pretrained( + model, original_model_path, is_trainable=True, torch_device=torch_device + ) + if int(os.environ.get("RANK", "0")) == 0: + model.print_trainable_parameters() + else: + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, + config=config, + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + torch_dtype=dtype, + low_cpu_mem_usage=True, + **kwargs, + ) + + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # --------------------------------------------------------------------------- + # Dataset: a tiny fixed subset, tokenized once via the chat template. + # --------------------------------------------------------------------------- + dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config, split=script_args.dataset_split) + dataset = dataset.select(range(script_args.num_examples)) + + def tokenize(example): + input_ids = tokenizer.apply_chat_template( + example["messages"], + tokenize=True, + return_dict=False, + truncation=True, + max_length=script_args.max_length, + ) + return {"input_ids": input_ids} + + dataset = dataset.map(tokenize, remove_columns=dataset.column_names) + + trainer = Trainer( + model=model, + args=training_args, + train_dataset=dataset, + data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False), + ) + + # Check that the model can be saved and reloaded correctly, even when sharded. + rank = int(os.environ.get("RANK", "0")) + trainer.save_model(training_args.output_dir) + + if rank == 0: + if model_args.use_lora: + original_state_dict = original_model.state_dict() + + unsharded_model = PeftModel.from_pretrained( + AutoModelForCausalLM.from_pretrained(training_args.output_dir, torch_dtype=dtype), + original_model_path, + ) + unsharded_state_dict = unsharded_model.state_dict() + else: + original_state_dict = original_model.state_dict() + + unsharded_model = AutoModelForCausalLM.from_pretrained(training_args.output_dir, torch_dtype=dtype) + unsharded_state_dict = unsharded_model.state_dict() + + mismatches = [] + missing_keys = set(original_state_dict) - set(unsharded_state_dict) + if missing_keys: + mismatches.append(f"missing from saved checkpoint: {sorted(missing_keys)}") + for key, expected_value in unsharded_state_dict.items(): + if key not in original_state_dict: + mismatches.append(f"{key}: unexpected key in saved checkpoint") + continue + try: + torch.testing.assert_close(expected_value, original_state_dict[key], rtol=0, atol=0) + except AssertionError as e: + mismatches.append(f"{key}: {e}") + if mismatches: + raise AssertionError("Save correctness check failed:\n" + "\n".join(mismatches)) + + print( + f"Save correctness check passed: {len(unsharded_state_dict)} parameters " + "match the unsharded checkpoint exactly." + ) + + # We can start training. + trainer.train() + + + +if __name__ == "__main__": + parser = HfArgumentParser((ScriptArguments, TrainingArguments, ModelArguments)) + script_args, training_args, model_args = parser.parse_args_into_dataclasses() + main(script_args, training_args, model_args) diff --git a/finetune_overfit.sh b/finetune_overfit.sh new file mode 100755 index 000000000000..d54efa787cf9 --- /dev/null +++ b/finetune_overfit.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# Minimalist TP/FSDP overfitting test on AWS Trainium or CUDA GPUs using the regular +# transformers Trainer (no SFTTrainer; LoRA optional via USE_LORA). Trains on a tiny fixed +# subset of examples for many steps to check that the loss goes to ~0, as a correctness check +# for TP or FSDP. + +set -euo pipefail + +# --------------------------------------------------------------------------- +# Parallelism: pick a variant with PARALLEL_MODE=tp|fsdp (default: fsdp). +# TP and FSDP2 are both enabled through the same mechanism -- `DistributedConfig`, which pre-shards +# the model into DTensor at from_pretrained time: +# PARALLEL_MODE=tp NUM_PROC=4 -> TP_SIZE=4 +# PARALLEL_MODE=fsdp NUM_PROC=4 -> FSDP_SIZE=4 +# --------------------------------------------------------------------------- +PARALLEL_MODE=${PARALLEL_MODE:-tp} +NUM_PROC=${NUM_PROC:-2} +USE_LORA=${USE_LORA:-true} + +case "$PARALLEL_MODE" in + tp) + export TP_SIZE=$NUM_PROC + export FSDP_SIZE=1 + ;; + fsdp) + export TP_SIZE=1 + export FSDP_SIZE=$NUM_PROC + ;; + *) + echo "Unknown PARALLEL_MODE: $PARALLEL_MODE (expected tp|fsdp)" >&2 + exit 1 + ;; +esac + +# --------------------------------------------------------------------------- +# Neuron runtime environment (Trainium only -- no-op, skipped entirely on CUDA) +# --------------------------------------------------------------------------- +if ! command -v nvidia-smi &> /dev/null || ! nvidia-smi &> /dev/null; then + export ON_NEURON_EAGER=1 + export NEURON_EAGER_MODEL_CACHE_SIZE=10000 + export OMP_NUM_THREADS=128 + export HF_DEACTIVATE_ASYNC_LOAD=1 + + export TORCH_NEURONX_ENABLE_HOST_CC=1 + export TORCH_NEURONX_ENABLE_ASYNC_NRT=1 + #export NEURON_RT_NUM_CORES=1 +fi + +# --------------------------------------------------------------------------- +# Model / data / hyperparameters +# --------------------------------------------------------------------------- +MODEL_NAME=Qwen/Qwen3-1.7B +DATASET_NAME=trl-lib/Capybara +LEARNING_RATE=5.0e-4 +NUM_TRAIN_EXAMPLES=16 +MAX_STEPS=50 +MAX_SEQ_LENGTH=1024 +BATCH_SIZE=4 + +LORA_SUFFIX="" +if [ "$USE_LORA" = "true" ]; then + LORA_SUFFIX="-lora" +fi +OUTPUT_DIR=Qwen3-1.7B-${PARALLEL_MODE}-Overfit${LORA_SUFFIX} + + +echo "==========================================" +echo "Plain Trainer parallelism overfitting test" +echo " Model: $MODEL_NAME" +echo " Dataset: $DATASET_NAME" +echo " PARALLEL_MODE: $PARALLEL_MODE" +echo " NUM_PROC: $NUM_PROC" +echo " TP_SIZE: $TP_SIZE" +echo " FSDP_SIZE: $FSDP_SIZE" +echo " USE_LORA: $USE_LORA" +echo " Num examples: $NUM_TRAIN_EXAMPLES" +echo " Max steps: $MAX_STEPS" +echo " Batch: $BATCH_SIZE" +echo " Max seq len: $MAX_SEQ_LENGTH" +echo " Output dir: $OUTPUT_DIR" +echo "==========================================" + +if [ "$NUM_PROC" -eq 1 ]; then + LAUNCHER="python" +else + LAUNCHER="torchrun --nproc_per_node=${NUM_PROC}" +fi + +$LAUNCHER \ + finetune_overfit.py \ + --model_name_or_path "$MODEL_NAME" \ + --use_lora $USE_LORA \ + --dataset_name "$DATASET_NAME" \ + --num_examples $NUM_TRAIN_EXAMPLES \ + --max_length $MAX_SEQ_LENGTH \ + --learning_rate $LEARNING_RATE \ + --gradient_checkpointing true \ + --bf16 true \ + --per_device_train_batch_size $BATCH_SIZE \ + --max_steps $MAX_STEPS \ + --eval_strategy no \ + --logging_steps 10 \ + --save_strategy no \ + --dataloader_num_workers 0 \ + --report_to trackio \ + --output_dir "$OUTPUT_DIR" diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index eb1bafd7355f..c3d2ff5c3027 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -196,6 +196,7 @@ def maybe_distribute_model( if isinstance(distributed_config.tp_plan, dict): model.tp_plan = distributed_config.tp_plan model = apply_tensor_parallelism(model, tp_mesh) + model._tp_size = distributed_config.tp_size elif distributed_config.fsdp_size > 1: fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh diff --git a/src/transformers/integrations/peft.py b/src/transformers/integrations/peft.py index 91ebf98791ce..219fe26c7c28 100644 --- a/src/transformers/integrations/peft.py +++ b/src/transformers/integrations/peft.py @@ -17,7 +17,7 @@ from dataclasses import replace from typing import TYPE_CHECKING, Any, Literal, Optional -from safetensors import safe_open +from transformers.utils.import_utils import is_peft_greater_or_equal from .._typing import PeftConfigLike from ..conversion_mapping import get_model_conversion_mapping @@ -155,9 +155,8 @@ def load_adapter( `find_adapter_config_file` method. """ from peft import PeftType - from peft.utils.save_and_load import _maybe_shard_state_dict_for_tp - from ..modeling_utils import LoadStateDictConfig, _get_resolved_checkpoint_files, load_state_dict + from ..modeling_utils import LoadStateDictConfig, _get_resolved_checkpoint_files if local_files_only: kwargs["local_files_only"] = True @@ -274,44 +273,22 @@ def is_adapter_key(key: str) -> bool: device_map = getattr(self, "hf_device_map", {"": self.device}) - # If the model is tensor parallel, we handle the sharding of the state dict here since the logic in `self._load_pretrained_model` - # is not compatible with the way PEFT adapter should be sharded. has_tp_adapters = False for module in self.modules(): + # Legacy, pre-DTensor TP integration: PEFT stamps a `_tp_info` marker on each TP-sharded LoRA module. tp_info = getattr(module, "_tp_info", None) if tp_info is not None: has_tp_adapters = True break + # DTensor TP integration: the base model itself carries the TP plan, so any adapter injected into it will be + # TP-sharded too; no per-module PEFT marker is needed to detect this. + has_tp_adapters = has_tp_adapters or bool(getattr(self, "_tp_plan", None)) - if has_tp_adapters: - all_pointer = set() - if adapter_state_dict is not None: - merged_state_dict = adapter_state_dict - elif ( - checkpoint_files is not None - and checkpoint_files[0].endswith(".safetensors") - and adapter_state_dict is None - ): - merged_state_dict = {} - for file in checkpoint_files: - file_pointer = safe_open(file, framework="pt", device="cpu") - all_pointer.add(file_pointer) - for k in file_pointer.keys(): - merged_state_dict[k] = file_pointer.get_tensor(k) - # Checkpoints are .bin - elif checkpoint_files is not None: - merged_state_dict = {} - for ckpt_file in checkpoint_files: - merged_state_dict.update(load_state_dict(ckpt_file)) - else: - raise ValueError("Neither a state dict nor checkpoint files were found.") - - adapter_state_dict = merged_state_dict - - if any(not isinstance(v, torch.Tensor) for v in adapter_state_dict.values()): - raise ValueError("Expected all values in the adapter state dict to be tensors.") - - _maybe_shard_state_dict_for_tp(self, adapter_state_dict, adapter_name) + if has_tp_adapters and not is_peft_greater_or_equal("0.20.1", accept_dev=True): + raise ValueError( + "Loading a tensor-parallel PEFT adapter requires peft >= 0.20.1, please upgrade your peft " + "installation." + ) load_config = replace( load_config, diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 7a5005146858..89050ac40286 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -49,6 +49,7 @@ from huggingface_hub import CommitInfo, ModelCard from packaging import version from torch import nn +from torch.distributed.tensor import DTensor from torch.utils.data import DataLoader, Dataset, IterableDataset, RandomSampler, SequentialSampler from . import __version__ @@ -1238,20 +1239,21 @@ def create_optimizer(self, model=None) -> torch.optim.Optimizer: if self.optimizer is None: decay_parameters = self.get_decay_parameter_names(opt_model) - optimizer_grouped_parameters = [ - { - "params": [ - p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad) - ], - "weight_decay": self.args.weight_decay, - }, - { - "params": [ - p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad) - ], - "weight_decay": 0.0, - }, - ] + optimizer_grouped_parameters = [] + for in_decay, weight_decay in ((True, self.args.weight_decay), (False, 0.0)): + params = [ + p + for n, p in opt_model.named_parameters() + if p.requires_grad and (n in decay_parameters) == in_decay + ] + # Fused/foreach optimizers batch every param of a group into a single op call, which + # errors out if the group mixes DTensor and plain Tensor params. + # Split each group by type so every batched call stays homogeneous. + dtensor_params = [p for p in params if isinstance(p, DTensor)] + plain_params = [p for p in params if not isinstance(p, DTensor)] + for group_params in (dtensor_params, plain_params): + if group_params: + optimizer_grouped_parameters.append({"params": group_params, "weight_decay": weight_decay}) if self.optimizer_cls_and_kwargs is not None: optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs @@ -1676,9 +1678,11 @@ def _prepare_for_training(self, max_steps, train_dataloader, resume_from_checkpo # wrapped (e.g. in DataParallel) on subsequent `train()` calls and avoid double wrapping. model = self._wrap_model(self.model_wrapped) + is_natively_fsdp_sharded = getattr(model, "_is_fsdp_managed_module", False) + # If the model is wrapped, don't use `accelerator.prepare` # this is for unhandled cases in accelerate such as FSDP-XLA, SageMaker MP/DP, DataParallel - use_accelerator_prepare = model is self.model + use_accelerator_prepare = model is self.model and not is_natively_fsdp_sharded # prepare using `accelerator` prepare if use_accelerator_prepare: @@ -3925,6 +3929,19 @@ def save_model(self, output_dir: str | None = None, _internal_call: bool = False remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME]) self.model_wrapped.save_checkpoint(output_dir) + elif getattr(self.model.config, "distributed_config", None) is not None: + os.makedirs(output_dir, exist_ok=True) + self.model.save_pretrained(output_dir) + if self.args.should_save: + if self.processing_class is not None: + self.processing_class.save_pretrained(output_dir) + elif ( + self.data_collator is not None + and hasattr(self.data_collator, "tokenizer") + and self.data_collator.tokenizer is not None + ): + self.data_collator.tokenizer.save_pretrained(output_dir) + torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) elif self.args.should_save: self._save(output_dir)