easydel.trainers.trainer_protocol

class easydel.trainers.trainer_protocol.BaseProgressBar

Bases: ABC

Abstract base class for progress bar implementations.

abstract close() → None

abstract reset() → None

abstract set_postfix(**kwargs) → None

abstract update(n: int = 1) → None

class easydel.trainers.trainer_protocol.BaseTrainerProtocol(arguments: TrainingArguments, model: EasyDeLBaseModule, dataset_train: Optional[Any] = None, dataset_eval: Optional[Any] = None, finetune: bool = True, checkpoint_path: Optional[Union[str, PathLike]] = None)

Bases: object

abstract apply_training_hooks(metrics: LossMetrics) → LossMetrics

Apply training hooks to the model.

arguments: TrainingArguments

checkpoint_manager: Any

checkpoint_path: Optional[Union[str, PathLike]]

abstract compile_aot() → bool

Compile the state ahead of time for faster execution.

config: EasyDeLBaseConfig

abstract configure_dataloaders() → TrainerConfigureDataloaderOutput

Configures the dataloaders for training and evaluation.

abstract configure_functions() → TrainerConfigureFunctionOutput

Configures and JIT-compiles the training and evaluation step functions.

abstract configure_model() → TrainerConfigureModelOutput

Configures the model, optimizer, scheduler, and configuration.

abstract static count_model_parameters(prm)

Prints the number of model parameters in billions.

abstract create_collect_function(max_sequence_length: int, truncation_mode: Literal['keep_end', 'keep_start']) → Callable

Creates a function to collect and process batches of data for training or evaluation.

abstract create_progress_bar(desc: str = '', disabled: bool = False) → BaseProgressBar

Create a progress bar of the specified type.

data_collator: Optional[Callable]

dataloader_eval: Optional[Iterator[ndarray]]

dataloader_train: Iterator[ndarray]

dataset_eval: Optional[Any]

dataset_train: Optional[Any]

dtype: Any

abstract eval(model_state: EasyDeLState) → Iterator[dict]

Evaluates using the provided model state.

evalu_tracker: CompilationTracker

abstract property evaluation_batch_size

finetune: bool

abstract finish()

Finalize the training process.

abstract get_runstage_flops(is_training: bool) → float

Return the total number of FLOPs for the model.

abstract initialize_trainer_utils()

Initializes all trainer utilities.

abstract log_metrics(metrics: Dict[str, float], pbar: BaseProgressBar, step: int, mode: str = 'train') → None

Log metrics and update progress bar.

abstract log_weight_distribution(state: EasyDeLState, step: int)

Log distribution of weights.

max_evaluation_steps: int

max_training_steps: int

memory_monitor: Any

abstract property mesh

abstract property model

model_state: EasyDeLState

abstract on_step_end(state: EasyDeLState, metrics: Any, step: int) → Tuple[EasyDeLState, Any]

hook process to call in start of the step.

abstract on_step_start(state: EasyDeLState, step: int) → EasyDeLState

hook process to call in start of the step.

param_dtype: Any

pruning_module: Any

abstract save_information(output_path: Union[str, Path]) → None

Save the generated information to a markdown file.

abstract save_pretrained(state: EasyDeLState, save_directory: Optional[str] = None, gather_fns: Optional[Union[Any, Mapping[str, Callable], dict[Callable]]] = None, to_torch: bool = False, base_hf_auto_class=None, easystate_to_huggingface_model_kwargs: Optional[dict] = None, torch_save_pretrained_kwargs: Optional[dict] = None)

Saves the model state as a checkpoint file or to a Torch compatible directory.

scheduler: Callable[[Union[Array, ndarray, bool, number, float, int]], Union[Array, ndarray, bool, number, float, int]]

sharded_evaluation_step_function: Callable

sharded_training_step_function: Callable

abstract specs_to_name_sharding(tree, mesh=None)

Convert specs to named sharding.

abstract start_evaluation_hook()

Hook to run before evaluation starts.

abstract start_training_hook()

Hook to run before training starts.

state: Any

state_named_sharding: Any

state_partition_spec: Any

state_shape: Any

timer: Timers

abstract train(model_parameters: Optional[FrozenDict] = None, state: Optional[EasyDeLState] = None) → Any

Train using the provided model state.

train_tracker: CompilationTracker

