Training module – unified trainers, results management, metrics, and callbacks.

Sub-modules

neuralops.trainer.base_trainer

neuralops.trainer.callbacks

neuralops.trainer.custom_loop_trainer

High-performance custom training loop trainer with DDP gradient accumulation synchronization safety.

neuralops.trainer.hf_trainer

neuralops.trainer.hyperopt

neuralops.trainer.losses

neuralops.trainer.methods

Training method plugins: LoRA, DoRA, DisLoRA, Fisher, GeLoRA, LLRD, Mixout, R‑Drop, R3F. All methods are registered via the METHOD_REGISTRY and can be …

neuralops.trainer.metrics

neuralops.trainer.optimizers

Optimizers package for NeuralOps.

neuralops.trainer.registry

neuralops.trainer.results_manager

neuralops.trainer.rl_trainer

RLTrainer – Reinforcement Learning and Preference Optimization Trainer. Supports TRL (Transformer Reinforcement Learning: PPO, DPO, SFT) when …

neuralops.trainer.schedulers

neuralops.trainer.train_lra

Train small Transformer from scratch on LRA tasks, with optional module injection. LRA requires training from scratch, NOT fine-tuning.

neuralops.trainer.trainer

SVDAwareTrainer – custom HuggingFace Trainer with SVD compression callbacks.

neuralops.trainer.trainer_factory

neuralops.trainer.trainer_svd

neuralops.trainer.unsloth_trainer

UnslothTrainer – Accelerated Fast-LoRA and Memory-Efficient Training Engine. Integrates Unsloth 2x faster cross-entropy and fast Triton kernels when …

neuralops.trainer.utils

High-Performance Model Training, Factorized Gradient Clipping & Evaluation Utilities.

Functions

def aggregate_results(results_list) ‑> Any

Expand source code

def aggregate_results(results_list) -> Any:
    """Aggregate results."""
    if not results_list:
        return {}
    aggregated = {}
    for key in results_list[0].keys():
        values = [r[key] for r in results_list if key in r]
        if values:
            aggregated[f"{key}_mean"] = np.mean(values)
            aggregated[f"{key}_std"] = np.std(values)
            aggregated[f"{key}_values"] = values
    return aggregated

Aggregate results.

def benchmark_model(model: nn.Module, data_loader: DataLoader, device: torch.device, warmup: int = 2) ‑> dict[str, float]

Expand source code

def benchmark_model(
    model: nn.Module, data_loader: DataLoader, device: torch.device, warmup: int = 2
) -> dict[str, float]:
    """Benchmark model forward pass latency and throughput."""
    model.eval()
    for i, batch in enumerate(data_loader):
        if i >= warmup:
            break
        if isinstance(batch, dict) and "input_ids" in batch:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch.get("attention_mask", torch.ones_like(input_ids)).to(device)
            with torch.no_grad():
                _ = model(input_ids=input_ids, attention_mask=attention_mask)

    memory_before = psutil.Process().memory_info().rss / 1024 / 1024
    start_time = time.time()
    batch_times: list[float] = []

    with torch.no_grad():
        for batch in data_loader:
            b_start = time.time()
            if isinstance(batch, dict) and "input_ids" in batch:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch.get("attention_mask", torch.ones_like(input_ids)).to(device)
                _ = model(input_ids=input_ids, attention_mask=attention_mask)
            if device.type == "cuda":
                torch.cuda.synchronize()
            batch_times.append(time.time() - b_start)

    total_time = time.time() - start_time
    memory_after = psutil.Process().memory_info().rss / 1024 / 1024
    num_samples = (
        len(data_loader.dataset)
        if hasattr(data_loader, "dataset")
        else len(batch_times) * data_loader.batch_size
    )

    return {
        "total_time": total_time,
        "avg_batch_time": float(np.mean(batch_times)) if batch_times else 0.0,
        "std_batch_time": float(np.std(batch_times)) if batch_times else 0.0,
        "samples_per_second": (float(num_samples / total_time) if total_time > 0 else 0.0),
        "memory_usage_mb": float(memory_after - memory_before),
        "peak_memory_mb": float(memory_after),
        "throughput": float(num_samples / total_time) if total_time > 0 else 0.0,
    }

Benchmark model forward pass latency and throughput.

def compute_accuracy(predictions, labels) ‑> Any

Expand source code

def compute_accuracy(predictions, labels) -> Any:
    """Compute and return accuracy."""
    return accuracy_score(labels, predictions)

Compute and return accuracy.

def compute_f1(predictions, labels, average='weighted') ‑> Any

Expand source code

def compute_f1(predictions, labels, average="weighted") -> Any:
    """Compute and return f1."""
    return f1_score(labels, predictions, average=average, zero_division=0)

Compute and return f1.

def compute_matthews_corrcoef(predictions, labels) ‑> Any

Expand source code

def compute_matthews_corrcoef(predictions, labels) -> Any:
    """Compute and return matthews corrcoef."""
    return matthews_corrcoef(labels, predictions)

Compute and return matthews corrcoef.

def compute_metrics(eval_pred: tuple[np.ndarray, np.ndarray], task: str = 'classification') ‑> dict[str, float]

Expand source code

def compute_metrics(
    eval_pred: tuple[np.ndarray, np.ndarray], task: str = "classification"
) -> dict[str, float]:
    """Compute and return metrics."""
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    metrics = {"accuracy": float(accuracy_score(labels, preds))}
    if task == "classification":
        metrics["f1_macro"] = float(f1_score(labels, preds, average="macro", zero_division=0))
    if len(np.unique(labels)) == 2:
        metrics["matthews_corrcoef"] = float(matthews_corrcoef(labels, preds))
    return metrics

Compute and return metrics.

def compute_pearson_spearman(predictions, labels) ‑> Any

Expand source code

def compute_pearson_spearman(predictions, labels) -> Any:
    """Compute and return pearson spearman."""
    pred_flat = np.squeeze(predictions)
    labels_flat = np.squeeze(labels)
    p = pearsonr(pred_flat, labels_flat)[0]
    s = spearmanr(pred_flat, labels_flat)[0]
    return {"pearson": p, "spearman": s, "pearson_spearman": (p + s) / 2}

Compute and return pearson spearman.

def count_parameters(model: torch.nn.modules.module.Module, trainable_only: bool = True) ‑> int

Expand source code

def count_parameters(model: torch.nn.Module, trainable_only: bool = True) -> int:
    """Count parameters."""
    if trainable_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())

Count parameters.

