diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index b4d15daefd2..11aa6c49585 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -60,40 +60,48 @@ OptimizerConfig, ParamKey, ParamPredicate, + ParamWithNamePredicate, SGDOptimizerConfig, ) logger = logging.getLogger(__name__) -def get_standard_config_overrides( - decoupled_lr: float | None = None, decoupled_min_lr: float | None = None -) -> Dict[ParamKey, ParamGroupOverride]: +def get_standard_config_overrides(config: OptimizerConfig) -> Dict[ParamKey, ParamGroupOverride]: """Get standard config overrides for the optimizer, handling decoupled LR and common wd skips. Args: - decoupled_lr (float | None): decoupled learning rate. - decoupled_min_lr (float | None): decoupled minimum learning rate. + config (OptimizerConfig): optimizer configuration object. Returns: Dict[ParamKey, ParamGroupOverride]: standard config overrides. """ config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = {} - if decoupled_lr is not None: - decoupled_lr_config: ParamGroupOverride = {"max_lr": decoupled_lr} - decoupled_param_key = ParamKey(attr="is_embedding_or_output_parameter") - if decoupled_min_lr is not None: - decoupled_lr_config["min_lr"] = decoupled_min_lr - config_overrides[decoupled_param_key] = decoupled_lr_config + # First, figure out how we are going to do wd skipping. The two main approaches are: + # 1. The classic megatron approach of skipping all len 1 and bias parameters. + # 2. The Qwen3-Next approach of doing 1, other than qk layernorm parameters. + if config.apply_wd_to_qk_layernorm: + shape_1_not_qkln_param = ParamWithNamePredicate( + name="s1_not_qkln", + fn=lambda param, name: (len(param.shape) == 1 or name.endswith(".bias")) + and not ("q_layernorm." in name or "k_layernorm." in name), + ) + param_wd_mult_key = ParamKey(with_name_predicate=shape_1_not_qkln_param) + else: + param_length_1_match = ParamPredicate( + name="param_len_1", fn=lambda param: len(param.shape) == 1 + ) + param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match) - # Next construct the standard param group overrides for no weight decay on bias parameters - # as well as any length 1 parameters. - param_length_1_match = ParamPredicate( - name="param_len_1", fn=lambda param: len(param.shape) == 1 - ) - param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match) config_overrides[param_wd_mult_key] = ParamGroupOverride(wd_mult=0.0) + if config.decoupled_lr is not None: + decoupled_lr_config: ParamGroupOverride = {"max_lr": config.decoupled_lr} + decoupled_param_key = ParamKey(attr="is_embedding_or_output_parameter") + if config.decoupled_min_lr is not None: + decoupled_lr_config["min_lr"] = config.decoupled_min_lr + config_overrides[decoupled_param_key] = decoupled_lr_config + return config_overrides @@ -132,7 +140,7 @@ def _get_param_groups( # the config_overrides argument by default lead to bias parameters and length 1 parameters. # We assume that users of decoupled LR already provide config overrides so will adapt # to the new API. - config_overrides = get_standard_config_overrides() + config_overrides = get_standard_config_overrides(config=config) for model_chunk in model_chunks: for name, param in model_chunk.named_parameters(): diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index f57c7aba95c..2d3e3ca08e0 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -33,6 +33,34 @@ def __call__(self, param: torch.nn.Parameter) -> bool: return self.fn(param) +@dataclass(frozen=True) +class ParamWithNamePredicate: + """Wraps a matching function to make it hashable for ParamKey. + Example: + >>> shape_1_not_qkln_param = ParamWithNamePredicate( + name="s1_not_qkln", + fn=lambda param, name: ( + len(param.shape) == 1 or name.endswith(".bias") + and not ("q_layernorm." in name or "k_layernorm." in name) + ) + ) + >>> shape_1_not_qkln_param(torch.empty(10), "interesting.bias") + True + >>> shape_1_not_qkln_param(torch.empty(10), "interesting.q_layernorm.bias") + False + + NOTE: + __hash__ and __eq__ are automatically generated by @dataclass(frozen=True) + based solely on 'name' because we set compare=False/hash=False on 'fn'. + """ + + name: str + fn: Callable[[torch.nn.Parameter, str], bool] = field(compare=False, hash=False) + + def __call__(self, param: torch.nn.Parameter, name: str) -> bool: + return self.fn(param, name) + + @dataclass(frozen=True, slots=True) class ParamKey: """Key to group parameters by. All such grouped parameters can share an @@ -49,6 +77,15 @@ class ParamKey: predicate: Union[ParamPredicate, Tuple[ParamPredicate]] = field(default_factory=tuple) """Predicate(s) to match parameters by. If multiple predicates are provided, any must match.""" + with_name_predicate: Union[ParamWithNamePredicate, Tuple[ParamWithNamePredicate]] = field( + default_factory=tuple + ) + """ + Predicate(s) to match parameters with their name. If multiple predicates are provided, + any must match. This is useful if you need to filter out some parameters from an otherwise + positive match by their name. + """ + def matches(self, param: torch.nn.Parameter, param_name: str) -> bool: """Returns true if passed-in parameter (with name) matches `param_key`. @@ -86,6 +123,15 @@ def matches(self, param: torch.nn.Parameter, param_name: str) -> bool: for predicate in self.predicate: if predicate(param): return True + + # Check if with_name_predicate matches. + if isinstance(self.with_name_predicate, ParamWithNamePredicate): + if self.with_name_predicate(param, param_name): + return True + else: + for predicate in self.with_name_predicate: + if predicate(param, param_name): + return True return False @@ -105,9 +151,20 @@ class OptimizerConfig: min_lr: Optional[float] = None """Minumum value for learning rate. The scheduler clip values below this threshold.""" + decoupled_lr: Optional[float] = None + """Separate learning rate for the input and output layer.""" + + decoupled_min_lr: Optional[float] = None + """Minimum value for learning rate for the input and output layer. The scheduler clip values + below this threshold. + """ + weight_decay: float = 0.01 """Weight decay coefficient for L2 regularization.""" + apply_wd_to_qk_layernorm: bool = False + """If true, apply weight decay to qk layernorm as a special case.""" + ############## # Precision ############## diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 7cb83a868a4..aba961b6b29 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -893,6 +893,17 @@ def validate_args(args, defaults={}): dc = torch.cuda.get_device_capability() assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels." + if args.no_weight_decay_cond_type is not None: + print_rank_0( + 'WARNING: --no-weight-decay-cond-type is deprecated. Please use --apply-wd-to-qk-layernorm instead.', + args.rank, + ) + if args.no_weight_decay_cond_type == "apply_wd_to_qk_layernorm": + args.apply_wd_to_qk_layernorm = True + else: + raise ValueError(f"Invalid no_weight_decay_cond_type: {args.no_weight_decay_cond_type}") + args.no_weight_decay_cond_type = None + if args.weight_decay_incr_style == 'constant': assert args.start_weight_decay is None assert args.end_weight_decay is None @@ -2038,6 +2049,8 @@ def _add_regularization_args(parser): help='Dropout probability for hidden state transformer.') group.add_argument('--weight-decay', type=float, default=0.01, help='Weight decay coefficient for L2 regularization.') + group.add_argument('--apply-wd-to-qk-layernorm', action='store_true', + help='Apply weight decay to qk layernorm as a special case.') group.add_argument('--clip-grad', type=float, default=1.0, help='Gradient clipping based on global L2 norm.') group.add_argument('--adam-beta1', type=float, default=0.9, @@ -2071,6 +2084,12 @@ def _add_regularization_args(parser): help='How to perform NS calculation for tensor model parallel weights') group.add_argument('--muon-extra-scale-factor', type=float, default=1.0, help='Additional scale factor for the muon update') + group.add_argument('--no-weight-decay-cond-type', type=str, choices=['apply_wd_to_qk_layernorm'], + help='Type of no weight decay condition. Choices: ' + 'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ' + 'None (default): apply weight decay to 1D weights and biases.' + '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' + 'qk layernorm as a special case.') return parser diff --git a/megatron/training/training.py b/megatron/training/training.py index b25d802f152..7398edbea3f 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1325,7 +1325,7 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig: # Construct the appropriate config_overrides object. This default handles many cases, but # can be added to as needed by the user, or replaced entirely with a custom override. - config_overrides = get_standard_config_overrides(args.decoupled_lr, args.decoupled_min_lr) + config_overrides = get_standard_config_overrides(config=config) return config, config_overrides diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 4f914b56f7c..1f5bbc3f14c 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -21,6 +21,7 @@ _get_param_groups, check_config_overrides_consistency, get_megatron_optimizer, + get_standard_config_overrides, ) from megatron.core.optimizer_param_scheduler import ParamGroupOverride from megatron.core.process_groups_config import ProcessGroupCollection @@ -45,7 +46,7 @@ class Net(nn.Module): - def __init__(self): + def __init__(self, add_layernorm=False): super().__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) @@ -53,6 +54,10 @@ def __init__(self): self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) + if add_layernorm: + self.q_layernorm = nn.LayerNorm(10, bias=False) + self.k_layernorm = nn.LayerNorm(10, bias=False) + self.layernorm = nn.LayerNorm(10, bias=False) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) @@ -206,6 +211,81 @@ def test_get_param_groups_overlapping_matches(mock_get_world_size): assert param_groups[2]['max_lr'] == 0.01 +@patch('torch.distributed.get_world_size', return_value=1) +@patch( + 'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj) +) +def test_get_param_groups_with_standard_config_overrides(apply_wd_to_qk_layernorm: bool): + """In this test, we see if the standard config overrides are applied correctly.""" + + # Initialize the model with layernorm + net = Net() + + config = OptimizerConfig(optimizer='adam', lr=0.01) + config_overrides = get_standard_config_overrides(config=config) + param_groups = _get_param_groups([net], config, config_overrides) + + assert len(param_groups) == 2 + p_set = set(net.parameters()) + + assert p_set == set(param_groups[0]['params']) | set(param_groups[1]['params']) + assert len(p_set) == len(param_groups[0]['params']) + len(param_groups[1]['params']) + assert param_groups[0]['wd_mult'] == 0.0 or param_groups[1]['wd_mult'] == 0.0 + assert param_groups[0]['wd_mult'] == 1.0 or param_groups[1]['wd_mult'] == 1.0 + assert len(param_groups[0]['params']) > 0 and len(param_groups[1]['params']) > 0 + + # Both param groups should have 5 parameters. + # Param group A (wd_mult=1.0): conv1.weight, conv2.weight, fc1.weight, fc2.weight, fc3.weight + # Param group B (wd_mult=0.0): conv1.bias, conv2.bias, fc1.bias, fc2.bias, fc3.bias + assert len(param_groups[0]['params']) == 5, ( + f"Expected 5 parameters in the first param group, " + f"but got {len(param_groups[0]['params'])}" + ) + assert len(param_groups[1]['params']) == 5, ( + f"Expected 5 parameters in the second param group, " + f"but got {len(param_groups[1]['params'])}" + ) + + +@patch('torch.distributed.get_world_size', return_value=1) +@patch( + 'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj) +) +def test_get_param_groups_appling_wd_to_qk_layernorm(apply_wd_to_qk_layernorm: bool): + """In this test, we see if the `apply_wd_to_qk_layernorm` config is applied correctly.""" + + # Initialize the model with layernorm + net = Net(add_layernorm=True) + + config = OptimizerConfig( + optimizer='adam', lr=0.01, apply_wd_to_qk_layernorm=apply_wd_to_qk_layernorm + ) + config_overrides = get_standard_config_overrides(config=config) + param_groups = _get_param_groups([net], config, config_overrides) + + assert len(param_groups) == 2 + p_set = set(net.parameters()) + + assert p_set == set(param_groups[0]['params']) | set(param_groups[1]['params']) + assert len(p_set) == len(param_groups[0]['params']) + len(param_groups[1]['params']) + assert param_groups[0]['wd_mult'] == 1.0 + assert param_groups[1]['wd_mult'] == 0.0 + + # There are two param groups, having 7, and 6 parameters respectively. + # Param group A (wd_mult=1.0): conv1.weight, conv2.weight, fc1.weight, fc2.weight, fc3.weight, + # q_layernorm.weight, k_layernorm.weight + # Param group B (wd_mult=0.0): conv1.bias, conv2.bias, fc1.bias, fc2.bias, fc3.bias, + # layernorm.weight + assert len(param_groups[0]['params']) == 7, ( + f"Expected 5 parameters in the first param group, " + f"but got {len(param_groups[0]['params'])}" + ) + assert len(param_groups[1]['params']) == 6, ( + f"Expected 6 parameters in the second param group, " + f"but got {len(param_groups[1]['params'])}" + ) + + def test_chained_optimizer(): net = Net() optimizer_1 = Adam(list(net.parameters())[:2], lr=0.01)