def revert_liger_kernel_to_mllama(model_config: MiniModelConfig, model_type: str = "causal_lm"):
    """
    Revert all Liger kernel patches applied to MLlama.
    """

    assert model_type in [
        "causal_lm",
        "conditional_generation",
    ], f'model_type must be "causal_lm" or "conditional_generation", Got: {model_type}'
    import torch.nn as nn

    from transformers.models.mllama import modeling_mllama

    importlib.reload(nn)
    importlib.reload(modeling_mllama)
    if model_type == "causal_lm":
        model_config.model_class = modeling_mllama.MllamaForCausalLM
    else:
        model_config.model_class = modeling_mllama.MllamaForConditionalGeneration

    print("Liger kernel patches have been reverted.")

    Revert all Liger kernel patches applied to MLlama.
    
test/utils.py