def create_adabelief(params: Any, lr: float = 0.001, weight_decay: float = 0.0, **kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("adabelief")
def create_adabelief(
    params: Any, lr: float = 1e-3, weight_decay: float = 0.0, **kwargs: Any
) -> Any:
    """Build and return AdaBelief."""
    from adabelief_pytorch import AdaBelief

    return AdaBelief(params, lr=lr, weight_decay=weight_decay, **kwargs)

Build and return AdaBelief.

def create_adam(params: Any, lr: float = 0.001, weight_decay: float = 0.0, **kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("adam")
def create_adam(params: Any, lr: float = 1e-3, weight_decay: float = 0.0, **kwargs: Any) -> Any:
    """Build and return Adam."""
    return optim.Adam(params, lr=lr, weight_decay=weight_decay, **kwargs)

Build and return Adam.

def create_adamw(params: Any, lr: float = 0.001, weight_decay: float = 0.01, **kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("adamw")
def create_adamw(params: Any, lr: float = 1e-3, weight_decay: float = 0.01, **kwargs: Any) -> Any:
    """Build and return AdamW."""
    return optim.AdamW(params, lr=lr, weight_decay=weight_decay, **kwargs)

Build and return AdamW.

def create_cosine(optimizer, T_max=10000, **kwargs) ‑> Any

Expand source code

@SCHEDULER_REGISTRY.register("cosine")
def create_cosine(optimizer, T_max=10000, **kwargs) -> Any:
    """Build and return cosine."""
    return lr_scheduler.CosineAnnealingLR(optimizer, T_max=T_max)

Build and return cosine.

def create_lamb(params: Any, lr: float = 0.001, weight_decay: float = 0.01, **kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("lamb")
def create_lamb(params: Any, lr: float = 1e-3, weight_decay: float = 0.01, **kwargs: Any) -> Any:
    """Build and return Lamb."""
    from torch_optimizer import Lamb

    return Lamb(params, lr=lr, weight_decay=weight_decay, **kwargs)

Build and return Lamb.

def create_lars(params: Any, lr: float = 0.001, weight_decay: float = 0.0, **kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("lars")
def create_lars(params: Any, lr: float = 1e-3, weight_decay: float = 0.0, **kwargs: Any) -> Any:
    """Build and return Lars."""
    from torch_optimizer import Lars

    return Lars(params, lr=lr, weight_decay=weight_decay, **kwargs)

Build and return Lars.

def create_linear(optimizer, num_warmup_steps=0, num_training_steps=10000, **kwargs) ‑> Any

Expand source code

@SCHEDULER_REGISTRY.register("linear")
def create_linear(optimizer, num_warmup_steps=0, num_training_steps=10000, **kwargs) -> Any:
    """Build and return linear."""
    return lr_scheduler.LinearLR(
        optimizer, start_factor=1.0, end_factor=0.0, total_iters=num_training_steps
    )

Build and return linear.

def create_muon(params: Any,
lr: float = 0.02,
momentum: float = 0.95,
weight_decay: float = 0.01,
**kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("muon")
def create_muon(
    params: Any,
    lr: float = 0.02,
    momentum: float = 0.95,
    weight_decay: float = 0.01,
    **kwargs: Any,
) -> Any:
    """Build and return Muon optimizer."""
    return Muon(params, lr=lr, momentum=momentum, weight_decay=weight_decay, **kwargs)

Build and return Muon optimizer.

def create_muon_adamw_reg(params: Any,
lr: float = 0.02,
lr_adam: float = 5e-05,
weight_decay: float = 0.01,
**kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("muon_adamw")
def create_muon_adamw_reg(
    params: Any,
    lr: float = 0.02,
    lr_adam: float = 5e-5,
    weight_decay: float = 0.01,
    **kwargs: Any,
) -> Any:
    """Build and return MuonAdamW hybrid optimizer."""
    return create_muon_adamw_optimizer(
        params, lr_muon=lr, lr_adam=lr_adam, weight_decay=weight_decay, **kwargs
    )

Build and return MuonAdamW hybrid optimizer.

def create_onecycle(optimizer, max_lr=0.01, steps_per_epoch=1000, epochs=10, **kwargs) ‑> Any

Expand source code

@SCHEDULER_REGISTRY.register("one_cycle")
def create_onecycle(optimizer, max_lr=0.01, steps_per_epoch=1000, epochs=10, **kwargs) -> Any:
    """Build and return onecycle."""
    return lr_scheduler.OneCycleLR(
        optimizer, max_lr=max_lr, steps_per_epoch=steps_per_epoch, epochs=epochs
    )

Build and return onecycle.

def create_poly(optimizer, total_iters=10000, power=1.0, **kwargs) ‑> Any

Expand source code

@SCHEDULER_REGISTRY.register("polynomial")
def create_poly(optimizer, total_iters=10000, power=1.0, **kwargs) -> Any:
    """Build and return poly."""
    return lr_scheduler.PolynomialLR(optimizer, total_iters=total_iters, power=power)

Build and return poly.

def create_ranger(params: Any, lr: float = 0.001, weight_decay: float = 0.0, **kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("ranger")
def create_ranger(params: Any, lr: float = 1e-3, weight_decay: float = 0.0, **kwargs: Any) -> Any:
    """Build and return Ranger."""
    from ranger import Ranger

    return Ranger(params, lr=lr, weight_decay=weight_decay, **kwargs)

Build and return Ranger.

def create_riemannian_stiefel(params: Any,
lr: float = 0.001,
lr_stiefel: float | None = None,
base_optimizer: str = 'adamw',
weight_decay: float = 0.01,
**kwargs: Any) ‑> RiemannianStiefelOptimizer

Expand source code

@OPTIMIZER_REGISTRY.register("riemannian_stiefel")
@OPTIMIZER_REGISTRY.register("stiefel")
@OPTIMIZER_REGISTRY.register("stiefel_optimizer")
def create_riemannian_stiefel(
    params: Any,
    lr: float = 1e-3,
    lr_stiefel: float | None = None,
    base_optimizer: str = "adamw",
    weight_decay: float = 0.01,
    **kwargs: Any,
) -> RiemannianStiefelOptimizer:
    """Build and return RiemannianStiefelOptimizer."""
    return RiemannianStiefelOptimizer(
        params,
        lr=lr,
        lr_stiefel=lr_stiefel,
        base_optimizer=base_optimizer,
        weight_decay=weight_decay,
        **kwargs,
    )

Build and return RiemannianStiefelOptimizer.

def create_rop(optimizer, factor=0.1, patience=10, **kwargs) ‑> Any

Expand source code

@SCHEDULER_REGISTRY.register("reduce_on_plateau")
def create_rop(optimizer, factor=0.1, patience=10, **kwargs) -> Any:
    """Build and return rop."""
    return lr_scheduler.ReduceLROnPlateau(optimizer, factor=factor, patience=patience)

Build and return rop.

def create_sgd(params: Any,
lr: float = 0.01,
momentum: float = 0.9,
weight_decay: float = 0.0,
**kwargs: Any) ‑> Any

Expand source code

@OPTIMIZER_REGISTRY.register("sgd")
def create_sgd(
    params: Any,
    lr: float = 0.01,
    momentum: float = 0.9,
    weight_decay: float = 0.0,
    **kwargs: Any,
) -> Any:
    """Build and return SGD."""
    return optim.SGD(params, lr=lr, momentum=momentum, weight_decay=weight_decay, **kwargs)

Build and return SGD.

def evaluate_model(model: nn.Module,
data_loader: DataLoader,
device: torch.device,
num_labels: int,
is_multichoice: bool = False,
compute_all_metrics: bool = False) ‑> float | dict[str, float]

Expand source code

def evaluate_model(
    model: nn.Module,
    data_loader: DataLoader,
    device: torch.device,
    num_labels: int,
    is_multichoice: bool = False,
    compute_all_metrics: bool = False,
) -> float | dict[str, float]:
    """Evaluate model with non-blocking GPU execution and memory safety."""
    cleanup_gpu_memory(force=False)
    model.eval()
    all_preds, all_labels = [], []

    with torch.no_grad():
        for batch in data_loader:
            if is_multichoice:
                input_ids = batch["input_ids"].to(device, non_blocking=True)
                attention_mask = batch["attention_mask"].to(device, non_blocking=True)
                true_labels = batch["labels"]
                bs, nc, sl = input_ids.shape
                flat_ids = input_ids.view(-1, sl)
                flat_mask = attention_mask.view(-1, sl)
                outputs = model(input_ids=flat_ids, attention_mask=flat_mask)
                logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
                avg_logits = logits.view(bs, nc, -1).mean(dim=1)
                preds = torch.argmax(avg_logits, dim=-1).cpu().numpy()
                all_preds.extend(preds)
                all_labels.extend(true_labels.numpy())
            else:
                input_ids = batch["input_ids"].to(device, non_blocking=True)
                attention_mask = batch["attention_mask"].to(device, non_blocking=True)
                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
                preds = torch.argmax(logits, dim=-1).cpu().numpy()
                all_preds.extend(preds)
                all_labels.extend(batch["labels"].numpy())
            del outputs

    cleanup_gpu_memory(force=False)
    if not all_preds:
        return 0.0 if not compute_all_metrics else {"accuracy": 0.0}

    acc = float(accuracy_score(all_labels, all_preds))
    if not compute_all_metrics:
        return acc

    return {
        "accuracy": acc,
        "f1_macro": float(f1_score(all_labels, all_preds, average="macro", zero_division=0)),
        "f1_weighted": float(f1_score(all_labels, all_preds, average="weighted", zero_division=0)),
        "matthews_corrcoef": float(matthews_corrcoef(all_labels, all_preds)),
    }

Evaluate model with non-blocking GPU execution and memory safety.

def get_task_metrics(task_name: str) ‑> Any

Expand source code

def get_task_metrics(task_name: str) -> Any:
    """Return a compute_metrics function for a GLUE task."""
    task_name = task_name.lower()
    if task_name == "stsb":

        def fn(p) -> Any:
            """Fn."""
            preds = np.squeeze(p.predictions)
            labels = p.label_ids
            return compute_pearson_spearman(preds, labels)

    elif task_name == "cola":

        def fn(p) -> Any:
            """Fn."""
            preds = np.argmax(p.predictions, axis=1)
            labels = p.label_ids
            return {
                "matthews_correlation": matthews_corrcoef(labels, preds),
                "accuracy": accuracy_score(labels, preds),
            }

    elif task_name in ("mrpc", "qqp"):

        def fn(p) -> Any:
            """Fn."""
            preds = np.argmax(p.predictions, axis=1)
            labels = p.label_ids
            acc = accuracy_score(labels, preds)
            f1 = f1_score(labels, preds)
            return {"accuracy": acc, "f1": f1, "accuracy_f1": (acc + f1) / 2}

    else:

        def fn(p) -> Any:
            """Fn."""
            preds = np.argmax(p.predictions, axis=1)
            labels = p.label_ids
            return {
                "accuracy": accuracy_score(labels, preds),
                "f1_macro": f1_score(labels, preds, average="macro", zero_division=0),
                "f1_weighted": f1_score(labels, preds, average="weighted", zero_division=0),
            }

    return fn

Return a compute_metrics function for a GLUE task.

def get_trainer(key: str = 'hf') ‑> Any

Expand source code

def get_trainer(key: str = "hf") -> Any:
    """Return a trainer class (not instance)."""
    trainer_cls = TRAINER_FACTORY_REGISTRY.get(key)
    if trainer_cls is None:
        return HuggingFaceTrainer  # default fallback
    return trainer_cls

Return a trainer class (not instance).

def plot_training_curves(history: list[dict[str, Any]], save_path: str) ‑> None

Expand source code

def plot_training_curves(history: list[dict[str, Any]], save_path: str) -> None:
    """Plot training curves."""
    import matplotlib.pyplot as plt

    losses = [h.get("loss") for h in history if "loss" in h]
    eval_losses = [h.get("eval_loss") for h in history if "eval_loss" in h]
    eval_acc = [h.get("eval_accuracy") for h in history if "eval_accuracy" in h]
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    if losses:
        axes[0].plot(range(len(losses)), losses, "b-", label="Train Loss")
    if eval_losses:
        axes[0].plot(range(len(eval_losses)), eval_losses, "r-", label="Val Loss")
    axes[0].set_xlabel("Step")
    axes[0].set_ylabel("Loss")
    axes[0].set_title("Loss Curves")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    if eval_acc:
        axes[1].plot(range(len(eval_acc)), eval_acc, "g-", marker="o", label="Val Accuracy")
    axes[1].set_xlabel("Epoch")
    axes[1].set_ylabel("Accuracy")
    axes[1].set_title("Accuracy Curves")
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)

Plot training curves.

def project_stiefel_gradient(W: torch.Tensor, G: torch.Tensor) ‑> torch.Tensor

Expand source code

def project_stiefel_gradient(W: torch.Tensor, G: torch.Tensor) -> torch.Tensor:
    """Project Euclidean gradient G onto the tangent space of the Stiefel manifold St(p, n).

    Args:
        W: Current orthogonal parameter tensor.
        G: Euclidean gradient tensor (same shape as W).

    Returns:
        Riemannian gradient in tangent space T_W St(p, n).
    """
    if W.dim() == 2:
        n, p = W.shape
        if n >= p:
            wt_g = torch.matmul(W.t(), G)
            sym = 0.5 * (wt_g + wt_g.t())
            grad_r = G - torch.matmul(W, sym)
        else:
            g_wt = torch.matmul(G, W.t())
            sym = 0.5 * (g_wt + g_wt.t())
            grad_r = G - torch.matmul(sym, W)
        return grad_r
    elif W.dim() > 2:
        n, p = W.shape[-2], W.shape[-1]
        if n >= p:
            wt = W.transpose(-2, -1)
            gt = G.transpose(-2, -1)
            wt_g = torch.matmul(wt, G)
            gt_w = torch.matmul(gt, W)
            sym = 0.5 * (wt_g + gt_w)
            grad_r = G - torch.matmul(W, sym)
        else:
            wt = W.transpose(-2, -1)
            gt = G.transpose(-2, -1)
            g_wt = torch.matmul(G, wt)
            w_gt = torch.matmul(W, gt)
            sym = 0.5 * (g_wt + w_gt)
            grad_r = G - torch.matmul(sym, W)
        return grad_r
    else:
        return G

Project Euclidean gradient G onto the tangent space of the Stiefel manifold St(p, n).

Args

-----=

W

Current orthogonal parameter tensor.

G

Euclidean gradient tensor (same shape as W).

Returns -----= Riemannian gradient in tangent space T_W St(p, n).

def quick_optimize(train_fn, eval_fn, n_trials=10, study_name='quick_opt') ‑> dict[str, typing.Any]

Expand source code

def quick_optimize(train_fn, eval_fn, n_trials=10, study_name="quick_opt") -> dict[str, Any]:
    """Quick optimize."""
    config_template = {
        "learning_rate": 2e-5,
        "batch_size": 16,
        "rank": 128,
        "weight_decay": 0.01,
        "warmup_ratio": 0.1,
    }
    optimizer = HyperparameterOptimizer(train_fn, eval_fn, config_template, study_name)
    return optimizer.run(n_trials=n_trials)

Quick optimize.

def save_training_history(history: list[dict[str, Any]], save_path: str) ‑> None

Expand source code

def save_training_history(history: list[dict[str, Any]], save_path: str) -> None:
    """Persist training history to disk."""
    Path(save_path).parent.mkdir(parents=True, exist_ok=True)
    with open(save_path, "w") as f:
        json.dump(history, f, indent=2)

Persist training history to disk.

def stiefel_qr_retraction(W: torch.Tensor, grad_R: torch.Tensor, lr: float) ‑> torch.Tensor

Expand source code

def stiefel_qr_retraction(W: torch.Tensor, grad_R: torch.Tensor, lr: float) -> torch.Tensor:
    """Apply QR retraction step to project back onto the Stiefel manifold.

    Args:
        W: Current orthogonal parameter tensor.
        grad_R: Riemannian gradient in tangent space.
        lr: Learning rate step size along tangent vector.

    Returns:
        Retracted orthogonal parameter tensor on Stiefel manifold.
    """
    w_step = W - lr * grad_R

    if W.dim() == 2:
        n, p = W.shape
        if n >= p:
            q, r = torch.linalg.qr(w_step, mode="reduced")
            d = torch.diagonal(r, dim1=-2, dim2=-1)
            ph = torch.sign(d)
            ph = torch.where(ph == 0, torch.ones_like(ph), ph)
            w_next = q * ph.unsqueeze(-2)
        else:
            q, r = torch.linalg.qr(w_step.t(), mode="reduced")
            d = torch.diagonal(r, dim1=-2, dim2=-1)
            ph = torch.sign(d)
            ph = torch.where(ph == 0, torch.ones_like(ph), ph)
            w_next = (q * ph.unsqueeze(-2)).t()
        return w_next
    elif W.dim() > 2:
        n, p = W.shape[-2], W.shape[-1]
        if n >= p:
            q, r = torch.linalg.qr(w_step, mode="reduced")
            d = torch.diagonal(r, dim1=-2, dim2=-1)
            ph = torch.sign(d)
            ph = torch.where(ph == 0, torch.ones_like(ph), ph)
            w_next = q * ph.unsqueeze(-2)
        else:
            w_step_t = w_step.transpose(-2, -1)
            q, r = torch.linalg.qr(w_step_t, mode="reduced")
            d = torch.diagonal(r, dim1=-2, dim2=-1)
            ph = torch.sign(d)
            ph = torch.where(ph == 0, torch.ones_like(ph), ph)
            w_next = (q * ph.unsqueeze(-2)).transpose(-2, -1)
        return w_next
    else:
        norm = torch.linalg.norm(w_step)
        return w_step / (norm + 1e-12)

Apply QR retraction step to project back onto the Stiefel manifold.

Args

-----=

W

Current orthogonal parameter tensor.

grad_R

Riemannian gradient in tangent space.

lr

Learning rate step size along tangent vector.

Returns -----= Retracted orthogonal parameter tensor on Stiefel manifold.

def train_model(model: nn.Module,
train_loader: DataLoader,
val_loader: DataLoader,
device: torch.device,
config: Any,
num_labels: int,
is_multichoice: bool = False) ‑> dict[str, float]

Expand source code

def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    device: torch.device,
    config: Any,
    num_labels: int,
    is_multichoice: bool = False,
) -> dict[str, float]:
    """
    Train model with maximum GPU compute utilization, AMP precision, DDP gradient accumulation,
    and adaptive factorized gradient clipping safety.
    """
    cleanup_gpu_memory(force=True)
    model.to(device)
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = True
    model.train()

    use_amp = bool(getattr(config, "fp16", False) or getattr(config, "amp", True))
    amp_dtype = (
        torch.bfloat16
        if (device.type == "cuda" and torch.cuda.is_bf16_supported())
        else (torch.float16 if device.type == "cuda" else torch.float32)
    )
    scaler = torch.cuda.amp.GradScaler(
        enabled=(use_amp and device.type == "cuda" and amp_dtype == torch.float16)
    )

    # Optimizer with fused AdamW option if available on CUDA
    optimizer_kwargs = {
        "lr": getattr(config, "learning_rate", 2e-5),
        "weight_decay": getattr(config, "weight_decay", 0.01),
    }
    if device.type == "cuda" and hasattr(torch.optim.AdamW, "fused"):
        optimizer_kwargs["fused"] = True
    optimizer = torch.optim.AdamW(model.parameters(), **optimizer_kwargs)

    epochs = getattr(config, "epochs", 3)
    max_grad_norm = getattr(config, "max_grad_norm", 1.0)
    gradient_accumulation_steps = max(1, int(getattr(config, "gradient_accumulation_steps", 1)))
    steps_per_epoch = max(1, len(train_loader))
    total_loss = 0.0

    for epoch in range(epochs):
        epoch_loss = 0.0
        optimizer.zero_grad(set_to_none=True)
        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}", leave=True)

        for step, batch in enumerate(pbar):
            # FIX: Determine whether this micro-batch is the synchronization step of the accumulation window
            is_sync_step = ((step + 1) % gradient_accumulation_steps == 0) or (
                step + 1 == len(train_loader)
            )

            # FIX: Wrap non-final micro-batches in model.no_sync() for DDP to eliminate redundant all-reduce synchronization
            sync_context = (
                contextlib.nullcontext()
                if (is_sync_step or not hasattr(model, "no_sync") or not callable(model.no_sync))
                else model.no_sync()
            )

            with sync_context:
                with torch.autocast(
                    device_type=device.type,
                    dtype=amp_dtype,
                    enabled=(use_amp and device.type in ["cuda", "mps"]),
                ):
                    if is_multichoice:
                        input_ids = batch["input_ids"].to(device, non_blocking=True)
                        attention_mask = batch["attention_mask"].to(device, non_blocking=True)
                        true_labels = batch["labels"].to(device, non_blocking=True)
                        bs, nc, sl = input_ids.shape
                        flat_ids = input_ids.view(-1, sl)
                        flat_mask = attention_mask.view(-1, sl)
                        outputs = model(input_ids=flat_ids, attention_mask=flat_mask)
                        logits = (
                            outputs["logits"]
                            if isinstance(outputs, dict)
                            else (outputs.logits if hasattr(outputs, "logits") else outputs)
                        )
                        combined_labels = true_labels.repeat_interleave(nc)
                        loss = nn.functional.cross_entropy(logits, combined_labels)
                    else:
                        input_ids = batch["input_ids"].to(device, non_blocking=True)
                        attention_mask = batch["attention_mask"].to(device, non_blocking=True)
                        labels = batch["labels"].to(device, non_blocking=True)
                        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                        logits = (
                            outputs["logits"]
                            if isinstance(outputs, dict)
                            else (outputs.logits if hasattr(outputs, "logits") else outputs)
                        )
                        loss = nn.functional.cross_entropy(logits, labels)

                    # FIX: Normalize loss for accumulation steps
                    if gradient_accumulation_steps > 1:
                        loss = loss / gradient_accumulation_steps

                if torch.isfinite(loss):
                    if scaler.is_enabled():
                        scaler.scale(loss).backward()
                    else:
                        loss.backward()

                    curr_loss = loss.detach().item() * gradient_accumulation_steps
                    epoch_loss += curr_loss
                    total_loss += curr_loss
                else:
                    curr_loss = 0.0

            # FIX: Gradient clipping and optimizer step on accumulation boundaries
            if is_sync_step:
                if scaler.is_enabled():
                    try:
                        scaler.unscale_(optimizer)
                    except RuntimeError:
                        pass
                    clip_grad_norm_factorized(model, max_norm=max_grad_norm)
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    clip_grad_norm_factorized(model, max_norm=max_grad_norm)
                    optimizer.step()
                optimizer.zero_grad(set_to_none=True)

            if (step + 1) % max(1, getattr(config, "logging_steps", 10)) == 0 or step == len(
                train_loader
            ) - 1:
                pbar.set_postfix({"loss": f"{curr_loss:.4f}"})

            del outputs, logits, loss

        cleanup_gpu_memory(force=False)

    val_acc = evaluate_model(model, val_loader, device, num_labels, is_multichoice)
    cleanup_gpu_memory(force=True)
    return {
        "final_accuracy": (val_acc if isinstance(val_acc, float) else val_acc.get("accuracy", 0.0)),
        "avg_loss": total_loss / (epochs * steps_per_epoch),
    }

Train model with maximum GPU compute utilization, AMP precision, DDP gradient accumulation, and adaptive factorized gradient clipping safety.

Classes

class BaseTrainer

Expand source code

class BaseTrainer(ABC):
    """Abstract base class for all trainers."""

    @abstractmethod
    def train(self) -> dict[str, Any]:
        """Fit the model on the provided dataset."""
        pass

    @abstractmethod
    def evaluate(self) -> dict[str, Any]:
        """Run evaluate and return results."""
        pass

Abstract base class for all trainers.

Ancestors

• abc.ABC

Subclasses

• CustomLoopTrainer
• HuggingFaceTrainer
• RLTrainer
• UnslothTrainer

Methods

def evaluate(self) ‑> dict[str, typing.Any]

Expand source code

@abstractmethod
def evaluate(self) -> dict[str, Any]:
    """Run evaluate and return results."""
    pass

Run evaluate and return results.

def train(self) ‑> dict[str, typing.Any]

Expand source code

@abstractmethod
def train(self) -> dict[str, Any]:
    """Fit the model on the provided dataset."""
    pass

Fit the model on the provided dataset.

class ContrastiveLoss (margin=1.0, **kwargs)

Expand source code

@LOSS_REGISTRY.register("contrastive")
class ContrastiveLoss(nn.Module):
    def __init__(self, margin=1.0, **kwargs) -> None:
        super().__init__()
        self.margin = margin

    def forward(self, embeddings, labels) -> Any:
        """Forward pass — compute module output from inputs."""
        dist = F.pairwise_distance(embeddings[0], embeddings[1])
        loss = (1 - labels) * 0.5 * dist.pow(2) + labels * 0.5 * torch.clamp(
            self.margin - dist, min=0
        ).pow(2)
        return loss.mean()

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, embeddings, labels) ‑> Any

Expand source code

def forward(self, embeddings, labels) -> Any:
    """Forward pass — compute module output from inputs."""
    dist = F.pairwise_distance(embeddings[0], embeddings[1])
    loss = (1 - labels) * 0.5 * dist.pow(2) + labels * 0.5 * torch.clamp(
        self.margin - dist, min=0
    ).pow(2)
    return loss.mean()

Forward pass — compute module output from inputs.

class CrossEntropyLoss (weight=None, ignore_index=-100, **kwargs)

Expand source code

@LOSS_REGISTRY.register("cross_entropy")
class CrossEntropyLoss(nn.Module):
    def __init__(self, weight=None, ignore_index=-100, **kwargs) -> None:
        super().__init__()
        self.loss = nn.CrossEntropyLoss(weight=weight, ignore_index=ignore_index)

    def forward(self, logits, labels) -> Any:
        """Forward pass — compute module output from inputs."""
        return self.loss(logits, labels)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, logits, labels) ‑> Any

Expand source code

def forward(self, logits, labels) -> Any:
    """Forward pass — compute module output from inputs."""
    return self.loss(logits, labels)

Forward pass — compute module output from inputs.

class CustomLoopTrainer (model: nn.Module,
train_dataset: Any,
eval_dataset: Any,
config: dict[str, Any],
results_manager: Any | None = None,
device: torch.device | None = None,
num_labels: int = 2,
is_multichoice: bool = False,
compute_metrics_fn: Callable | None = None,
**kwargs: Any)

Expand source code

@TrainerRegistry.register("custom_loop")
class CustomLoopTrainer(BaseTrainer):
    """
    Custom loop trainer optimized for maximum GPU compute throughput, mixed-precision training,
    and distributed data parallel (DDP) gradient accumulation without redundant all-reduce overhead.
    """

    def __init__(
        self,
        model: nn.Module,
        train_dataset: Any,
        eval_dataset: Any,
        config: dict[str, Any],
        results_manager: Any | None = None,
        device: torch.device | None = None,
        num_labels: int = 2,
        is_multichoice: bool = False,
        compute_metrics_fn: Callable | None = None,
        **kwargs: Any,
    ) -> None:
        self.model = model
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.config = config or {}
        self.results_manager = results_manager
        self.device = device or (
            torch.device(
                "cuda"
                if torch.cuda.is_available()
                else (
                    "mps"
                    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
                    else "cpu"
                )
            )
        )
        self.num_labels = num_labels
        self.is_multichoice = is_multichoice
        self.compute_metrics_fn = compute_metrics_fn or get_task_metrics(
            self.config.get("task_name", "mrpc")
        )
        self.kwargs = kwargs

    def _make_loader(self, dataset: Any, shuffle: bool = False) -> DataLoader:
        """Make DataLoader optimized for high-throughput GPU training."""
        pin = self.device.type in ["cuda", "mps"]
        num_workers = int(self.config.get("num_workers", 4 if self.device.type == "cuda" else 0))
        batch_size = int(self.config.get("batch_size", 128))
        loader_kwargs = {
            "batch_size": batch_size,
            "shuffle": shuffle,
            "pin_memory": pin,
            "num_workers": num_workers,
            "collate_fn": (dataset.collate_fn if hasattr(dataset, "collate_fn") else None),
        }
        if num_workers > 0:
            loader_kwargs["prefetch_factor"] = 2
        if torch.cuda.is_available():
            torch.backends.cudnn.benchmark = True
        return DataLoader(dataset, **loader_kwargs)

    def train_epoch(
        self,
        train_loader: DataLoader,
        optimizer: torch.optim.Optimizer,
        scaler: torch.cuda.amp.GradScaler | None,
        gradient_accumulation_steps: int = 1,
        max_grad_norm: float = 1.0,
        amp_dtype: torch.dtype = torch.float32,
        use_amp: bool = False,
        pbar: Any | None = None,
    ) -> float:
        """
        Execute a single training epoch with DDP gradient accumulation synchronization optimization.

        # FIX: Non-final micro-batches are wrapped in model.no_sync() or contextlib.nullcontext()
        # to eliminate redundant all-reduce synchronization across distributed workers during gradient accumulation.
        """
        self.model.train()
        total_epoch_loss = 0.0
        num_batches = len(train_loader)
        accum_steps = max(1, gradient_accumulation_steps)

        for step, batch in enumerate(train_loader):
            # FIX: Determine whether this micro-batch is the final step of the accumulation window or epoch
            is_sync_step = ((step + 1) % accum_steps == 0) or (step + 1 == num_batches)

            # FIX: Wrap non-final micro-batches in model.no_sync() for DDP to eliminate redundant all-reduce synchronization
            sync_context = (
                contextlib.nullcontext()
                if (
                    is_sync_step
                    or not hasattr(self.model, "no_sync")
                    or not callable(self.model.no_sync)
                )
                else self.model.no_sync()
            )

            with sync_context:
                with torch.autocast(
                    device_type=self.device.type,
                    dtype=amp_dtype,
                    enabled=(use_amp and self.device.type in ["cuda", "mps"]),
                ):
                    if self.is_multichoice:
                        input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                        attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                        true_labels = batch["labels"].to(self.device, non_blocking=True)
                        bs, nc, sl = input_ids.shape
                        flat_ids = input_ids.view(-1, sl)
                        flat_mask = attention_mask.view(-1, sl)
                        outputs = self.model(input_ids=flat_ids, attention_mask=flat_mask)
                        logits = (
                            outputs["logits"]
                            if isinstance(outputs, dict)
                            else (outputs.logits if hasattr(outputs, "logits") else outputs)
                        )
                        combined_labels = true_labels.repeat_interleave(nc)
                        loss = nn.functional.cross_entropy(logits, combined_labels)
                    else:
                        input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                        attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                        labels = batch["labels"].to(self.device, non_blocking=True)
                        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                        logits = (
                            outputs["logits"]
                            if isinstance(outputs, dict)
                            else (outputs.logits if hasattr(outputs, "logits") else outputs)
                        )
                        loss = nn.functional.cross_entropy(logits, labels)

                    # FIX: Scale loss by accumulation factor so gradient magnitudes match the full effective batch size
                    if accum_steps > 1:
                        loss = loss / accum_steps

                if torch.isfinite(loss):
                    if scaler is not None and scaler.is_enabled():
                        scaler.scale(loss).backward()
                    else:
                        loss.backward()

                    curr_loss = loss.detach().item() * accum_steps
                    total_epoch_loss += curr_loss
                else:
                    curr_loss = 0.0

                # FIX: Only unscale, clip gradients, and step optimizer when synchronized on accumulation boundary
                if scaler is not None and scaler.is_enabled():
                    try:
                        scaler.unscale_(optimizer)
                    except RuntimeError:
                        pass
                    clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                    optimizer.step()
                optimizer.zero_grad(set_to_none=True)

            if pbar is not None and (
                (step + 1) % max(1, int(self.config.get("logging_steps", 10))) == 0
                or step == num_batches - 1
            ):
                pbar.set_postfix({"loss": f"{curr_loss:.4f}"})

        return total_epoch_loss / max(1, num_batches)

    def train(self) -> dict[str, Any]:
        """Fit the model on the provided dataset with full DDP and accumulation safety."""
        train_loader = self._make_loader(self.train_dataset, shuffle=True)
        val_loader = self._make_loader(self.eval_dataset, shuffle=False)

        class AttrDict:
            pass

        cfg = AttrDict()
        cfg.learning_rate = float(self.config.get("learning_rate", 2e-5))
        cfg.epochs = int(self.config.get("epochs", 3))
        cfg.max_grad_norm = float(self.config.get("max_grad_norm", 1.0))
        cfg.gradient_accumulation_steps = int(self.config.get("gradient_accumulation_steps", 1))
        cfg.fp16 = bool(self.config.get("fp16", False))
        cfg.amp = bool(self.config.get("amp", True))
        cfg.logging_steps = int(self.config.get("logging_steps", 10))

        result = train_model(
            self.model,
            train_loader,
            val_loader,
            self.device,
            cfg,
            self.num_labels,
            self.is_multichoice,
        )

        if self.results_manager is not None:
            if hasattr(self.results_manager, "save_final_model"):
                self.results_manager.save_final_model(self.model)
            if hasattr(self.results_manager, "save_plots"):
                self.results_manager.save_plots()

        return result

    def evaluate(self) -> dict[str, Any]:
        """Run evaluate and return results."""
        val_loader = self._make_loader(self.eval_dataset, shuffle=False)
        return evaluate_model(
            self.model,
            val_loader,
            self.device,
            self.num_labels,
            self.is_multichoice,
            compute_all_metrics=True,
        )

Custom loop trainer optimized for maximum GPU compute throughput, mixed-precision training, and distributed data parallel (DDP) gradient accumulation without redundant all-reduce overhead.

Ancestors

• BaseTrainer
• abc.ABC

Methods

def train(self) ‑> dict[str, typing.Any]

Expand source code

def train(self) -> dict[str, Any]:
    """Fit the model on the provided dataset with full DDP and accumulation safety."""
    train_loader = self._make_loader(self.train_dataset, shuffle=True)
    val_loader = self._make_loader(self.eval_dataset, shuffle=False)

    class AttrDict:
        pass

    cfg = AttrDict()
    cfg.learning_rate = float(self.config.get("learning_rate", 2e-5))
    cfg.epochs = int(self.config.get("epochs", 3))
    cfg.max_grad_norm = float(self.config.get("max_grad_norm", 1.0))
    cfg.gradient_accumulation_steps = int(self.config.get("gradient_accumulation_steps", 1))
    cfg.fp16 = bool(self.config.get("fp16", False))
    cfg.amp = bool(self.config.get("amp", True))
    cfg.logging_steps = int(self.config.get("logging_steps", 10))

    result = train_model(
        self.model,
        train_loader,
        val_loader,
        self.device,
        cfg,
        self.num_labels,
        self.is_multichoice,
    )

    if self.results_manager is not None:
        if hasattr(self.results_manager, "save_final_model"):
            self.results_manager.save_final_model(self.model)
        if hasattr(self.results_manager, "save_plots"):
            self.results_manager.save_plots()

    return result

Fit the model on the provided dataset with full DDP and accumulation safety.

def train_epoch(self,
train_loader: DataLoader,
optimizer: torch.optim.Optimizer,
scaler: torch.cuda.amp.GradScaler | None,
gradient_accumulation_steps: int = 1,
max_grad_norm: float = 1.0,
amp_dtype: torch.dtype = torch.float32,
use_amp: bool = False,
pbar: Any | None = None) ‑> float

Expand source code

def train_epoch(
    self,
    train_loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scaler: torch.cuda.amp.GradScaler | None,
    gradient_accumulation_steps: int = 1,
    max_grad_norm: float = 1.0,
    amp_dtype: torch.dtype = torch.float32,
    use_amp: bool = False,
    pbar: Any | None = None,
) -> float:
    """
    Execute a single training epoch with DDP gradient accumulation synchronization optimization.

    # FIX: Non-final micro-batches are wrapped in model.no_sync() or contextlib.nullcontext()
    # to eliminate redundant all-reduce synchronization across distributed workers during gradient accumulation.
    """
    self.model.train()
    total_epoch_loss = 0.0
    num_batches = len(train_loader)
    accum_steps = max(1, gradient_accumulation_steps)

    for step, batch in enumerate(train_loader):
        # FIX: Determine whether this micro-batch is the final step of the accumulation window or epoch
        is_sync_step = ((step + 1) % accum_steps == 0) or (step + 1 == num_batches)

        # FIX: Wrap non-final micro-batches in model.no_sync() for DDP to eliminate redundant all-reduce synchronization
        sync_context = (
            contextlib.nullcontext()
            if (
                is_sync_step
                or not hasattr(self.model, "no_sync")
                or not callable(self.model.no_sync)
            )
            else self.model.no_sync()
        )

        with sync_context:
            with torch.autocast(
                device_type=self.device.type,
                dtype=amp_dtype,
                enabled=(use_amp and self.device.type in ["cuda", "mps"]),
            ):
                if self.is_multichoice:
                    input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                    attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                    true_labels = batch["labels"].to(self.device, non_blocking=True)
                    bs, nc, sl = input_ids.shape
                    flat_ids = input_ids.view(-1, sl)
                    flat_mask = attention_mask.view(-1, sl)
                    outputs = self.model(input_ids=flat_ids, attention_mask=flat_mask)
                    logits = (
                        outputs["logits"]
                        if isinstance(outputs, dict)
                        else (outputs.logits if hasattr(outputs, "logits") else outputs)
                    )
                    combined_labels = true_labels.repeat_interleave(nc)
                    loss = nn.functional.cross_entropy(logits, combined_labels)
                else:
                    input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                    attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                    labels = batch["labels"].to(self.device, non_blocking=True)
                    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                    logits = (
                        outputs["logits"]
                        if isinstance(outputs, dict)
                        else (outputs.logits if hasattr(outputs, "logits") else outputs)
                    )
                    loss = nn.functional.cross_entropy(logits, labels)

                # FIX: Scale loss by accumulation factor so gradient magnitudes match the full effective batch size
                if accum_steps > 1:
                    loss = loss / accum_steps

            if torch.isfinite(loss):
                if scaler is not None and scaler.is_enabled():
                    scaler.scale(loss).backward()
                else:
                    loss.backward()

                curr_loss = loss.detach().item() * accum_steps
                total_epoch_loss += curr_loss
            else:
                curr_loss = 0.0

            # FIX: Only unscale, clip gradients, and step optimizer when synchronized on accumulation boundary
            if scaler is not None and scaler.is_enabled():
                try:
                    scaler.unscale_(optimizer)
                except RuntimeError:
                    pass
                clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                scaler.step(optimizer)
                scaler.update()
            else:
                clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                optimizer.step()
            optimizer.zero_grad(set_to_none=True)

        if pbar is not None and (
            (step + 1) % max(1, int(self.config.get("logging_steps", 10))) == 0
            or step == num_batches - 1
        ):
            pbar.set_postfix({"loss": f"{curr_loss:.4f}"})

    return total_epoch_loss / max(1, num_batches)

Execute a single training epoch with DDP gradient accumulation synchronization optimization.

FIX: Non-final micro-batches are wrapped in model.no_sync() or contextlib.nullcontext()

to eliminate redundant all-reduce synchronization across distributed workers during gradient accumulation.

Inherited members

• BaseTrainer:
  • evaluate

class DiceLoss (smooth=1.0, **kwargs)

Expand source code

@LOSS_REGISTRY.register("dice")
class DiceLoss(nn.Module):
    def __init__(self, smooth=1.0, **kwargs) -> None:
        super().__init__()
        self.smooth = smooth

    def forward(self, logits, labels) -> Any:
        """Forward pass — compute module output from inputs."""
        probs = F.softmax(logits, dim=-1)
        one_hot = F.one_hot(labels, num_classes=probs.shape[-1]).float()
        intersection = (probs * one_hot).sum(dim=0)
        union = probs.sum(dim=0) + one_hot.sum(dim=0)
        dice = (2.0 * intersection + self.smooth) / (union + self.smooth)
        return 1 - dice.mean()

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, logits, labels) ‑> Any

Expand source code

def forward(self, logits, labels) -> Any:
    """Forward pass — compute module output from inputs."""
    probs = F.softmax(logits, dim=-1)
    one_hot = F.one_hot(labels, num_classes=probs.shape[-1]).float()
    intersection = (probs * one_hot).sum(dim=0)
    union = probs.sum(dim=0) + one_hot.sum(dim=0)
    dice = (2.0 * intersection + self.smooth) / (union + self.smooth)
    return 1 - dice.mean()

Forward pass — compute module output from inputs.

class FocalLoss (alpha=0.25, gamma=2.0, **kwargs)

Expand source code

@LOSS_REGISTRY.register("focal")
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0, **kwargs) -> None:
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, logits, labels) -> Any:
        """Forward pass — compute module output from inputs."""
        ce_loss = F.cross_entropy(logits, labels, reduction="none")
        p_t = torch.exp(-ce_loss)
        return (self.alpha * (1 - p_t) ** self.gamma * ce_loss).mean()

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, logits, labels) ‑> Any

