How to set different learning rates for different parameters in the model?
How can I set different learning rates for different parameters in the model? I have rewritten the optimizers and separately set the learning rate for theact_fnin the model, but during training, I found that it doesn’t work:if optimizer_grouped_parameters is None:
        # Default parameter groups
        decay_parameters = Trainer.get_decay_parameter_names(None, model)
        optimizer_grouped_parameters = [
            {
                'params': [p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad and
                                                                       'act_fn' not in n)],
                'weight_decay': args.weight_decay,
            },
            {
                'params': [p for n, p in model.named_parameters() if (n not in decay_parameters and p.requires_grad)],
                'weight_decay': 0.0,
            },
            {
                "params": [
                    p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad and
                                                                   'act_fn' in n)
                ],
                "weight_decay": 0.0,
                'lr': 0.5
            },
        ]
    optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(args)
2024-12-16T10:22:21Z
This is updated version of your code.decay_parameters = Trainer.get_decay_parameter_names(None, model)
optimizer_grouped_parameters = [
    {
        'params': [p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad and
                                                               'act_fn' not in n)],
        'weight_decay': args.weight_decay,
        'lr': args.learning_rate,  # Default learning rate
    },
    {
        'params': [p for n, p in model.named_parameters() if (n not in decay_parameters and p.requires_grad)],
        'weight_decay': 0.0,
        'lr': args.learning_rate,  # Default learning rate
    },
    {
        'params': [p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad and
                                                               'act_fn' in n)],
        'weight_decay': 0.0,
        'lr': 0.5,  # Custom learning rate for act_fn
    },
]

optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(args)
optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

# Debugging optimizer parameter groups
for i, param_group in enumerate(optimizer.param_groups):
    print(f"Param group {i}: lr={param_group.get('lr', args.learning_rate)}, "
          f"weight_decay={param_group['weight_decay']}")
2024-12-16T11:58:16Z
I printed the relevant parameters:Param group 0: lr=5e-05, weight_decay=0.1
Param group 1: lr=5e-05, weight_decay=0.0
Param group 2: lr=0.5, weight_decay=0.0However, intransformer.trainer, afterself.optimizer.step(), I also checked it with:self.optimizer.step()
for i, param_group in enumerate(self.optimizer.optimizer.param_groups):
    print(f"Param group {i}: lr={param_group.get('lr', args.learning_rate)}, "
          f"weight_decay={param_group['weight_decay']}")The output is:Param group 0: lr=5e-05, weight_decay=0.1This is strange; there are no Param group 1 and 2. I am using DeepSpeed’s Zero3. Does this change the Param group?
2024-12-16T14:26:43Z
The issue is caused by DeepSpeed. When using DeepSpeed, it results in parameter groups being merged into a single group. I am not sure how to configure it to prevent this merging.
2024-12-16T14:29:31Z
I have some experience about that.The issue with DeepSpeed merging parameter groups into a single group is related to how it handles optimization parameter groups. DeepSpeed, by default, optimizes parameters more efficiently by merging them into a single group to streamline operations like gradient updates and memory management. However, if you want to maintain separate parameter groups for different learning rates or other configurations, you need to adjust DeepSpeed’s configuration.Solutions:1. Usezero_allow_untested_optimizerin the DeepSpeed ConfigDeepSpeed’s ZeRO optimizer merges parameter groups by default for memory efficiency. You can disable this behavior using thezero_allow_untested_optimizerflag in the DeepSpeed configuration file. For example:jsonCopy code{
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 0.001
    }
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "zero_allow_untested_optimizer": true
  }
}This flag prevents DeepSpeed from enforcing its internal parameter group merging.2. Define Custom Parameter GroupsWhen initializing the optimizer in your code, explicitly define parameter groups before passing them to DeepSpeed. For example:pythonCopy codeoptimizer_grouped_parameters = [
    {"params": model.base_parameters, "lr": 1e-3},
    {"params": model.special_parameters, "lr": 1e-4}
]

optimizer = torch.optim.AdamW(optimizer_grouped_parameters)
model, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config_params=deepspeed_config
)Ensurezero_allow_untested_optimizeris enabled if you’re using ZeRO optimization.
2024-12-16T14:46:59Z
It works! I’m truly so grateful to you!
2024-12-16T14:48:59Z
I am so happy to work your code.There is another solution for you codewith deepspeed.zero.Init(config=deepspeed_config):
    model = MyModel()

optimizer_grouped_parameters = [
    {"params": model.base_parameters, "lr": 1e-3},
    {"params": model.special_parameters, "lr": 1e-4}
]

optimizer = torch.optim.AdamW(optimizer_grouped_parameters)
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config_params=deepspeed_config
)
2024-12-17T02:49:50Z
This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.