abstract property training_batch_size

tx: GradientTransformation

wandb_runtime: Any

class easydel.trainers.trainer_protocol.JSONProgressBar(desc='')

Bases: BaseProgressBar

Wrapper for JSON

close() → None

reset() → None

set_postfix(**kwargs) → None

update(n: int = 1) → None

class easydel.trainers.trainer_protocol.MetricsColumn(metrics_to_show=None)

Bases: ProgressColumn

A custom progress column for displaying metrics.

render(task: Task) → Text

Render the metrics in a organized way.

class easydel.trainers.trainer_protocol.MetricsTracker

Bases: object

Tracks and aggregates training metrics over time.

reset(step)

Reset tracked metrics.

update(loss, accuracy, step)

Update tracked metrics with new values.

class easydel.trainers.trainer_protocol.NullProgressBar

Bases: BaseProgressBar

Dummy progress bar that does nothing - useful for multiprocessing.

close() → None

reset() → None

set_postfix(**kwargs) → None

update(n: int = 1) → None

class easydel.trainers.trainer_protocol.RichProgressBar(progress: Progress, task_id: TaskID)

Bases: BaseProgressBar

Wrapper for rich progress bar.

close() → None

reset() → None

set_postfix(**kwargs) → None

update(n: int = 1) → None

class easydel.trainers.trainer_protocol.StepMetrics(arguments)

Bases: object

Handles calculation and tracking of training metrics.

calculate(metrics: LossMetrics, current_step: int, epoch: int, flops: float, batch_size: int, seq_length: int, learning_rate: float, mode: Optional[Literal['eval', 'train']] = None, **extras) → Dict[str, float]

Calculate comprehensive metrics for the training step.

start_step()

Mark the start of a training step.

class easydel.trainers.trainer_protocol.TqdmProgressBar(pbar: tqdm)

Bases: BaseProgressBar

Wrapper for tqdm progress bar.

close() → None

reset() → None

set_postfix(**kwargs) → None

update(n: int = 1) → None

class easydel.trainers.trainer_protocol.TrainerConfigureDataloaderOutput(dataloader_train: 'tp.Iterator[np.ndarray]', max_training_steps: 'int', dataloader_eval: 'tp.Optional[tp.Iterator[np.ndarray]]' = None, max_evaluation_steps: 'tp.Optional[int]' = None)

Bases: object

dataloader_eval: Optional[Iterator[ndarray]] = None

dataloader_train: Iterator[ndarray]

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

max_evaluation_steps: Optional[int] = None

max_training_steps: int

replace(**kwargs)

Creates a new instance with specified fields replaced.

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

class easydel.trainers.trainer_protocol.TrainerConfigureFunctionOutput(sharded_training_step_function: 'JitWrapped', mesh: 'Mesh', checkpoint_manager: 'CheckpointManager', sharded_evaluation_step_function: 'tp.Optional[JitWrapped]' = None)

Bases: object

checkpoint_manager: CheckpointManager

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

mesh: Mesh

replace(**kwargs)

Creates a new instance with specified fields replaced.

sharded_evaluation_step_function: Optional[Callable] = None

sharded_training_step_function: Callable

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

class easydel.trainers.trainer_protocol.TrainerConfigureModelOutput(model: 'EasyDeLBaseModule', tx: 'GradientTransformation', scheduler: 'Schedule', config: 'tp.Optional[EasyDeLBaseConfig]' = None)

Bases: object

config: Optional[EasyDeLBaseConfig] = None

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

model: EasyDeLBaseModule

replace(**kwargs)

Creates a new instance with specified fields replaced.

scheduler: Callable[[Union[Array, ndarray, bool, number, float, int]], Union[Array, ndarray, bool, number, float, int]]

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

tx: GradientTransformation

class easydel.trainers.trainer_protocol.TrainerOutput(state: 'EasyDeLState', mesh: 'tp.Optional[jax.sharding.Mesh]', last_save_file_name: 'tp.Optional[str]' = None, checkpoint_path: 'tp.Optional[str]' = None)

Bases: object

checkpoint_path: Optional[str] = None

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

last_save_file_name: Optional[str] = None

mesh: Optional[Mesh]

replace(**kwargs)

Creates a new instance with specified fields replaced.

state: EasyDeLState

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.