It is possible to finetune the model freezing some layers?

[davide97l]

I would like to fine-tune the Moirai model by freezing some of its layers. Specifically, I want to freeze all layers except for:

• The top MLP layer
• The embedding layers

Is there an existing method or configuration within the Moirai framework that supports this selective layer freezing for fine-tuning? If not, could you provide guidance or potential solutions for implementing this?

Thank you for your assistance!

[davide97l]

Solved by modifying train.py in the following way

@hydra.main(version_base="1.3", config_name="default.yaml")
def main(cfg: DictConfig):
    if cfg.tf32:
        assert cfg.trainer.precision == 32
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    model: L.LightningModule = instantiate(cfg.model, _convert_="all")

    #for name, module in model.named_children():
    #    print(f"Module name: {name}")
    #    print(module)

    learnable_layers = []
    frozen_layers = []
    learnable_layers_keywords = ["param_proj", "in_proj"]
    for name, param in model.named_parameters():
        if any(k in name for k in learnable_layers_keywords):
            param.requires_grad = True
            learnable_layers.append(name)
        else:
            param.requires_grad = False
            frozen_layers.append(name)
    print('Learnable layers:')
    for l in learnable_layers:
        print(l)
    print('Frozen layers:')
    for l in frozen_layers:
        print(l)

    weight_tensor = model.module.in_proj.weight
    weight_sum = torch.sum(weight_tensor)
    print("Sum of weights (in_proj):", weight_sum)

    if cfg.compile:
        model.module.compile(mode=cfg.compile)
    trainer: L.Trainer = instantiate(cfg.trainer)
    train_dataset: Dataset = instantiate(cfg.data).load_dataset(
        model.train_transform_map
    )
    val_dataset: Optional[Dataset | list[Dataset]] = (
        tree_map(
            lambda ds: ds.load_dataset(model.val_transform_map),
            instantiate(cfg.val_data, _convert_="all"),
        )
        if "val_data" in cfg
        else None
    )
    L.seed_everything(cfg.seed + trainer.logger.version, workers=True)
    trainer.fit(
        model,
        datamodule=DataModule(cfg, train_dataset, val_dataset),
        ckpt_path=cfg.ckpt_path,
    )

if __name__ == "__main__":
    main()