from transformers import AutoModelForCausalLM

def enable_grad_only_every_nth(model, n):
    """
    This function configures the specified model to enable gradient calculations exclusively for every nth layer, starting 
    from the first layer (0-indexed), to accommodate newly added blocks for training. Concurrently, it freezes the gradients 
    for all other components of the model, including the embedding layers and the model's head. This setup is particularly 
    useful for fine-tuning processes where only a subset of layers are targeted for updates, ensuring efficient training and 
    adaptation of newly integrated layers while maintaining the pre-trained behavior of other model components.
    """

    # Freeze embeddings.
    for param in model.model.embed_tokens.parameters():
        param.requires_grad = False

    # Freeze lm_head.
    for param in model.lm_head.parameters():
        param.requires_grad = False

    # Enable gradients for every nth layer
    layers = model.model.layers  # Access the ModuleList containing the layers

    for index, layer in enumerate(layers):

        if (index + 1) % n == 0:  # Enables gradients for every nth layer, starting from the layer after the 0th
            for param in layer.parameters():
                param.requires_grad = True
        else:
            for param in layer.parameters():
                param.requires_grad = False

model = transformers.AutoModelForCausalLM.from_pretrained(
    "arcee-ai/Mistral-7B-Instruct-v0.2-expanded"
    )
# Update layer gradients, specify the correct value for n based on your model's architecture
n =5
enable_grad_only_every_nth(model, n)