SmolFactory / docs /TRACKIO_UPDATE_FIX.md
Tonic's picture
adds update attribute for trl compatibility bug fix
5fe0328 verified
|
raw
history blame
4.53 kB

TrackioConfig Update Method Fix

Problem Description

The error 'TrackioConfig' object has no attribute 'update' occurred because the TRL library (specifically SFTTrainer) expects the Trackio configuration object to have an update method, but our custom TrackioConfig class didn't implement it.

Additionally, TRL calls the update method with keyword arguments like allow_val_change, which our initial implementation didn't support.

Root Cause

Based on the Trackio documentation, Trackio is designed to be API compatible with wandb.init, wandb.log, and wandb.finish. However, the TRL library has additional expectations for the configuration object, including an update method that allows dynamic configuration updates with both dictionary and keyword arguments.

Solution Implementation

1. Enhanced Update Method for TrackioConfig

Modified src/trackio.py to add a flexible update method that handles both dictionary and keyword arguments:

class TrackioConfig:
    """Configuration class for trackio (TRL compatibility)"""
    
    def __init__(self):
        self.project_name = os.environ.get('EXPERIMENT_NAME', 'smollm3_experiment')
        self.experiment_name = os.environ.get('EXPERIMENT_NAME', 'smollm3_experiment')
        self.trackio_url = os.environ.get('TRACKIO_URL')
        self.trackio_token = os.environ.get('TRACKIO_TOKEN')
        self.hf_token = os.environ.get('HF_TOKEN')
        self.dataset_repo = os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
    
    def update(self, config_dict: Dict[str, Any] = None, **kwargs):
        """
        Update configuration with new values (TRL compatibility)
        
        Args:
            config_dict: Dictionary of configuration values to update (optional)
            **kwargs: Additional configuration values to update
        """
        # Handle both dictionary and keyword arguments
        if config_dict is not None:
            for key, value in config_dict.items():
                if hasattr(self, key):
                    setattr(self, key, value)
                else:
                    # Add new attributes dynamically
                    setattr(self, key, value)
        
        # Handle keyword arguments
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                # Add new attributes dynamically
                setattr(self, key, value)

2. Key Features of the Enhanced Fix

  • Flexible Argument Handling: Supports both dictionary and keyword arguments
  • TRL Compatibility: Handles TRL's allow_val_change and other keyword arguments
  • Dynamic Attribute Updates: Can update existing attributes and add new ones dynamically
  • Backward Compatibility: Doesn't break existing functionality
  • Future-Proof: Supports additional TRL requirements

3. Usage Examples

Dictionary-based updates:

import trackio

config = trackio.config
config.update({
    'project_name': 'my_experiment',
    'experiment_name': 'test_run_1',
    'custom_setting': 'value'
})

Keyword argument updates (TRL style):

config.update(allow_val_change=True, project_name="test_project")

Mixed updates:

config.update({'experiment_name': 'test'}, allow_val_change=True, new_attr='value')

Verification

The enhanced fix has been verified to work correctly:

  1. Import Test: import trackio works without errors
  2. Config Access: trackio.config is available
  3. Update Method: trackio.config.update() method exists and works
  4. Keyword Arguments: Handles TRL's allow_val_change and other kwargs
  5. TRL Compatibility: All TRL-expected methods are available

Benefits

  1. Resolves Training Error: Fixes both 'TrackioConfig' object has no attribute 'update' and 'TrackioConfig.update() got an unexpected keyword argument 'allow_val_change' errors
  2. Maintains TRL Compatibility: Ensures SFTTrainer can use Trackio for logging with any argument style
  3. Dynamic Configuration: Allows runtime configuration updates via multiple methods
  4. Future-Proof: Supports additional TRL requirements and argument patterns

Related Documentation