fishy.experiments.deep_training

Training orchestrator for the deep learning pipeline.

Classes

class fishy.experiments.deep_training.MixedOptimizer(optimizers: List[Optimizer])

Bases: Optimizer

Wrapper that combines Muon for matrices and another optimizer for vectors.

__init__(optimizers: List[Optimizer])

step(closure=None)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

zero_grad(set_to_none: bool = True)

Reset the gradients of all optimized torch.Tensor s.

Parameters:

set_to_none (bool, optional) –

Instead of setting to zero, set the grads to None. Default: True

This will in general have lower memory footprint, and can modestly improve performance. However, it changes certain behaviors. For example:

1. When the user tries to access a gradient and perform manual ops on it, a None attribute or a Tensor full of 0s will behave differently.

2. If the user requests zero_grad(set_to_none=True) followed by a backward pass, .grads are guaranteed to be None for params that did not receive a gradient.

3. torch.optim optimizers have a different behavior if the gradient is 0 or None (in one case it does the step with a gradient of 0 and in the other it skips the step altogether).

class fishy.experiments.deep_training.ModelTrainer(config: TrainingConfig, wandb_run: Any | None = None, ctx: RunContext | None = None)

Bases: object

Orchestrates deep learning experiments.

__init__(config: TrainingConfig, wandb_run: Any | None = None, ctx: RunContext | None = None)

create_optimizer(model: Module) → Optimizer

Create the requested optimizer(s) for the model.

create_scheduler(optimizer: Optimizer, train_loader: DataLoader) → Tuple[LRScheduler | None, bool]

Create a learning rate scheduler with warmup if configured.

pre_train() → Module | None

train(pre_trained_model: Module | None = None) → Tuple[Module, Dict[str, Any]]

Functions

fishy.experiments.deep_training.run_training_pipeline(config: TrainingConfig, wandb_run: Any | None = None, ctx: RunContext | None = None) → Dict[str, Any]

s