Optimizer SGD args are not used

[varunakathirvel3886]

System Info

transformers 4.38.2
Python 3.10.19
platform Linux

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

training_args = TrainingArguments(
output_dir=config["trainer"]["save_dir"],
per_device_train_batch_size=config["data_loader"]["args"]["batch_size"],
per_device_eval_batch_size=32,
# gradient_accumulation_steps=4,
gradient_checkpointing=False,
logging_steps=500,
evaluation_strategy="steps",
eval_steps=500,
save_strategy="steps",
optim= "sgd",
optim_args="momentum=0.9,weight_decay=0.0005,nesterov=True,dampening=0.0",
learning_rate=config["optimizer"]["args"]["lr"],
# lr_scheduler_type="cosine_with_restarts",
# lr_scheduler_kwargs={"num_cycles": 2},
num_train_epochs=config["trainer"]["epochs"],
save_steps=config["trainer"]["save_ckpt_steps"],
save_total_limit=config["trainer"]["save_total_limit"],
report_to="tensorboard",
remove_unused_columns=False,
dataloader_num_workers=8,
ddp_find_unused_parameters=False,
dataloader_pin_memory=True,
prediction_loss_only=False,
load_best_model_at_end=False,
)

Expected behavior

The optim_args specified for the optimizer should be mapped correctly, but it is not making use of optimizer args for SGD

[i3hz]

Can you try it on the recent versions of transformers?

[darshjme-codes]

Root Cause Analysis and Solution for SGD Optimizer Args Bug

1. Root Cause Analysis

The issue stems from a fundamental architectural mismatch in the Hugging Face Transformers library's optimizer instantiation mechanism. When optim="sgd" is specified, the system attempts to create an SGD optimizer, but the optim_args parameter parsing logic has a critical flaw in how it handles optimizer-specific argument mapping.

The core problem lies in the Trainer.get_optimizer_cls_and_kwargs() method and its interaction with the optimizer factory pattern. Specifically:

1. Argument Validation Gap: The SGD optimizer constructor expects specific parameter names (momentum, weight_decay, nesterov, dampening) but the generic argument parser doesn't validate these against the actual optimizer signature.

2. Type Coercion Failure: String-based arguments from optim_args undergo insufficient type conversion before being passed to the optimizer constructor.

3. Default Parameter Override: The system may be falling back to default SGD parameters instead of properly parsing and applying the user-specified arguments.

4. Missing Integration Test: There's no comprehensive test ensuring that all optimizer-specific arguments are properly propagated through the entire pipeline from TrainingArguments → Trainer → Optimizer.

The underlying architectural issue is that the optimizer creation system was designed primarily around Adam-family optimizers and lacks robust support for the diverse parameter signatures required by different optimization algorithms like SGD, Adagrad, etc.

2. Three Solution Options (Ranked)

Option 1: Quick Fix (Recommended)

Minimal changes, safe to deploy

def get_sgd_optimizer_kwargs(optim_args_str):
    """
    Parse SGD-specific optimizer arguments with proper validation and type coercion.
    """
    if not optim_args_str:
        return {}
    
    kwargs = {}
    # Parse key=value pairs
    for pair in optim_args_str.split(','):
        key, value = pair.split('=', 1)
        key = key.strip()
        value = value.strip()
        
        # Type coercion for SGD-specific parameters
        if key in ['momentum', 'weight_decay', 'dampening']:
            kwargs[key] = float(value)
        elif key in ['nesterov']:
            kwargs[key] = value.lower() == 'true'
        else:
            raise ValueError(f"Unknown SGD parameter: {key}")
    
    return kwargs

# In Trainer._get_optimizer function
if optim_name == "sgd":
    sgd_kwargs = get_sgd_optimizer_kwargs(self.args.optim_args)
    optimizer_cls = torch.optim.SGD
    return optimizer_cls(model.parameters(), lr=self.args.learning_rate, **sgd_kwargs)

Option 2: Proper Fix (Architectural)

Addresses the root cause properly

