Trainer.save_model saves config with FSDP-wrapped architectures instead of base model
Author: JdRion
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.56.0
- Platform: Linux-5.15.92-2.el8.navix.ncc.x86_64-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.34.4
- Safetensors version: 0.5.2
- Accelerate version: 1.10.1
- Accelerate config:
```
distributed_type: FSDP
mixed_precision: bf16

fsdp_config:
  fsdp_version: 2
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_activation_checkpointing: True
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_transformer_layer_cls_to_wrap: Gemma3DecoderLayer
  fsdp_state_dict_type: FULL_STATE_DICT

parallelism_config:
  parallelism_config_cp_size: 1
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 96
  parallelism_config_tp_size: 1
```
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
- 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?: 
```
pc = ParallelismConfig()` , 
...
training_args = GRPOConfig(parallelism_config=pc, ...)


accelerate launch \
                    --num_processes ${NUM_PROCESSES} \
                    --num_machines ${NUM_MACHINES} \
                    --machine_rank ${MACHINE_RANK} \
                    --main_process_ip ${MASTER_ADDR} \
                    --main_process_port ${MASTER_PORT} \
                    --config_file fsdp.yaml \
                    dapo.py --config mg --max-completion-len 8192
```
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

- trainer: @zach-huggingface @SunMarc

### Information

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

### Tasks

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

### Reproduction

When saving a model trained with ParallelismConfig(FSDP), the saved config.json includes an architectures entry with the FSDP-wrapped class name (e.g., "FSDPGemma3ForCausalLM") instead of the original base model ("Gemma3ForCausalLM").

Custom Trainer (to avoid OOM on save)
To avoid OOM when saving in FSDP + ND-parallel mode, I overrode save_model in my custom trainer:
```
class SafeGRPOTrainer(GRPOTrainer):
    def save_model(
        self, output_dir: Optional[str] = None, _internal_call: bool = False
    ):
        if output_dir is None:
            output_dir = self.args.output_dir

        out_dir = Path(output_dir)
        out_dir.mkdir(parents=True, exist_ok=True)

        # N-D parallelism branch
        if getattr(self.accelerator, "parallelism_config", None) is not None:
            if self.accelerator.should_save_model:
                # Added this line to avoid OOM:
                state_dict = self.accelerator.get_state_dict(self.model)
                self._save(output_dir, state_dict=state_dict)

        elif self.is_fsdp_enabled:
            if (
                "FULL_STATE_DICT"
                in str(self.accelerator.state.fsdp_plugin.state_dict_type)
            ) and (version.parse(accelerate_version) > version.parse("0.24.1")):
                state_dict = self.accelerator.get_state_dict(self.model)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)

        elif self.args.should_save:
            self._save(output_dir)

        if self.args.push_to_hub and not _internal_call:
            self.push_to_hub(
                commit_message="Model save", revision=self.args.hub_revision
            )
```
- Code Example
```
import torch
from accelerate.utils import ParallelismConfig
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GRPOConfig
from util.safe_trainer import SafeGRPOTrainer  # <-- custom override with get_state_dict

def main():
    # Parallelism config (trigger N-D parallelism branch)
    pc = ParallelismConfig()

    model_name = "google/medgemma-27b-text-it"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Minimal training args
    training_args = GRPOConfig(
        parallelism_config=pc,
        output_dir="./ckpt_test",
        per_device_train_batch_size=1,
        num_train_epochs=1,
        save_strategy="steps",
        save_steps=1,
        save_total_limit=1,
        logging_steps=1,
        report_to="none",
    )

    # Init trainer
    trainer = SafeGRPOTrainer(
        model=model,
        processing_class=tokenizer,
        args=training_args,
        train_dataset=train_ds,
        reward_funcs=[...],  # not needed for repro
    )

    # Only run save to trigger the bug
    trainer.save_model("./ckpt_test")

if __name__ == "__main__":
    main()
```

### Expected behavior

- Actual result (config.json)
```
{
  "architectures": [
    "FSDPGemma3ForCausalLM"
  ],
  ...
}
```
- Expected result (config.json)
```
{
  "architectures": [
    "Gemma3ForCausalLM"
  ],
  ...
}
```

Comments:
cc @S1ro1 

Will revisit, thought this PR was already merged. However this is only cosmetical change, you can just rename the architecture in the saved file yourself for now - the architecture is the same, just class name is different. https://github.com/huggingface/transformers/pull/40207

@S1ro1 Hi, thanks for the work on N-D parallelism 👏

In the current implementation, when using FSDP with FULL_STATE_DICT, Trainer.save_model already calls
```
state_dict = self.accelerator.get_state_dict(self.model)
```
which avoids OOM and works well.

However, in the N-D parallelism branch, it still falls back to model.state_dict().
I had to override save_model and add get_state_dict manually to prevent OOM.

Question:
Since FULL_STATE_DICT already relies on get_state_dict, is there any plan to apply the same logic to the N-D parallelism branch as well?

hi @JdRion Thank you for the nice words! Yes, I'll be checking the saving with N-D parallelism. It should for sure use `accelerator.get_state_dict` to work properly.

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.