from enum import Enum
from time import perf_counter
import torch
import numpy as np
from pathlib import Path
from typing import Dict, Iterator, List, NamedTuple, Any, Optional
from experimaestro import (
Task,
Config,
Param,
pathgenerator,
Annotated,
tqdm,
field,
Meta,
)
import lightning as L
from lightning.fabric.strategies.strategy import Strategy as l_Strategy
from xpm_torch import Random
from xpm_torch.configuration import FabricConfiguration
from xpm_torch.metrics import Metrics, ScalarMetric
from .batchers import RecoverableOOMError
from .optim import (
Module,
ModuleLoader,
ParameterOptimizer,
ScheduledOptimizer,
OptimizationHook,
)
from xpm_torch.context import Hook, InitializationHook
from xpm_torch.utils.logging import EasyLogger
from xpm_torch.trainers.context import (
StepTrainingHook,
TrainState,
TrainerContext,
)
from xpm_torch.trainers import Trainer
from xpm_torch.module import fallback_fa2_if_incompatible_precision
import logging
logger = logging.getLogger(__name__)
[docs]
class Strategy(Config, l_Strategy):
"""A Lightning strategy"""
pass
class LearnerListenerStatus(Enum):
NO_DECISION = 0
STOP = 1
DONT_STOP = 2
def update(self, other: "LearnerListenerStatus") -> "LearnerListenerStatus":
return LearnerListenerStatus(max(self.value, other.value))
class UpdateModelHook(InitializationHook):
"""
A generalized hook that executes exactly once after the Learner has wrapped its model.
It traverses a target object tree and replaces all occurrences of the unwrapped model
with the true Fabric-wrapped model. Optionally, it can also inject the Fabric instance.
"""
def __init__(self, target_obj, unwrapped_model, fabric=None):
self.target_obj = target_obj
self.unwrapped_model = unwrapped_model
self.fabric = fabric
def _swap_model(self, obj, wrapped_model, visited=None):
if visited is None:
visited = set()
# Prevent infinite recursion on cyclic references
if id(obj) in visited:
return
visited.add(id(obj))
# 1. Update direct references to the model
for attr in ["scorer", "model", "retriever"]:
if getattr(obj, attr, None) is self.unwrapped_model:
setattr(obj, attr, wrapped_model)
if self.fabric is not None:
# Inject fabric so components like TwoStageRetriever can autocast
try:
setattr(obj, "fabric", self.fabric)
except Exception:
pass
# 2. Recurse down children to support nested component trees
if hasattr(obj, "named_children"):
for _, child in obj.named_children():
self._swap_model(child, wrapped_model, visited)
# Also check experimaestro config values if it's a Config object
if hasattr(obj, "__fields__"):
for field_name in obj.__fields__:
field_val = getattr(obj, field_name, None)
if field_val is not None and not isinstance(field_val, (str, int, float, bool, dict, list)):
self._swap_model(field_val, wrapped_model, visited)
def after(self, context):
# Fetch the definitively wrapped model from the training state
wrapped_model = getattr(context.state, "_fabric_model", getattr(context.state, "model", None))
if wrapped_model is not None and wrapped_model is not self.unwrapped_model:
self._swap_model(self.target_obj, wrapped_model)
[docs]
class CheckpointSettings(Config):
"""Settings for a checkpoint-specific ModuleLoader."""
epoch: Param[Optional[int]] = field(default=None, ignore_default=True)
"""The epoch of the checkpoint"""
[docs]
class LearnerListener(Config):
"""Hook for learner
Performs some operations after a learning epoch"""
id: Meta[str]
"""Unique ID to identify the listener (ignored for signature)"""
def initialize(self, learner: "Learner", context: TrainerContext):
self.learner = learner
self.context = context
def __call__(self, state: TrainState) -> LearnerListenerStatus:
"""Process and returns whether the training process should stop"""
return LearnerListenerStatus.NO_DECISION
def update_metrics(self, metrics: Dict[str, float]):
"""Add metrics"""
pass
def init_task(self, learner: "Learner", dep, add_action):
"""Returns the initialization task that loads the associated checkpoint
:param learner: The learner object
:param dep: The function that adds a dependency
:param add_action: Function to register an action
"""
return None
class LearnerOutput(NamedTuple):
"""The data structure for the output of a learner. It contains a dictionary
where the key is the name of the listener and the value is the output of
that listener. It also allows to access the checkpoints saved during
the training"""
listeners: Dict[str, Any]
learned_model: ModuleLoader
checkpoints: Dict[str, Any]
[docs]
class Learner(Task, EasyLogger):
"""Model Learner
The learner task is generic, and takes two main arguments: (1) the model
defines the model (e.g. DRMM), and (2) the trainer defines how the model
should be trained (e.g. pointwise, pairwise, etc.)
When submitted, it returns a dictionary based on the `listeners`
"""
# Training
random: Param[Random]
"""The random generator"""
trainer: Param[Trainer]
"""Specifies how to train the model"""
model: Param[Module]
"""Defines the model to be learned. If multiple models are used, one can use
MultipleModel.
"""
max_epochs: Param[int] = field(default=1000, ignore_default=True)
"""Maximum number of epochs"""
steps_per_epoch: Param[int] = field(default=128, ignore_default=True)
"""Number of steps for one epoch (after each epoch results are logged)"""
optimizers: Param[List[ParameterOptimizer]]
"""The list of parameter optimizers"""
listeners: Param[List[LearnerListener]]
"""Listeners are in charge of handling the validation of the model, and
saving the relevant checkpoints"""
checkpoint_interval: Meta[int] = field(default=1, ignore_default=True)
"""Number of epochs between each checkpoint"""
logpath: Annotated[Path, pathgenerator("runs")]
"""The path to the tensorboard logs"""
checkpointspath: Annotated[Path, pathgenerator("checkpoints")]
"""The path to the checkpoints"""
hooks: Param[List[Hook]] = field(default=[], ignore_default=True)
"""Global learning hooks
:class:`Initialization hooks <xpm_torch.context.InitializationHook>` are called
before and after the initialization of the trainer and listeners.
"""
fabric_config: Param[FabricConfiguration] = field(default_factory=FabricConfiguration.C)
"""Runtime configuration, managed by Fabric"""
def __validate__(self):
assert self.optimizers, "At least one optimizer should be defined"
assert len(set(listener.id for listener in self.listeners)) == len(
self.listeners
), "IDs of listeners should be unique"
return super().__validate__()
def __submit__(self, dep, add_action):
"""Submit the learner task and register export actions."""
learned_model = dep(
self.model.loader_config(
self.last_checkpoint_path / TrainState.MODEL_DIR
)
)
# Register export action — the model controls the action type
add_action(self.model.export_action(learned_model, default_name="last"))
return LearnerOutput(
listeners={
listener.id: listener.init_task(self, dep, add_action=add_action)
for listener in self.listeners
},
learned_model=learned_model,
checkpoints={
interval: dep(
self.model.loader_config(
TrainerContext.get_checkpoint_path(
self.checkpointspath, interval
)
/ TrainState.MODEL_DIR,
settings=CheckpointSettings.C(epoch=interval),
)
)
for interval in range(0, self.max_epochs, self.checkpoint_interval)
},
)
@property
def last_checkpoint_path(self):
return self.checkpointspath / "last"
def execute(self):
"""Main training loop, executed using the fabric context.
The training process is stopped either by the listeners
or when max_epoch is reached.
"""
# 1. Launch Fabric
fabric = self.fabric_config.get_fabric()
fabric.launch()
self.optimizer = ScheduledOptimizer()
self.only_cached = False
self.context = TrainerContext(
self.logpath,
self.checkpointspath,
self.max_epochs,
self.steps_per_epoch,
self.trainer,
self.model,
self.optimizer,
fabric=fabric,
)
for hook in self.hooks:
self.context.add_hook(hook)
# Call init hooks
for hook in self.context.hooks(InitializationHook):
hook.before(self.context)
# Sets the random seed
# WARNING - will still not be fully deterministic unless using (lot slower):
# - torch.use_deterministic_algorithms(True) (PyTorch ≥1.8).
# - torch.backends.cudnn.deterministic = True
# - torch.backends.cudnn.benchmark = False.
# can also use fabric.seed_everything(self.random.state.randint((2**32) - 1))
seed = self.random.state.randint((2**32) - 1)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# Initialize the model and trainer
with fabric.init_module():
self.trainer.initialize(self.random.state, self.context)
# Wrap dataloader with Fabric for device placement (if using new path)
if hasattr(self.trainer, "dataloader") and self.trainer.dataloader is not None:
self.trainer.dataloader = fabric.setup_dataloaders(
self.trainer.dataloader
)
for listener in self.listeners:
listener.initialize(self, self.context)
num_training_steps = self.max_epochs * self.steps_per_epoch
self.optimizer.initialize(
self.optimizers,
num_training_steps,
self.model,
use_scaler=False,
hooks=[hook for hook in self.hooks if isinstance(hook, OptimizationHook)],
trainer_context=self.context,
)
# Trainer hook for one-time setup that needs the model on the device but
# NOT yet DDP-wrapped (e.g. predictive batcher probing on rank 0). Must
# run after the optimizer is initialized and before fabric.setup(model).
self.trainer.pre_train_setup(fabric)
# Setup model and optimizers with Fabric
self.model = self.model.setup_with_fabric(fabric)
self.logger.info(
f"[Learner] Model setup via setup_with_fabric completed (type: {type(self.model).__name__})"
)
if self.optimizer.optimizers:
self.optimizer.optimizers = list(fabric.setup_optimizers(*self.optimizer.optimizers))
else:
self.logger.error("No optimizers were initialized; training will not update model parameters.")
# Propagate the Fabric-wrapped model back to the state and trainer
# so that training loops automatically use the wrapped forward pass.
# We use a private attribute `_fabric_model` to prevent Experimaestro
# `Param` descriptors from stripping the PyTorch proxy wrapper!
self.context.state._fabric_model = self.model
if self.trainer is not None:
self.trainer._fabric_model = self.model
# Also attempt standard assignment for compatibility
self.context.state.model = self.model
if getattr(self.trainer, "model", None) is not None:
self.trainer.model = self.model
self.logger.info(
f"Model is on device {self.model.device} using dtype {next(self.model.parameters()).dtype}"
)
if torch.cuda.is_available():
# This is the definitive check for BF16 support
supports_bf16 = torch.cuda.is_bf16_supported()
self.logger.info(f"Hardware supports BF16: {supports_bf16}")
if not supports_bf16 and "bf16" in fabric.precision.precision:
self.logger.error(
"CRITICAL: You are forcing BF16 on incompatible hardware!"
)
for hook in self.context.hooks(InitializationHook):
hook.after(self.context)
self.logger.info("Starting to train")
current = 0
state = None
with tqdm(
total=self.max_epochs, desc=f"Training ({self.max_epochs} epochs)"
) as tqdm_epochs:
for state in self.iter_train(fabric):
# Report progress
tqdm_epochs.update(state.epoch - current)
current = state.epoch
if state.epoch >= 0 and not self.only_cached:
message = f"epoch {state.epoch}"
if state.cached:
self.logger.debug(f"[train] [cached] {message}")
else:
self.logger.debug(f"[train] {message}")
if state.epoch == -1:
continue
if not state.cached and state.epoch % self.checkpoint_interval == 0:
# Save checkpoint on rank 0 and barrier all ranks
if fabric.is_global_zero:
self.context.save_checkpoint()
self.context.copy(self.last_checkpoint_path)
fabric.barrier()
# Call listeners
decision = LearnerListenerStatus.NO_DECISION
for listener in self.listeners:
# listener.__call__ returns LearnerListenerStatus
decision = decision.update(listener(state))
# Synchronize decision across DDP ranks so all ranks reach identical stopping decision
if fabric.world_size > 1:
decision_val = torch.tensor([decision.value], device=fabric.device)
fabric.broadcast(decision_val, src=0)
decision = LearnerListenerStatus(decision_val.item())
fabric.barrier()
if decision == LearnerListenerStatus.STOP:
self.logger.warning(
"stopping after epoch {epoch} ({early_stop} epochs) since "
"all listeners asked for it"
)
break
# Stop if max epoch is reached
if self.context.epoch >= self.max_epochs:
self.logger.warning(
"stopping after epoch {max_epochs} (max_epoch)".format(
**self.__dict__
)
)
break
# End of the learning process
if state is not None and not state.cached:
if fabric.is_global_zero:
# Set the hyper-parameters
metrics = {}
for listener in self.listeners:
listener.update_metrics(metrics)
if self.context.writer is not None:
self.context.writer.add_hparams(getattr(self, "__tags__", {}), metrics)
fabric.barrier()
def iter_train(self, fabric: L.Fabric) -> Iterator[TrainState]:
"""Infinite generator of training states: one per epoch, containing self.steps_per_epoch steps
"""
# Try to load a checkpoint
if self.context.load_bestcheckpoint(self.max_epochs):
yield self.context.state
# Get an iterator over batches
batch_iter = self.trainer.iter_batches()
while True:
# Step to the next epoch
self.context.nextepoch()
# Train for an epoch
with tqdm(
leave=False,
total=self.steps_per_epoch,
ncols=100,
desc=f"Train - epoch #{self.context.epoch}",
) as pbar:
# Put the model into training mode (just in case)
self.context.state.model.train()
# Epoch: loop over batches
metrics = Metrics()
start = perf_counter()
for b in range(self.steps_per_epoch):
# Get the next batch, recreate iterator on exhaustion
try:
batch = next(batch_iter)
except StopIteration:
batch_iter = self.trainer.iter_batches()
batch = next(batch_iter)
self.context.nextbatch()
while True:
try:
# Computes the gradient, takes a step and collect metrics
with self.context.step(metrics):
# Call epoch hooks
for hook in self.context.hooks(StepTrainingHook):
hook.before(self.context)
# Computes the gradient
self.trainer.process_batch(batch)
# Update metrics and counter
pbar.update(1)
break
except RecoverableOOMError:
logger.warning(
"Recoverable OOM detected"
" - re-running the training step"
)
continue
for hook in self.context.hooks(StepTrainingHook):
hook.after(self.context)
metrics.add(
ScalarMetric("iter_per_seconds", self.steps_per_epoch / (perf_counter() - start), 1)
)
# Yields the current state (after one epoch)
# -> allows listeners to process it and decide whether to stop or not
yield self.context.state
# Report metrics over the epoch, and log them in tensorboard
# Note that this is done after the listeners are called, so that they can update the metrics if needed (e.g., with validation results)
metrics.report(
self.context.state.step,
self.context.writer,
"train",
)