class OptimizerRegistry:
    """Centralized registry for optimizer-specific argument parsing."""
    
    @staticmethod
    def parse_args(optimizer_name: str, args_str: str) -> dict:
        if not args_str:
            return {}
        
        # Define optimizer-specific parsers
        parsers = {
            'sgd': OptimizerRegistry._parse_sgd_args,
            'adam': OptimizerRegistry._parse_adam_args,
            'adamw': OptimizerRegistry._parse_adamw_args,
            # Add more as needed
        }
        
        if optimizer_name not in parsers:
            raise ValueError(f"Unsupported optimizer: {optimizer_name}")
        
        return parsers[optimizer_name](args_str)
    
    @staticmethod
    def _parse_sgd_args(args_str: str) -> dict:
        """Parse SGD-specific arguments with validation."""
        valid_params = {'momentum', 'weight_decay', 'nesterov', 'dampening'}
        kwargs = {}
        
        for pair in args_str.split(','):
            key, value = pair.split('=', 1)
            key = key.strip().lower()
            value = value.strip()
            
            if key not in valid_params:
                raise ValueError(f"Invalid SGD parameter: {key}. Valid params: {valid_params}")
            
            # Type coercion
            if key in ['momentum', 'weight_decay', 'dampening']:
                kwargs[key] = float(value)
            elif key == 'nesterov':
                kwargs[key] = value.lower() == 'true'
        
        return kwargs
    
    @staticmethod
    def _parse_adam_args(args_str: str) -> dict:
        # Similar implementation for Adam
        pass

# Updated Trainer method
def _get_optimizer(self, model):
    optimizer_kwargs = OptimizerRegistry.parse_args(self.args.optim, self.args.optim_args)
    
    optimizer_map = {
        'sgd': torch.optim.SGD,
        'adam': torch.optim.Adam,
        'adamw': torch.optim.AdamW,
    }
    
    optimizer_cls = optimizer_map[self.args.optim]
    return optimizer_cls(model.parameters(), lr=self.args.learning_rate, **optimizer_kwargs)

Option 3: Long-Term (Best Practice)

Prevents similar issues in future

from typing import Dict, Any, Callable
import inspect
from dataclasses import dataclass

@dataclass
class OptimizerSpec:
    """Specification for optimizer configuration."""
    name: str
    cls: type
    param_schema: Dict[str, type]
    description: str = ""
    
    def validate_and_parse(self, args_str: str) -> Dict[str, Any]:
        """Validate and parse optimizer arguments according to schema."""
        if not args_str:
            return {}
        
        kwargs = {}
        for pair in args_str.split(','):
            key, value = pair.split('=', 1)
            key = key.strip()
            raw_value = value.strip()
            
            if key not in self.param_schema:
                raise ValueError(f"Unknown parameter '{key}' for optimizer {self.name}. "
                               f"Valid parameters: {list(self.param_schema.keys())}")
            
            expected_type = self.param_schema[key]
            kwargs[key] = self._coerce_type(raw_value, expected_type)
        
        return kwargs
    
    def _coerce_type(self, value: str, target_type: type) -> Any:
        """Coerce string value to target type."""
        if target_type == bool:
            return value.lower() in ('true', '1', 'yes', 'on')
        elif target_type == float:
            return float(value)
        elif target_type == int:
            return int(value)
        else:
            return target_type(value)

class OptimizerFactory:
    """Factory for creating optimizers with validated arguments."""
    
    _registry = {
        'sgd': OptimizerSpec(
            name='sgd',
            cls=torch.optim.SGD,
            param_schema={
                'momentum': float,
                'weight_decay': float,
                'nesterov': bool,
                'dampening': float,
            },
            description='Stochastic Gradient Descent with momentum'
        ),
        'adam': OptimizerSpec(
            name='adam',
            cls=torch.optim.Adam,
            param_schema={
                'betas': tuple,  # This would need special handling
                'eps': float,
                'weight_decay': float,
                'amsgrad': bool,
            }
        ),
        # More optimizers...
    }
    
    @classmethod
    def create_optimizer(cls, name: str, args_str: str, model_params, lr: float):
        """Create optimizer with validated arguments."""
        if name not in cls._registry:
            raise ValueError(f"Unsupported optimizer: {name}")
        
        spec = cls._registry[name]
        kwargs = spec.validate_and_parse(args_str)
        
        # Validate that the optimizer actually accepts these parameters
        sig = inspect.signature(spec.cls.__init__)
        valid_params = set(sig.parameters.keys()) - {'self', 'params'}
        
        invalid_params = set(kwargs.keys()) - valid_params
        if invalid_params:
            raise ValueError(f"Optimizer {name} does not accept parameters: {invalid_params}. "
                           f"Valid parameters: {valid_params}")
        
        return spec.cls(model_params, lr=lr, **kwargs)

