Checkpointing broken for classifier training multi-gpu
Author: ojh31
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.52.4
- Platform: Linux-5.15.0-141-generic-x86_64-with-glibc2.35
- Python version: 3.11.11
- Huggingface_hub version: 0.33.0
- Safetensors version: 0.5.3
- Accelerate version: 1.7.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes
- Using GPU in script?: yes
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@zach-huggingface @SunMarc

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Script:
```
"""Simple IMDB binary classification training script using Qwen3-14B and HF Trainer."""

import torch
from datasets import DatasetDict, load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
)
from transformers.training_args import TrainingArguments
from transformers.trainer import Trainer
from transformers.data.data_collator import DataCollatorWithPadding
from peft import LoraConfig, get_peft_model, TaskType


def tokenize_function(examples, tokenizer, max_length=512):
    """Tokenize the text data."""
    return tokenizer(
        examples["text"],
        truncation=True,
        padding=False,
        max_length=max_length,
    )


def main():
    # Model and tokenizer setup
    model_name = "Qwen/Qwen3-14B"

    print(f"Loading tokenizer from {model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

    # Set pad token if not present
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"Loading model from {model_name}...")
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    )

    # Set pad_token_id in model config
    model.config.pad_token_id = tokenizer.pad_token_id

    # Configure LoRA
    print("Adding LoRA adapters...")
    lora_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.1,
        bias="none",
        task_type=TaskType.SEQ_CLS,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # Load IMDB dataset
    print("Loading IMDB dataset...")
    dataset = load_dataset("imdb")

    # Tokenize dataset
    print("Tokenizing dataset...")
    tokenized_dataset = dataset.map(
        lambda examples: tokenize_function(examples, tokenizer),
        batched=True,
        remove_columns=["text"],
    )
    assert isinstance(tokenized_dataset, DatasetDict)

    # Data collator
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    # Training arguments
    training_args = TrainingArguments(
        output_dir="./results_imdb_qwen14b",
        num_train_epochs=1,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=4,
        gradient_accumulation_steps=16,
        log_level="debug",
        save_strategy="steps",
        save_steps=1,
        bf16=True,
    )

    # Initialize trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset["train"],
        eval_dataset=tokenized_dataset["test"],
        processing_class=tokenizer,
        data_collator=data_collator,
    )

    # Train the model
    print("Starting training...")
    trainer.train()


if __name__ == "__main__":
    main()

```
Launch with `accelerate launch --num_processes=2 --config_file=accelerate_config.yaml script.py`
```
# accelerate config
compute_environment: LOCAL_MACHINE
debug: false
# We want FSDP to shard model parameters between devices.
distributed_type: FSDP
downcast_bf16: "no"
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
# Previously we used "fp16" since that is how Pythia models were trained,
# but decided to turn off mixed precision for simplicity.
mixed_precision: "no"
num_machines: 1
# We overwrite this with a CLI argument
num_processes: 1
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

### Expected behavior

The trainer should save a checkpoint to the local `results_imdb_qwen14b` folder after a single step, but instead it hangs and eventually there is an NCCL timeout. The code never gets past the first call to save_fsdp_model -> _get_model_state_dict -> get_peft_model_state_dict -> model.state_dict()

Comments:
Variations which DO work:

- Removing peft (i.e. using full parameter fine-tuning)
- Switching SequenceClassification to CausalLM (and changing the data collator as necessary)
- Using a single gpu (and a smaller model as necessary)

If I don't save checkpoints, then I can train on a few datapoints, but if the code runs for too long then I hit an NCCL timeout

Ran into similar issues with multi-gpu + LoRA fine-tuning that hangs at the end of training (eventually time out..). But this issue only occurs for `transformers>4.51.0`. 

Following this thread.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.