Expand source code

def forward(self, logits, labels) -> Any:
    """Forward pass — compute module output from inputs."""
    ce_loss = F.cross_entropy(logits, labels, reduction="none")
    p_t = torch.exp(-ce_loss)
    return (self.alpha * (1 - p_t) ** self.gamma * ce_loss).mean()

Forward pass — compute module output from inputs.

class GradientNormCallback (log_every_n_steps: int = 100)

Expand source code

class GradientNormCallback(TrainerCallback):
    """Logs gradient norms."""

    def __init__(self, log_every_n_steps: int = 100) -> None:
        self.log_every_n_steps = log_every_n_steps

    def on_step_end(self, args, state, control, model=None, **kwargs) -> None:
        """On step end."""
        if model is None or state.global_step % self.log_every_n_steps != 0:
            return
        total_norm = 0.0
        for p in model.parameters():
            if p.grad is not None:
                param_norm = p.grad.data.norm(2)
                total_norm += param_norm.item() ** 2
        total_norm = total_norm**0.5
        # logs will be passed to on_log? Actually we can just log to the trainer
        if hasattr(args, "logging_dir"):
            pass  # We can't easily log here; the SVDAwareTrainer logs it automatically
        # For simplicity we just print
        print(f"Step {state.global_step}: grad norm = {total_norm:.4f}")