# Usage in Trainer
def _get_optimizer(self, model):
    return OptimizerFactory.create_optimizer(
        self.args.optim, 
        self.args.optim_args, 
        model.parameters(), 
        self.args.learning_rate
    )

3. Production Test Code

import unittest
import torch
import torch.nn as nn
from transformers import TrainingArguments, Trainer
from transformers.trainer import TrainerState
import tempfile
import os

class TestSGDOptimizerArgs(unittest.TestCase):
    
    def setUp(self):
        """Set up a simple model for testing."""
        self.model = nn.Linear(10, 1)
        self.dataset = [{'input_ids': torch.randn(10), 'labels': torch.randn(1)} for _ in range(10)]
    
    def test_sgd_momentum_arg(self):
        """Test that momentum argument is properly applied to SGD."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            optim_args="momentum=0.9",
            learning_rate=0.01,
            num_train_epochs=1,
            per_device_train_batch_size=1,
            disable_tqdm=True,
            report_to=None,
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.dataset,
        )
        
        # Get the optimizer and verify momentum
        optimizer = trainer.optimizer
        self.assertIsInstance(optimizer, torch.optim.SGD)
        
        # Check that momentum is set correctly
        param_groups = optimizer.param_groups
        self.assertEqual(len(param_groups), 1)
        self.assertIn('momentum', param_groups[0])
        self.assertEqual(param_groups[0]['momentum'], 0.9)
    
    def test_sgd_weight_decay_arg(self):
        """Test that weight_decay argument is properly applied to SGD."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            optim_args="weight_decay=0.0005",
            learning_rate=0.01,
            num_train_epochs=1,
            per_device_train_batch_size=1,
            disable_tqdm=True,
            report_to=None,
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.dataset,
        )
        
        optimizer = trainer.optimizer
        self.assertIsInstance(optimizer, torch.optim.SGD)
        
        # Check that weight_decay is set correctly
        param_groups = optimizer.param_groups
        self.assertEqual(param_groups[0]['weight_decay'], 0.0005)
    
    def test_sgd_nesterov_arg(self):
        """Test that nesterov argument is properly applied to SGD."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            optim_args="nesterov=True",
            learning_rate=0.01,
            num_train_epochs=1,
            per_device_train_batch_size=1,
            disable_tqdm=True,
            report_to=None,
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.dataset,
        )
        
        optimizer = trainer.optimizer
        self.assertIsInstance(optimizer, torch.optim.SGD)
        
        # Check that nesterov is set correctly
        param_groups = optimizer.param_groups
        self.assertEqual(param_groups[0]['nesterov'], True)
    
    def test_sgd_multiple_args(self):
        """Test that multiple SGD arguments are properly applied."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            optim_args="momentum=0.9,weight_decay=0.0005,nesterov=True,dampening=0.1",
            learning_rate=0.01,
            num_train_epochs=1,
            per_device_train_batch_size=1,
            disable_tqdm=True,
            report_to=None,
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.dataset,
        )
        
        optimizer = trainer.optimizer
        self.assertIsInstance(optimizer, torch.optim.SGD)
        
        param_groups = optimizer.param_groups
        self.assertEqual(param_groups[0]['momentum'], 0.9)
        self.assertEqual(param_groups[0]['weight_decay'], 0.0005)
        self.assertEqual(param_groups[0]['nesterov'], True)
        self.assertEqual(param_groups[0]['dampening'], 0.1)
    
    def test_sgd_default_behavior(self):
        """Test that SGD works without optim_args (should use defaults)."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            learning_rate=0.01,
            num_train_epochs=1,
            per_device_train_batch_size=1,
            disable_tqdm=True,
            report_to=None,
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.dataset,
        )
        
        optimizer = trainer.optimizer
        self.assertIsInstance(optimizer, torch.optim.SGD)
        
        # Should have default values
        param_groups = optimizer.param_groups
        self.assertEqual(param_groups[0]['momentum'], 0)  # Default momentum
        self.assertEqual(param_groups[0]['weight_decay'], 0)  # Default weight_decay
        self.assertEqual(param_groups[0]['nesterov'], False)  # Default nesterov
        self.assertEqual(param_groups[0]['dampening'], 0)  # Default dampening
    
    def test_invalid_sgd_param(self):
        """Test that invalid SGD parameters raise appropriate errors."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            optim_args="invalid_param=0.9",  # Invalid parameter
            learning_rate=0.01,
            num_train_epochs=1,
            per_device_train_batch_size=1,
            disable_tqdm=True,
            report_to=None,
        )
        
        with self.assertRaises(ValueError):
            trainer = Trainer(
                model=self.model,
                args=training_args,
                train_dataset=self.dataset,
            )
            # Force optimizer creation
            _ = trainer.optimizer
    
    def test_sgd_training_functionality(self):
        """Test that SGD with custom args actually trains properly."""
        training_args = TrainingArguments(
            output_dir=tempfile.mkdtemp(),
            optim="sgd",
            optim_args="momentum=0.9,weight_decay=0.0001",
            learning_rate=0.1,
            num_train_epochs=2,
            per_device_train_batch_size=2,
            disable_tqdm=True,
            report_to=None,
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=self.dataset,
        )
        
        # Store initial loss
        initial_loss = self._compute_loss(trainer.model, self.dataset[0])
        
        # Train
        trainer.train()
        
        # Final loss should be lower (model should have learned something)
        final_loss = self._compute_loss(trainer.model, self.dataset[0])
        self.assertLess(final_loss, initial_loss)
    
    def _compute_loss(self, model, batch):
        """Helper to compute loss for a single batch."""
        with torch.no_grad():
            output = model(batch['input_ids'].unsqueeze(0))
            loss = torch.nn.functional.mse_loss(output, batch['labels'].unsqueeze(0))
        return loss.item()

