How to Configure Weight Decay with Parameter Exclusions in Nanotron

Nanotron enables global weight decay with selective exclusions for specific parameter types via the weight_decay_exclude_named_params field, automatically partitioning model parameters into optimizer groups that apply zero decay to excluded names while maintaining the global rate for others.

Configuring weight decay with parameter exclusions in Nanotron is critical for modern transformer training, where standard practice dictates that bias vectors and normalization parameters should not be regularized. The Hugging Face Nanotron library simplifies this through declarative configuration options and internal helper utilities that automatically construct the appropriate optimizer parameter groups without requiring manual group management in your training scripts.

Configuration Options for Weight Decay Exclusions

The foundation of weight decay configuration resides in the OptimizerArgs class within src/nanotron/config/config.py. This configuration dataclass exposes two key fields that control regularization behavior:

• weight_decay: The global decay coefficient applied to all parameters by default
• weight_decay_exclude_named_params: A list of name substrings that trigger zero decay for matching parameters

When you define these fields in your training configuration, Nanotron automatically handles the complexity of parameter group construction during trainer initialization.

The Three-Step Process for Building Parameter Groups

Nanotron implements weight decay exclusion through a systematic three-step pipeline defined in src/nanotron/helpers.py. This process transforms simple configuration declarations into valid PyTorch optimizer parameter groups.

Step 1: Define Global Decay and Exclusions

First, specify your regularization strategy in the configuration file or object. The weight_decay_exclude_named_params field accepts a list of substrings; any parameter whose name contains at least one of these substrings receives zero weight decay.

from nanotron.config import OptimizerArgs

optimizer_args = OptimizerArgs(
    lr=6e-4,
    weight_decay=0.01,
    weight_decay_exclude_named_params=["bias", "LayerNorm.weight", "norm"]
)

Step 2: Generate Custom Weight Decay Groups

The helper function get_custom_weight_decay_for_named_parameters in src/nanotron/helpers.py (lines 186–229) iterates over all model parameters and partitions them based on the exclusion list. It returns a list of dictionaries, each containing a named_params list and its corresponding weight_decay value (either the global value or 0.0).

from nanotron.helpers import get_custom_weight_decay_for_named_parameters

wd_groups = get_custom_weight_decay_for_named_parameters(
    named_params=list(model.named_parameters()),
    weight_decay=0.01,
    exclude_named_params=["bias", "LayerNorm.weight"]
)

This function ensures that excluded parameters are isolated into separate groups where weight_decay=0.0, while standard parameters remain in groups with the global decay value.

Step 3: Merge with Learning Rate Groups

Since optimizer groups must specify both learning rate and weight decay for every parameter, Nanotron uses merge_named_param_groups (lines 291–311 in src/nanotron/helpers.py) to combine the weight decay groups with learning rate groups. This utility creates the final parameter groups that include both lr and weight_decay for each subset of parameters.

The NanotronTrainer class in src/nanotron/trainer.py orchestrates this entire workflow internally, ensuring that by the time the optimizer is instantiated, every parameter belongs to exactly one group with the correct hyperparameters.

Practical Implementation Examples

You can configure weight decay exclusions through declarative YAML files or programmatically for advanced use cases.

YAML Configuration Approach

The simplest method uses the standard Nanotron configuration file. Define your exclusions in the optimizer_args section:

optimizer_args:
  lr: 6e-4
  weight_decay: 0.01
  weight_decay_exclude_named_params:
    - "bias"
    - "LayerNorm.weight"
    - "ln"

When you initialize NanotronTrainer with this configuration, the trainer automatically applies zero weight decay to all bias terms and LayerNorm weights without additional code.

Advanced Programmatic Usage

For custom training loops outside the standard trainer, import the helper functions directly to construct optimizer groups:

from nanotron.helpers import (
    get_custom_weight_decay_for_named_parameters,
    merge_named_param_groups,
    get_named_param_groups_with_lr,
)
import torch

named_params = list(model.named_parameters())

# 1. Create learning rate groups

lr_groups = [{"named_params": named_params, "lr": 6e-4}]

# 2. Create weight decay groups with exclusions

wd_groups = get_custom_weight_decay_for_named_parameters(
    named_params=named_params,
    weight_decay=0.01,
    exclude_named_params=["bias", "LayerNorm.weight"],
)

# 3. Merge groups: each entry now has both 'lr' and 'weight_decay'

optimizer_groups = merge_named_param_groups(lr_groups, wd_groups)

# 4. Initialize optimizer with the merged groups

optimizer = torch.optim.AdamW(optimizer_groups)

The resulting optimizer_groups list contains entries structured like:

{
    "named_params": [("layers.0.attention.bias", Tensor(...))],
    "lr": 6e-4,
    "weight_decay": 0.0  # Excluded from decay

}

End-to-End Training Integration

For standard supervised fine-tuning or pre-training workflows, no manual group construction is necessary:

from nanotron.config import Config
from nanotron.trainer import NanotronTrainer

# Load configuration containing optimizer_args with weight_decay_exclude_named_params

config = Config.from_yaml("config.yaml")

# Initialize trainer - handles all parameter group construction internally

trainer = NanotronTrainer(config=config, model=model, data_collator=collator)
trainer.train()

The trainer internally calls get_custom_weight_decay_for_named_parameters and merge_named_param_groups during setup, applying your exclusion rules across all training steps.

Summary

• Set global decay using optimizer_args.weight_decay to define the default regularization strength for all parameters
• Exclude specific parameters by providing a list of name substrings in weight_decay_exclude_named_params; matches receive zero decay automatically
• Automatic group construction occurs via get_custom_weight_decay_for_named_parameters in src/nanotron/helpers.py, which partitions parameters based on exclusion rules
• Group merging happens through merge_named_param_groups, ensuring every optimizer group contains both learning rate and weight decay specifications
• Broad compatibility with PyTorch optimizers including AdamW, since Nanotron generates standard parameter group dictionaries recognized by torch.optim classes

Frequently Asked Questions

Which parameters should typically be excluded from weight decay in Nanotron?

Bias terms and normalization layer weights (such as LayerNorm, GroupNorm, or RMSNorm parameters) are standard exclusions in transformer training. These parameters benefit from different optimization dynamics, and applying weight decay to them often degrades model performance. The Nanotron configuration accepts any substring match, so you can exclude "bias", "LayerNorm.weight", "norm", or specific module prefixes as needed.

How does Nanotron match parameter names for exclusion?

The get_custom_weight_decay_for_named_parameters function uses simple substring matching against the full parameter name. If a parameter name contains any string in the exclude_named_params list, it is assigned to a group with weight_decay=0.0. This approach matches PyTorch naming conventions, allowing you to target broad categories (like all "bias" parameters) or specific layers (like "transformer.h.0.attn.bias") using partial string matches.

Can weight decay exclusions coexist with learning rate scheduling?

Yes. The merge_named_param_groups function explicitly combines learning rate groups with weight decay groups, ensuring that every final optimizer group contains both lr and weight_decay fields. This merged structure is fully compatible with PyTorch learning rate schedulers, which adjust the learning rate across all groups while preserving the fixed weight decay values assigned during initialization.

Where is the weight decay exclusion logic tested in the Nanotron repository?

The test suite in tests/test_optimizer_params_groups.py validates that weight decay exclusions are applied correctly. These tests verify that parameters matching exclusion substrings receive weight_decay=0.0, that non-matching parameters retain the global decay value, and that the merging logic correctly combines LR and weight decay specifications without parameter duplication or omission.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}