Logs gradient norms.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_step_end(self, args, state, control, model=None, **kwargs) ‑> None

Expand source code

def on_step_end(self, args, state, control, model=None, **kwargs) -> None:
    """On step end."""
    if model is None or state.global_step % self.log_every_n_steps != 0:
        return
    total_norm = 0.0
    for p in model.parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm**0.5
    # logs will be passed to on_log? Actually we can just log to the trainer
    if hasattr(args, "logging_dir"):
        pass  # We can't easily log here; the SVDAwareTrainer logs it automatically
    # For simplicity we just print
    print(f"Step {state.global_step}: grad norm = {total_norm:.4f}")

On step end.

class HuggingFaceTrainer (model,
train_dataset,
eval_dataset,
config: dict[str, typing.Any],
results_manager,
compute_metrics=None,
**kwargs)

Expand source code

@TrainerRegistry.register("hf")
class HuggingFaceTrainer(BaseTrainer):
    def __init__(
        self,
        model,
        train_dataset,
        eval_dataset,
        config: dict[str, Any],
        results_manager,
        compute_metrics=None,
        **kwargs,
    ) -> None:
        self.model = model
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.config = config
        self.results_manager = results_manager
        self.compute_metrics = compute_metrics
        self.kwargs = kwargs
        self._trainer = None
        self._callbacks: list[Any] = []

    def add_callback(self, callback: Any) -> None:
        """Add a callback to the trainer."""
        if callback is not None:
            self._callbacks.append(callback)
            if self._trainer is not None:
                self._trainer.add_callback(callback)

    @property
    def args(self) -> Any:
        """Return training arguments."""
        if self._trainer is None:
            self._build()
        return self._trainer.args

    @property
    def state(self) -> Any:
        """Return trainer state."""
        if self._trainer is not None:
            return self._trainer.state
        return None

    def _build(self) -> Any:
        """Build underlying HuggingFace Trainer."""
        output_dir = self.results_manager.get_checkpoint_dir()
        training_args_dict = {
            "output_dir": output_dir,
            "learning_rate": self.config.get("learning_rate", 2e-5),
            "per_device_train_batch_size": self.config.get("batch_size", 16),
            "per_device_eval_batch_size": self.config.get("batch_size", 16),
            "num_train_epochs": self.config.get("epochs", 3),
            "warmup_ratio": self.config.get("warmup_ratio", 0.1),
            "max_grad_norm": self.config.get("max_grad_norm", 1.0),
            "weight_decay": self.config.get("weight_decay", 0.01),
            "eval_strategy": self.config.get("eval_strategy", "epoch"),
            "save_strategy": self.config.get("save_strategy", "epoch"),
            "save_total_limit": self.config.get("save_total_limit", 2),
            "load_best_model_at_end": (
                self.config.get("load_best_model_at_end", True)
                if self.eval_dataset is not None
                else False
            ),
            "metric_for_best_model": self.config.get("metric_for_best_model", "eval_loss"),
            "fp16": self.config.get("fp16", False),
            "logging_dir": str(self.results_manager.log_dir),
            "logging_steps": self.config.get("logging_steps", 10),
            "report_to": "none",
            # Disable safetensors to avoid RuntimeError from shared-memory tensors in custom attention modules
            "save_safetensors": False,
            # GPU throughput
            "dataloader_pin_memory": True,
            "dataloader_num_workers": self.config.get("num_workers", 0),
            "remove_unused_columns": False,
        }
        self._trainer = create_trainer(
            model=self.model,
            train_dataset=self.train_dataset,
            eval_dataset=self.eval_dataset,
            compute_metrics=self.compute_metrics,
            **training_args_dict,
            **self.kwargs,
        )
        self._trainer = self.results_manager.wrap_trainer(self._trainer)
        for cb in self._callbacks:
            self._trainer.add_callback(cb)

        original_train = self._trainer.train

        def train_and_save(*args, **train_kwargs) -> Any:
            """Train, persist artifacts, and generate structured summary reports."""
            result = original_train(*args, **train_kwargs)
            self.results_manager.save_final_model(self.model)
            self.results_manager.save_plots()
            last_metrics = (
                getattr(self._trainer.state, "log_history", [])[-1]
                if getattr(self._trainer.state, "log_history", [])
                else {}
            )
            self.results_manager.generate_summary_report(final_metrics=last_metrics)
            return result

        self._trainer.train = train_and_save

    def train(self) -> Any:
        """Fit the model on the provided dataset."""
        if self._trainer is None:
            self._build()
        return self._trainer.train()

    def evaluate(self) -> Any:
        """Run evaluate and return results."""
        if self._trainer is None:
            self._build()
        return self._trainer.evaluate()

