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

    Wrap any model for FSDP which follows one of the 3 existing conventions from
    HuggingFace for decoder-only LLMs.
    """
    causal_base_model = hf_get_causal_base_model(model)  # type: ignore

    # OPT and olmo have an extra layer of wrapping, so special case here
    if isinstance(causal_base_model, OPTDecoder):
        underlying_model = maybe_get_underlying_model(model)
        underlying_model.model._fsdp_wrap = False  # type: ignore
    model_block = hf_get_hidden_layers(causal_base_model)
    lm_head = model.get_output_embeddings()  # type: ignore
    # Try to get input embeddings from the transformer backbone
    # and then from the XXXForCausalLM
    try:
        tied_embeddings = causal_base_model.get_input_embeddings()
    except:
        tied_embeddings = model.get_input_embeddings()  # type: ignore

    modules = {
        'base_model': causal_base_model,
        'model_block': model_block,
        'lm_head': lm_head,
        'tied_embeddings': tied_embeddings,
    }

    for mod_name, module in modules.items():
        if module is None:
            raise ValueError(
                f'Unable to FSDP-wrap this model! `{mod_name}` does not ' +
                'follow common layer/weight naming conventions.',
            )
    block_type = type(model_block[0])

    # When using the HF LM models,
    # the weights of the self.lm_head and self.transformer.wte are tied.
    # This tying occurs inside the `self.post_init()` function.
    # This is a hurdle for FSDP because they need to be in the same FSDP block
    # These lines ensures that both modules stay together in the top-most block when
    # the model has this tying enabled (almost all do; this property defaults to True)
    if model.config.tie_word_embeddings:  # type: ignore
        causal_base_model._fsdp_wrap = False
        tied_embeddings._fsdp_wrap = False
        lm_head._fsdp_wrap = False

    # PEFT layers should be individually wrapped
    # TODO: Revisit this if we enforce use_orig_params=True, which seems to support
    # mixed frozen/unfrozen FSDP modules
    if hasattr(model, 'peft_type') and model.peft_type is not None:
        from peft import PeftModel
        assert isinstance(model, PeftModel)
        peft_type = model.peft_type.lower()  # type: ignore
        active_adapters = [
            adapter.lower() for adapter in model.active_adapters  # type: ignore
        ]  # type: ignore
        for name, module in model.named_modules():
            if peft_type in name.lower() and any(
                adapter in name.lower() for adapter in active_adapters
            ):
                has_parameters = next(module.parameters(), None) is not None
                has_buffers = next(module.buffers(), None) is not None
                if has_parameters or has_buffers:
                    module._fsdp_wrap = True

    # FSDP Wrap and Activation Checkpoint every model block
    model.fsdp_wrap_fn = lambda module: isinstance(module, block_type)
    model.activation_checkpointing_fn = lambda module: isinstance(
        module,
        block_type,
    )
FSDP wrap a HuggingFace decoder.

    Wrap any model for FSDP which follows one of the 3 existing conventions from
    HuggingFace for decoder-only LLMs.
    
llmfoundry/models/hf/hf_fsdp.py