expert_usage

expert_usage

Expert usage logging callback for mixture-of-experts models.

ExpertUsageCallback

ExpertUsageCallback(log_every_n_epochs=1)

Bases: Callback

Callback to log expert usage statistics for mixture-of-experts models.

Logs expert usage ratios as bar plots to WandB at the end of each epoch.

Initialize expert usage callback.

Parameters:

Name	Type	Description	Default
log_every_n_epochs	int	Log expert usage every N epochs	1

Source code in src/versatil/training/callbacks/expert_usage.py

def __init__(self, log_every_n_epochs: int = 1):
    """Initialize expert usage callback.

    Args:
        log_every_n_epochs: Log expert usage every N epochs
    """
    super().__init__()
    self.log_every_n_epochs = log_every_n_epochs

on_train_epoch_end

on_train_epoch_end(trainer, pl_module)

Log training expert usage at end of epoch.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/expert_usage.py

def on_train_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Log training expert usage at end of epoch.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if trainer.current_epoch % self.log_every_n_epochs != 0:
        return

    expert_usages = pl_module.train_metrics.compute_expert_usage()
    if expert_usages is not None:
        for key, expert_usage in expert_usages.items():
            fig = self._create_expert_usage_figure(expert_usage, f"Train {key}")
            if trainer.logger is not None:
                wandb_image = figure_to_wandb_image(fig)
                trainer.logger.log_metrics(
                    {f"train_{key}": wandb_image},
                    step=trainer.current_epoch,
                )
            plt.close(fig)

on_validation_epoch_end

on_validation_epoch_end(trainer, pl_module)

Log validation expert usage picture at end of epoch.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/expert_usage.py

def on_validation_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Log validation expert usage picture at end of epoch.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if trainer.sanity_checking:
        return
    if trainer.current_epoch % self.log_every_n_epochs != 0:
        return
    expert_usages = pl_module.val_metrics.compute_expert_usage()
    if expert_usages is not None:
        for key, expert_usage in expert_usages.items():
            fig = self._create_expert_usage_figure(expert_usage, f"Val {key}")
            if trainer.logger is not None:
                wandb_image = figure_to_wandb_image(fig)
                trainer.logger.log_metrics(
                    {f"val_{key}": wandb_image},
                    step=trainer.current_epoch,
                )
            plt.close(fig)