huggingface/transformers · error · ValueError

save_pretrained(..., distributed_checkpoint=True) is only su

Error message

save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models.

What it means

The distributed checkpoint save path only knows how to serialize FSDP2-managed (fully_shard-wrapped) modules via is_fsdp_managed_module. If the model was never wrapped — because fsdp_size was 1 or the distributed config was not applied — the writer cannot map sharded parameters and raises ValueError.

Source

Thrown at src/transformers/distributed/mixin.py:235

    def save_distributed_checkpoint(
        self,
        model_to_save,
        save_directory: str | os.PathLike,
        *,
        push_to_hub: bool = False,
        save_on_this_rank: bool = True,
        repo_id: str | None = None,
        files_timestamps: dict | None = None,
        commit_message: str | None = None,
        token: str | bool | None = None,
        create_pr: bool = False,
    ) -> None:
        """Save an FSDP-wrapped model via DCP and optionally push to the Hub."""
        if not is_torch_greater_or_equal("2.7"):
            raise OSError("save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7.")
        if not is_fsdp_managed_module(model_to_save):
            raise ValueError(
                "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models."
            )
        if getattr(model_to_save, "_device_mesh", None) is None:
            raise ValueError(
                "save_pretrained(..., distributed_checkpoint=True) requires the model to have been "
                "initialized with a distributed_config (_device_mesh is None)."
            )
        save_model_checkpoint_distributed(model_to_save, save_directory)

        if push_to_hub and save_on_this_rank:
            model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
            model_card.save(os.path.join(save_directory, "README.md"))
            self._upload_modified_files(
                save_directory,
                repo_id,
                files_timestamps,
                commit_message=commit_message,
                token=token,

View on GitHub (pinned to a597f97485)

Solutions

  1. Initialize the model with distributed_config={'fsdp_size': N>1, ...} so from_pretrained wraps it with fully_shard, then save.
  2. If the model is intentionally not FSDP-sharded, use the standard save_pretrained (distributed_checkpoint omitted/False).
  3. Verify wrapping before saving: assert any(is_fsdp_managed_module(m) for m in model.modules()).

Example fix

# before
model = AutoModelForCausalLM.from_pretrained(model_id)  # no fsdp
model.save_pretrained(out_dir, distributed_checkpoint=True)  # raises

# after
model = AutoModelForCausalLM.from_pretrained(model_id, distributed_config={"fsdp_size": world_size})
model.save_pretrained(out_dir, distributed_checkpoint=True)
Defensive patterns

Strategy: validation

Validate before calling

from transformers.distributed.utils import is_fsdp_managed_module

def can_distributed_save(model) -> bool:
    return any(is_fsdp_managed_module(m) for m in model.modules())

Type guard

def is_fsdp_model(model) -> bool:
    return any(is_fsdp_managed_module(m) for m in model.modules())

Try / catch

try:
    model.save_pretrained(out, distributed_checkpoint=True)
except ValueError as e:
    if "only supported for FSDP-wrapped" in str(e):
        model.save_pretrained(out)  # plain save for non-FSDP models
    else:
        raise

Prevention

When it happens

Trigger: Calling save_pretrained(..., distributed_checkpoint=True) on a plain model: loaded with from_pretrained without distributed_config, or with a config where fsdp_size=1 (TP-only or no sharding).

Common situations: Trying the new save flag on a TP-only model; saving after manually unwrapping FSDP; calling the FSDP save helper on a checkpoint-gathered model.

Related errors


AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/d0e2660331df6079. Report an issue: GitHub.