ViperEkura
0546331637
- GradientCheckpointingCallback.on_train_begin returns early on empty module list - previously logged "Gradient checkpointing enabled" even when checkpointing was inactive, misleading profiling
378 lines
12 KiB
Python
378 lines
12 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.utils.checkpoint import checkpoint as torch_checkpoint
|
|
from tqdm import tqdm
|
|
|
|
from astrai.factory import BaseFactory
|
|
from astrai.parallel import only_on_rank
|
|
from astrai.parallel.setup import get_current_device
|
|
from astrai.serialization import Checkpoint
|
|
from astrai.trainer.metric_util import (
|
|
ctx_get_grad_norm,
|
|
ctx_get_grad_snr,
|
|
ctx_get_loss,
|
|
ctx_get_lr,
|
|
ctx_get_moe_metric,
|
|
ctx_get_val_loss,
|
|
)
|
|
from astrai.trainer.train_context import TrainContext
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@runtime_checkable
|
|
class TrainCallback(Protocol):
|
|
"""
|
|
Callback interface for trainer.
|
|
"""
|
|
|
|
def on_train_begin(self, context: TrainContext):
|
|
"""Called at the beginning of training."""
|
|
|
|
def on_train_end(self, context: TrainContext):
|
|
"""Called at the end of training."""
|
|
|
|
def on_epoch_begin(self, context: TrainContext):
|
|
"""Called at the beginning of each epoch."""
|
|
|
|
def on_epoch_end(self, context: TrainContext):
|
|
"""Called at the end of each epoch."""
|
|
|
|
def on_batch_begin(self, context: TrainContext):
|
|
"""Called at the beginning of each batch."""
|
|
|
|
def on_batch_end(self, context: TrainContext):
|
|
"""Called at the end of each batch."""
|
|
|
|
def on_optimizer_step(self, context: TrainContext):
|
|
"""Called on every optimizer step (sync step only)."""
|
|
|
|
def on_error(self, context: TrainContext):
|
|
"""Called when an error occurs during training."""
|
|
|
|
|
|
class CallbackFactory(BaseFactory[TrainCallback]):
|
|
"""Factory for registering and creating training callbacks.
|
|
|
|
Example:
|
|
@CallbackFactory.register("my_callback")
|
|
class MyCallback(TrainCallback):
|
|
...
|
|
|
|
callback = CallbackFactory.create("my_callback", **kwargs)
|
|
"""
|
|
|
|
|
|
@CallbackFactory.register("gradient_clipping")
|
|
class GradientClippingCallback(TrainCallback):
|
|
"""
|
|
Gradient clipping callback for trainer.
|
|
"""
|
|
|
|
def __init__(self, max_grad_norm: float):
|
|
self.max_grad_norm = max_grad_norm
|
|
|
|
def on_optimizer_step(self, context: TrainContext):
|
|
context.grad_norm = context.executor.clip_grad_norm(
|
|
context.model, self.max_grad_norm
|
|
)
|
|
|
|
|
|
@CallbackFactory.register("gradient_checkpointing")
|
|
class GradientCheckpointingCallback(TrainCallback):
|
|
"""
|
|
Activation checkpointing callback — trades compute for memory
|
|
by recomputing specified module activations during the backward pass.
|
|
|
|
Args:
|
|
modules: Module types to apply checkpointing to.
|
|
"""
|
|
|
|
def __init__(self, modules: Optional[List[type]] = None):
|
|
self.modules = tuple(modules) if modules else ()
|
|
|
|
def _enable(self, module: nn.Module):
|
|
if self.modules and isinstance(module, self.modules):
|
|
fn = module.forward
|
|
module._original_forward = fn
|
|
module.forward = lambda *a, **kw: torch_checkpoint(
|
|
fn, *a, use_reentrant=False, **kw
|
|
)
|
|
|
|
@staticmethod
|
|
def _disable(module: nn.Module):
|
|
if hasattr(module, "_original_forward"):
|
|
module.forward = module._original_forward
|
|
del module._original_forward
|
|
|
|
def on_train_begin(self, context: TrainContext):
|
|
if not self.modules:
|
|
return
|
|
context.model.apply(self._enable)
|
|
logger.info("Gradient checkpointing enabled")
|
|
|
|
def on_train_end(self, context: TrainContext):
|
|
context.model.apply(self._disable)
|
|
|
|
|
|
@CallbackFactory.register("checkpoint")
|
|
class CheckpointCallback(TrainCallback):
|
|
"""
|
|
Checkpoint callback for trainer.
|
|
"""
|
|
|
|
extra_keys = ("optimizer", "scheduler")
|
|
|
|
def __init__(
|
|
self,
|
|
save_dir: str,
|
|
interval: int,
|
|
weight_only: bool = False,
|
|
save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
|
|
):
|
|
self.save_dir = save_dir
|
|
self.interval = interval
|
|
self.weight_only = weight_only
|
|
self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
|
|
self.last_ckpt_step = None
|
|
|
|
def on_train_begin(self, context: TrainContext):
|
|
self.last_ckpt_step = context.optimizer_step
|
|
|
|
def _save_checkpoint(self, context: TrainContext):
|
|
self.last_ckpt_step = context.optimizer_step
|
|
|
|
with context.executor.checkpoint_context(context.model) as state_dict:
|
|
if state_dict is not None:
|
|
save_path = os.path.join(
|
|
self.save_dir,
|
|
f"epoch_{context.epoch}_step_{context.optimizer_step}",
|
|
)
|
|
extra = self.save_extra_fn(context)
|
|
meta = context.config.to_dict()
|
|
context.checkpoint = Checkpoint(
|
|
state_dict=state_dict,
|
|
epoch=context.epoch,
|
|
consumed_samples=context.consumed_samples,
|
|
config=context.model_config,
|
|
extra=extra,
|
|
meta=meta,
|
|
)
|
|
context.checkpoint.save(save_path)
|
|
|
|
def on_batch_end(self, context: TrainContext):
|
|
if context.optimizer_step - self.last_ckpt_step >= self.interval:
|
|
self._save_checkpoint(context)
|
|
|
|
def on_train_end(self, context: TrainContext):
|
|
if context.optimizer_step != self.last_ckpt_step:
|
|
self._save_checkpoint(context)
|
|
|
|
def on_error(self, context: TrainContext):
|
|
self._save_checkpoint(context)
|
|
|
|
@staticmethod
|
|
def save_extra(context: TrainContext) -> dict:
|
|
extra = {}
|
|
for name in CheckpointCallback.extra_keys:
|
|
obj = getattr(context, name, None)
|
|
if obj:
|
|
extra[name] = obj.state_dict()
|
|
return extra
|
|
|
|
|
|
@CallbackFactory.register("progress_bar")
|
|
class ProgressBarCallback(TrainCallback):
|
|
"""
|
|
Progress bar callback for trainer.
|
|
"""
|
|
|
|
def __init__(
|
|
self, num_epoch: int, log_interval: int = 100, file: Optional[IO[str]] = None
|
|
):
|
|
self.num_epoch = num_epoch
|
|
self.log_interval = log_interval
|
|
self.file = file
|
|
self.progress_bar: tqdm = None
|
|
|
|
@only_on_rank(0)
|
|
def on_epoch_begin(self, context: TrainContext):
|
|
total_steps = len(context.dataloader) // context.executor.grad_accum_steps
|
|
self.progress_bar = tqdm(
|
|
total=total_steps,
|
|
desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
|
|
dynamic_ncols=True,
|
|
file=self.file or sys.stdout,
|
|
)
|
|
|
|
@only_on_rank(0)
|
|
def on_optimizer_step(self, context: TrainContext):
|
|
postfix = {
|
|
"step": f"{context.optimizer_step:d}",
|
|
"loss": f"{context.loss:.4f}",
|
|
"lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
|
|
}
|
|
if context.grad_norm is not None:
|
|
postfix["grad_norm"] = f"{context.grad_norm:.2f}"
|
|
if context.val_loss is not None:
|
|
postfix["val_loss"] = f"{context.val_loss:.4f}"
|
|
self.progress_bar.set_postfix(postfix)
|
|
self.progress_bar.update(1)
|
|
|
|
@only_on_rank(0)
|
|
def on_epoch_end(self, context: TrainContext):
|
|
_ = context
|
|
if self.progress_bar:
|
|
self.progress_bar.close()
|
|
|
|
|
|
@CallbackFactory.register("metric")
|
|
class MetricCallback(TrainCallback):
|
|
def __init__(
|
|
self,
|
|
ckpt_dir: str,
|
|
save_interval: int,
|
|
metrics: List[str] = None,
|
|
val_step: int = 0,
|
|
):
|
|
self.last_log_flush_step = None
|
|
self.save_interval = save_interval
|
|
self.metrics = metrics or ["loss", "lr"]
|
|
self.val_step = val_step
|
|
self._next_val_step = 0
|
|
|
|
self.ckpt_dir = Path(ckpt_dir) if ckpt_dir else Path.cwd() / "checkpoint"
|
|
|
|
self.log_cache = []
|
|
|
|
self._metric_funcs = {
|
|
"loss": ctx_get_loss,
|
|
"lr": ctx_get_lr,
|
|
"val_loss": ctx_get_val_loss,
|
|
"grad_norm": ctx_get_grad_norm,
|
|
"grad_snr": ctx_get_grad_snr,
|
|
"moe_aux_loss": partial(ctx_get_moe_metric, key="aux_loss"),
|
|
"router_entropy": partial(ctx_get_moe_metric, key="router_entropy"),
|
|
"dead_expert_fraction": partial(
|
|
ctx_get_moe_metric, key="dead_expert_fraction"
|
|
),
|
|
"load_imbalance_mean": partial(
|
|
ctx_get_moe_metric, key="load_imbalance_mean"
|
|
),
|
|
"load_imbalance_max": partial(ctx_get_moe_metric, key="load_imbalance_max"),
|
|
}
|
|
|
|
def _metrics(self, context: TrainContext, names):
|
|
metrics = dict(context.metrics)
|
|
for name in names:
|
|
metric_fn = self._metric_funcs.get(name)
|
|
if metric_fn is None:
|
|
continue
|
|
value = metric_fn(context)
|
|
if value is not None:
|
|
metrics[name] = value
|
|
selected = set(context.metrics) | set(names)
|
|
selected.discard("*")
|
|
result = {name: metrics[name] for name in selected if name in metrics}
|
|
if context.world_size > 1 and dist.is_initialized() and result:
|
|
metric_names = sorted(result)
|
|
values = torch.tensor(
|
|
[result[name] for name in metric_names],
|
|
dtype=torch.float32,
|
|
device=get_current_device(),
|
|
)
|
|
dist.all_reduce(values, op=dist.ReduceOp.SUM)
|
|
values /= context.world_size
|
|
result.update(zip(metric_names, values.tolist()))
|
|
return result
|
|
|
|
@only_on_rank(0)
|
|
def _append(self, event_type: str, context: TrainContext, **extra):
|
|
entry = {
|
|
"type": event_type,
|
|
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
|
|
"epoch": context.epoch,
|
|
"step": context.optimizer_step,
|
|
"consumed_samples": context.consumed_samples,
|
|
**extra,
|
|
}
|
|
self.log_cache.append(entry)
|
|
|
|
def _run_validation(self, context: TrainContext) -> float:
|
|
context.model.eval()
|
|
|
|
total_loss = 0.0
|
|
num_batches = 0
|
|
|
|
with torch.no_grad():
|
|
for batch in context.val_dataloader:
|
|
loss_output = context.strategy(batch)
|
|
total_loss += loss_output["loss"].item()
|
|
num_batches += 1
|
|
|
|
if context.world_size > 1 and dist.is_initialized():
|
|
stats = torch.tensor(
|
|
[total_loss, float(num_batches)], device=get_current_device()
|
|
)
|
|
dist.all_reduce(stats, op=dist.ReduceOp.SUM)
|
|
avg_loss = (stats[0] / stats[1]).item()
|
|
else:
|
|
avg_loss = total_loss / max(num_batches, 1)
|
|
|
|
context.model.train()
|
|
return avg_loss
|
|
|
|
def on_train_begin(self, context: TrainContext):
|
|
self.last_log_flush_step = context.optimizer_step
|
|
|
|
@only_on_rank(0)
|
|
def _flush(self, epoch, step):
|
|
log_file = self.ckpt_dir / f"epoch_{epoch}_step_{step}" / "metric.jsonl"
|
|
log_file.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(log_file, "w") as f:
|
|
for log in self.log_cache:
|
|
f.write(json.dumps(log) + "\n")
|
|
|
|
def on_optimizer_step(self, context):
|
|
context.grad_snr_tracker.update(context.model)
|
|
|
|
if (
|
|
context.val_dataloader is not None
|
|
and self.val_step > 0
|
|
and context.optimizer_step >= self._next_val_step
|
|
):
|
|
context.val_loss = self._run_validation(context)
|
|
self._next_val_step = context.optimizer_step + self.val_step
|
|
self._append("validation", context, val_loss=context.val_loss)
|
|
|
|
step_metrics = [m for m in self.metrics if m != "val_loss"]
|
|
self._append("step", context, **self._metrics(context, step_metrics))
|
|
|
|
if context.optimizer_step - self.last_log_flush_step >= self.save_interval:
|
|
self._flush(context.epoch, context.optimizer_step)
|
|
self.last_log_flush_step = context.optimizer_step
|
|
|
|
def on_epoch_end(self, context):
|
|
self._append("epoch", context)
|
|
|
|
def on_train_end(self, context):
|
|
if (
|
|
self.last_log_flush_step is None
|
|
or context.optimizer_step != self.last_log_flush_step
|
|
):
|
|
self._flush(context.epoch, context.optimizer_step)
|
|
self.last_log_flush_step = context.optimizer_step
|
|
|
|
def on_error(self, context):
|
|
self._flush(context.epoch, context.optimizer_step)
|