Abstract base class for all trainers.

Ancestors

• BaseTrainer
• abc.ABC

Instance variables

prop args : Any

Expand source code

@property
def args(self) -> Any:
    """Return training arguments."""
    if self._trainer is None:
        self._build()
    return self._trainer.args

Return training arguments.

prop state : Any

Expand source code

@property
def state(self) -> Any:
    """Return trainer state."""
    if self._trainer is not None:
        return self._trainer.state
    return None

Return trainer state.

Methods

def add_callback(self, callback: Any) ‑> None

Expand source code

def add_callback(self, callback: Any) -> None:
    """Add a callback to the trainer."""
    if callback is not None:
        self._callbacks.append(callback)
        if self._trainer is not None:
            self._trainer.add_callback(callback)

Add a callback to the trainer.

Inherited members

• BaseTrainer:
  • evaluate
  • train

class HyperparameterOptimizer (train_fn: ,
eval_fn: ,
config_template: dict[str, typing.Any],
study_name: str = 'svd_attention_opt',
storage: str | None = None,
direction: str = 'maximize')

Expand source code

class HyperparameterOptimizer:
    def __init__(
        self,
        train_fn: Callable,
        eval_fn: Callable,
        config_template: dict[str, Any],
        study_name: str = "svd_attention_opt",
        storage: str | None = None,
        direction: str = "maximize",
    ) -> None:
        if not _HAS_OPTUNA:
            raise ImportError("optuna is required. pip install optuna")
        self.train_fn = train_fn
        self.eval_fn = eval_fn
        self.config_template = config_template
        self.study_name = study_name
        self.storage = storage
        self.direction = direction
        self.study = None
        self.best_params = None

    def objective(self, trial) -> float:
        """Objective."""
        config = self.config_template.copy()
        config["learning_rate"] = trial.suggest_float("learning_rate", 1e-6, 5e-5, log=True)
        config["batch_size"] = trial.suggest_categorical("batch_size", [8, 16, 32])
        config["rank"] = trial.suggest_int("rank", 32, 256, step=32)
        config["weight_decay"] = trial.suggest_float("weight_decay", 0.0, 0.1)
        config["warmup_ratio"] = trial.suggest_float("warmup_ratio", 0.0, 0.3)
        if "compression_method" in config:
            config["compression_method"] = trial.suggest_categorical(
                "compression_method",
                ["standard_svd", "randomized_svd", "quantized_svd"],
            )
        if "energy_threshold" in config:
            config["energy_threshold"] = trial.suggest_float("energy_threshold", 0.85, 0.99)
        self.train_fn(config)
        metrics = self.eval_fn(config)
        return metrics.get("accuracy", metrics.get("matthews_correlation", 0.0))

    def run(
        self,
        n_trials: int = 20,
        timeout: int | None = None,
        n_jobs: int = 1,
        show_progress_bar: bool = True,
    ) -> dict[str, Any]:
        """Execute the run operation and return results."""
        self.study = optuna.create_study(
            direction=self.direction,
            study_name=self.study_name,
            storage=self.storage,
            load_if_exists=True,
        )
        self.study.optimize(
            self.objective,
            n_trials=n_trials,
            timeout=timeout,
            n_jobs=n_jobs,
            show_progress_bar=show_progress_bar,
        )
        self.best_params = self.study.best_params
        with open("best_params.json", "w") as f:
            json.dump(self.best_params, f, indent=2)
        df = self.study.trials_dataframe()
        df.to_csv("optuna_trials.csv", index=False)
        print(f"\n✅ Optimization complete! Best value: {self.study.best_value:.4f}")
        print(f"Best params: {self.best_params}")
        return self.best_params

    def get_trial_dataframe(self) -> pd.DataFrame:
        """Return the trial dataframe."""
        if self.study is None:
            raise ValueError("Run optimization first")
        return self.study.trials_dataframe()

    def get_importance(self) -> dict[str, float]:
        """Return the importance."""
        if self.study is None:
            raise ValueError("Run optimization first")
        return optuna.importance.get_param_importances(self.study)

    def plot_optimization_history(self, save_path=None) -> Any:
        """Plot optimization history."""
        if self.study is None:
            raise ValueError("Run optimization first")
        fig = optuna.visualization.plot_optimization_history(self.study)
        if save_path:
            fig.write_image(save_path)
        return fig

    def plot_parallel_coordinate(self, save_path=None) -> Any:
        """Plot parallel coordinate."""
        if self.study is None:
            raise ValueError("Run optimization first")
        fig = optuna.visualization.plot_parallel_coordinate(self.study)
        if save_path:
            fig.write_image(save_path)
        return fig

    def plot_param_importances(self, save_path=None) -> Any:
        """Plot param importances."""
        if self.study is None:
            raise ValueError("Run optimization first")
        fig = optuna.visualization.plot_param_importances(self.study)
        if save_path:
            fig.write_image(save_path)
        return fig

Subclasses

• PrunedHyperparameterOptimizer

Methods

def get_importance(self) ‑> dict[str, float]

Expand source code

def get_importance(self) -> dict[str, float]:
    """Return the importance."""
    if self.study is None:
        raise ValueError("Run optimization first")
    return optuna.importance.get_param_importances(self.study)

Return the importance.

def get_trial_dataframe(self) ‑> pandas.DataFrame

Expand source code

def get_trial_dataframe(self) -> pd.DataFrame:
    """Return the trial dataframe."""
    if self.study is None:
        raise ValueError("Run optimization first")
    return self.study.trials_dataframe()

Return the trial dataframe.

def objective(self, trial) ‑> float

Expand source code

def objective(self, trial) -> float:
    """Objective."""
    config = self.config_template.copy()
    config["learning_rate"] = trial.suggest_float("learning_rate", 1e-6, 5e-5, log=True)
    config["batch_size"] = trial.suggest_categorical("batch_size", [8, 16, 32])
    config["rank"] = trial.suggest_int("rank", 32, 256, step=32)
    config["weight_decay"] = trial.suggest_float("weight_decay", 0.0, 0.1)
    config["warmup_ratio"] = trial.suggest_float("warmup_ratio", 0.0, 0.3)
    if "compression_method" in config:
        config["compression_method"] = trial.suggest_categorical(
            "compression_method",
            ["standard_svd", "randomized_svd", "quantized_svd"],
        )
    if "energy_threshold" in config:
        config["energy_threshold"] = trial.suggest_float("energy_threshold", 0.85, 0.99)
    self.train_fn(config)
    metrics = self.eval_fn(config)
    return metrics.get("accuracy", metrics.get("matthews_correlation", 0.0))

Objective.

def plot_optimization_history(self, save_path=None) ‑> Any

Expand source code

def plot_optimization_history(self, save_path=None) -> Any:
    """Plot optimization history."""
    if self.study is None:
        raise ValueError("Run optimization first")
    fig = optuna.visualization.plot_optimization_history(self.study)
    if save_path:
        fig.write_image(save_path)
    return fig

Plot optimization history.

def plot_parallel_coordinate(self, save_path=None) ‑> Any

Expand source code

def plot_parallel_coordinate(self, save_path=None) -> Any:
    """Plot parallel coordinate."""
    if self.study is None:
        raise ValueError("Run optimization first")
    fig = optuna.visualization.plot_parallel_coordinate(self.study)
    if save_path:
        fig.write_image(save_path)
    return fig

Plot parallel coordinate.

def plot_param_importances(self, save_path=None) ‑> Any

Expand source code

def plot_param_importances(self, save_path=None) -> Any:
    """Plot param importances."""
    if self.study is None:
        raise ValueError("Run optimization first")
    fig = optuna.visualization.plot_param_importances(self.study)
    if save_path:
        fig.write_image(save_path)
    return fig

Plot param importances.

def run(self,
n_trials: int = 20,
timeout: int | None = None,
n_jobs: int = 1,
show_progress_bar: bool = True) ‑> dict[str, typing.Any]

Expand source code