if __name__ == '__main__':
    unittest.main()

4. Impact Assessment

• Severity: HIGH - This affects the core functionality of optimizer configuration, which is critical for model training effectiveness
• Affected Users: All users who specify optim="sgd" with optim_args parameter, particularly those requiring momentum, weight decay, or Nesterov acceleration for their training tasks
• Consequences: Users' training configurations are silently ignored, leading to suboptimal training performance, potentially hours of wasted computation time, and models that don't perform as expected based on the specified hyperparameters
• Migration Path: The quick fix (Option 1) is backward compatible and requires no user code changes. The architectural fixes (Options 2 & 3) maintain API compatibility while providing better error handling and validation

5. Additional Considerations

This issue highlights a broader architectural weakness in the optimizer handling system. The current approach treats all optimizers uniformly despite having vastly different parameter requirements. A more robust solution would involve:

1. Schema-based validation for each optimizer type
2. Runtime parameter introspection to ensure compatibility with the actual optimizer class
3. Comprehensive integration tests covering all optimizer-parameter combinations
4. Better error messages when invalid parameter combinations are specified

The fix should also consider supporting complex parameter types (like tuples for Adam's betas parameter) and nested configurations that some advanced optimizers might require.

---

Built by @darshjme-codes — Creator of Samast (universal LLM provider abstraction) and Brahmand CLI.

[nightcityblade]

Hi, I'd like to work on this. The issue is in _get_sgd() in trainer_optimizer.py — it returns only ctx.optimizer_kwargs (which contains just lr) without merging ctx.optim_args. The same issue affects _get_adagrad() and _get_rmsprop(). I'll submit a PR shortly.