def prepare_hf_model_for_fsdp(
    model: Union[PreTrainedModel, 'PeftModel'],
    init_device: Optional[str],
) -> None:
    """FSDP wrap a HuggingFace model.

    Call specific functions
    """
    if model.config.is_encoder_decoder:  # type: ignore
        prepare_hf_enc_dec_model_for_fsdp(model, init_device)
    else:
        # many common decoder-only model do not set the flag
        # model.config.is_decoder, so we can't trust it
        prepare_hf_causal_lm_model_for_fsdp(model, init_device)
FSDP wrap a HuggingFace model.

    Call specific functions
    
llmfoundry/models/hf/hf_fsdp.py