def run(
    self,
    n_trials: int = 20,
    timeout: int | None = None,
    n_jobs: int = 1,
    show_progress_bar: bool = True,
) -> dict[str, Any]:
    """Execute the run operation and return results."""
    self.study = optuna.create_study(
        direction=self.direction,
        study_name=self.study_name,
        storage=self.storage,
        load_if_exists=True,
    )
    self.study.optimize(
        self.objective,
        n_trials=n_trials,
        timeout=timeout,
        n_jobs=n_jobs,
        show_progress_bar=show_progress_bar,
    )
    self.best_params = self.study.best_params
    with open("best_params.json", "w") as f:
        json.dump(self.best_params, f, indent=2)
    df = self.study.trials_dataframe()
    df.to_csv("optuna_trials.csv", index=False)
    print(f"\n✅ Optimization complete! Best value: {self.study.best_value:.4f}")
    print(f"Best params: {self.best_params}")
    return self.best_params

Execute the run operation and return results.

class LabelSmoothingLoss (smoothing=0.1, **kwargs)

Expand source code

@LOSS_REGISTRY.register("label_smoothing")
class LabelSmoothingLoss(nn.Module):
    def __init__(self, smoothing=0.1, **kwargs) -> None:
        super().__init__()
        self.smoothing = smoothing
        self.confidence = 1.0 - smoothing

    def forward(self, logits, labels) -> Any:
        """Forward pass — compute module output from inputs."""
        log_probs = F.log_softmax(logits, dim=-1)
        labels = labels.to(log_probs.device)
        nll = -log_probs.gather(dim=-1, index=labels.unsqueeze(1))
        smooth = -log_probs.mean(dim=-1)
        return (self.confidence * nll + self.smoothing * smooth).mean()

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, logits, labels) ‑> Any

Expand source code

def forward(self, logits, labels) -> Any:
    """Forward pass — compute module output from inputs."""
    log_probs = F.log_softmax(logits, dim=-1)
    labels = labels.to(log_probs.device)
    nll = -log_probs.gather(dim=-1, index=labels.unsqueeze(1))
    smooth = -log_probs.mean(dim=-1)
    return (self.confidence * nll + self.smoothing * smooth).mean()

Forward pass — compute module output from inputs.

class MemoryCallback

Expand source code

class MemoryCallback(TrainerCallback):
    """Logs GPU/MPS memory usage."""

    def on_log(self, args, state, control, logs=None, **kwargs) -> None:
        """On log."""
        if logs is None:
            return
        if torch.cuda.is_available():
            logs["gpu_memory_allocated"] = torch.cuda.memory_allocated() / 1024 / 1024
            logs["gpu_memory_reserved"] = torch.cuda.memory_reserved() / 1024 / 1024
        if torch.backends.mps.is_available():
            logs["mps_memory"] = torch.mps.current_allocated_memory() / 1024 / 1024

Logs GPU/MPS memory usage.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_log(self, args, state, control, logs=None, **kwargs) ‑> None

Expand source code

def on_log(self, args, state, control, logs=None, **kwargs) -> None:
    """On log."""
    if logs is None:
        return
    if torch.cuda.is_available():
        logs["gpu_memory_allocated"] = torch.cuda.memory_allocated() / 1024 / 1024
        logs["gpu_memory_reserved"] = torch.cuda.memory_reserved() / 1024 / 1024
    if torch.backends.mps.is_available():
        logs["mps_memory"] = torch.mps.current_allocated_memory() / 1024 / 1024

On log.

class PolyLoss (epsilon=1.0, **kwargs)

Expand source code

@LOSS_REGISTRY.register("poly")
class PolyLoss(nn.Module):
    def __init__(self, epsilon=1.0, **kwargs) -> None:
        super().__init__()
        self.epsilon = epsilon

    def forward(self, logits, labels) -> Any:
        """Forward pass — compute module output from inputs."""
        ce = F.cross_entropy(logits, labels)
        pt = torch.exp(-ce)
        return ce + self.epsilon * (1 - pt)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, logits, labels) ‑> Any

Expand source code

def forward(self, logits, labels) -> Any:
    """Forward pass — compute module output from inputs."""
    ce = F.cross_entropy(logits, labels)
    pt = torch.exp(-ce)
    return ce + self.epsilon * (1 - pt)

Forward pass — compute module output from inputs.

class PreemptionCallback (preempt_flag_dir='/tmp/preempt_flags', job_id=None)

Expand source code

class PreemptionCallback(TrainerCallback):
    """Stops training when a preemption flag file is present."""

    def __init__(self, preempt_flag_dir="/tmp/preempt_flags", job_id=None) -> None:
        self.preempt_flag_dir = preempt_flag_dir
        self.job_id = job_id

    def on_step_end(self, args, state, control, **kwargs) -> None:
        """On step end."""
        if self.job_id and os.path.exists(f"{self.preempt_flag_dir}/{self.job_id}"):
            print("Preemption requested, stopping training.")
            control.should_training_stop = True

Stops training when a preemption flag file is present.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_step_end(self, args, state, control, **kwargs) ‑> None

Expand source code

def on_step_end(self, args, state, control, **kwargs) -> None:
    """On step end."""
    if self.job_id and os.path.exists(f"{self.preempt_flag_dir}/{self.job_id}"):
        print("Preemption requested, stopping training.")
        control.should_training_stop = True

On step end.

class PrunedHyperparameterOptimizer (*args, pruner='median', **kwargs)

Expand source code

class PrunedHyperparameterOptimizer(HyperparameterOptimizer):
    def __init__(self, *args, pruner="median", **kwargs) -> None:
        super().__init__(*args, **kwargs)
        if pruner == "median":
            self.pruner = optuna.pruners.MedianPruner()
        elif pruner == "halving":
            self.pruner = optuna.pruners.HalvingPruner()
        elif pruner == "hyperband":
            self.pruner = optuna.pruners.HyperbandPruner()
        else:
            self.pruner = optuna.pruners.MedianPruner()

Ancestors

• HyperparameterOptimizer

Inherited members

• HyperparameterOptimizer:
  • get_importance
  • get_trial_dataframe
  • objective
  • plot_optimization_history
  • plot_parallel_coordinate
  • plot_param_importances
  • run

class ResultsManager (experiment_name: str,
base_dir: str = 'results',
config: dict[str, typing.Any] | None = None,
log_to_file: bool = True)

Expand source code

class ResultsManager:
    """Manages per‑run result directories, logging, metrics, and trainer wrapping."""

    def __init__(
        self,
        experiment_name: str,
        base_dir: str = "results",
        config: dict[str, Any] | None = None,
        log_to_file: bool = True,
    ) -> None:
        self.experiment_name = experiment_name
        self.base_dir = Path(base_dir)
        self.config = config or {}
        self.log_to_file = log_to_file

        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        self.run_dir = self.base_dir / f"{experiment_name}_{timestamp}"
        self.run_dir.mkdir(parents=True, exist_ok=True)

        self.checkpoint_dir = self.run_dir / "checkpoints"
        self.log_dir = self.run_dir / "logs"
        self.final_model_dir = self.run_dir / "final_model"
        for d in [self.checkpoint_dir, self.log_dir, self.final_model_dir]:
            d.mkdir(parents=True, exist_ok=True)

        self.logger = self._setup_logger()
        self.metrics_path = self.run_dir / "metrics.jsonl"
        self.history: list[dict[str, Any]] = []

        if config:
            with open(self.run_dir / "config.yaml", "w") as f:
                yaml.dump(config, f, default_flow_style=False)

        self.logger.info(f"Results directory: {self.run_dir}")

    def _setup_logger(self) -> logging.Logger:
        """setup logger."""
        logger = logging.getLogger(f"training.{self.experiment_name}")
        logger.setLevel(logging.INFO)
        logger.handlers.clear()
        ch = logging.StreamHandler()
        ch.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
        logger.addHandler(ch)
        if self.log_to_file:
            fh = logging.FileHandler(self.log_dir / "run.log")
            fh.setFormatter(
                logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
            )
            logger.addHandler(fh)
        return logger

    def log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:
        """Log metrics."""
        record = {"step": step} if step is not None else {}
        record.update(metrics)
        record["timestamp"] = datetime.now().isoformat()
        record = _sanitize_json_value(record)
        self.history.append(record)
        with open(self.metrics_path, "a") as f:
            f.write(json.dumps(record, allow_nan=False) + "\n")

    def save_plots(self) -> None:
        """Persist plots to disk."""
        if self.history:
            plot_path = self.run_dir / "training_history.png"
            plot_training_curves(self.history, str(plot_path))
            self.logger.info(f"Saved training plot to {plot_path}")

    def save_final_model(self, model, tokenizer=None) -> None:
        """Persist final model to disk."""
        model.save_pretrained(self.final_model_dir)
        if tokenizer:
            tokenizer.save_pretrained(self.final_model_dir)
        self.logger.info(f"Final model saved to {self.final_model_dir}")

    def get_checkpoint_dir(self) -> str:
        """Return the checkpoint dir."""
        return str(self.checkpoint_dir)

    def push_to_hf(self, repo_id: str | None = None) -> None:
        """Push full run results (models, plots, metrics) to HuggingFace Hub."""
        from neuralops.results_saver.backends import huggingface as hf_backend

        cfg = {"hf_options": {"repo_id": repo_id}} if repo_id else self.config
        hf_backend.push(self.experiment_name, self.run_dir, cfg)

    def generate_summary_report(self, final_metrics: dict[str, Any] | None = None) -> None:
        """Generate structured markdown and json summary reports for the run."""

        def _make_serializable(obj):
            if hasattr(obj, "to_dict") and callable(obj.to_dict):
                return obj.to_dict()
            if hasattr(obj, "__dict__"):
                return {
                    k: _make_serializable(v)
                    for k, v in obj.__dict__.items()
                    if not k.startswith("_")
                }
            if isinstance(obj, dict):
                return {str(k): _make_serializable(v) for k, v in obj.items()}
            if isinstance(obj, (list, tuple, set)):
                return [_make_serializable(item) for item in obj]
            if isinstance(obj, (datetime, Path)):
                return str(obj)
            if hasattr(obj, "value"):
                return obj.value
            try:
                json.dumps(obj)
                return obj
            except Exception:
                return str(obj)

        summary = {
            "experiment_name": self.experiment_name,
            "run_dir": str(self.run_dir),
            "total_steps": len(self.history),
            "config": _make_serializable(self.config),
            "final_metrics": _sanitize_json_value(
                _make_serializable(final_metrics or (self.history[-1] if self.history else {}))
            ),
            "timestamp": datetime.now().isoformat(),
        }
        summary_path = self.run_dir / "run_summary.json"
        with open(summary_path, "w") as f:
            json.dump(summary, f, indent=2, allow_nan=False)

        card_path = self.run_dir / "run_summary.md"
        with open(card_path, "w") as f:
            f.write(f"# Run Summary: {self.experiment_name}\n\n")
            f.write(f"- **Timestamp**: {summary['timestamp']}\n")
            f.write(f"- **Total Logging Steps**: {summary['total_steps']}\n")
            f.write(f"- **Run Directory**: `{self.run_dir}`\n\n")
            f.write("## Final Metrics\n```json\n")
            f.write(json.dumps(summary["final_metrics"], indent=2))
            f.write("\n```\n")
        self.logger.info(f"Generated structured run summary report at {card_path}")

    def wrap_trainer(self, trainer) -> Any:
        """Add standard callbacks and metrics logging to a Hugging Face Trainer."""
        try:
            trainer.add_callback(MemoryCallback())
            trainer.add_callback(TimeCallback())
            trainer.add_callback(GradientNormCallback(log_every_n_steps=100))
        except Exception as e:
            self.logger.warning(f"Could not add callbacks: {e}")

        original_log = trainer.log

        def custom_log(logs) -> None:
            """Custom log."""
            original_log(logs)
            self.log_metrics(logs)

        trainer.log = custom_log
        return trainer

Manages per‑run result directories, logging, metrics, and trainer wrapping.

Methods

def generate_summary_report(self, final_metrics: dict[str, typing.Any] | None = None) ‑> None

Expand source code

