Scheduler Factory

The SchedulerFactory is responsible for generating the scheduler used in training from the optimizer created by the arctic_training.optimizer.factory.OptimizerFactory. A Scheduler Factory can be created by inheriting from the SchedulerFactory class and implementing the create_scheduler() method.

class arctic_training.scheduler.factory.SchedulerFactory(trainer, scheduler_config=None)[source]

Bases: ABC, CallbackMixin

Base class for all scheduler factories.

Parameters:
name: str

The name of the scheduler factory. This is used to identify the scheduler factory in the registry.

config: SchedulerConfig

The configuration class for the scheduler factory. This is used to validate the configuration passed to the factory.

property trainer: Trainer
property device: str
property world_size: int
property global_rank: int
property optimizer: Any
abstract create_scheduler(optimizer)[source]

Create a scheduler from the optimizer.

Return type:

Any

Parameters:

optimizer (Any)

Attributes

Similar to other Factory classes in ArcticTraining, the SchedulerFactory class must have a name attribute that is used to identify the factory when registering it with ArcticTraining and a config attribute type hint that is used to validate the config object passed to the factory.

Properties

A SchedulerFactory has several attributes that can be used to access information about the trainer and distributed state at runtime, including device, trainer, optimizer, world_size, and global_rank.

Methods

SchedulerFactories have just one method that must be defined: create_scheduler(). This method should return the scheduler object created using the optimizer object passed to the factory.

Huggingface Scheduler Factory

A custom scheduler factory can be created from the SchedulerFactory, but ArcticTraining comes with a HFSchedulerFactory implementation that can be used out of the box.