def generate_summary_report(self, final_metrics: dict[str, Any] | None = None) -> None:
    """Generate structured markdown and json summary reports for the run."""

    def _make_serializable(obj):
        if hasattr(obj, "to_dict") and callable(obj.to_dict):
            return obj.to_dict()
        if hasattr(obj, "__dict__"):
            return {
                k: _make_serializable(v)
                for k, v in obj.__dict__.items()
                if not k.startswith("_")
            }
        if isinstance(obj, dict):
            return {str(k): _make_serializable(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple, set)):
            return [_make_serializable(item) for item in obj]
        if isinstance(obj, (datetime, Path)):
            return str(obj)
        if hasattr(obj, "value"):
            return obj.value
        try:
            json.dumps(obj)
            return obj
        except Exception:
            return str(obj)

    summary = {
        "experiment_name": self.experiment_name,
        "run_dir": str(self.run_dir),
        "total_steps": len(self.history),
        "config": _make_serializable(self.config),
        "final_metrics": _sanitize_json_value(
            _make_serializable(final_metrics or (self.history[-1] if self.history else {}))
        ),
        "timestamp": datetime.now().isoformat(),
    }
    summary_path = self.run_dir / "run_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2, allow_nan=False)

    card_path = self.run_dir / "run_summary.md"
    with open(card_path, "w") as f:
        f.write(f"# Run Summary: {self.experiment_name}\n\n")
        f.write(f"- **Timestamp**: {summary['timestamp']}\n")
        f.write(f"- **Total Logging Steps**: {summary['total_steps']}\n")
        f.write(f"- **Run Directory**: `{self.run_dir}`\n\n")
        f.write("## Final Metrics\n```json\n")
        f.write(json.dumps(summary["final_metrics"], indent=2))
        f.write("\n```\n")
    self.logger.info(f"Generated structured run summary report at {card_path}")

Generate structured markdown and json summary reports for the run.

def get_checkpoint_dir(self) ‑> str

Expand source code

def get_checkpoint_dir(self) -> str:
    """Return the checkpoint dir."""
    return str(self.checkpoint_dir)

Return the checkpoint dir.

def log_metrics(self, metrics: dict[str, typing.Any], step: int | None = None) ‑> None

Expand source code

def log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:
    """Log metrics."""
    record = {"step": step} if step is not None else {}
    record.update(metrics)
    record["timestamp"] = datetime.now().isoformat()
    record = _sanitize_json_value(record)
    self.history.append(record)
    with open(self.metrics_path, "a") as f:
        f.write(json.dumps(record, allow_nan=False) + "\n")

Log metrics.

def push_to_hf(self, repo_id: str | None = None) ‑> None

Expand source code

def push_to_hf(self, repo_id: str | None = None) -> None:
    """Push full run results (models, plots, metrics) to HuggingFace Hub."""
    from neuralops.results_saver.backends import huggingface as hf_backend

    cfg = {"hf_options": {"repo_id": repo_id}} if repo_id else self.config
    hf_backend.push(self.experiment_name, self.run_dir, cfg)

Push full run results (models, plots, metrics) to HuggingFace Hub.

def save_final_model(self, model, tokenizer=None) ‑> None

Expand source code

def save_final_model(self, model, tokenizer=None) -> None:
    """Persist final model to disk."""
    model.save_pretrained(self.final_model_dir)
    if tokenizer:
        tokenizer.save_pretrained(self.final_model_dir)
    self.logger.info(f"Final model saved to {self.final_model_dir}")

Persist final model to disk.

def save_plots(self) ‑> None

Expand source code

def save_plots(self) -> None:
    """Persist plots to disk."""
    if self.history:
        plot_path = self.run_dir / "training_history.png"
        plot_training_curves(self.history, str(plot_path))
        self.logger.info(f"Saved training plot to {plot_path}")

Persist plots to disk.

def wrap_trainer(self, trainer) ‑> Any

Expand source code

def wrap_trainer(self, trainer) -> Any:
    """Add standard callbacks and metrics logging to a Hugging Face Trainer."""
    try:
        trainer.add_callback(MemoryCallback())
        trainer.add_callback(TimeCallback())
        trainer.add_callback(GradientNormCallback(log_every_n_steps=100))
    except Exception as e:
        self.logger.warning(f"Could not add callbacks: {e}")

    original_log = trainer.log

    def custom_log(logs) -> None:
        """Custom log."""
        original_log(logs)
        self.log_metrics(logs)

    trainer.log = custom_log
    return trainer

Add standard callbacks and metrics logging to a Hugging Face Trainer.

class RiemannianStiefelOptimizer (params: Iterable[torch.nn.Parameter] | Iterable[dict[str, Any]] | nn.Module | list[tuple[str, torch.nn.Parameter]],
lr: float = 0.001,
lr_stiefel: float | None = None,
base_optimizer: str = 'adamw',
betas: tuple[float, float] = (0.9, 0.999),
eps: float = 1e-08,
weight_decay: float = 0.01,
momentum: float = 0.9,
stiefel_momentum: float = 0.0,
dampening: float = 0.0,
nesterov: bool = False,
orthogonality_threshold: float = 0.15,
**kwargs: Any)

Expand source code

class RiemannianStiefelOptimizer(Optimizer):
    """Riemannian Stiefel Manifold Optimizer with Retraction-based Updates."""

    def __init__(
        self,
        params: (
            Iterable[torch.nn.Parameter]
            | Iterable[dict[str, Any]]
            | nn.Module
            | list[tuple[str, torch.nn.Parameter]]
        ),
        lr: float = 1e-3,
        lr_stiefel: float | None = None,
        base_optimizer: str = "adamw",
        betas: tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0.01,
        momentum: float = 0.9,
        stiefel_momentum: float = 0.0,
        dampening: float = 0.0,
        nesterov: bool = False,
        orthogonality_threshold: float = 0.15,
        **kwargs: Any,
    ) -> None:
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if lr_stiefel is not None and lr_stiefel < 0.0:
            raise ValueError(f"Invalid Stiefel learning rate: {lr_stiefel}")
        if eps < 0.0:
            raise ValueError(f"Invalid epsilon value: {eps}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

        effective_lr_stiefel = lr_stiefel if lr_stiefel is not None else lr

        defaults = dict(
            lr=lr,
            lr_stiefel=effective_lr_stiefel,
            base_optimizer=base_optimizer,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
            momentum=momentum,
            stiefel_momentum=stiefel_momentum,
            dampening=dampening,
            nesterov=nesterov,
            orthogonality_threshold=orthogonality_threshold,
            **kwargs,
        )

        param_groups: list[dict[str, Any]] = []

        if isinstance(params, nn.Module):
            named_params = list(params.named_parameters())
            stiefel_p: list[torch.nn.Parameter] = []
            euclidean_p: list[torch.nn.Parameter] = []
            for name, p in named_params:
                if not p.requires_grad:
                    continue
                if is_stiefel_parameter(name, p, orthogonality_threshold):
                    stiefel_p.append(p)
                else:
                    euclidean_p.append(p)

            if stiefel_p:
                param_groups.append(
                    {
                        "params": stiefel_p,
                        "manifold": "stiefel",
                        "lr": effective_lr_stiefel,
                        "momentum": stiefel_momentum,
                        "weight_decay": 0.0,
                    }
                )
            if euclidean_p:
                param_groups.append(
                    {
                        "params": euclidean_p,
                        "manifold": "euclidean",
                        "lr": lr,
                        "weight_decay": weight_decay,
                        "betas": betas,
                        "eps": eps,
                        "base_optimizer": base_optimizer,
                        "momentum": momentum,
                        "dampening": dampening,
                        "nesterov": nesterov,
                    }
                )
        elif isinstance(params, list) and len(params) > 0 and isinstance(params[0], dict):
            for group in params:
                if "manifold" in group:
                    param_groups.append(group)
                else:
                    stiefel_p = []
                    euclidean_p = []
                    grp_params = group.get("params", [])
                    if isinstance(grp_params, torch.Tensor):
                        grp_params = [grp_params]
                    for p in grp_params:
                        if not getattr(p, "requires_grad", True):
                            continue
                        name = getattr(p, "_name", "")
                        if is_stiefel_parameter(name, p, orthogonality_threshold):
                            stiefel_p.append(p)
                        else:
                            euclidean_p.append(p)

                    if stiefel_p:
                        st_group = dict(group)
                        st_group["params"] = stiefel_p
                        st_group["manifold"] = "stiefel"
                        st_group["lr"] = group.get("lr_stiefel", effective_lr_stiefel)
                        st_group["momentum"] = group.get("stiefel_momentum", stiefel_momentum)
                        st_group["weight_decay"] = 0.0
                        param_groups.append(st_group)

                    if euclidean_p:
                        euc_group = dict(group)
                        euc_group["params"] = euclidean_p
                        euc_group["manifold"] = "euclidean"
                        param_groups.append(euc_group)
        else:
            stiefel_p = []
            euclidean_p = []
            param_list = list(params)
            for item in param_list:
                if isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str):
                    name, p = item
                else:
                    name = getattr(item, "_name", "")
                    p = item
                if not getattr(p, "requires_grad", True):
                    continue
                if is_stiefel_parameter(name, p, orthogonality_threshold):
                    stiefel_p.append(p)
                else:
                    euclidean_p.append(p)

            if stiefel_p:
                param_groups.append(
                    {
                        "params": stiefel_p,
                        "manifold": "stiefel",
                        "lr": effective_lr_stiefel,
                        "momentum": stiefel_momentum,
                        "weight_decay": 0.0,
                    }
                )
            if euclidean_p:
                param_groups.append(
                    {
                        "params": euclidean_p,
                        "manifold": "euclidean",
                        "lr": lr,
                        "weight_decay": weight_decay,
                        "betas": betas,
                        "eps": eps,
                        "base_optimizer": base_optimizer,
                        "momentum": momentum,
                        "dampening": dampening,
                        "nesterov": nesterov,
                    }
                )

        if not param_groups:
            param_groups = [{"params": [], "manifold": "euclidean", "lr": lr}]

        super().__init__(param_groups, defaults)

    @staticmethod
    def project_gradient(W: torch.Tensor, G: torch.Tensor) -> torch.Tensor:
        """Riemannian gradient projection onto Stiefel tangent space."""
        return project_stiefel_gradient(W, G)

    @staticmethod
    def retract(W: torch.Tensor, grad_R: torch.Tensor, lr: float) -> torch.Tensor:
        """QR retraction onto Stiefel manifold."""
        return stiefel_qr_retraction(W, grad_R, lr)

    def check_orthogonality_error(
        self, param: torch.Tensor | None = None
    ) -> float | dict[int, float]:
        """Compute Frobenius norm orthogonality error for manifold parameters."""
        if param is not None:
            w = param.detach()
            n, p = w.shape[-2], w.shape[-1]
            if n >= p:
                gram = torch.matmul(w.transpose(-2, -1), w)
                eye = torch.eye(p, device=w.device, dtype=w.dtype)
            else:
                gram = torch.matmul(w, w.transpose(-2, -1))
                eye = torch.eye(n, device=w.device, dtype=w.dtype)
            return float(torch.linalg.norm(gram - eye).item())

        errors: dict[int, float] = {}
        idx = 0
        for group in self.param_groups:
            if group.get("manifold") == "stiefel":
                for p in group["params"]:
                    w = p.detach()
                    if w.dim() >= 2:
                        n, dim_k = w.shape[-2], w.shape[-1]
                        if n >= dim_k:
                            gram = torch.matmul(w.transpose(-2, -1), w)
                            eye = torch.eye(dim_k, device=w.device, dtype=w.dtype)
                        else:
                            gram = torch.matmul(w, w.transpose(-2, -1))
                            eye = torch.eye(n, device=w.device, dtype=w.dtype)
                        errors[idx] = float(torch.linalg.norm(gram - eye).item())
                        idx += 1
        return errors

    @torch.no_grad()
    def step(self, closure: Callable[[], float] | None = None) -> float | None:
        """Execute single Riemannian optimization step."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            manifold = group.get("manifold", "euclidean")
            lr = group["lr"]

            if manifold == "stiefel":
                momentum = group.get("momentum", 0.0)
                for p in group["params"]:
                    if p.grad is None:
                        continue
                    grad = p.grad
                    if grad.is_sparse:
                        grad = grad.to_dense()

                    grad_r = project_stiefel_gradient(p.data, grad)

                    state = self.state[p]
                    if len(state) == 0:
                        state["step"] = 0
                        if momentum > 0:
                            state["momentum_buffer"] = torch.zeros_like(p.data)

                    state["step"] += 1

                    if momentum > 0:
                        buf = state["momentum_buffer"]
                        buf.mul_(momentum).add_(grad_r, alpha=1.0 - momentum)
                        tangent_update = project_stiefel_gradient(p.data, buf)
                    else:
                        tangent_update = grad_r

                    w_next = stiefel_qr_retraction(p.data, tangent_update, lr)
                    p.data.copy_(w_next)

            else:
                base_opt = group.get("base_optimizer", "adamw").lower()
                weight_decay = group.get("weight_decay", 0.0)

                if base_opt == "sgd":
                    momentum = group.get("momentum", 0.0)
                    dampening = group.get("dampening", 0.0)
                    nesterov = group.get("nesterov", False)

                    for p in group["params"]:
                        if p.grad is None:
                            continue
                        d_p = p.grad
                        if d_p.is_sparse:
                            d_p = d_p.to_dense()

                        state = self.state[p]
                        if len(state) == 0:
                            state["step"] = 0

                        state["step"] += 1

                        if weight_decay != 0.0:
                            d_p = d_p.add(p.data, alpha=weight_decay)

                        if momentum != 0.0:
                            if "momentum_buffer" not in state:
                                buf = state["momentum_buffer"] = torch.clone(d_p).detach()
                            else:
                                buf = state["momentum_buffer"]
                                buf.mul_(momentum).add_(d_p, alpha=1.0 - dampening)
                            if nesterov:
                                d_p = d_p.add(buf, alpha=momentum)
                            else:
                                d_p = buf

                        p.data.add_(d_p, alpha=-lr)

                else:
                    betas = group.get("betas", (0.9, 0.999))
                    beta1, beta2 = betas
                    eps = group.get("eps", 1e-8)

                    for p in group["params"]:
                        if p.grad is None:
                            continue
                        grad = p.grad
                        if grad.is_sparse:
                            grad = grad.to_dense()

                        state = self.state[p]
                        if len(state) == 0:
                            state["step"] = 0
                            state["exp_avg"] = torch.zeros_like(p.data)
                            state["exp_avg_sq"] = torch.zeros_like(p.data)

                        exp_avg = state["exp_avg"]
                        exp_avg_sq = state["exp_avg_sq"]
                        state["step"] += 1

                        if weight_decay != 0.0:
                            p.data.mul_(1.0 - lr * weight_decay)

                        exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)

                        bias_correction1 = 1.0 - beta1 ** state["step"]
                        bias_correction2 = 1.0 - beta2 ** state["step"]
                        step_size = lr / bias_correction1
                        denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)

                        p.data.addcdiv_(exp_avg, denom, value=-step_size)

        return loss

Riemannian Stiefel Manifold Optimizer with Retraction-based Updates.

Ancestors

• torch.optim.optimizer.Optimizer

Static methods

def project_gradient(W: torch.Tensor, G: torch.Tensor) ‑> torch.Tensor

Expand source code

@staticmethod
def project_gradient(W: torch.Tensor, G: torch.Tensor) -> torch.Tensor:
    """Riemannian gradient projection onto Stiefel tangent space."""
    return project_stiefel_gradient(W, G)

Riemannian gradient projection onto Stiefel tangent space.

def retract(W: torch.Tensor, grad_R: torch.Tensor, lr: float) ‑> torch.Tensor

Expand source code

@staticmethod
def retract(W: torch.Tensor, grad_R: torch.Tensor, lr: float) -> torch.Tensor:
    """QR retraction onto Stiefel manifold."""
    return stiefel_qr_retraction(W, grad_R, lr)

QR retraction onto Stiefel manifold.

Methods

def check_orthogonality_error(self, param: torch.Tensor | None = None) ‑> float | dict[int, float]

Expand source code

def check_orthogonality_error(
    self, param: torch.Tensor | None = None
) -> float | dict[int, float]:
    """Compute Frobenius norm orthogonality error for manifold parameters."""
    if param is not None:
        w = param.detach()
        n, p = w.shape[-2], w.shape[-1]
        if n >= p:
            gram = torch.matmul(w.transpose(-2, -1), w)
            eye = torch.eye(p, device=w.device, dtype=w.dtype)
        else:
            gram = torch.matmul(w, w.transpose(-2, -1))
            eye = torch.eye(n, device=w.device, dtype=w.dtype)
        return float(torch.linalg.norm(gram - eye).item())

    errors: dict[int, float] = {}
    idx = 0
    for group in self.param_groups:
        if group.get("manifold") == "stiefel":
            for p in group["params"]:
                w = p.detach()
                if w.dim() >= 2:
                    n, dim_k = w.shape[-2], w.shape[-1]
                    if n >= dim_k:
                        gram = torch.matmul(w.transpose(-2, -1), w)
                        eye = torch.eye(dim_k, device=w.device, dtype=w.dtype)
                    else:
                        gram = torch.matmul(w, w.transpose(-2, -1))
                        eye = torch.eye(n, device=w.device, dtype=w.dtype)
                    errors[idx] = float(torch.linalg.norm(gram - eye).item())
                    idx += 1
    return errors

Compute Frobenius norm orthogonality error for manifold parameters.

def step(self, closure: Callable[[], float] | None = None) ‑> float | None

Expand source code

@torch.no_grad()
def step(self, closure: Callable[[], float] | None = None) -> float | None:
    """Execute single Riemannian optimization step."""
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()

    for group in self.param_groups:
        manifold = group.get("manifold", "euclidean")
        lr = group["lr"]

        if manifold == "stiefel":
            momentum = group.get("momentum", 0.0)
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.is_sparse:
                    grad = grad.to_dense()

                grad_r = project_stiefel_gradient(p.data, grad)

                state = self.state[p]
                if len(state) == 0:
                    state["step"] = 0
                    if momentum > 0:
                        state["momentum_buffer"] = torch.zeros_like(p.data)

                state["step"] += 1

                if momentum > 0:
                    buf = state["momentum_buffer"]
                    buf.mul_(momentum).add_(grad_r, alpha=1.0 - momentum)
                    tangent_update = project_stiefel_gradient(p.data, buf)
                else:
                    tangent_update = grad_r

                w_next = stiefel_qr_retraction(p.data, tangent_update, lr)
                p.data.copy_(w_next)

        else:
            base_opt = group.get("base_optimizer", "adamw").lower()
            weight_decay = group.get("weight_decay", 0.0)

            if base_opt == "sgd":
                momentum = group.get("momentum", 0.0)
                dampening = group.get("dampening", 0.0)
                nesterov = group.get("nesterov", False)

                for p in group["params"]:
                    if p.grad is None:
                        continue
                    d_p = p.grad
                    if d_p.is_sparse:
                        d_p = d_p.to_dense()

                    state = self.state[p]
                    if len(state) == 0:
                        state["step"] = 0

                    state["step"] += 1

                    if weight_decay != 0.0:
                        d_p = d_p.add(p.data, alpha=weight_decay)

                    if momentum != 0.0:
                        if "momentum_buffer" not in state:
                            buf = state["momentum_buffer"] = torch.clone(d_p).detach()
                        else:
                            buf = state["momentum_buffer"]
                            buf.mul_(momentum).add_(d_p, alpha=1.0 - dampening)
                        if nesterov:
                            d_p = d_p.add(buf, alpha=momentum)
                        else:
                            d_p = buf

                    p.data.add_(d_p, alpha=-lr)

            else:
                betas = group.get("betas", (0.9, 0.999))
                beta1, beta2 = betas
                eps = group.get("eps", 1e-8)

                for p in group["params"]:
                    if p.grad is None:
                        continue
                    grad = p.grad
                    if grad.is_sparse:
                        grad = grad.to_dense()

                    state = self.state[p]
                    if len(state) == 0:
                        state["step"] = 0
                        state["exp_avg"] = torch.zeros_like(p.data)
                        state["exp_avg_sq"] = torch.zeros_like(p.data)

                    exp_avg = state["exp_avg"]
                    exp_avg_sq = state["exp_avg_sq"]
                    state["step"] += 1

                    if weight_decay != 0.0:
                        p.data.mul_(1.0 - lr * weight_decay)

                    exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)

                    bias_correction1 = 1.0 - beta1 ** state["step"]
                    bias_correction2 = 1.0 - beta2 ** state["step"]
                    step_size = lr / bias_correction1
                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)

                    p.data.addcdiv_(exp_avg, denom, value=-step_size)

    return loss

Execute single Riemannian optimization step.

class SVDAttentionCallback (log_attention_stats: bool = True)

Expand source code

class SVDAttentionCallback(TrainerCallback):
    """Logs SVD compression statistics."""

    def __init__(self, log_attention_stats: bool = True) -> None:
        self.log_attention_stats = log_attention_stats

    def on_log(self, args, state, control, logs=None, model=None, **kwargs) -> None:
        """On log."""
        if not self.log_attention_stats or logs is None or model is None:
            return
        svd_layers = 0
        compressed_params = 0
        total_params = 0
        for module in model.modules():
            total_params += sum(p.numel() for p in module.parameters() if p.requires_grad)
            name = module.__class__.__name__
            if name in [
                "SVDLinearLayer",
                "CompressedAttentionLayer",
                "FactorisedLinear",
            ]:
                svd_layers += 1
                compressed_params += sum(p.numel() for p in module.parameters())
        if svd_layers > 0:
            logs["svd_layers"] = svd_layers
            logs["compression_ratio"] = compressed_params / total_params if total_params > 0 else 0

Logs SVD compression statistics.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_log(self, args, state, control, logs=None, model=None, **kwargs) ‑> None

Expand source code

def on_log(self, args, state, control, logs=None, model=None, **kwargs) -> None:
    """On log."""
    if not self.log_attention_stats or logs is None or model is None:
        return
    svd_layers = 0
    compressed_params = 0
    total_params = 0
    for module in model.modules():
        total_params += sum(p.numel() for p in module.parameters() if p.requires_grad)
        name = module.__class__.__name__
        if name in [
            "SVDLinearLayer",
            "CompressedAttentionLayer",
            "FactorisedLinear",
        ]:
            svd_layers += 1
            compressed_params += sum(p.numel() for p in module.parameters())
    if svd_layers > 0:
        logs["svd_layers"] = svd_layers
        logs["compression_ratio"] = compressed_params / total_params if total_params > 0 else 0

On log.

class TimeCallback

Expand source code

class TimeCallback(TrainerCallback):
    """Logs elapsed time."""

    def __init__(self) -> None:
        self.start_time = None

    def on_train_begin(self, args, state, control, **kwargs) -> None:
        """On train begin."""
        self.start_time = time.time()

    def on_log(self, args, state, control, logs=None, **kwargs) -> None:
        """On log."""
        if logs is not None and self.start_time is not None:
            logs["elapsed_time"] = time.time() - self.start_time

Logs elapsed time.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_log(self, args, state, control, logs=None, **kwargs) ‑> None

Expand source code

def on_log(self, args, state, control, logs=None, **kwargs) -> None:
    """On log."""
    if logs is not None and self.start_time is not None:
        logs["elapsed_time"] = time.time() - self.start_time

On log.

def on_train_begin(self, args, state, control, **kwargs) ‑> None

Expand source code

def on_train_begin(self, args, state, control, **kwargs) -> None:
    """On train begin."""
    self.start_time = time.time()

On train begin.

class TrainerRegistry

Expand source code

class TrainerRegistry:
    """Registry for trainer classes."""

    _registry: dict[str, type[BaseTrainer]] = {}

    @classmethod
    def register(cls, name: str) -> Any:
        """Register."""

        def decorator(klass: type[BaseTrainer]) -> type[BaseTrainer]:
            """Decorator."""
            if not issubclass(klass, BaseTrainer):
                raise TypeError(f"{klass.__name__} must inherit from BaseTrainer")
            cls._registry[name] = klass
            logger.info(f"Registered trainer: {name} -> {klass.__name__}")
            return klass

        return decorator

    @classmethod
    def get(cls, name: str) -> type[BaseTrainer]:
        """Get."""
        if name not in cls._registry:
            raise ValueError(f"Unknown trainer '{name}'. Available: {list(cls._registry.keys())}")
        return cls._registry[name]

    @classmethod
    def list(cls) -> list:
        """List."""
        return list(cls._registry.keys())

    @classmethod
    def create(cls, name: str, **kwargs) -> BaseTrainer:
        """Create."""
        return cls.get(name)(**kwargs)

Registry for trainer classes.

Static methods

def create(name: str, **kwargs) ‑> BaseTrainer

Create.

def get(name: str) ‑> type[BaseTrainer]

Get.

def list() ‑> TrainerRegistry.list() at 0x12f024bf0>)>

List.

def register(name: str) ‑> Any

Register.